微分萬物:深度學習的啟示

|作者：王 磊1,2,† 劉金國3

(1 中國科學院物理研究所) 

(2 松山湖材料實驗室) 

(3 哈佛大學物理系) 

本文選自《物理》2021年第2期

摘要   深度學習教會了人們一種新的和計算機打交道的方式：將一些可微分的計算單元組合形成一段程序，再通過梯度優化的方法調整程序參數，使其達成期望的目標。這就是微分編程的思想。深度學習技術的快速發展為微分編程提供了趁手的工具，也為計算物理開闢了一番新天地。文章介紹微分編程的基本概念，並舉例說明它在建模、優化、控制、反向設計等物理問題中的應用。

關鍵詞  微分編程，自動微分，計算物理

深度學習在做什麼？對於這個問題，人們的第一反應往往是「訓練人工神經網絡」。這個解釋讓深度學習聽起來既神奇又神秘，卻沒有正面回答問題。讓我們再追問一句：訓練神經網絡又是在做什麼？純粹從編程的角度看，深度學習將一些可微分的計算單元搭建成為一段程序，再使用梯度優化的方法調節程序參數，使其達成所需的功能。其實，人工神經網絡是一類特別簡單的電腦程式。大多數神經網絡中主要涉及的是張量運算，例如張量收縮和元素非線性變換等。而所謂「深層」的神經網絡，只不過意味著這類張量計算程序比較長而已。在這個意義下，人工神經網絡只是一小類可以被訓練的程序。推而廣之，我們可以對一般的電腦程式求導，利用梯度信息調節程序的行為，這就是微分編程。

以Yann LeCun(Fackbook的首席AI科學家)和 Andrej Karpathy(Tesla的AI部門主管)為代表的一些機器學習專家認為微分編程是深度學習的未來[1]。深度學習中所發明的各式各樣的神經網絡結構，例如卷積、循環、注意力機制等，都是針對特定任務的可微分模塊。設計這類模塊的指導思想是儘可能地利用關於問題本身的先驗知識，例如對稱性、局域性、層級結構等等。微分編程通過學習一般的電腦程式，可以更加充分地利用問題相關的特定領域的先驗知識，同時還能保證模型的可解釋性。另一方面，越來越多其他領域的科學家們也認識到微分編程的價值，認為它是聯繫深度學習與科學計算的一條紐帶。

讓我們通過幾個例子來直觀地認識微分編程。圖1(a)定義了這樣一個問題：如何設計合適的外勢場，使得其中的自由電子的基態密度分布達到期望的形式。這個問題的正向求解是相對容易的，只需要根據外勢場構造哈密頓量，再求解相應的本徵值問題就能得到基態的電子密度。按照微分編程的思路解決反向問題，我們首先定義正向計算的密度與目標密度的差別作為目標函數，再穿過正向計算過程求導得到目標函數對於外勢場的梯度。之後，我們可以利用梯度信息優化外勢場，使其中的電子密度分布逐漸逼近目標分布。這類反向求解自由電子勢函數的問題是密度泛函理論研究中的基本問題[4]，也是一個有代表性的哈密頓量反向設計問題。微分編程是解決這類問題簡單而通用的方法。讀者可以參考筆者編寫的演示程序，使用深度學習框架PyTorch解決一維薛丁格方程的逆問題[5]。

圖1 三個微分編程實例 (a)逆薛丁格方程；(b)憤怒的小鳥；(c)圖像識別

圖1(b)示意了「憤怒的小鳥」遊戲的基本設置。我們希望通過調整小紅飛出的初始速度，使得它能夠命中目標。與反覆的試錯相比，一種更加高效的調整方式需要用到和目標的距離對初始速度的導數信息。這要求能夠穿過整個飛行軌跡求導。微分編程可以幫助我們做到這一點。更妙的是，為了對運動方程的積分計算求導，我們只需要沿著時間反向，對另一套運動方程積分[6]。

圖1(c)是一個典型的深度學習任務：圖像識別。我們希望調節神經網絡的參數，使得網絡輸出正確的圖片標籤。基於網絡參數梯度的優化方案是目前最成功的神經網絡訓練方法。關於這點有一個簡單的理解：多元函數的函數值只是一個標量的信息，而梯度是高維的矢量。從信息的角度，越是高維的函數，梯度優化相對於非梯度優化方法的優勢就越明顯。現代的深度學習應用甚至可以訓練含有一千多億個參數的神經網絡。如果不使用梯度的信息，實在難以想像如何處理這麼高維的優化問題。

表1總結了以上三個例子的基本特徵。可以看出，無論是從問題性質、學習對象、計算過程，還是目標函數的角度，這三個問題看起來都毫無關聯。然而，它們其實都是微分編程的特例：我們穿過整個計算過程求導，得到目標函數相對於學習對象的梯度，從而解決問題。

表1 三個微分編程實例的基本特徵

和通常的深度學習問題不同，圖1(a，b)這兩個例子中都包含針對特定物理領域的計算過程(求解薛丁格方程和牛頓方程)。另外，讀者也許注意到了，在這兩個例子中甚至都沒有談到訓練數據。事實上，有不少科學問題都是如此。我們並不能夠、也不需要事先準備大量的訓練數據。對於科學問題，「一個好的模型勝過一千個數據」[7]。而微分編程將模型表達成電腦程式，再利用梯度信息學習其中的未知成分。這背後的關鍵既不是大數據，也不是神經網絡，而是自動微分。

顧名思義，自動微分就是自動化地計算程序的輸出相對於輸入的導數。這項技術的本質是機械地對計算過程使用鏈式法則，從而得到數值精確的導數信息。我們可以用計算圖來形象化地表達程序運行和自動微分的過程。計算圖是一個有向無環圖(directed acyclic graph)，其中的節點和箭頭代表程序運行時變量之間的依賴關係。圖2展示了一類典型的計算圖。其中變量 x1 代表程序輸入，x2，x3 是計算中間結果，θ1，θ2 是程序的參數。最後，計算圖末端的節點L 代表程序輸出的標量損失函數。圖1中的三個例子所對應的計算圖都具有這樣的結構[8]。

為了考慮自動微分，我們定義計算圖中某一變量 x 的梯度為其伴隨變量(adjoint variable)，記為= ∂L/∂x。自動微分的目的是把計算圖最末端 = 1 的信息反向傳遞迴所有需要計算梯度的節點。這可以使用反向傳播算法實現：中間節點收集下遊節點的伴隨變量信息，根據鏈式法則得到自己的伴隨變量，再往它的上遊節點傳遞這個信息。在反向傳播的每一步，我們只需要計算下遊傳來的伴隨變量與當前計算步驟雅可比矩陣的乘積。遵循這種信息傳遞的方法，計算機可自動地對任意複雜的計算過程微分。理論保證了使用這個辦法求導與程序前向運行的計算複雜度是一樣的。此外，因為梯度的計算過程也可以表達為計算圖，我們甚至可以迭代地使用自動微分方法計算程序的高階導數。

圖2 計算圖表達計算過程中變量之間的依賴關係。自動微分方法回溯計算圖，得到程序輸出相對於所有變量的梯度

自動微分相對於使用數值差分的方法計算梯度有不少優點。首先，它可以保證計算的梯度是數值精確的，而不依賴於額外的差分參數。其次 ，圖2所展示的自動微分方式，僅通過向前和向後遍歷計算圖兩遍，就能高效地計算出所有變量的梯度。這遠比數值差分高效：因為數值差分需要對參數的每一個分量分別做微擾，再反覆地執行前向計算來得到參數在各個分量上的梯度。由於圖2所示的計算方法有一個從後向前回溯計算圖的過程，這被稱為反向模式(reverse mode)的自動微分。實現反向模式的自動微分往往需要付出額外的存儲代價，以記錄前向計算時的中間變量信息。還有一類前向模式(forward mode)的自動微分：通過計算雅可比矩陣和變量梯度的乘積，向前傳遞對於程序輸入或參數的擾動。前向模式求導的方向與程序的運行方向相同，因此不需要付出額外的存儲代價。但是前向模式不如反向模式高效：它的計算複雜度和數值差分是一樣的，因此通常不被用於含有大量參數的優化問題。

值得強調的是，自動微分的對象不見得非要是數學函數。微分編程完全可以應用於包含循環和控制流等指令的一般程序。在程序的運行過程中，取決於不同的輸入數據與分支條件，動態生成的計算圖可以具有完全不同的拓撲結構與深度。一旦確定了計算圖的結構，在自動微分計算的每一步，僅需要關心計算單元局部的信息傳遞。這裡的計算單元可以是基本的四則運算，也可以是一系列基本運算的複合，例如本徵方程、微分方程、自洽方程求解等等。只要我們有辦法指定計算單元的伴隨變量回傳規則，就可不必關心其內部計算的細節，也不需要求導進入計算單元內部。利用這樣的方式控制自動微分的顆粒度，可以保證數值穩定性、提升運行效率、降低反向模式計算的內存消耗。如何針對具體問題設計合適的計算單元，正是微分編程的藝術。模塊化和複合性是深度學習的名片，也是微分編程的關鍵。

微分編程並不是一個從天而降的嶄新方法。科學家早在20世紀50、60年代就開始探索對於電腦程式的自動微分[2,3]。在控制和反向設計等領域，也早已存在與微分編程密切相關的伴隨變量方法。而物理學家熟知的微擾論和線性響應理論，可以說就是穿過物理過程的微分計算。如今，在基於張量運算的自動微分應用——深度學習——成為一門顯學的時代，微分編程獲得了長足的進步和更普遍的關注。它也給物理學中的建模、優化、控制和反向設計問題提供了一個統一而靈活的研究方法。下面我們分別舉例說明。  

建模是指根據觀測數據推斷實際物理過程，從而預測新的情形。建模的目標是能夠舉一反三，亦即從少量的觀測數據出發做出正確的預測。用機器學習的術語來說，就是訓練出具有泛化能力的模型。圖1(c)所示的圖像識別任務就是一個建模問題。然而，不同於常見的圖像識別任務，我們通常難以為待解決的物理問題準備大量準確的訓練數據。僅僅依賴於少量數據建模，容易發生過擬合現象。這就需要精心地設計模型，最好讓它本身就契合待研究的物理過程。微分編程是實現這個目標的一種方式。

蛋白質摺疊是一個公認的難題。從胺基酸序列出發預測蛋白質的結構，中間涉及複雜的物理過程。完全基於第一性原理的計算恐怕是極其困難的。因此，人們寄希望於從已知的序列——蛋白質結構數據中尋找規律，直接預測蛋白質結構。DeepMind的AlphaFold憑著精妙的模型設計與充沛的算力，已在這方面取得了了不起的成果[11]。但這並沒有妨礙人們另闢蹊徑，探索不一樣的技術路線。文獻[9]將胺基酸序列到蛋白質結構的映射建模成一段分子動力學模擬，並從現有的序列——結構數據中學習分子模擬所需的能量函數。利用微分編程技術，可以穿過如圖3(a)所示的整個模擬過程學習模型參數，以預測未知的蛋白質結構。

圖3 微分編程應用於蛋白質摺疊和仿星器設計 (a)從胺基酸序列到蛋白質結構的計算圖；(b)從仿星器的線圈參數到目標函數的計算圖(摘自文獻[9]和[10])

對於量子效應更加顯著的化學分子，預測其結構需要快速準確地求解量子多體問題的基態。Google Research 的李力等[12]根據密度泛函理論建立模型，將 Kohn—Sham 密度泛函計算當作可微分的程序，穿過整個自洽計算過程來學習交換關聯勢。他們僅僅需要非常少的分子基態數據就可以成功地訓練模型，預測一維氫分子隨著原子間距的整條勢能曲線。

以上兩個工作的共同特點是模型中包含了特定的物理計算(分子動力學模擬、Kohn—Sham方程)，因此不再需要從零開始學習關於物理問題的基本規律，而只需填補成熟的理論框架中空缺的部分(分子模擬的能量函數、密度泛函理論的交換關聯勢)。這樣的建模方式有利於從小樣本中學習到具有足夠泛化能力的模型。但是由於涉及到比神經網絡複雜得多的計算，這些工作目前還只停留在小規模概念驗證的階段。時間將會告訴我們，基於物理機制的微分編程是否還能帶我們走得更遠。

優化問題在物理學中無處不在。自然界偏愛低能的狀態。更一般地看，甚至可以說最小作用量原理支配了整個物理學。在具體的物理學研究中，可以利用自動微分計算梯度幫助求解優化問題。這是微分編程非常自然的一類應用。

一個有代表性的例子是全變分的Hartree—Fock方法[13]。在傳統的Hartree—Fork計算中，人們通常經驗性地選取一些基函數展開分子的哈密頓量，再通過自洽計算優化這些基函數上的組合係數來近似體系的基態。而在全變分的計算中，可以利用自動微分穿過自洽循環，計算Hartree—Fock能量相對於基函數中經驗參數的導數，從而徹底優化基函數的位置和形狀。因為優化空間變大了，全變分的Hartree—Fock方法可以得到比傳統方法更低的變分能量。

另一個例子是我們與合作者們發展的張量網絡微分編程[14]。張量網絡方法經過多年的發展已經成為經典統計、量子多體物理，乃至機器學習等領域重要的理論和計算工具。然而，針對張量網絡的優化仍是一個長期存在的難點。傳統的優化方法要麼是基於一些近似的方案，要麼依賴於繁瑣的人工解析推導梯度。這限制了張量網絡方法在複雜問題中的廣泛應用。採用微分編程的方法穿過整個張量網絡重正化計算進行優化，不僅大大降低了程序實現的門檻，還更加高效和可靠。例如，作者對於二維無窮大方格子上的反鐵磁海森伯模型得到了最低的變分能量。對於更有挑戰的阻挫磁性問題[15](如籠目格(Kagome)格子上的反鐵磁海森伯模型等)張量網絡微分編程也都給出比之前更低的變分能量。特別的，對於六角格子上的Kitaev模型，通過對張量網絡進行梯度優化可以直接得到物理上定性正確的非磁性基態。這些結果給了我們更強的信心：假以時日，對張量網絡的微分編程可以帶給我們新的物理！

上面舉的兩個例子涉及到對於Hartree—Fock自洽方程或張量網絡重正化這類迭代計算的求導。可以簡單直接地將迭代過程展開，然後利用現成的自動微分工具穿過整個計算過程求導。除此以外，還可以把達到迭代收斂條件視作基本的計算單元，利用隱函數定理得出伴隨變量的回傳規則[16]。後面這種做法更加契合微分編程的思想，可以讓程序更加高效和省內存。

控制問題是指通過調整體系的初始條件和外場等，以實現對於含時物理過程的調控。圖1(b)所示的憤怒的小鳥遊戲就是一個控制問題。通過對計算機模擬的經典物理過程求導，微分編程已經在計算機圖形學、機器人控制、分子模擬等領域彰顯威力。例如，文獻[17]利用可微分的分子模擬手段調整膠體顆粒的相互作用形式，以控制其自組裝過程中的動力學行為。

量子計算中也有類似的控制問題。例如，如何調控微波或光脈衝以實現特定的量子門？從微分編程的視角審視現有的量子控制方法，會發現它們有的使用了不夠高效的前向微分模式，有的依賴於不夠靈活的離散控制方案。利用圖1(b)中我們見過的穿過微分方程求導的方法[6]，可以類似地控制含時薛丁格方程的演化[18]。這類微分編程方法不僅高效，還適用於任意的控制手段、限制條件和目標函數，是研究量子最優控制問題的一般方法。另外，在實際的控制問題中，計算機模擬未必能夠完全精確地描述物理過程。因此還需要將基於微分方程的最優控制和上文提到的建模方法相結合[7]。一方面根據實際觀測的結果對物理過程建立準確的模型，另一方面利用學習到的模型做出正確的控制預測。和徹底黑箱化的控制手段(例如無模型的強化學習方案)相比，微分編程可以充分利用關於控制過程的物理認識，因此優化效率要高很多。

反向設計是指通過設計合適的物理結構以實現特定的力場、光場、電磁場等。圖1(a)所示的逆薛丁格問題就是一個反向設計問題。一個類似的例子是光子晶體的反向設計。光子晶體通過人為加工的微結構來調製光的傳播。因為計算機模擬的光學現象可以幾乎完美地在樣品中實現，這是一個特別適合於反向設計的領域。史丹福大學的Momchil Minkov等人[19]通過對麥克斯韋方程組求解過程求導，以優化光子晶體的色散關係和品質因子。他們體會到，自動微分方法將光子晶體領域裡常用的伴隨變量方法拓展到了任意的計算圖，可以輔助設計更加靈活多樣的光學器件。

另一個值得一提的應用是使用自動微分設計仿星器[10]。仿星器是可控核聚變的途徑之一。如圖3(b)所示，為了實現它，需要設計合適的線圈以產生磁場來約束幾億度的等離子體。這樣的極端條件對於設計提出了極高的要求。微分編程的方法可以統一地考慮產生磁約束的物理過程以及各種現實的工程限制。深度學習專家David Duvenaud讀到這篇文章之後，興奮地在Twitter上發文說：「這增強了我對人類文明的信念。當我聽說仿星器時，我就想『希望他用上了自動微分，但我擔心他們其實還沒有。我希望有人告訴他們怎麼做。』 」。誠然，已經充分享受微分編程成果和樂趣的深度學習專家，自然希望它能夠被更廣泛地應用於各個領域，特別是那些與人類命運休戚相關的問題上。

科幻作家Arthur Clarke有一句名言：「任何充分先進的技術看起來都無異於魔法」。技術進步可以解放人們的思想，去探索更有創意的工作。微分編程的技術與應用都還在快速的發展中。例如，微分編程與可逆計算相結合，可以緩解反向模式自動微分的內存瓶頸[30]。此外，微分編程的應用範圍甚至還逐漸擴展到對於離散參數或是結構信息等「不可微分」對象的學習上。因此，本文所介紹的例子還僅僅是微分編程牛刀小試，希望可以讓更多的讀者了解和喜歡上微分編程，讓它在現實問題中大顯身手。

本文中列舉的應用涉及到對多種複雜運算的求導，例如非線性方程求解、快速傅立葉變換、矩陣分解等等。因為這些應用的推動，人們開發了越來越多的微分編程基本單元和工具。這些進展進一步豐富了微分編程的生態。下方表格介紹了一些微分編程工具。此外，www.autodiff.org網站還列舉了更多的微分編程工具。

基於這些編程工具，科學家們打造了針對具體領域的微分編程軟體。這包括張量網絡[20]、分子模擬[21]、量子線路[22—24]、系外行星[25]等。此外還有面向可微分的科學計算工具：ξ-torch[26]和 ADCME[27]等。SciML[28]和dolfimadjoint[29]對基於微分方程的可微分科學計算有豐富的支持。一般來說，如果程序中主要包含的是線性代數與張量計算，可以直接使用現有的深度學習框架，因為這些框架的本質是支持硬體加速的可微分張量計算庫。而如果程序中主要包含的是微分方程求解、自洽迭代等計算，則需要根據具體情況選用專門的工具以充分發揮微分編程的性能。

致 謝  感謝與謝浩、張潘、李偉、謝志遠、廖海軍和向濤的合作與討論。

[1] 其實，也可以說微分編程是深度學習的過去，見文獻2，3 

[2] Nolan J F. Analytical differentiation on a digital computer. Massachusetts Institute of Technology，1953

[3] Wengert R E. Communications of the ACM，1964，7：463

[4] Jensen D S，Wasserman A. International Journal of Quantum Chemistry，2018，118(1)：e25425

[5] https://github.com/QuantumBFS/SSSS/blob/master/1_deep_learning/schrodinger.py

[6] Chen R T Q，Rubanova Y，Bettencourt J et al. 2018，arXiv:1806.07366

[7] Rackauckas C et al. 2020，arXiv:2001.04385

[8] 其實逆薛丁格方程和憤怒的小鳥例子中的計算圖更簡單：它們沒有參數節點，是純鏈狀結構的。 

[9] Ingraham J et al. Learning Protein Structure with a Differentiable Simulator. International Conference on Learning Representations，2019

[10] McGreivy N，Hudson S R，Zhu C X. 2020，arXiv:2009.00196

[11] AlphaFold: a solution to a 50-year-old grand challenge in biology. https://deepmind.com/blog/article/alphafold-a-solution-to-a-50-year-old-grand-challenge-in-biology

[12] Li L et al. Phys. Rev. Lett.，2021，126：036401

[13] Tamayo-Mendoza T et al. ACS Cent. Sci.，2018，4(5)：559

[14] Liao H J，Liu J G，Wang L et al. Phys. Rev. X，2019，9：031041

[15] Liao H J. Talk at Computational Approaches to Quantum Many-body Problems. ISSP，2019. https://www.issp.u- tokyo.ac.jp/public/caqmp2019/slides/726S_Liao.pdf

[16] Kolter Z，Duvenaud D，Johnson M. NeurIPS 2020 tutorial，Deep Implicit Layers- Neural ODEs，Deep Equilibirum Models，and Beyond. https://implicit-layers-tutorial.org/

[17] Goodrich C P et al. 2020，arXiv: 2010.15175

[18] Güngördü U，Kestner J P. 2020，arXiv:2011.02512

[19] Minkov M et al. ACS Photonics，2020，7( 7)：1729

[20] https://github.com/jurajHasik/peps-torch

[21] https://github.com/google/jax-md

[22] Bergholm V et al. 2018，arXiv:1811.04968

[23] Luo X Z，Liu J G，Zhang P et al. Quantum，2020，4：341

[24] Broughton M et al. 2020，arXiv:2003.02989

[25] Morvan M et al. 2020，arXiv:2011.02030

[26] Kasim M F，Vinko S M. 2020，arXiv:2010.01921

[27] Xu K L. Automatic Differentiation Library for Computational and Mathematical Engineering. https://github.com/kailaix/ADCME.jl

[28] https://sciml.ai/

[29] http://www.dolfin-adjoint.org/en/latest/

[30] Liu J G，Zhao T. 2020，arXiv:2003.04617