從概念到實踐，我們該如何構建自動微分庫

像 PyTorch 或 TensorFlow 這樣通用的自動微分框架是非常有用和高效的，而且在大多數情況下，幾乎不需要再寫一些更專門化的東西。然而本文作者構建了一個自動微分庫，以高效地計算小批量數據上的訓練。此外，作者還詳細描述了在構建自動微分庫中的過程與思考，是理解自動微分理念的優秀博文。

我最近開始寫自己的 autodiff 程序包。這篇博客文章記錄了我一路以來學到的東西，並把它當成 Julia Evans 的「窮人版」博客文章。

因為有許多博客文章都解釋了自動微分的機制，比我解釋的要好得多，所以這裡我跳過了解釋。此外，在構建神經網路結構方面還有其他一些有趣的文章，因此，雖然我的庫遵循非常相似的模式（靜態計算圖和依賴類型），但我並沒有過多地關注類型系統。

最後，如果你想直接跳轉到代碼部分，最終的結果在以下 GitHub 地址，同時還有一個基於神經網路的 FizzBuzz 解決方案。

自動微分代碼：https://github.com/maciejkula/wyrm

FizzBuzz：https://github.com/maciejkula/fizzbuzz

動機

關於為什麼我想要有自己的 autodiff/backprop 框架，而不是直接用 PyTorch 或者 TensorFlow 的原因有以下幾點。

1.PyTorch 和 TF 在擬合每個 x 小批量所需計算量很少的模型時非常慢。因為在計算機視覺問題中，每個小批量處理的計算量非常大，以至於框架開銷幾乎不成問題。但這在矩陣分解式的模型中卻不可忽略，這種模型在推薦系統中是有用的。且即使在 GPU 上，擬合這些模型也很慢。

2. 我希望能夠用我的 autodiff 庫像 Python 包那樣以最小的依賴關係來編寫和構造模型。這個庫能夠生成一個相當小的、且獨立的二進位文件，這是相對於繁瑣的 TF 和 PyTorch 依賴的優勢。

3. 這是一個有趣的學習經驗，並且讓我更詳細地了解神經網路庫的內部工作機制。

受到對推薦模型（也可能是 NLP) 有效的輕量級解決方案需求的啟發，我編寫了一系列的設計約束條件（design constraints）。

1. 我希望框架能夠自然地支持稀疏梯度：即絕大多數梯度都為零的情況。這在 NLP 和使用大型嵌入層的推薦模型中非常常見。在任何給定的小批量中，只有很小一部分嵌入層被使用，其餘記錄的梯度均為零。在執行梯度更新時能夠跳過零對於快速創建這些模型非常重要。

2. 我希望除實際計算之外，框架有最小的開銷。因為我主要想要擬合小的、稀疏的模型，所以開銷是關鍵。在 PyTorch 中，此類模型的運行時間以 Python 中的循環為主要開銷。為了避免這種情況，我的庫必須在它的擬合循環中放棄 Python，並且需要完全用編譯語言編寫以充分利用編譯器優化的性質。

3. 模型圖必須逐個定義，就像 Chainer 或者 PyTorch 一樣。這種方法的可用性和可調試性對我來說是非常有價值的，以至於我甚至不想回到 TensorFlow 的處理方式。同時，我很高興圖形一旦被定義就是靜態的。這有助於保持較小的開銷：我可以分配一次中間計算緩衝區並繼續使用它們，而不是寫一個複雜的緩衝池系統（或者，更糟糕的是，在每次傳遞的時候不斷地分配和釋放內存）。

4. 我希望性能可以與可用 CPU 內核的數量大致呈線性關係。這意味著在整個圖形的層次上進行並行化，而不是對單獨的操作。每個計算線程將有它自己的計算圖副本，但在更新時寫入共享參數緩衝區。這實際上是 Hogwild! 方法，這個方法中多個計算線程同時更新共享參數緩衝區而沒有任何鎖定。只要梯度相對稀疏，就可以在模型質量下降很少的情況下進行近線性的縮放。

這裡還列出了一些我現在不想添加或不太關心的事情：

1.GPU 支持。我主要想要擬合小型模型（或者至少有很多參數但每個小批量的計算很少的模型）。

2.CNNs，或者，實際上具有兩個維度以上的張量。

考慮到需求（和非需求）列表，我們就能自然地得出一些設計決策。

1. 整個事情將用一種編譯語言（compiled language）編寫，這種編譯語言能夠生成沒有運行時間的本地共享對象，模型也將用相同的語言來定義。

2. 這個語言將會是 Rust，這是一門令人驚嘆的語言，而且非常適合這種任務。因此下面的許多東西都有 Rust 的味道。然而，我所描述的設計權衡在 C++、其他靜態類型和 AOT 編譯的編程語言中是相同的。

3. 我將會使用反向模式自動微分。這樣，我可以很容易地通過多輸入的任意（靜態）計算圖進行反向傳播。

在編寫庫時，我經常想到 API，我希望能夠將這個微分庫公開並獲得社區的幫助。在這種情況下，我想寫如下內容：

let slope = Parameter::new(1.0);

let intercept = Parameter::new(0.0);

let x = Input::new(3.0);

let y = Input::new(2.0 * 3.0 + 1.0);

let loss = (y?—?(slope * x + intercept)).square();

loss.backward();

並讓它工作。

準備工作完成之後，我們可以進入有趣的部分：弄清楚如何實現計算圖。

表示計算圖

我們選擇什麼樣的數據結構來表示計算圖？我了解有以下兩種方案：

1. 基於向量：所有計算節點都被連續地存儲在一個向量中，並使用索引來定址它們的父節點。例如，在創建輸入節點時，對象 InputNode 被壓入向量，且索引為 0。如果隨後將該節點平方，SquareNode 將被壓入索引為 1 的分量，並知道它的父節點是索引 0。在正向傳播過程中，SquareNode 將使用該索引來獲取其輸入的值。

2. 基於圖形。節點被放置在內存中的任意位置，並用指向其父節點的索引來維護計算圖的結構。(向量表示可以看作是圖模型的線性化。)

基於矢量的方法有很多優點。

1. 所有的節點都在同一個地方。他們連續地儲存在內存中，可能會減少內存的定址問題。

2. 他們的所有權很容易解釋。這使得克隆計算圖圖非常簡單：只需克隆節點向量即可。這一點很重要，因為我依靠於為我的並行處理方法提供多個圖的副本。

3. 節點按拓撲順序排列。我們可以通過簡單地沿著向量向前迭代來正確地執行前向傳播，且沒有重複的工作。

但是它也有缺點。

我們在節點向量中存儲了什麼類型的對象是不清楚的。所有的節點類型都不一樣（不同的大小），但向量都是同質的類型。Rust 為這種問題提供了兩種解決方案，但是都不是特別令人滿意。

第一種是枚舉（sum 類型；ADTs； tagged unions）。我們定義一個 Node 類型作為所有可能的節點類型的集合，並將其儲存在節點向量中。這樣，所有的節點就具有相同的類型了。但我們仍然需要將 Node 的方法從封裝的 Node 類型分配到所包含的內部節點。這可以通過模式匹配（聯合類型標籤上的 switch 語句）完成；有 Rust 對模式識別和宏的支持，編寫必要的代碼是輕而易舉的。

但是，這種做法會增加運行時間成本。每次我們使用一個節點，我們需要經過一個 switch 語句來解決內部類型問題。原則上，優化編譯器會將這種代碼編譯成跳轉表（jump tables）。實際上，在我的實驗中為分配代碼生成的程序集僅僅是對所有可能性的線性掃描，強加了與框架支持的具體節點類型數量呈線性關係的分配成本。更糟的是，編譯器不願意內聯 switch 本身和被調用的函數。前者是因為它增加了分支預測的失誤，後者增加了函數調用的開銷。（最近的分值預測攻擊加劇了這個問題： compiler mitigations 可能會導致像這樣的間接指令更加昂貴。）

對節點向量使用 sum 類型的最後一個缺點是它會導致一個封閉的系統（類似於 Scala『s 的 封閉特性）：庫的下游用戶不能添加新的節點類型。

另一種方法是用 Rust 的運行時多態機制（polymorphism mechanism）： trait objects。trait objects 是對目標具體類型進行抽象的一種方法：我們將他們隱藏在指向數據的指針和他們方法表的後面，而不是將結構存儲在內聯中。調用方法時，我們跳轉到 vtable，找到函數並執行。通過使用 trait ojbects，我們將這些 fat pointers 放到節點向量中而不是節點自身裡面。

然而，這種解決方案恰恰引入了我們開始時想要避免的那種間接性。此外，它完全否認了編譯器在內聯方面做的努力：被調用的函數直到運行時才知道。

那麼基於圖的設計呢？在這裡，每個節點都在內存中被放置在自己的位置，並且可以通過索引指向其祖先。因為每個節點可以重複使用任意次數，我用 Rust 中的 Rc相當於 C++中的 shared_ptr。

這種方法的一個直接缺點是模糊了圖的所有權結構，使克隆和序列化/反序列化變得困難：因為節點可以被重複利用，單純的克隆/反序列化將導致創建相同節點的多個副本。

第二個缺點是缺少一個容易獲得的拓撲排序：前向和後向傳遞都遞歸地完成，而且必須小心地避免重複計算共享子圖的值。

使用圖形表達的優點是在編譯時已知任何節點的父節點類型。每一個節點在其父節點類型上是（遞歸地）通用的：添加兩個 InputNodes 將會產生一個 AddNode。將其添加到另一個輸入節點會產生 AddNode，InputNode>等等。除了一個在類型系統中表現更好的設計，這給了我分配和內聯的靜態方法。

結果

使用一些非正式的基準，基於圖的方法比基於向量的方法快大約 30%。最後的結果可以在我很普通的雙核筆記本上，20 毫秒內在 Movielens 100K 數據集上完整地運行一個 BPR 學習-排序分解模型。此外，它的性能會隨著處理器內核的增加而線性增長。

除了底層的圖形結構之後，這裡還利用了很多優化。

1. 我用 Rust 的 SIMD 內在函數進行了很多操作，如向量點積和標量加法。

2. 對於大多數操作，我假定 C 為連續矩陣並直接在底層數據上迭代，而不是用 ndarrays 迭代方法。事實證明，這樣做要快得多，大概是因為它允許 LLVM 自動對向量實現向量化。

3. 事實證明，LLVM 足夠智能，能夠自動向量化大部分不涉及縮減步驟（主要是賦值）的數值循環。與（2）結合起來看，這種方法使得很多的數值循環以最小的優化努力獲得更高的效率。

仍然有很多方法可以使計算速度更快。

1. 此時，代碼在正向傳遞中不會緩存任何子圖的結果：如果一個節點在正向傳遞中被用了兩次，所有依賴它的計算將會執行兩次。這可以通過一個簡單的拓撲排序演算法很容易的解決，一旦評估了它們的值，就將該節點標記為已評估。

2. 類似地，在後向傳遞中梯度被直接傳遞給參數節點。如果一個節點被多次使用，這意味著在逐步向下傳遞梯度時做了不必要的工作。累積所有的梯度並且只遞歸一次將節省這項工作。

3. 對輸入有一些不必要的複製，在可能的情況下更好的利用索引應該會產生一些小的性能收益。

下一步是什麼

我寫了（並繼續維護）很多的開源 Python ML 包。這些模型是在 Cython 中手工編寫的，儘管它們表現的很好，但是擴展它們是困難的。部分是因為 Cython 的局限性，另一部分的原因在於手動派生更新規則所需的努力。

我希望這個庫（或它的一些變體）可以使這個任務變得簡單一些，並且可以讓我更輕鬆地實現複雜模型以將它們作為獨立的 Python 包發布出去。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！