上海交大全光AI芯片登Science！比A100快百倍

芯東西（公眾號：aichip001）
作者 陳駿達(dá)
編輯 心緣

全球首個(gè)大規(guī)模全光生成式AI芯片，來自中國！

芯東西12月19日報(bào)道，今天，國際頂級學(xué)術(shù)期刊《科學(xué)（Science）》刊登了一項(xiàng)來自上海交通大學(xué)的最新研究——“All-optical synthesis chip for large-scale intelligent semantic vision generation（大規(guī)模智能語義視覺生成全光芯片）”。這項(xiàng)研究是新一代算力芯片領(lǐng)域的重大突破，首次實(shí)現(xiàn)了支持大規(guī)模語義媒體生成模型的全光計(jì)算芯片。這篇論文被Science官方選為高光論文重點(diǎn)報(bào)道。

LightGen突破了光子計(jì)算領(lǐng)域的三大難題。其團(tuán)隊(duì)基于多模光學(xué)建立了光學(xué)潛空間（Optical Latent Space, OLS），實(shí)現(xiàn)了全光學(xué)條件下的維度變化，并提出了一種不依賴標(biāo)簽的訓(xùn)練算法。同時(shí)，該團(tuán)隊(duì)設(shè)計(jì)了高度集成的超表面結(jié)構(gòu)，通過三維封裝，在單芯片上集成了超過210萬個(gè)光子神經(jīng)元。

LightGen在實(shí)驗(yàn)中實(shí)現(xiàn)了多種先進(jìn)生成式AI任務(wù)，包括無需圖像分塊的512×512分辨率語義圖像生成、三維生成、視頻語義操控、風(fēng)格遷移以及去噪處理等。

當(dāng)在特定任務(wù)中達(dá)到與Stable Diffusion、StyleGAN、NeRF和VGG-19相當(dāng)?shù)男阅軙r(shí)，LightGen的端到端計(jì)算速度、能效和計(jì)算密度分別達(dá)到了3.57×10? TOPS、6.64×102 TOPS/W和 2.62×102 TOPS/mm2，這三項(xiàng)指標(biāo)均比英偉達(dá)A100芯片高出兩個(gè)數(shù)量級。

▲LightGen與英偉達(dá)A100等芯片性能對比（圖源：論文補(bǔ)充材料）

上海交通大學(xué)集成電路學(xué)院（信息與電子工程學(xué)院）陳一彤助理教授為這篇論文的第一作者及通訊作者。她長期致力于光計(jì)算領(lǐng)域的研究，其團(tuán)隊(duì)曾提出全模擬光電芯片ACCEL，國際首次實(shí)測驗(yàn)證了復(fù)雜智能任務(wù)中光計(jì)算的系統(tǒng)級算力優(yōu)越性。2023年，她提出的PED光計(jì)算架構(gòu)，被Science子刊稱為“首個(gè)全光生成網(wǎng)絡(luò)”。

陳一彤2019年本科畢業(yè)于清華大學(xué)錢學(xué)森班，2024年獲清華大學(xué)自動化系博士學(xué)位。

▲陳一彤（圖源：上海交通大學(xué)微信公眾號）

論文鏈接：

https://www.science.org/doi/10.1126/science.adv7434

一、三大挑戰(zhàn)困住光子計(jì)算：集成規(guī)模、維度變化與訓(xùn)練算法

學(xué)界認(rèn)為，光子計(jì)算是重塑傳統(tǒng)計(jì)算范式最有潛力的方向之一。然而，目前大多數(shù)端到端的光子計(jì)算成果仍局限于判決類任務(wù)，例如圖像分類和元音識別等等，未能擴(kuò)展到生成式AI模型領(lǐng)域。

是什么限制了光子計(jì)算在生成式AI中的應(yīng)用？LightGen團(tuán)隊(duì)總結(jié)了三大原因：

首先，先進(jìn)的生成式AI模型通常需要數(shù)百萬個(gè)神經(jīng)元，以學(xué)習(xí)并嵌入足夠的信息用于高維生成。而現(xiàn)有集成光子計(jì)算芯片通常僅能集成數(shù)千甚至數(shù)十個(gè)神經(jīng)元。

如果通過數(shù)字電路對光子芯片進(jìn)行互連和復(fù)用，則會引入顯著的延遲和能耗，功耗甚至可能提升數(shù)個(gè)數(shù)量級，從而削弱光計(jì)算的優(yōu)勢。

另一個(gè)關(guān)鍵挑戰(zhàn)是模擬光學(xué)信號的維度變化。生成模型需要合適的維度變化才能學(xué)到有用特征，但在純光學(xué)系統(tǒng)中信息是以連續(xù)光斑的形式存在的，不像電子信號那樣有明確的維度。

即使改變光子神經(jīng)元的尺寸，也只是改了調(diào)制方式，光斑本身仍然是連續(xù)的、理論上無限維的信號，因此在全光學(xué)系統(tǒng)里很難精確定義或改變維度。

第三個(gè)挑戰(zhàn)是訓(xùn)練算法。現(xiàn)有的分類型光子芯片通常依賴人工定義的輸出作為真實(shí)標(biāo)簽進(jìn)行優(yōu)化。然而，生成式AI的目標(biāo)通常是輸出此前并不存在的數(shù)據(jù)，因此光子生成芯片亟需一種不依賴預(yù)定義真實(shí)標(biāo)簽的新型訓(xùn)練算法。

二、光學(xué)潛空間實(shí)現(xiàn)維度轉(zhuǎn)換，提出新算法BOGT告別標(biāo)簽依賴

為解決上述問題，陳一彤領(lǐng)銜的研究團(tuán)隊(duì)提出了LightGen。LightGen由光子編碼器、光學(xué)潛空間和光子生成器組成。

▲LightGen架構(gòu)

LightGen的整體思路其實(shí)和電子生成模型很像，但成功把“編碼—潛空間—生成”這套流程完全搬到了光學(xué)域里。

光子編碼器由一系列集成衍射超表面構(gòu)成，面積小于35mm2，用于從高維圖像域中提取低維特征。

編碼后的光場通過耦合進(jìn)入單模光纖陣列進(jìn)行特征子采樣，通過同時(shí)利用光的幅度和相位進(jìn)行編碼，使得光纖陣列中的光學(xué)潛空間含有豐富的信息。

▲LightGen實(shí)物圖

由于單模光纖只允許基模傳播，除基模外的其他模式無法耦合進(jìn)入光學(xué)潛空間，因此，光學(xué)潛空間在全光學(xué)條件下完成了模擬光斑維度的轉(zhuǎn)換，而不是像傳統(tǒng)光子系統(tǒng)那樣只是在無限維光場上做調(diào)制卻無法控制維度。

LightGen利用光斑所攜帶的復(fù)數(shù)信息來表征高維潛變量。光纖陣列中每一個(gè)單模光纖對應(yīng)一個(gè)光斑，形成具有不同幅度和相位的高斯光斑陣列。

由此，LightGen在圖像域與光學(xué)潛域之間建立了橋梁。下圖展示了在以狗臉數(shù)據(jù)訓(xùn)練的光學(xué)潛空間中均勻采樣得到的輸出結(jié)果。相近的光學(xué)潛空間值生成語義相似的圖像，而相距較遠(yuǎn)的光學(xué)潛空間值生成語義差異明顯的對象，表明光學(xué)潛空間滿足連續(xù)性和自洽性，這是數(shù)學(xué)意義上潛空間進(jìn)行合理生成所必需的特性。

光子編碼器將原始圖像嵌入到100維光學(xué)潛空間中，并通過t-SNE進(jìn)行降維可視化，其分布如下圖所示。不同毛色的狗被成功聚類，說明光學(xué)潛空間能有效嵌入圖像中的細(xì)粒度特征。除前景外，背景特征（如草地與白色背景）也能被嵌入并區(qū)分。實(shí)驗(yàn)結(jié)果表明，光學(xué)潛空間的功能與電子變分自編碼器（VAE）及β-VAE相當(dāng)。

在訓(xùn)練方面，LightGen團(tuán)隊(duì)提出了一種基于貝葉斯思想的無監(jiān)督訓(xùn)練方法（BOGT），通過引入對語義數(shù)據(jù)概率分布的建模來提升生成質(zhì)量，告別了對標(biāo)簽的依賴。

三、集成超210萬個(gè)光子神經(jīng)元，計(jì)算速度比A100高兩個(gè)數(shù)量級

論文稱，LightGen能在高分辨率彩色圖像上執(zhí)行語義生成、去噪和風(fēng)格遷移等多種任務(wù)。

下圖展示了針對動物類別的實(shí)驗(yàn)生成結(jié)果。LightGen可生成512×512分辨率的動物圖像，涵蓋不同類別、顏色、表情和背景。放大區(qū)域顯示了豐富細(xì)節(jié)，如毛發(fā)紋理和眼睛反射光。

研究團(tuán)隊(duì)還對生成圖像的特征進(jìn)行了定量分析，結(jié)果顯示其分類準(zhǔn)確率與真實(shí)測試數(shù)據(jù)相當(dāng)，表明LightGen生成的數(shù)據(jù)在整體和細(xì)節(jié)特征上均接近自然數(shù)據(jù)。

LightGen在語義去噪方面同樣表現(xiàn)突出。自由空間模式與單模光纖基模之間的信息容量差異，使得LightGen在壞匹配像素比例（PBMP）較高的前提下，仍能有效去噪。

LightGen的另一大優(yōu)勢是無需將高分辨率圖像切分為小塊進(jìn)行處理。通過三維封裝，LightGen集成了超過210萬個(gè)光子神經(jīng)元，可直接處理≥512×512分辨率的圖像。

在風(fēng)格遷移任務(wù)中，LightGen成功實(shí)現(xiàn)了印象派（梵高風(fēng)格）、金屬風(fēng)格（馬列維奇風(fēng)格）和馬賽克風(fēng)格等多種轉(zhuǎn)換，其定性與定量性能均可與VGG-19、StyleGAN和StyleID等相媲美。

與基于分塊的方法相比，LightGen在處理全局結(jié)構(gòu)和連續(xù)曲線方面具有顯著優(yōu)勢，避免了塊間不連續(xù)和結(jié)構(gòu)破壞的問題。

LightGen還可從二維擴(kuò)展至三維生成與語義操控，用于生成高分辨率圖像和視頻。

在室內(nèi)場景中，LightGen可生成包含多種家具的256×256分辨率圖像。通過二維投影訓(xùn)練，LightGen無監(jiān)督地學(xué)習(xí)了物體的三維結(jié)構(gòu)及其在光學(xué)潛空間中的表示，性能可與NeRF相當(dāng)。通過操控光學(xué)潛空間，LightGen可生成不同風(fēng)格和視角的三維家具，在定性和PSNR指標(biāo)上均與NeRF相當(dāng)。

此外，光學(xué)潛空間中的潛變量直接對應(yīng)生成對象的物理特征，使得語義操控具有良好的可解釋性。實(shí)驗(yàn)表明，LightGen能在不干擾其他部分的情況下，精確操控目標(biāo)特征。

研究團(tuán)隊(duì)將LightGen與最先進(jìn)的電子和光子芯片在生成任務(wù)中的性能進(jìn)行了對比。在保守估計(jì)下，LightGen在實(shí)驗(yàn)中實(shí)現(xiàn)了3.57×10? TOPS的計(jì)算速度、6.64×102 TOPS/W的能效和2.62×102 TOPS/mm2的計(jì)算密度，均比NVIDIA A100高出兩個(gè)數(shù)量級以上。若采用更高速的空間光調(diào)制器，其理論計(jì)算速度可達(dá)5.69×10? TOPS。

結(jié)語：解決光子生成芯片三大瓶頸，未來可擴(kuò)展至更高分辨率

LightGen同時(shí)突破了光子生成芯片的三個(gè)關(guān)鍵瓶頸：集成規(guī)模、全光學(xué)維度變換和訓(xùn)練算法。三維封裝方案解決了集成規(guī)模問題，光學(xué)潛空間實(shí)現(xiàn)了全光學(xué)維度變換，而BOGT解決了對標(biāo)簽的依賴。

研究團(tuán)隊(duì)認(rèn)為，通過直接增加超表面的像素?cái)?shù)和層數(shù)、提高光學(xué)潛空間的維度或以分塊生成的方式，LightGen可以進(jìn)一步擴(kuò)展到更高的處理分辨率。增加超表面規(guī)模通常會提升系統(tǒng)整體計(jì)算速度，因?yàn)榭偛僮鲾?shù)增加，但這也需要更多的入射能量來維持足夠的計(jì)算精度。