天天看點

【自動微分原理】自動微分的具體實作

第一篇自動微分原理文章中我們大概初步談了談從手動微分到自動微分的過程,第二篇自動微分正反模式中深入了自動微分的正反向模式具體公式和推導。

實際上第二章了解到正反向模式隻是自動微分的原理模式,在實際代碼實作的過程,正方向模式隻是提供一個原理性的指導,在真正編碼過程會有很多細節需要打開,例如如何解析表達式,如何記錄反向求導表達式的操作等等。這一篇文章中,ZOMI希望通過介紹目前比較熱門的方法給大家普及一下自動微分的具體實作。

  • 【自動微分原理】01. 原理介紹
  • 【自動微分原理】02. 正反模式
  • 【自動微分原理】03. 具體實作

自動微分實作

了解自動微分的不同實作方式非常有用。在這裡呢,我們将介紹主要的自動微分實作方法。在上一篇的文章中,我們介紹了自動微分的基本數學原理。可以總結自動微分的關鍵步驟為:

  • 分解程式為一系列已知微分規則的基礎表達式的組合
  • 根據已知微分規則給出各基礎表達式的微分結果
  • 根據基礎表達式間的資料依賴關系,使用鍊式法則将微分結果組合完成程式的微分結果

雖然自動微分的數學原理已經明确,包括正向和反向的數學邏輯和模式。但具體的實作方法則可以有很大的差異,2018 年,Siskind 等學者在其綜述論文Automatic Differentiation in Machine Learning: a Survey [1] 中對自動微分實作方案劃分為三類:

基本表達式:基本表達式或者稱元素庫(Elemental Libraries),基于元素庫中封裝一系列基本的表達式(如:加減乘除等)及其對應的微分結果表達式,作為庫函數。使用者通過調用庫函數建構需要被微分的程式。而封裝後的庫函數在運作時會記錄所有的基本表達式和相應的組合關系,最後使用鍊式法則對上述基本表達式的微分結果進行組合完成自動微分。

操作符重載:操作符重載或者稱運算重載(Operator Overloading,OO),利用現代語言的多态特性(例如C++/JAVA/Python等進階語言),使用操作符重載對語言中基本運算表達式的微分規則進行封裝。同樣,重載後的操作符在運作時會記錄所有的操作符和相應的組合關系,最後使用鍊式法則對上述基本表達式的微分結果進行組合完成自動微分。

源代碼變換:源代碼變換或者叫做源碼轉換(Source Code Transformation,SCT)則是通過對語言預處理器、編譯器或解釋器的擴充,将其中程式表達(如:源碼、AST抽象文法樹 或 編譯過程中的中間表達 IR)的基本表達式微分規則進行預定義,再對程式表達進行分析得到基本表達式的組合關系,最後使用鍊式法則對上述基本表達式的微分結果進行組合生成對應微分結果的新程式表達,完成自動微分。

任何 AD 實作中的一個主要考慮因素是 AD 運算時候引入的性能開銷。就計算複雜性而言,AD 需要保證算術量增加不超過一個小的常數因子。另一方面,如果不小心管理 AD 算法,可能會帶來很大的開銷。例如,簡單的配置設定資料結構來儲存對偶數(正向運算和反向求導),将涉及每個算術運算的記憶體通路和配置設定,這通常比現代計算機上的算術運算更昂貴。同樣,使用運算符重載可能會引入伴随成本的方法分派,與原始函數的原始數值計算相比,這很容易導緻一個數量級的減速。

下面這個圖是論文作者回顧了一些比較通用的 AD 實作。

【自動微分原理】自動微分的具體實作

基本表達式

基本表達式法也叫做元素庫(Elemental Libraries),程式中實作構成自動微分中計算的最基本的類别或者表達式,并通過調用自動微分中的庫,來代替數學邏輯運算來工作。然後在函數定義中使用庫公開的方法,這意味着在編寫代碼時,手動将任何函數分解為基本操作。

這個方法呢從自動微分剛出現的時候就已經被廣泛地使用,典型的例子是 Lawson (1971) 的 WCOMP 和 UCOMP 庫,Neidinger (1989) 的 APL 庫,以及 Hinkins (1994) 的工作。同樣,Rich 和 Hill (1992) 使用基本表達式法在 MATLAB 中制定了他們的自動微分實作。

以公式為例子:

使用者首先需要手動将公式1中的各個操作,或者叫做子函數,分解為庫函數中基本表達式組合:

t1 = log(x)
t3 = sin(x)
t2 = x1 * x2
t4 = x1 + x2
t5 = x1 - x2
           

使用給定的庫函數,完成上述函數的程式設計:

// 參數為變量 x,y,t 和對應的導數變量 dx,dy,dt
def ADAdd(x, y, dx, dy, t, dt)

// 同理對上面的公式實作對應的函數
def ADSub(x, y, dx, dy, t, dt)
def ADMul(x, y, dx, dy, t, dt)
def ADLog(x, dx, t, dt)
def ADSin(x, dx, t, dt)
           

而庫函數中則定義了對應表達式的數學微分規則,和對應的鍊式法則:

// 參數為變量 x,y,t 和對應的導數變量 dx,dy,dt
def ADAdd(x, y, dx, dy, t, dt):
    t = x + y
    dt = dy + dx

// 參數為變量 x,y,t 和對應的導數變量 dx,dy,dt
def ADSub(x, y, dx, dy, t, dt):
    t = x - y
    dt = dy - dx

// ... 以此類推
           

針對公式1中基本表達式法,可以按照下面示例代碼來實作正向的推理功能,反向其實也是一樣,不過調用代碼更複雜一點:

x1 = xxx
x2 = xxx
t1 = ADlog(x1)
t2 = ADSin(x2)
t3 = ADMul(x1, x2)
t4 = ADAdd(t1, t3)
t5 = ADSub(t4, t2)
           

基本表達式法的優點可以總結如下:

  • 實作簡單,基本可在任意語言中快速地實作為庫

基本表達式法的缺點可以總結如下:

  • 使用者必須使用庫函數進行程式設計,而無法使用語言原生的運算表達式;
  • 另外實作邏輯和代碼也會備援較長,依賴于開發人員較強的數學背景

基本表達式法在沒有操作符重載AD的80到90年代初期,仍然是計算機中實作自動微分功能最簡單和快捷的政策啦。

操作符重載

在具有多态特性的現代程式設計語言中,運算符重載提供了實作自動微分的最直接方式,利用了程式設計語言的第一特性(first class feature),重新定義了微分基本操作語義的能力。

在 C++ 中使用運算符重載實作的流行工具是 ADOL-C(Walther 和 Griewank,2012)。 ADOL-C 要求對變量使用啟用 AD 的類型,并在 Tape 資料結構中記錄變量的算術運算,随後可以在反向模式 AD 計算期間“回放”。 Mxyzptlk 庫 (Michelotti, 1990) 是 C++ 能夠通過前向傳播計算任意階偏導數的另一個例子。 FADBAD++ 庫(Bendtsen 和 Stauning,1996 年)使用模闆和運算符重載為 C++ 實作自動微分。對于 Python 語言來說,autograd 提供正向和反向模式自動微分,支援高階導數。

在機器學習 ML 或者深度學習 DL 領域,目前AI架構中使用操作符重載的一個典型代表是 Pytroch,其中使用資料結構 Tape 來記錄計算流程,在反向模式求解梯度的過程中進行replay Operator。

  1. 操作符重載來實作自動微分的功能裡面,很重要的是利用進階語言的特性。下面簡單看看僞代碼,這裡面我們定義一個特殊的資料結構 Variable,然後基于 Variable 重載一系列的操作如 __mul__ 代替 * 操作。
class Variable:
    def __init__(self, value):
        self.value = value

    def __mul__(self, other):
        return ops_mul(self, other)

   # 同樣重載各種不同的基礎操作
   def __add__(self, other)
   def __sub__(self, other)
   def __div__(self, other)
           
  1. 實作操作符重載後的計算。
def ops_mul(self, other):
    x = Variable(self.value * other.value)
           
  1. 接着通過一個 Tape 的資料結構,來記錄每次 Variable

    執行計算的順序,Tape 這裡面主要是記錄正向的計算,把輸入、輸出和執行運算的操作符記錄下來。

class Tape(NamedTuple):
    inputs : []
    outputs : []
    propagate : (inputs, outpus)
           
  1. 因為大部分 ML 系統或者 AI 架構采用的是反向模式,是以最後會逆向周遊 Tape 裡面的資料(相當于反向傳播或者反向模式的過程),然後累積反向計算的梯度。
# 反向求導的過程,類似于 Pytroch 的 backward 接口
def grad(l, results):

    # 通過 reversed 操作把帶有梯度資訊的 tape 逆向周遊
    for entry in reversed(gradient_tape):
        # 進行梯度累積,反向傳播給上一次的操作計算
        dl_d[input] += dl_dinput
           

當然啦,我們會在下一節當中帶着大家親自通過操作符重載實作一個前向的自動微分和後向的自動微分。下面總結一下操作符重載的一個基本流程:

  • 預定義了特定的資料結構,并對該資料結構重載了相應的基本運算操作符
  • 程式在實際執行時會将相應表達式的操作類型和輸入輸出資訊記錄至特殊資料結構
  • 得到特殊資料結構後,将對資料結構進行周遊并對其中記錄的基本運算操作進行微分
  • 把結果通過鍊式法則進行組合,完成自動微分

操作符重載法的優點可以總結如下:

  • 實作簡單,隻要求語言提供多态的特性能力
  • 易用性高,重載操作符後跟使用原生語言的程式設計方式類似

操作符重載法的缺點可以總結如下:

  • 需要顯式的構造特殊資料結構和對特殊資料結構進行大量讀寫、周遊操作,這些額外資料結構和操作的引入不利于高階微分的實作
  • 對于一些類似 if,while 等控制流表達式,難以通過操作符重載進行微分規則定義。對于這些操作的處理會退化成基本表達式方法中特定函數封裝的方式,難以使用語言原生的控制流表達式

源代碼轉換

源碼轉換(Source Code Transformation,SCT)是最複雜的,實作起來也是非常具有挑戰性。

源碼轉換的實作提供了對程式設計語言的擴充,可自動将算法分解為支援自動微分的基本操作。通常作為預處理器執行,以将擴充語言的輸入轉換為原始語言。簡單來說就是利用源語言來實作領域擴充語言 DSL 的操作方式。

源代碼轉換的經典執行個體包括 Fortran 預處理器 GRESS(Horwedel 等人,1988 年)和 PADRE2(Kubo 和 Iri,1990 年),在編譯之前将啟用 AD 的 Fortran 變體轉換為标準 Fortran。類似地,ADIFOR 工具 (Bischof et al., 1996) 給定一個 Fortran 源代碼,生成一個增強代碼,其中除了原始結果之外還計算所有指定的偏導數。對于以 ANSI C 編碼的過程,ADIC 工具(Bischof 等人,1997)在指定因變量和自變量之後将 AD 實作為源代碼轉換。 Tapenade(Pascual 和 Hasco¨et,2008 年;Hasco¨et 和 Pascual,2013 年)是過去10年終 SCT 的流行工具,它為 Fortran 和 C 程式實作正向和反向模式 AD。

除了通過源代碼轉換進行語言擴充外,還有一些實作通過專用編譯器或解釋器引入了具有緊密內建的 AD 功能的新語言。一些最早的 AD 工具,例如 SLANG (Adamson and Winant, 1969) 和 PROSE (Pfeiffer, 1987) 屬于這一類。 NAGWare Fortran 編譯器 (Naumann and Riehme, 2005) 是一個較新的示例,其中使用與 AD 相關的擴充會在編譯時觸發衍生代碼的自動生成。

作為基于解釋器的實作的一個例子,代數模組化語言 AMPL (Fourer et al., 2002) 可以用數學符号表示目标和限制,系統從中推導出活動變量并安排必要的 AD 計算。此類别中的其他示例包括基于類似 Algol 的 DIFALG 語言的 FM/FAD 包 (Mazourik, 1991),以及類似于 Pascal 的面向對象的 COZY 語言 (Berz et al., 1996)。

而華為全場景AI架構 MindSpore 則是基于 Python 語言使用源代碼轉換實作 AD 的正反向模式,并采用了函數式程式設計的風格,該機制可以用控制流表示複雜的組合。函數被轉換成函數中間表達(Intermediate Representation,IR),中間表達構造出一個能夠在不同裝置上解析和執行的計算圖。在執行前,計算圖上應用了多種軟硬體協同優化技術,以提升端、邊、雲等不同場景下的性能和效率。

其主要流程是:分析獲得源程式的 AST 表達形式;然後基于 AST 完成基本表達式的分解和微分操作;再通過周遊 AST 得到基本表達式間的依賴關系,進而應用鍊式法則完成自動微分。

因為源碼轉換涉及到底層的抽象文法樹、編譯執行等細節,是以這裡就不給出僞代碼了(實在太難了給不出來),我們通過下面這張圖來簡單了解下 SCT 的一般性過程。

【自動微分原理】自動微分的具體實作

從圖中可以看到源碼轉換的整體流程分為編譯時間和執行時間,以 MindSpore 為例,其在運作之前的第一個 epoch 會等待一段時間,是因為需要對源碼進行編譯轉換解析等一系列的操作。然後再 run time 運作時則會比較順暢,直接對資料和代碼不斷地按照計算機指令來高速執行。

編譯階段呢,在 Initialization 過程中會對源碼進行 Parse 轉換成為抽象文法樹 AST,接着轉換為基于圖表示的中間表達 IR,這個基于圖的IR從概念上了解可以了解為計算圖,神經網絡層數的表示通過圖表示會比較直覺。

接着對 Graph base IR進行一些初級的類型推導,特别是針對 Tensor/List/Str 等不同的基礎資料表示,然後進行宏展開,還有語言單态化,最後再對變量或者自變量進行類型推導。可以從圖中看到,很多地方出現了不同形式的 IR,IR 其實是編譯器中常用的一個中間表達概念,在編譯的 Pass 中會有很多處理流程,每一步處理流程産生一個 IR,交給下一個Pass進行處理。

最後通過 LLVM 或者其他等不同的底層編譯器,最後把 IR 編譯成機器碼,然後就可以真正地在runtime執行起來。

源碼轉換法的優點可以總結如下:

  • 支援更多的資料類型(原生和使用者自定義的資料類型) + 原生語言操作(基本數學運算操作和控制流操作)
  • 高階微分中實作容易,不用每次使用 Tape 來記錄高階的微分中産生的大量變量,而是統一通過編譯器進行額外變量優化和重計算等優化
  • 進一步提升性能,沒有産生額外的 tape 資料結構和 tape 讀寫操作,除了利于實作高階微分以外,還能夠對計算表達式進行統一的編譯優化

源碼轉換法的缺點可以總結如下:

  • 實作複雜,需要擴充語言的預處理器、編譯器或解釋器,深入計算機體系和底層編譯
  • 支援更多資料類型和操作,使用者自由度雖然更高,但同時更容易寫出不支援的代碼導緻錯誤
  • 微分結果是以代碼的形式存在,在執行計算的過程當中,特别是深度學習中大量使用for循環過程中間錯誤了,或者是資料處理流程中出現錯誤,并不利于深度調試

參考

[1] Automatic Differentiation in Machine Learning: a Survey: https://arxiv.org/abs/1502.05767

[2] Rirchard L. Burden and J. Douglas Faires. Numerical Analysis. Brooks/Cole, 2001.

[3] Max E. Jerrell. Automatic differentiation and interval arithmetic for estimation of disequilibrium models. Computational Economics, 10(3):295–316, 1997.

[4] Johannes Grabmeier and Erich Kaltofen. Computer Algebra Handbook: Foundations, Applications, Systems. Springer, 2003.

[5] G. W. Leibniz. Machina arithmetica in qua non additio tantum et subtractio sed et multiplicatio nullo, diviso vero paene nullo animi labore peragantur. Hannover, 1685.

[6] George F. Corliss. Application of differentiation arithmetic, volume 19 of Perspectives in Computing, pages 127–48. Academic Press, Boston, 1988.

[7] Arun Verma. An introduction to automatic differentiation. Current Science, 78(7):804–7, 2000.

[8] Andreas Griewank and Andrea Walther. Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation. Society for Industrial and Applied Mathematics, Philadelphia, 2008. doi: 10.1137/1.9780898717761.

[9] John F. Nolan. Analytical differentiation on a digital computer. Master’s thesis, Massachusetts Institute of Technology, 1953.

[10] L. M. Beda, L. N. Korolev, N. V. Sukkikh, and T. S. Frolova. Programs for automatic differentiation for the machine BESM (in Russian). Technical report, Institute for Precise Mechanics and Computation Techniques, Academy of Science, Moscow, USSR, 1959.

[11] Robert E. Wengert. A simple automatic derivative evaluation program. Communications of the ACM, 7:463–4, 1964.

[12] Andreas Griewank. On automatic differentiation. pages 83–108, 1989.

[13] Hascoet, Laurent, and Valérie Pascual. "The Tapenade automatic differentiation tool: principles, model, and specification." ACM Transactions on Mathematical Software (TOMS) 39.3 (2013): 1-43.

[14] 知乎專欄:自動微分(Automatic Differentiation): https://zhuanlan.zhihu.com/p/61103504

[15] 知乎專欄:數值計算方法 第六章 數值積分和數值微分: https://zhuanlan.zhihu.com/p/14

[16] 知乎專欄:技術分享 | 從自動微分到可微程式設計語言設計 https://zhuanlan.zhihu.com/p/39