作者｜澤南、亞洲

?來源｜機器之心

麻省理工學院（MIT）提出的自動微分框架 Enzyme 在 NeurIPS 2020 大會上引起過不少人的興趣，現在人們已經可以嘗試使用這套工具了。

當前，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 上都有程序包，可讓開發者便捷訪問具有最新性能的外來代碼梯度，從而使外來代碼可直接合并到現有的機器學習工作流程中。

論文標題：

Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradient

論文鏈接：

https://arxiv.org/abs/2010.01709

項目網頁：

https://enzyme.mit.edu/

項目地址：

https://github.com/wsmoses/Enzyme

作者表示，Enzyme 可提供這些工具和能力：

• 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 項目概述

Enzyme 項目是一個用于可靜態分析 LLVM IR 的反向模式自動微分（AD）工具。它允許開發者可以自動創建基于源代碼的梯度，而無需更多額外工作。

double foo(double);double grad_foo(double x) { return __enzyme_autodiff(foo, x); }

通過優化微分后的代碼，Enzyme 可以比現有的優化工具提供更快的微分速度：

組件

Enzyme 工具主要由四部分組成：

• 可選的預處理階段，該階段執行對 AD 有用的較小轉換。

• 一種新的過程間類型分析，可推斷出內存位置的基礎類型。

• 活動分析，確定哪些指令或值會影響導數計算（在現有 AD 系統中很常見）。

• 優化遍歷可創建任何必需的派生函數，用生成的函數替換對__enzyme_autodiff 的調用。

更多詳細介紹，可查看 MIT 研究者們提交的 NeurIPS 2020 論文。

????

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

關于PaperWeekly

PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

總結

以上是生活随笔為你收集整理的MIT开源高性能自动微分框架Enzyme：速度提升4.5倍的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。