當前，PyTorch、TensorFlow 等機器學習框架已經成為了人們開發的重要工具。計算反向傳播、貝葉斯推理、不確定性量化和機率程式設計等演算法的梯度時，我們需要把所有的程式碼以微分型寫入框架內。這對於將機器學習引入新領域帶來了問題：在物理模擬、遊戲引擎、氣候模型中，原領域元件不是由機器學習框架的特定領域語言（DSL）編寫的。因此在將機器學習引入科學計算時，重寫需求成為了一個挑戰。為了解決這一問題，現在的發展趨勢包含構建新的 DSL，讓重寫過程變得簡單，或者在程式設計時直接進行構建。這些方法可以讓我們獲得有效的梯度，但是仍然需要使用 DSL 或可微分的程式語言進行重寫。開發者們自然會希望使用機器學習框架以外的程式碼重用已有工具，或在一種語言裡寫好損失函式，然後為其用例提供更簡單的抽象。儘管目前已經出現了一些針對不同語言的反向自動微分框架（AD），但我們很難在 ML 框架外的程式碼上自動使用它們，因為其仍需要重寫，且對於跨語言 AD 和庫的支援有限。為了方便開發者，來自 MIT 的研究者開源了 Enzyme，一種用於 LLVM 編譯器框架的高效能自動微分（AD）編譯器外掛。該外掛能夠合成以 LLVM 中間表示（IR）表示的靜態可分析程式的梯度。Enzyme 能夠合成任何以面向 LLVM IR 編譯器為語言編寫的程式的梯度，包括 C、C ++、Fortran、Julia、Rust、Swift、MLIR 等，從而提供這些語言的本機 AD 功能。據作者介紹，與傳統的源到源和 operator-overloading 工具不同，Enzyme 在最佳化的 IR 上執行 AD。在包括微軟 ADBench 在內的以機器學習為重點的基準套件上，經過最佳化的 IR 上的 AD 的幾何平均速度比未經過最佳化的 IR 上 AD 的幾何平均速度提高了 4.5 倍，這使得 Enzyme 達到了最高的效能。此外，Enzyme 方便使用，在 PyTorch 和 TensorFlow 上都有程式包，可讓開發者便捷訪問具有最新效能的外來程式碼梯度，從而使外來程式碼可直接合併到現有的機器學習工作流程中。Enzyme，一種用於 LLVM 的編譯器外掛，可以合成可靜態微分的 LLVM IR 的快速梯度。包括 C、C ++、Fortran、Rust、Swift 等編譯器前端生成的中間表示（IR）。
PyTorch-Enzyme/TensorFlow-Enzyme，一個外部功能介面，允許機器學習研究者使用 PyTorch 和 TensorFlow 使用以 LLVM 編譯語言編寫的外部程式碼。
Enzyme.jl，一個 Julia 包，透過動態高階語言編寫的程式碼，僅使用低層資訊獲得梯度。
透過連結時最佳化（LTO）支援了多來源 AD 和靜態庫支援。
研究表明最佳化後執行 AD，在標準機器學習基準測試上可獲得顯著效能提升，並達到 SOTA 水平。
上圖為 relu(pow(x,3)) 的梯度合成示例。左側為 LLVM IR 上的原始計算。左側註釋中展示了將新增到前向傳遞中的活動變數的影子分配。右側則是 Enzyme 將生成的反向傳遞。完整的合成梯度函式將結合使用這些函式（新增影子分配），將 if.end 中的返回替換為 reverse_if.end 的分支。Enzyme 專案是一個用於可靜態分析 LLVM IR 的反向模式自動微分（AD）工具。它允許開發者可以自動建立基於原始碼的梯度，而無需更多額外工作。double foo(double);double grad_foo(double x) { return __enzyme_autodiff(foo, x);}
透過最佳化微分後的程式碼，Enzyme 可以比現有的最佳化工具提供更快的微分速度：
可選的預處理階段，該階段執行對 AD 有用的較小轉換。
一種新的過程間型別分析，可推斷出記憶體位置的基礎型別。
活動分析，確定哪些指令或值會影響導數計算（在現有 AD 系統中很常見）。
最佳化遍歷可建立任何必需的派生函式，用生成的函式替換對__enzyme_autodiff 的呼叫。
更多詳細介紹，可檢視 MIT 研究者們提交的 NeurIPS 2020 論文：
論文地址：https://arxiv.org/pdf/2010.01709.pdfPowered by Froala Editor