我最近開始寫自己的 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<T>相當於 C++中的 shared_ptr。
這種方法的一個直接缺點是模糊了圖的所有權結構，使克隆和序列化/反序列化變得困難：因為節點可以被重複利用，單純的克隆/反序列化將導致建立相同節點的多個副本。
第二個缺點是缺少一個容易獲得的拓撲排序：前向和後向傳遞都遞迴地完成，而且必須小心地避免重複計算共享子圖的值。
使用圖形表達的優點是在編譯時已知任何節點的父節點型別。每一個節點在其父節點型別上是（遞迴地）通用的：新增兩個 InputNodes 將會產生一個 AddNode<InputNode, InputNode>。將其新增到另一個輸入節點會產生 AddNode<AddNode<InputNode, InputNode>，InputNode>等等。除了一個在型別系統中表現更好的設計，這給了我分配和內聯的靜態方法。
結果
使用一些非正式的基準，基於圖的方法比基於向量的方法快大約 30%。最後的結果可以在我很普通的雙核筆記本上，20 毫秒內在 Movielens 100K 資料集上完整地執行一個 BPR 學習-排序分解模型。此外，它的效能會隨著處理器核心的增加而線性增長。
除了底層的圖形結構之後，這裡還利用了很多最佳化。
1. 我用 Rust 的 SIMD 內在函式進行了很多操作，如向量點積和標量加法。
2. 對於大多數操作，我假定 C 為連續矩陣並直接在底層資料上迭代，而不是用 ndarrays 迭代方法。事實證明，這樣做要快得多，大概是因為它允許 LLVM 自動對向量實現向量化。
3. 事實證明，LLVM 足夠智慧，能夠自動向量化大部分不涉及縮減步驟（主要是賦值）的數值迴圈。與（2）結合起來看，這種方法使得很多的數值迴圈以最小的最佳化努力獲得更高的效率。
仍然有很多方法可以使計算速度更快。
1. 此時，程式碼在正向傳遞中不會快取任何子圖的結果：如果一個節點在正向傳遞中被用了兩次，所有依賴它的計算將會執行兩次。這可以透過一個簡單的拓撲排序演算法很容易的解決，一旦評估了它們的值，就將該節點標記為已評估。
2. 類似地，在後向傳遞中梯度被直接傳遞給引數節點。如果一個節點被多次使用，這意味著在逐步向下傳遞梯度時做了不必要的工作。累積所有的梯度並且只遞迴一次將節省這項工作。
3. 對輸入有一些不必要的複製，在可能的情況下更好的利用索引應該會產生一些小的效能收益。
下一步是什麼
我寫了（並繼續維護）很多的開源 Python ML 包。這些模型是在 Cython 中手工編寫的，儘管它們表現的很好，但是擴充套件它們是困難的。部分是因為 Cython 的侷限性，另一部分的原因在於手動派生更新規則所需的努力。
我希望這個庫（或它的一些變體）可以使這個任務變得簡單一些，並且可以讓我更輕鬆地實現複雜模型以將它們作為獨立的 Python 包釋出出去。
附錄
結果表明，當圖形表達應用到遞迴神經網路時有一些問題：在遞迴的每一步，結果型別的複雜度增加，導致了相當奇怪的型別：
Variable<nodes::LogNode<nodes::SoftmaxNode<nodes::DotNode<layers::recurrent::LSTMCellHidden<layers::recurrent::LSTMCellState<layers::recurrent::LSTMCellSt
ate<layers::recurrent::LSTMCellState<nodes::InputNode, nodes::InputNode, nodes::IndexNode<nodes::ParameterNode>>, layers::recurrent::LSTMCellHidden<nodes::InputNode, nodes::InputNode, nodes::IndexNode<nodes::Par
ameterNode>>, nodes::IndexNode<nodes::ParameterNode>>, layers::recurrent::LSTMCellHidden<layers::recurrent::LSTMCellState<nodes::InputNode, nodes::InputNode, nodes::IndexNode<nodes::ParameterNode>>, layers::recu
rrent::LSTMCellHidden<nodes::InputNode, nodes::InputNode, nodes::IndexNode<nodes::ParameterNode>>, nodes::IndexNode<nodes::ParameterNode>>, nodes::IndexNode<nodes::ParameterNode>>, layers::recurrent::LSTMCellHid
den<layers::recurrent::LSTMCellState<layers::recurrent::LSTMCellState<nodes::InputNode, nodes::InputNode, nodes::IndexNode<nodes::ParameterNode>>, layers::recurrent::LSTMCellHidden<nodes::InputNode, nodes::Input
Node, nodes::IndexNode<nodes::ParameterNode>>, nodes::IndexNode<nodes::ParameterNode>>, layers::recurrent::LSTMCellHidden<layers::recurrent::LSTMCellState<nodes::InputNode, nodes::InputNode, nodes::IndexNode<nod
es::ParameterNode>>, layers::recurrent::LSTMCellHidden<nodes::InputNode, nodes::InputNode, nodes::IndexNode<nodes::ParameterNode>>, nodes::IndexNode<nodes::ParameterNode>>, nodes::IndexNode<nodes::ParameterNode>
>, nodes::IndexNode<nodes::ParameterNode>>, nodes::ParameterNode>>>>
不用說，在經過一些迭代步驟後，編譯器放棄了。這可以透過實現一個融合的 LSTM 單元來解決，而不是將其從更簡單的操作中組裝起來，或者選擇透過 trait objects 選擇性擦除。目前為止，我已經使用了第二種方案：透過將每個 LSTM 單元的輸出值裝入 trait object 來將其具體型別刪除。
原文連結：https://medium.com/@maciejkula/building-an-autodifferentiation-library-9ccf32c7a658