劉金國 許開來

1) (哈佛大學(xué)物理系, 坎布里奇 02138)

2) (斯坦福大學(xué), 斯坦福 94305)

自動(dòng)微分是利用計(jì)算機(jī)自動(dòng)化求導(dǎo)的技術(shù), 最近幾十年因?yàn)樗跈C(jī)器學(xué)習(xí)研究中的應(yīng)用而被很多人了解. 如今越來越多的科學(xué)工作者意識(shí)到, 高效、自動(dòng)化的求導(dǎo)可以為很多科學(xué)問題的求解提供新的思路, 其中自動(dòng)微分在物理模擬中的應(yīng)用尤為重要, 而且具有挑戰(zhàn)性. 物理系統(tǒng)的可微分模擬可以幫助解決混沌理論、電磁學(xué)、地震學(xué)、海洋學(xué)等領(lǐng)域中的很多重要問題, 但又因?yàn)槠鋵τ?jì)算時(shí)間和空間的苛刻要求而對自動(dòng)微分技術(shù)本身提出了挑戰(zhàn). 本文回顧了將自動(dòng)微分技術(shù)運(yùn)用到物理模擬中的幾種方法, 并橫向?qū)Ρ人鼈冊谟?jì)算時(shí)間、空間和精度方面的優(yōu)勢和劣勢. 這些自動(dòng)微分技術(shù)包括伴隨狀態(tài)法, 前向自動(dòng)微分, 后向自動(dòng)微分, 以及可逆計(jì)算自動(dòng)微分.

1 引 言

自動(dòng)微分[1]是指自動(dòng)獲取一段計(jì)算機(jī)程序?qū)?shù)的技術(shù), 很多人對它的了解源自它在機(jī)器學(xué)習(xí)中的成功應(yīng)用, 人們可以用它優(yōu)化包含數(shù)千億參數(shù)的神經(jīng)元網(wǎng)絡(luò)[2]. 與很多人印象不同的是, 自動(dòng)微分其實(shí)是個(gè)很古老的技術(shù). Nolan[3]早在他1953年的博士論文中就提出過計(jì)算機(jī)自動(dòng)化求導(dǎo)的構(gòu)想, 后來針對這一構(gòu)想又出現(xiàn)了兩種不同的實(shí)踐, 分別是1964年由 Wengert[4]實(shí)現(xiàn)的前向自動(dòng)微分和1970年 Linnainmaa[5]實(shí)現(xiàn)的后向自動(dòng)微分. 而最近十幾年, 由于后向自動(dòng)微分在機(jī)器學(xué)習(xí)中的廣泛應(yīng)用, 相關(guān)技術(shù)越來越成熟并在科學(xué)計(jì)算中也得到了越來越多的應(yīng)用. 科學(xué)家們利用方便的、自動(dòng)化的計(jì)算機(jī)輔助求導(dǎo)解決了很多重要的物理問題, 其中包括幫助變分蒙特卡羅設(shè)計(jì)更加通用的多體物理波函數(shù)[6-9], 加速變分量子算法的模擬[10], 拓展變分張量網(wǎng)絡(luò)算法[11], 以及求解自旋玻璃基態(tài)構(gòu)型[12]等.

對物理模擬過程的自動(dòng)微分是自動(dòng)微分重要的應(yīng)用之一, 其中最大的技術(shù)挑戰(zhàn)是對電磁學(xué)、海洋學(xué)[13]和地震學(xué)[14,15]等問題中最核心的微分方程求解過程的自動(dòng)微分. 這些微分方程常見的求解方法是先將問題的時(shí)空坐標(biāo)離散化, 并以數(shù)值積分的形式完成求解. 要得到精確的結(jié)果, 離散化后的網(wǎng)格需要設(shè)計(jì)得很稠密, 從而對存儲(chǔ)空間和計(jì)算時(shí)間的需求巨大. 與此同時(shí), 在自動(dòng)微分計(jì)算過程中,機(jī)器學(xué)習(xí)中主流的后向自動(dòng)微分庫, 如 Jax[16]和PyTorch[17], 需要很多額外的空間來緩存中間結(jié)果, 導(dǎo)致人們經(jīng)常無法直接用它們對具有一定規(guī)模的實(shí)際物理模擬問題求導(dǎo). 一種傳統(tǒng)的解決常微分方程求導(dǎo)的方案叫作伴隨狀態(tài)法[18,19], 它假設(shè)了在較短時(shí)間內(nèi)積分器可逆, 并通過逆向積分來幫助自動(dòng)微分回溯狀態(tài). 事實(shí)上, 除了 Leapfrog 等少數(shù)辛積分器外, 大多數(shù)的積分器并不能保證嚴(yán)格的時(shí)間反演對稱, 所以伴隨狀態(tài)法往往存在與積分步長有關(guān)的系統(tǒng)性誤差. 后來, 有人把機(jī)器學(xué)習(xí)中的最優(yōu)檢查點(diǎn)算法帶到了物理模擬的反向自動(dòng)微分中[14],到了嚴(yán)格求導(dǎo). 也有一些人直接利用包含最優(yōu)檢查點(diǎn)算法[20,21]的自動(dòng)微分編譯器, 如 Tapenade[22],OpenAD[23]等, 編譯模擬代碼實(shí)現(xiàn)類似的時(shí)間和空間的權(quán)衡. 在最近幾年, 可逆計(jì)算開始被用作一種新的通用的微分的方案[24], 也可以做到對數(shù)的額外空間開銷. 可逆計(jì)算提供了傳統(tǒng)自動(dòng)微分框架所不具有的內(nèi)存控制的靈活性, 因此可以被用來微分 GPU 核函數(shù)、復(fù)雜的控制流等.

本文將會(huì)回顧共軛態(tài)方法, 前向自動(dòng)微分以及基于最優(yōu)檢查點(diǎn)算法和可逆計(jì)算的后向自動(dòng)微分等自動(dòng)微分方法在物理模擬問題中的應(yīng)用, 并對比不同方法的優(yōu)劣以及適用的場景. 第2節(jié)介紹了幾種自動(dòng)微分方法的基本理論. 第3節(jié)介紹了不同自動(dòng)微分技術(shù)在波的傳播模擬中的應(yīng)用. 關(guān)于如何在后向自動(dòng)微分中權(quán)衡程序的運(yùn)行時(shí)間和空間的理論, 雖然重要, 但受篇幅限制我們將其放在附錄.

2 自動(dòng)微分方法

物理模擬過程的常見求解方案是先將偏微分方程的空間部分離散化并做差分處理[26], 從而將其轉(zhuǎn)換為對時(shí)間的常微分方程

其中s為狀態(tài),t為時(shí)間,θ為控制參數(shù). 一個(gè)常微分方程的數(shù)值求解器會(huì)把時(shí)間維度離散化, 做n步疊代求解, 每步僅做從時(shí)刻ti到時(shí)刻ti+1=ti+Δt的演化,

其中si為完成第i步積分后的狀態(tài). 為了簡便, 下文將單步運(yùn)算簡記為 O DEStep(si). 最后我們還會(huì)定義一個(gè)損失函數(shù)L=loss(sn). 自動(dòng)微分的目標(biāo)則是求解損失量對輸入狀態(tài)和控制參數(shù)的導(dǎo)數(shù)和

2.1 共軛態(tài)方法

共軛態(tài)方法[18,19]是專門針對常微分方程的反向傳播方法. 在研究中, 人們發(fā)現(xiàn)積分過程的導(dǎo)數(shù)的反向傳播同樣是一個(gè)積分過程, 只不過方向相反. 于是人們通過構(gòu)造一個(gè)可以同時(shí)回溯函數(shù)值和回傳導(dǎo)數(shù)的拓展函數(shù), 以對拓展函數(shù)的逆向積分的形式完成導(dǎo)數(shù)的計(jì)算, 如表1中算法1所示. 該算法的描述來自文獻(xiàn)[19], 其中可以找到詳細(xì)的推導(dǎo)過程, 這里對原算法中的符號(hào)做了替換以方便讀者理解. 算法中的為局域?qū)?shù), 它們可以通過手動(dòng)推導(dǎo)或借助其他自動(dòng)微分庫來實(shí)現(xiàn).該方案在積分器嚴(yán)格可逆的時(shí)候梯度也嚴(yán)格, 而當(dāng)積分器反向積分誤差不可忽略時(shí), 則需要額外的處理保證誤差在可控范圍, 這一點(diǎn)我們會(huì)在隨后的例子中涉及.

表1 伴隨狀態(tài)法(算法1)Table 1. Continuous adjoint state method (Algorithm 1).

2.2 前向自動(dòng)微分

顧名思義, 前向自動(dòng)微分是指向前(指與程序運(yùn)行方向相同)傳播導(dǎo)數(shù). 這里導(dǎo)數(shù)傳播的規(guī)則和數(shù)學(xué)分析中的無窮小量的運(yùn)算規(guī)則有關(guān). 數(shù)學(xué)中,在對一個(gè)輸入變量p求導(dǎo)時(shí), 會(huì)讓它攜帶一個(gè)無窮小量 dp, 并通過對這個(gè)無窮小量的運(yùn)算完成對程序的求導(dǎo). 例如當(dāng)作用函數(shù)f時(shí), 會(huì)有如下鏈?zhǔn)椒▌t:

其中→x為函數(shù)f的輸入?yún)?shù)的集合, 可以包括p本身.為局域雅可比矩陣, 前向傳播中我們將它與梯度矢量相乘得到新的梯度矢量. 實(shí)際程序?qū)崿F(xiàn)中, 這個(gè)局域雅可比矩陣并不需要構(gòu)造出來. 考慮到任何程序都具有可拆分為基礎(chǔ)指令的特點(diǎn), 人們把程序拆解為基礎(chǔ)標(biāo)量指令, 并在這些基礎(chǔ)指令上通過代碼變換或算符重載的方式實(shí)現(xiàn)梯度矢量的計(jì)算. 以標(biāo)量的乘法為例, 變量會(huì)同時(shí)記錄它的數(shù)值和一階小量的系數(shù) (v,v˙) , 其中人們重新定義它的基本運(yùn)算規(guī)則如下:

使得其在計(jì)算同時(shí)更新一階小量的系數(shù), 而自動(dòng)微分所要做的就是將所有的運(yùn)算指令都引入小量的系數(shù)并做上述運(yùn)算規(guī)則的替換. 簡單的運(yùn)算規(guī)則的替換對于人類來說尚可手動(dòng), 但真實(shí)的程序可能會(huì)包含數(shù)以億計(jì)的這樣的基礎(chǔ)操作, 雖然結(jié)果依然是解析的, 但是人們很難再通過人力得到具體的導(dǎo)數(shù)表達(dá)式, 而計(jì)算機(jī)恰恰很擅長這樣繁瑣但是規(guī)則簡單的任務(wù). 如圖1(a)所示, 在求解常微分方程中,單步前向自動(dòng)微分可以形式化的表達(dá)為

圖1 物理模擬程序中對狀態(tài)變量梯度的(a)前向傳播和(b)后向傳播過程. 其中線條上的箭頭代表運(yùn)算的方向, 圓圈為后向自動(dòng)微分中緩存的變量Fig. 1. The (a) forward mode and (b) backward mode gradient propagation for state variables in a physics simulation. The arrows are computation directions, and circles are cached states in the backward mode automatic differentiation.

這里為了簡潔略去了時(shí)間等常數(shù)參量. 由于狀態(tài)si,si+1和控制參數(shù)θ均可包含多個(gè)變量, 上述偏微分均解釋為雅可比行列式. 一般前向自動(dòng)微分在一次前向運(yùn)行中只對一個(gè)或若干個(gè)變量求導(dǎo), 因此計(jì)算空間會(huì)隨著一次求導(dǎo)的變量數(shù)目線性增加. 無論是一次求導(dǎo)多少個(gè)變量, 前向自動(dòng)微分求導(dǎo)的時(shí)間都會(huì)隨著需求導(dǎo)的變量的數(shù)目線性增長, 這是限制前向自動(dòng)微分應(yīng)用場景的最主要因素.

2.3 后向自動(dòng)微分

后向自動(dòng)微分與前向自動(dòng)微分梯度傳播方向相反, 它解決了前向自動(dòng)微分中計(jì)算開銷隨著需要求導(dǎo)的變量數(shù)目線性增長的問題. 后向自動(dòng)微分包括正向計(jì)算和梯度反向傳播兩個(gè)過程. 正向計(jì)算過程中, 程序進(jìn)行普通的計(jì)算并獲取所需的運(yùn)行時(shí)信息, 最后得到一個(gè)標(biāo)量損失L. 梯度回傳的過程是計(jì)算導(dǎo)數(shù)的過程, 可表示為更新變量的梯度的過程. 從出發(fā), 梯度反向傳播對應(yīng)如下鏈?zhǔn)椒▌t:

為了實(shí)現(xiàn)該鏈?zhǔn)椒▌t, 軟件包設(shè)計(jì)者們一般會(huì)在軟件中規(guī)定一類可微分的函數(shù)集合, 叫作原子函數(shù){fp}, 并定義了梯度在原子函數(shù)上的回傳規(guī)則為伴隨變量的更新規(guī)則,

這里也同樣略去了時(shí)間等參量, push和pop分別是對全局堆棧 Σ 的入棧和出棧操作. 雖然導(dǎo)數(shù)回傳的計(jì)算復(fù)雜度與需要求導(dǎo)的變量數(shù)目無關(guān), 但后向自動(dòng)微分向堆棧中存儲(chǔ)數(shù)據(jù)帶來了正比于計(jì)算步驟數(shù)的額外空間開銷. 如何設(shè)計(jì)一個(gè)可以對計(jì)算狀態(tài)逆序的訪問的算法, 使得緩存帶來的時(shí)間和空間開銷更小, 是反向自動(dòng)微分設(shè)計(jì)復(fù)雜性的來源, 也被稱作“時(shí)間與空間的權(quán)衡”問題. 實(shí)際應(yīng)用中, 人們可以用檢查點(diǎn)[20]或者反向計(jì)算[24]技巧來避免緩存全部中間結(jié)果. 如圖2(a)所示的檢查點(diǎn)方案中, 正向計(jì)算中程序可以選擇性的緩存(黑色圓點(diǎn))或者不緩存部分狀態(tài)(空心圓點(diǎn)), 這些被緩存的狀態(tài)被稱作檢查點(diǎn). 在圖中下方的反向傳播過程中, 當(dāng)需要的數(shù)據(jù)沒有被緩存時(shí), 程序會(huì)從最近的檢查點(diǎn)出發(fā)重新計(jì)算該數(shù)據(jù)(下方步驟2). 圖2(b)所示的反計(jì)算方案是可逆計(jì)算中常用的避免緩存的方案. 當(dāng)運(yùn)算在數(shù)學(xué)層面不可逆, 可逆計(jì)算要求保留輸入狀態(tài)以使其在程序?qū)用婵赡? 因此默認(rèn)也需要緩存每步的結(jié)果. 而反計(jì)算通過反向運(yùn)行(上方步驟5)清零已分配的內(nèi)存的方式并將空間資源歸還給系統(tǒng)控. 在反向傳播過程中, 執(zhí)行順序和前向鏡像對稱, 每個(gè)指令變成原指令的逆, 因此可以自然地獲得運(yùn)行時(shí)的狀態(tài)信息. 不論是檢查點(diǎn)方案還是可逆計(jì)算方案, 節(jié)省內(nèi)存都需要消耗更多的時(shí)間, 那么如何權(quán)衡兩者的額外開銷才是最好呢?我們在附錄中引入一個(gè)簡單的理論模型——鵝卵石游戲并詳細(xì)討論了如何在檢查點(diǎn)方案中用Treeverse算法[20](也叫作最優(yōu)檢查點(diǎn)算法)實(shí)現(xiàn)僅需對數(shù)的額外時(shí)間和空間逆向遍歷狀態(tài), 以及如何在可逆計(jì)算中用Bennett算法[25]實(shí)現(xiàn)對數(shù)的額外空間和多項(xiàng)式的額外時(shí)間實(shí)現(xiàn)逆向遍歷狀態(tài), 這里僅把基本結(jié)論列于表2中.

表2 不同自動(dòng)微分方案中, 時(shí)間與空間復(fù)雜度的關(guān)系. 這里伴隨狀態(tài)法沒有考慮了緩存部分中間結(jié)果以保證反向積分的正確性, 考慮之后應(yīng)為 O (TS).前向自動(dòng)微分時(shí)間復(fù)雜度中的N代表了需要求導(dǎo)的參數(shù)個(gè)數(shù). 可逆計(jì)算中的多項(xiàng)式時(shí)間復(fù)雜度中的?＞0Table 2. The time and space complexities in different automatic differentiation schemes. Here, the adjoint state method does not consider the effect of round off error into consideration. The parameter N in the forward mode autodiff represent the number of differentiable input parameters. The ? in the reverse computing is greater than 0.

圖2 使用(a)檢查點(diǎn)方案和(b)反計(jì)算方案避免緩存全部中間狀態(tài). 圖中黑箭頭為正常計(jì)算過程, 紅箭頭為梯度回傳過程, 藍(lán)箭頭為梯度回傳和反向計(jì)算的復(fù)合過程, 箭頭上的數(shù)字代表了執(zhí)行順序. 黑色和白色的圓點(diǎn)為被緩存和未被緩存 (或通過反計(jì)算消除) 的狀態(tài)Fig. 2. Using (a) checkpointing scheme and (b) reverse computing scheme to avoid caching all intermediate states. The black arrows are regular forward computing, red arrows are gradient back propagation, and blue arrows are reverse computing. The numbers above the arrows are the execution order. Black and white circles represent cached states and not cached states (or those states deallocated in reverse computing) respecitively.

3 自動(dòng)微分在物理模擬中的應(yīng)用

本節(jié)有兩個(gè)案例, 其一是用前向自動(dòng)微分和伴隨狀態(tài)法求解參數(shù)較少的洛倫茨系統(tǒng)模擬的導(dǎo)數(shù),其二是用最基于最優(yōu)檢查點(diǎn)算法和可逆計(jì)算的后向微分對步驟較多且內(nèi)存消耗巨大的地震學(xué)模擬的求導(dǎo). 這些例子的源代碼可在Github倉庫中找到: https://github.com/GiggleLiu/WuLiXueBao.

3.1 洛倫茨方程求解

洛倫茨系統(tǒng)[27]是人們研究混沌問題的經(jīng)典模型, 它描述了一個(gè)定義在三維空間中的動(dòng)力學(xué)

其中,σ,ρ和β為三個(gè)控制參數(shù). 該系統(tǒng)的含時(shí)演化的曲線如圖3(b)所示.ρ＞1 時(shí), 系統(tǒng)有兩個(gè)吸引子[28], 但僅當(dāng)時(shí)才會(huì)出現(xiàn)圖3(b)中橘色曲線所示的粒子穩(wěn)定地圍繞其中一個(gè)吸引子運(yùn)動(dòng)的情況, 這時(shí)系統(tǒng)較為穩(wěn)定并表現(xiàn)出對初值較為不敏感的特點(diǎn). 末了位置坐標(biāo)對控制參數(shù)和初始坐標(biāo)的導(dǎo)數(shù)反映了末態(tài)對控制參數(shù)和初始位置的敏感度, 在一定程度上反映了系統(tǒng)出現(xiàn)了混沌現(xiàn)象. 在數(shù)值模擬中, 我們用4階龍格庫塔方法對時(shí)間積分并得到末了位置, 模擬中固定初始位置為 (x0,y0,z0)=(1,0,0) , 控制參數(shù)為β=8/3 , 積分時(shí)間區(qū)間為 [ 0,T=30] 以及積分步長為 3×10-3.由于該過程所含參數(shù)僅有6個(gè), 包括初始位置的三個(gè)坐標(biāo) (x0,y0,z0) 和三個(gè)控制參數(shù) (σ,ρ,β) , 因此用前向自動(dòng)微分工具 ForwardDiff[29]求導(dǎo)比后向自動(dòng)微分有很大優(yōu)勢. 我們把導(dǎo)數(shù)的絕對值的平均與初始ρ,σ的關(guān)系畫在圖3中. 可以看出只有在理論預(yù)言的黑線下方才會(huì)有較小的導(dǎo)數(shù), 說明穩(wěn)定吸引子參數(shù)下系統(tǒng)動(dòng)力學(xué)的確對初值依賴性較低.表3對比了不同方法的時(shí)間和空間的消耗, 可以看到前向自動(dòng)微分所用的時(shí)間僅為原代碼的不到2倍. 這里的高效來自 ForwardDiff 中允許一次對多個(gè)變量求導(dǎo), 代價(jià)是用了正比于參數(shù)數(shù)目倍數(shù)的空間. 該技術(shù)雖然沒有改變隨著輸入變量數(shù)目增加, 計(jì)算復(fù)雜度線性增加的本質(zhì), 但是減少了線性部分的系數(shù), 對處理實(shí)際問題很有幫助. 基于可逆計(jì)算的后向自動(dòng)微分庫NiLang[24]需要的計(jì)算時(shí)間為原計(jì)算的約3.5倍, 其中包含了前向計(jì)算和反向傳播過程, 因此這個(gè)速度并不算差. 但是由于龍格庫塔積分器不可逆, 在不利用額外的計(jì)算時(shí)間交換空間的情況下, 需要緩存每一步的計(jì)算以保證可逆性, 因此要求有 1 04倍于狀態(tài)大小的緩存空間.好在該問題的單個(gè)狀態(tài)空間僅有三個(gè)維度, 不做任何緩存求導(dǎo)仍然很容易. 表中最后一列的伴隨狀態(tài)法的單步求導(dǎo)利用了 NiLang 的自動(dòng)微分, 雖然伴隨狀態(tài)法在理論上可以做到無額外內(nèi)存消耗, 但是會(huì)引入由于積分器不可逆帶來的系統(tǒng)誤差, 這對于研究混沌問題非常致命. 我們以前向自動(dòng)微分的導(dǎo)數(shù)是嚴(yán)格的導(dǎo)數(shù)作為基準(zhǔn), 把伴隨狀態(tài)法求得的導(dǎo)數(shù)的l2相對誤差與積分步長的關(guān)系繪制于圖4(a)中, 可以看出相對誤差隨著積分步長指數(shù)增加. 因此, 我們需要每隔一段積分, 就設(shè)置一個(gè)檢查點(diǎn)重新加載正確的坐標(biāo). 圖4(b)顯示檢查點(diǎn)越密, 誤差越小, 消耗的額外空間也越多. 在表格中的模擬中,我們選擇了檢查點(diǎn)步長為200, 對應(yīng)檢查點(diǎn)數(shù)目為 50, 相對誤差約為10-5.

圖4 (a) 利用共軛態(tài)方法對 ρ =27,σ=10,β=8/3 的洛倫茨系統(tǒng)求導(dǎo)時(shí), 相對 l2 誤差與積分步數(shù)的關(guān)系. 其中一個(gè)點(diǎn)代表了在該步長下, 對 1 00 個(gè)隨機(jī)初始點(diǎn)計(jì)算得到的中位數(shù), 缺失的數(shù)據(jù)代表該處出現(xiàn)數(shù)值溢出. (b) 每隔固定步長設(shè)置檢查點(diǎn)后,104步積分的誤差Fig. 4. (a) Relative l2 error as a function of the integration step in differentiating the Lorenz equation solver at parameters ρ=27,σ=10,β=8/3 using the adjoint state method. Each data point in the graph represents a median of 1 00 random samples. Some data are missing due to the numerical overflow. (b) The same error curve of computing 1 04 steps with checkpointing to decrease the round off error.

表3 不同方法對 1 04 步洛倫茨系統(tǒng)積分過程的微分所需時(shí)間和空間, 其中空間部分以狀態(tài)數(shù)目為單位, 時(shí)間的估計(jì)通過BenchmarkTools.jl軟件包多次統(tǒng)計(jì)得出. 這里, NeuralODE中單步計(jì)算的微分由NiLang完成, 檢查點(diǎn)的步長為200步Table 3. The time and space to simulate a104 step Lorenz integration, where the space is measured in unit of single state size, and the time is obtained by repeated sample using the Julia package BenchmarkTools.jl. Here, the single step back propagation in NeuralODE is obtained with NiLang,and the checkpoint step size is 200.

圖3 (a) 洛倫茨系統(tǒng)中固定參數(shù) β =8/3 , 梯度大小與參數(shù) ρ 和 σ 的關(guān)系. 圖中顏色代表了平均梯度大小, 黑線是理論上會(huì)不會(huì)存在穩(wěn)定吸引子分界線. (b) 中的藍(lán)色和黃色的線分別對應(yīng) (a) 圖標(biāo)識(shí)的藍(lán)點(diǎn)處參數(shù) ( σ=10,ρ=27) 和黃點(diǎn)處參數(shù)(σ=10,ρ=15)對應(yīng)的動(dòng)力學(xué)模擬Fig. 3. (a) The magnitude of gradients as a function of parameters ρ and σ in a Lorenz system with β fixed to 8 /3. The black line is a theoretical predicted stable attractor phase transition line. The blue and yellow curves in (b) are simulations at parameters corresponding to the blue ( σ=10,ρ=27) and yellow ( σ=10,ρ=15) dots in (a).

3.2 波的傳播方程的微分

考慮由如下方程決定的波函數(shù)u(x1,x2,t) 在非均勻二維介質(zhì)中的傳播過程:

其中c為波在介質(zhì)中的傳播速度. 理想匹配層(PML)[30-32]是模擬波在介質(zhì)中運(yùn)動(dòng)的一種準(zhǔn)確可靠的方案, 為了在有限尺寸模擬該動(dòng)力學(xué), PML方法引入了吸收層防止邊界反射的影響. 在引入輔助場并對空間和時(shí)間進(jìn)行離散化后, 上述方程可變形為如下數(shù)值計(jì)算過程:

這里的第一項(xiàng)為近似, 因?yàn)樗雎粤嗽街薪橘|(zhì)傳播速度c的空間梯度項(xiàng)的貢獻(xiàn).ζ1和ζ2分別是x和y方向的衰減系數(shù),φx和φy分別為引入的輔助場的x和y方向的分量, Δx, Δy和 Δt分別是 空間和 時(shí)間的離散化參數(shù). 該方程的詳細(xì)推導(dǎo)可參考文獻(xiàn)[26]. PML 模擬的自動(dòng)微分在地震學(xué)中有著重要的應(yīng)用[15], 而且人們很早就認(rèn)識(shí)到檢查點(diǎn)方案可以用于地震波模擬中[14], 從而使回溯中間狀態(tài)的內(nèi)存需求大大減少.

在數(shù)值模擬中, 我們用雙精度浮點(diǎn)數(shù)模擬了1000×1000的二維格點(diǎn)上的PML方程求解, 每個(gè)狀態(tài)要存儲(chǔ)4個(gè)矩陣, 占用存儲(chǔ)空間為32 MB. 雖然前向自動(dòng)微分可以僅用常數(shù)倍空間開銷微分該程序, 但由于僅傳播速度c就包含 1 06個(gè)參數(shù), 前向自動(dòng)微分帶來的線性時(shí)間復(fù)雜度的增加是不可接受的. 同時(shí), 若不做任何內(nèi)存優(yōu)化對該程序后向自動(dòng)微分, 積分 1 04步需要存儲(chǔ)空間至少為320 G, 遠(yuǎn)超出了普通GPU的存儲(chǔ)能力. 此時(shí), 我們需要用附錄中描述的 Bennett 算法和 Treeverse 算法來節(jié)省后向自動(dòng)微分的緩存.圖5中展示了這兩種時(shí)間空間交換方案下, 實(shí)際程序在 GPU 上運(yùn)行的時(shí)間與空間的關(guān)系, 計(jì)算設(shè)備為 Nvidia Tesla V100. 純可逆計(jì)算實(shí)現(xiàn)的 Bennett 算法中, 計(jì)算梯度的部分隨著前向計(jì)算的步驟數(shù)的增加而增加, 額外開銷和理論模型幾乎一致.Bennett 算法在可逆計(jì)算的意義下是最優(yōu)的時(shí)空交換策略, 但對普通硬件并不是最優(yōu)的. Treeverse+NiLang 的方案是指用可逆計(jì)算處理單步運(yùn)算的微分, 同時(shí)用 Treeverse 算法處理步驟間的微分. 這里隨著檢查點(diǎn)數(shù)目的減少, 計(jì)算時(shí)間的減少并不明顯. 這是因?yàn)樵黾拥挠?jì)算時(shí)間是前向計(jì)算, 而這里單步后向計(jì)算梯度的時(shí)間是前向時(shí)間的20多倍,因此即使僅用5個(gè)檢查點(diǎn), 額外的增加的時(shí)間也不到1倍. 這里單步計(jì)算梯度的之所比前向計(jì)算慢很多, 是因?yàn)楫?dāng)使用 NiLang 微分并行運(yùn)行的GPU核函數(shù), 必須要避免變量的共享讀取以避免程序在后向計(jì)算中同時(shí)并行更新同一塊內(nèi)存的梯度, 這些額外的設(shè)計(jì)帶來了很多性能的損失. 與此對比, 在單線程版本的CPU上, 前向和后向的單步運(yùn)算時(shí)間差距在4倍以內(nèi). 此外, 雖然 Treeverse算法在可以做到高效的時(shí)間和空間的交換, 但由于內(nèi)存的管理需要用到系統(tǒng)內(nèi)存中的全局堆棧, 無法直接用于微分 GPU 的核函數(shù). 而可逆計(jì)算則非常適合微分這種非線性且具有一定可逆性的程序, 這也是為什么這里選擇用可逆計(jì)算對單步運(yùn)算求導(dǎo)來避免手動(dòng)求導(dǎo) GPU 核函數(shù)的麻煩.

圖5 Bennett算法和Treeverse算法應(yīng)用于PML求解過程的微分的時(shí)間與空間開銷. 其中圖中標(biāo)記的數(shù)字為函數(shù)的實(shí)際前向運(yùn)行步驟數(shù)與原始模擬步驟數(shù) ( 1 04 ) 的比值.縱軸的總時(shí)間是前向計(jì)算和后向傳播時(shí)間的總和, 橫軸的空間的數(shù)值的實(shí)際含義為檢查點(diǎn)的數(shù)目或可逆計(jì)算中的最大狀態(tài)數(shù). 在Bennett算法中, 后向運(yùn)行次數(shù)和前向運(yùn)行次數(shù)一樣, 而Treeverse算法中, 反向傳播的次數(shù)固定為104Fig. 5. The time and space cost of using Bennett algorithm and Treeverse algorithm to solve the PML equation. The annotated numbers are relative forward overhead defined as the ratio between actual forward computing steps and the original computing steps. The time in the y axis is the actual time that includes both forward computing and back propagation. The space in the x axis is defined as the maximum number of checkpoints in the checkpointing scheme or the maximum intermediate states in the reverse computing scheme. In the Bennett's algorithm, the number of forward and backward steps are the same, while in Treeverse algorithm, the number of back propagation steps is fixed to 104.

感謝王磊老師的討論以及計(jì)算設(shè)備的支持. 感謝彩云天氣CTO苑明理老師關(guān)于自動(dòng)微分關(guān)于自動(dòng)微分在天氣預(yù)報(bào)中應(yīng)用的討論. 感謝QuEra公司提供的研究贊助.

附錄A1 名詞表

伴隨變量 adjoint

核函數(shù) kernel function

原子函數(shù) primitive function

可逆計(jì)算 reversible computing

可逆編程 reversible programming

鵝卵石游戲 pebble game

伴隨狀態(tài)法 adjoint state method

自動(dòng)微分 automatic differentiation (AD)

前向自動(dòng)微分 forward mode AD

后向自動(dòng)微分 reverse mode AD

地震學(xué) seismic

檢查點(diǎn) checkpoint

拓展動(dòng)力學(xué) augmented dynamics

損失 loss

控制流 control flow

辛積分器 symplectic integrator

洛倫茨 Lorenz

龍格庫塔 Runge-Kutta

附錄A2 時(shí)間與空間的交換, 鵝卵石游戲

鵝卵石游戲是一個(gè)定義在一維格子上的單人游戲. 人們最初提出這個(gè)模型是為了描述可逆計(jì)算中的時(shí)間與空間的交換關(guān)系. 游戲開始時(shí), 玩家擁有一堆鵝卵石以及一個(gè)一維排布的n個(gè)格子, 標(biāo)記為 0,1,2···n, 并且在 0 號(hào)格子上有一個(gè)預(yù)先布置的鵝卵石. 游戲規(guī)則見表A1.

表 A1 鵝卵石游戲-可逆計(jì)算版本Table A1. Pebble game rules—the reversible computing version.

這里一個(gè)鵝卵石代表了一個(gè)單位的內(nèi)存, 而放置和取回鵝卵石的過程分別代表了計(jì)算和反計(jì)算, 因此均需要一個(gè)步驟數(shù), 對應(yīng)計(jì)算中的一個(gè)單位的運(yùn)算時(shí)間. 這里對應(yīng)回收規(guī)則中要求前一個(gè)格點(diǎn)中存在鵝卵石, 對應(yīng)可逆計(jì)算在釋放內(nèi)存時(shí), 要求其前置狀態(tài)存在以保證反計(jì)算可行. 當(dāng)鵝卵石數(shù)目充足 (S≥n), 我們用n個(gè)鵝卵石依次鋪至終點(diǎn)格子, 此時(shí)時(shí)間復(fù)雜度和空間復(fù)雜度均為O(n). 最少的鵝卵石數(shù)目的玩法則需要用到可逆計(jì)算框架下時(shí)間和空間最優(yōu)交換方案Bennett算法.

如表A2中算法(圖A1(b)) 所示, Bennett 算法將格子均勻的分割為k≥2 等份, 先是像前執(zhí)子k個(gè)區(qū)塊使最后一個(gè)區(qū)塊的末尾存在鵝卵石, 然后從最后第k-1 個(gè)區(qū)塊開始反向執(zhí)子收回中間k-1 個(gè)鵝卵石到自由堆中. 每個(gè)區(qū)塊又遞歸的均勻分割為k個(gè)子分塊做同樣的放置鵝卵石-保留最后的鵝卵石-取回鵝卵石的操作, 直到格子無法再分割.假設(shè)次過程的遞歸次數(shù)為l, 我們可以得到步驟數(shù)和鵝卵石數(shù)為

表 A2 Bennett算法Table A2. The Bennett's algorithm.

其中,k與l滿足總格子數(shù)n=kl. 可以看出可逆計(jì)算的時(shí)間復(fù)雜度和原時(shí)間為多項(xiàng)式關(guān)系. 同時(shí)可以看出k越小,使用的總鵝卵石數(shù)目越小, 因此最省空間的鵝卵石游戲解法對應(yīng)k=2. 作為例子, 圖A2(b)展示了n=16 ,k=2(l=4 ) 時(shí)候的游戲解法, 對應(yīng)步驟數(shù)為 (T+1)/2=41 ,這里的實(shí)際操作數(shù)少了大約一半, 是因?yàn)樽钔鈱拥腂ennett 過程不需要取回鵝卵石.

圖 A1 (a) Treeverse算法[20]示意圖, 其中(b) Bennett算法對應(yīng) k =3 的示意圖[33,25], 其中P和P-1分別代表了計(jì)算和反計(jì)算Fig. A1. (a) An illustration of the Treeverse algorithm[20], where (b) An illustration of the Bennett's algorithm for k =3 [33,25], where P and P -1 are forward computing and reverse computing respectively.

圖 A2 (a) Treeverse算 法 ( τ =3 , δ =3 ) 和 (b) Bennett算法 ( k =2 , n =4 ) 對應(yīng)的最優(yōu)時(shí)間空間交換策略下的鵝卵石游戲解法, 橫向是一維棋盤的格子, 縱向是步驟. 其中“○”為在這一步中收回的鵝卵石, “●”為在這一步中放上的鵝卵石, 而顏色稍淡的“”則對應(yīng)之前步驟中遺留在棋盤上未收回的鵝卵石. (a) 中的紅色格子代表已被涂鴉.(b) 中帶旗幟的格點(diǎn)代表終點(diǎn)Fig. A2. (a) The Treeverse algorithm( τ =3 , δ =3 ) and(b) Bennett's algorithm ( k =2 , n =4 ) solutions to the pebble game, the x direction is the grid layout and the y direction is the number of steps. Here, a “○” represents a pebble returned to the free pool in current step, a “●” represents the pebble added to the board in current step, and a“” represents a pebbles left on the board. In (a), red grids are painted grids. In (b), the grid with a flag sign is the goal.

我們稍微修改游戲設(shè)定可以得到檢查點(diǎn)版本的鵝卵石游戲. 該版本中用戶多了一支畫筆可用于涂鴉格點(diǎn), 且初始狀態(tài)的n號(hào)格子已經(jīng)被涂鴉, 新的規(guī)則描述見表A3.

檢查點(diǎn)版本的鵝卵石游戲中, 鵝卵石可以被隨時(shí)取下,代表不可逆的內(nèi)存擦除. 涂鴉過程則代表了梯度反向傳播的過程, 它要求前置計(jì)算狀態(tài)和后置梯度都存在. 在鵝卵石充足的情況下, 最節(jié)省步驟數(shù)的解法和可逆計(jì)算版本一樣,即計(jì)算過程中不取下任何鵝卵石. 而用最少鵝卵石的解法則僅需要兩枚鵝卵石交替使用, 每當(dāng)我們需要涂鴉一個(gè)格子i, 我們總是從初始鵝卵石開始掃描(依次放置一個(gè)鵝卵石并取下前一個(gè)鵝卵石)i步至指定格子, 因此總步驟k≤n(n-1)/2. 鵝卵石數(shù)目為 2＜δ＜n的最優(yōu)解最難, 需要用到如表A4中算法(圖A1(a)) 所示的 Treeverse算法. 該算法在完成第一遍從格子 0 開始的掃描會(huì)在格子上留下δ-1=S-2 個(gè)鵝卵石 (不包括初始鵝卵石), 把格子分割成δ個(gè)區(qū)塊. 我們把這些沒有被取下的鵝卵石稱為檢查點(diǎn), 總可以從任意一個(gè)檢查點(diǎn)出發(fā)掃描后方的格子. 區(qū)塊的大小由二項(xiàng)分布函數(shù)η(τ-1,δ),···,η(τ-1,2),η(τ-1,1) 確定, 其中τ的取值滿足η(τ,δ)=n. 為了涂鴉n-1 號(hào)格點(diǎn), 我們從離n-1 號(hào)格點(diǎn)最近的檢查點(diǎn) (最后一個(gè)區(qū)塊的開始處) 出發(fā)掃描至該點(diǎn), 重復(fù)該過程直至涂鴉至最后一個(gè)區(qū)塊的開始處. 由于最后一個(gè)區(qū)塊尺寸最小,僅為τ-1 , 因此我們并不擔(dān)心這樣的掃描會(huì)使得步驟數(shù)增加太多. 當(dāng)完成了最后一個(gè)區(qū)塊的涂鴉, 便可把格子上用于標(biāo)記最后一個(gè)區(qū)塊起點(diǎn)的鵝卵石取下以便重復(fù)利用. 為了涂鴉倒數(shù)第二個(gè)區(qū)塊, 先掃描這個(gè)區(qū)塊, 并用剛回收的鵝卵石將其分割為大小分別為η(τ-2,2) 和η(τ-2,1) 的兩個(gè)子區(qū)間. 用同樣的方式計(jì)算最后一個(gè)區(qū)間并遞歸分割前一個(gè)子區(qū)間直至區(qū)塊大小為1而無法繼續(xù)分割. 整個(gè)算法的步驟數(shù)和鵝卵石數(shù)目的關(guān)系是

表 A3 鵝卵石游戲-檢查點(diǎn)版本Table A3. Pebble game rules—the checkpointing version.

表 A4 算法A2: Treeverse算法Table A4. Algorithm A2: The Treeverse algorithm.

由二項(xiàng)分布的性質(zhì)可得,τ和δ的大小可以都正比于log(n). 圖A2(a)展示了 Treeverse 算法僅用4個(gè)鵝卵石,46個(gè)步驟數(shù)涂鴉完所有20個(gè)格子.①這里的46步并不是嚴(yán)格的最優(yōu)解, 因?yàn)閳D中掃描過程的最后一步并不需要立即釋放內(nèi)存從而可以減少步驟數(shù).

鵝卵石游戲是對程序的時(shí)間和空間的交換關(guān)系的非常理想化的描述, 它恰巧非常合適用于描述常微分方程求解這樣的不可逆線性程序. 圖A3展示了在固定額外時(shí)間開銷下, 程序應(yīng)用 Bennett 算法和 Treeverse 算法得到的最優(yōu)的空間開銷. 可以看出, 可逆計(jì)算整體上需要更多的空間開銷. 尤其是當(dāng)步驟數(shù)更多, 或是允許的時(shí)間的額外開銷更大時(shí), 該差別愈加明顯. 當(dāng)程序具有一定結(jié)構(gòu), 可逆計(jì)算也有不錯(cuò)的優(yōu)點(diǎn), 例如可以利用可逆性節(jié)省內(nèi)存. 另外由于可逆計(jì)算不需要對程序自動(dòng)設(shè)置檢查點(diǎn), 因此不需要借助全局堆棧, 對于微分運(yùn)行在GPU設(shè)備上的核函數(shù), 避免全局堆棧操作是必要的.

圖 A3 為了回溯中間狀態(tài), 時(shí)間和空間在兩種最優(yōu)時(shí)間-空間交換策略下的關(guān)系 (a) 固定橫軸為狀態(tài)回溯的計(jì)算時(shí)間與原函數(shù)計(jì)算時(shí)間的比值, 對比再允許固定時(shí)間開銷下, 內(nèi)存的額外開銷. 其中Bennett算法代表了可逆計(jì)算下的最優(yōu)策略, 而Treeverse則是傳統(tǒng)計(jì)算允許的最優(yōu)策略.(b) 對比Bennett算法與Treeverse算法空間開銷的比值Fig. A3. The time and space to trace back states in two time-space tradeoff strategies. (a) Space overheads versus time overheads at different numbers of simulation steps.The Bennett's algorithm is the optimal time-space tradeoff scheme in reversible computing, while the Treeverse algorithm is the optimal time-space tradeoff scheme in regular irreversible computing. (b) The ratio between the space overheads in the Bennett's algorithm and that in the Treeverse algorithm.