芯東西(公眾號(hào):aichip001)
作者 | 陳駿達(dá)
編輯 | 心緣

全球首個(gè)大規(guī)模全光生成式AI芯片,來(lái)自中國(guó)!

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

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

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),通過(guò)三維封裝,在單芯片上集成了超過(guò)210萬(wàn)個(gè)光子神經(jīng)元

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

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

當(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ù)量級(jí)。

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

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

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

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

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

▲陳一彤(圖源:上海交通大學(xué)微信公眾號(hào))

論文鏈接:

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ù),例如圖像分類和元音識(shí)別等等,未能擴(kuò)展到生成式AI模型領(lǐng)域。

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

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

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

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

即使改變光子神經(jīng)元的尺寸,也只是改了調(diào)制方式,光斑本身仍然是連續(xù)的、理論上無(wú)限維的信號(hào),因此在全光學(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)簽依賴

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

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

▲LightGen架構(gòu)

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

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

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

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

▲LightGen實(shí)物圖

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

LightGen的另一大優(yōu)勢(shì)是無(wú)需將高分辨率圖像切分為小塊進(jìn)行處理。通過(guò)三維封裝,LightGen集成了超過(guò)210萬(wàn)個(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等相媲美。

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

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

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

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

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

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

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

研究團(tuán)隊(duì)將LightGen與最先進(jìn)的電子和光子芯片在生成任務(wù)中的性能進(jìn)行了對(duì)比。在保守估計(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ù)量級(jí)以上。若采用更高速的空間光調(diào)制器,其理論計(jì)算速度可達(dá)5.69×10? TOPS。

結(jié)語(yǔ):解決光子生成芯片三大瓶頸,未來(lái)可擴(kuò)展至更高分辨率

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

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