實(shí)施用于機(jī)器學(xué)習(xí)的光子張量核心

喬治華盛頓大學(xué)的研究人員報(bào)告了一種構(gòu)建光子張量核心的方法，該方法利用相變光子存儲(chǔ)器來(lái)實(shí)現(xiàn)神經(jīng)網(wǎng)絡(luò)(NN)。他們的新穎架構(gòu)已于上周在《 AIP應(yīng)用物理評(píng)論》(AIP Applied Physics Review)上在線報(bào)道，與傳統(tǒng)的GPU和其他張量核心設(shè)備相比，有望實(shí)現(xiàn)性能提升和功耗優(yōu)勢(shì)。雖然已經(jīng)探索了幾種光子神經(jīng)網(wǎng)絡(luò)設(shè)計(jì)，但用于執(zhí)行張量操作的光子張量核心尚未實(shí)現(xiàn)。

光子學(xué)[i]涵蓋了廣泛的光技術(shù)，涵蓋了通過(guò)混合光電技術(shù)跨越光纖的各種技術(shù)。例如，光學(xué)互連技術(shù)是提高內(nèi)存到處理器乃至處理器到處理器帶寬的重要研究領(lǐng)域。高帶寬和低功耗是光子學(xué)的吸引力。

在他們的 論文中， 用于機(jī)器學(xué)習(xí)的光子張量核心，Mario Miscuglio和Volker Sorger認(rèn)為，在異構(gòu)計(jì)算時(shí)代，基于光子的專(zhuān)用處理器具有增強(qiáng)電子系統(tǒng)的巨大潛力，并且可能在網(wǎng)絡(luò)邊緣設(shè)備以及其他設(shè)備中表現(xiàn)出色5G通信。例如，用于推理的預(yù)訓(xùn)練光子張量核心神經(jīng)網(wǎng)絡(luò)將消耗很少的功率。

Miscuglio告訴 HPCwire：“除了可以直接在光域中工作而增加的速度和帶寬之外，還可以利用光纖中傳輸?shù)男盘?hào)的固有光學(xué)特性，使用光子架構(gòu)的優(yōu)勢(shì)是執(zhí)行推理的功耗更低這對(duì)于智能光學(xué)低功耗傳感器很有用。”

廣義上講，神經(jīng)網(wǎng)絡(luò)大量使用矩陣向量乘法。毫不奇怪，在這種計(jì)算方式下，最新的GPU和TPU比CPU更好。研究人員在論文中很好地總結(jié)了挑戰(zhàn)：

“對(duì)于提供高計(jì)算靈活性的通用處理器，這些矩陣運(yùn)算是串行(即一次一次)進(jìn)行的，同時(shí)需要連續(xù)訪問(wèn)高速緩存，因此產(chǎn)生了所謂的”馮·諾依曼瓶頸”。已經(jīng)設(shè)計(jì)出了用于神經(jīng)網(wǎng)絡(luò)的專(zhuān)用架構(gòu)，例如圖形處理單元(GPU)和張量處理單元(TPU)，以減少馮·諾依曼瓶頸的影響，從而實(shí)現(xiàn)尖端的機(jī)器學(xué)習(xí)模型。這些架構(gòu)的范例是提供域特定性，例如對(duì)卷積進(jìn)行優(yōu)化或與CPU不同的是執(zhí)行并行運(yùn)算的矩陣向量乘法(MVM)，從而并行部署 脈動(dòng)算法。

“ GPU具有針對(duì)矩陣數(shù)學(xué)運(yùn)算而優(yōu)化的數(shù)千個(gè)處理核心，可提供數(shù)十至數(shù)百個(gè)TFLOPS(層級(jí)浮點(diǎn)運(yùn)算)的性能，這使GPU成為基于深度NN的AI和ML應(yīng)用程序的顯而易見(jiàn)的計(jì)算平臺(tái)。GPU和TPU相對(duì)于CPU尤其有利，但是當(dāng)用于對(duì)大型二維數(shù)據(jù)集(例如圖像)實(shí)施深度NN執(zhí)行推理時(shí)，它們非常耗電，并且需要更長(zhǎng)的計(jì)算時(shí)間(>幾十毫秒)。此外，用于不太復(fù)雜的推理任務(wù)的較小矩陣乘法[例如，標(biāo)準(zhǔn)與技術(shù)研究院數(shù)據(jù)庫(kù)(MNIST)的手寫(xiě)數(shù)字的分類(lèi)]仍然受到不可忽略的延遲的挑戰(zhàn)，

他們提出了一種在光子學(xué)中實(shí)現(xiàn)的張量核心單元，它依賴(lài)于光子復(fù)用(WDM，波分復(fù)用)??信號(hào)，“經(jīng)過(guò)濾波后，使用基于在 波導(dǎo)上構(gòu)圖的Ge 2 Sb 2 Se 5導(dǎo)線的工程化多態(tài)光子存儲(chǔ)器進(jìn)行加權(quán)后加權(quán)。通過(guò)使用鎢電極引起的焦耳熱的電熱切換，通過(guò)有選擇地改變導(dǎo)線的相(非晶/晶體)來(lái)對(duì)光子存儲(chǔ)器進(jìn)行重新編程。如果需要，可以并行(幾微秒)實(shí)現(xiàn)光子存儲(chǔ)器編程，或者，該光子張量核心可以作為具有預(yù)設(shè)置內(nèi)核矩陣的無(wú)源系統(tǒng)來(lái)運(yùn)行。”

Miscuglio說(shuō)，相變存儲(chǔ)技術(shù)是一項(xiàng)至關(guān)重要的進(jìn)步，“我們大腦中的每個(gè)神經(jīng)元同時(shí)存儲(chǔ)和處理數(shù)據(jù)。同樣，在我們的體系結(jié)構(gòu)中，我們使用的存儲(chǔ)單元可以通過(guò)電子方式寫(xiě)入，并且可以存儲(chǔ)多位權(quán)重，并且可以通過(guò)簡(jiǎn)單地讓光線與材料相互作用來(lái)光學(xué)讀取。我們的光子存儲(chǔ)器依賴(lài)寬帶透明相變材料，與其他基于更成熟的GST(鍺-銻-碲)的實(shí)現(xiàn)方式不同，其特點(diǎn)是在電信波長(zhǎng)下的非晶態(tài)損耗可忽略不計(jì)。”

“這很重要，因?yàn)樗С指畹募軜?gòu)，無(wú)需使用額外的激光源或放大器就可以解決更復(fù)雜的任務(wù)。我們還提出了一種多狀態(tài)光子存儲(chǔ)器(4位)架構(gòu)，該架構(gòu)可以使用電熱加熱器輕松擦除并寫(xiě)在芯片上。與依賴(lài)于片上或片外的繁瑣的光學(xué)寫(xiě)入/擦除的其他實(shí)現(xiàn)方式不同，所有存儲(chǔ)器均具有專(zhuān)用電路并可以并行寫(xiě)入。

Miscuglio說(shuō)，該體系結(jié)構(gòu)未映射特定的網(wǎng)絡(luò)體系結(jié)構(gòu)，而是神經(jīng)網(wǎng)絡(luò)的更通用的加速器。利用其模塊化架構(gòu)，可以“直接將光子TPU用于一系列操作，包括但不限于矩陣矩陣乘法，例如矢量矩陣乘法，卷積。這些代數(shù)運(yùn)算是許多復(fù)雜的科學(xué)和社會(huì)問(wèn)題的關(guān)鍵運(yùn)算。”

“我們認(rèn)為，從長(zhǎng)期來(lái)看，數(shù)據(jù)中心將從這種架構(gòu)中受益匪淺，因?yàn)樗鼈兲幚淼脑S多信息已經(jīng)在光域中。我們認(rèn)為它不會(huì)取代超級(jí)計(jì)算機(jī)，但可用作預(yù)處理單元，可與超級(jí)計(jì)算機(jī)協(xié)同工作，處理靠近網(wǎng)絡(luò)邊緣的數(shù)據(jù)，從而對(duì)信號(hào)進(jìn)行分類(lèi)和關(guān)聯(lián)，以尋找特定的數(shù)據(jù)或模式塊，從而減少數(shù)據(jù)流量。”

在撰寫(xiě)本文時(shí)，他們已經(jīng)測(cè)試了多態(tài)低損耗光子存儲(chǔ)設(shè)備，“顯示出的性能與仿真非常吻合”。Miscuglio說(shuō)：“我們開(kāi)發(fā)了執(zhí)行4×4矩陣乘法的單光子核心架構(gòu)，目前正在開(kāi)發(fā)第一代光子張量核心。關(guān)于時(shí)間表，我們計(jì)劃在六個(gè)月至一年內(nèi)對(duì)單核進(jìn)行實(shí)驗(yàn)演示，并在未來(lái)幾年內(nèi)對(duì)功能齊全的多核張量處理器進(jìn)行演示。”