近日，Julia Computing 團隊發表論文表示他們構建了一種可微編程系統，它能將自動微分內嵌于 Julia 語言，從而將其作為第一級的語言特性。也就是說，我們以后直接用 Julia 語言及可微編程就能寫模型了？都不需要再調用 TensorFlow 或 PyTorch 這樣的框架了？

如果我們將可微編程系統視為編程語言最重要的特性之一，那么不論是機器學習還是其它科學計算都將方便不少，這樣的語言也將是科學計算最好的語言。Y Combinator Research 研究者 Michael Nielsen 對此也非常興奮，他非常贊同 Andrej Karpathy 所說的「梯度下降是更好的程序員」。

Karpathy 同時也回復到：「我們正向前移動了一點，與原來對程序有完整的定義不同，我們現在只是寫一個大致的架構，這樣的架構會通過權重參數化。如果我們有一個好的評估標準，那么最優化算法就能幫我們找到更好的解。」

這里 Karpathy 說的就是機器學習和編程的區別，ML 會通過梯度下降等最優化方法自動搜索最優解。但這里有個問題，模型需要梯度才能向著最優前進，因此模型的很多部分都要求是可微的。鑒于這一點，很多人也就將 ML 稱呼為可微編程了。

但是可微編程只能用于機器學習嗎？它能不能擴展到其它領域，甚至成為編程語言的基本特性？答案是可以的，這就是 Julia 團隊及 MIT 等其他研究機構正在嘗試的。

近年來，機器學習模型越來越精妙，展現出了很多科學計算的特性，側面凸顯了機器學習框架的強大能力。研究者表示，由于廣泛的科學計算和機器學習領域在底層結構上都需要線性代數的支持，因此有可能以可微編程的形式，創造一種新的計算基礎設施。

• 論文地址：https://arxiv.org/pdf/1907.07587.pdf

• GitHub項目地址：https://github.com/MikeInnes/zygote-paper

可微編程與 DL 框架又有什么不同

在論文中，研究人員提出了一種可微編程系統，它能在 Julia 語言中完成梯度計算，并成為 Julia 語言的第一級特性。這使得構建深度學習模型變得更直觀。更重要的是，這令用戶可以使用已有的 Julia 科學計算包去構建深度學習模型，并高效實現梯度計算。

那么可微編程與 TensorFlow 或 PyTorch 等 DL 框架又有什么不同，它的速度難道還要快于這些構建細算圖的系統？研究者表示，可微編程會執行 source-to-source 的轉換，自動微分轉換基本上沒有運行時開銷，因此它要比反向傳播的實際計算成本更有優勢。

在 Reddit 上也有開發者討論到了這一點：

 一位名為「Coconut_island」開發者表示，Zygote 和現有的機器學習框架的不同之處在于，Zygote 將源代碼視為計算圖，省略了中間語言的過程。更重要的是，通過控制流進行子梯度計算，這種梯度計算的方法可以應用在任何一種控制流分支上，具有代碼的抽象性。

科學計算和機器學習的聯姻：可微編程

咋一看，科學計算和機器學習分屬不同的領域。現代機器學習在神經網絡上的突破使得整個學科取得了極大的進步。神經網絡可以解決很多不同的計算機科學問題，使得人們開始設計新的硬件和軟件，以便于提升性能，處理極大量的標注數據，并能夠同時將訓練好的模型部署在設備上。科學計算與之相反，它擁有非常悠久的歷史，且希望使用一系列從物理現象中獲取的建模技術。

和典型的機器學習研究者不同，許多科學計算學家在數據量較少、但計算復雜度更高和范圍更廣的數據上進行研究。但是觀察來看，科學計算和機器學習有共同點。兩個領域都傾向于使用動態編程語言進行計算，如 Python、R 和 Julia。通常來說，Python 和 R 的關鍵性能模塊都使用 C++和 Fortran，而 Julia 中則較少。而且兩者的核心計算流程都是基于線性代數的，并且有硬件專門用于加速這種計算。

由于機器學習和科學計算都依賴線性代數，且已有大量硬件專門優化用于加速運算，因此可微編程具有進一步融合兩者的潛力。使用何種語言完成這項工作呢？由于編程工作量的問題，研究人員僅選擇在 Julia 語言上增加可微編程能力。選擇這門語言的一個原因是，Julia 語言已有大量的機器學習和科學計算包，都是純粹的 Julia 語言實現。這可以幫助他們在相對較多的應用中進行測試。

研究人員表示，已有的 Julia 程序包可以直接使用他們的系統，包括處理用戶定義類、基于狀態的控制流，以及通過「source-to-source」自動微分進行大量標量運算。論文展示了一些使用可微編程的案例。

用 Zygote 對正弦函數進行微分求導

研究人員以正弦函數為例，解釋可微編程的方法。如圖為正弦函數 sin(X) 的泰勒展開：

這個公式可以使用 Julia 語言進行定義，為了運行，請安裝 Julia 語言 1.1 以上版本，并安裝 Zygote.jl 和 ForwardDiff.jl 包，代碼如下所示：

using Pkg 
Pkg.add("Zygote") 
Pkg.add("ForwardDiff") 
using Zygote, ForwardDiff 
function s(x) 
  t = 0.0 
  sign = -1.0 
  for i in 1:19 
    if isodd(i) 
      newterm = x^i/factorial(i) 
      abs(newterm)<1e-8 && return t 
      println("i=",i) 
      sign = -sign 
      t += sign * newterm 
     end  
   end 
   return t 
end 

盡管用 Julia 語言可以更好的表示正弦函數，但是為了展示計算過程，這里使用了一個循環、一個條件語句、「isodd」和「factorial」函數，這些都是 Julia 的原生實現。在微分計算中，自動微分直接工作。如下為 x = 1.0 開始，i 的變化情況，計算得到的梯度，以及何時這個梯度和 cos(1.0) 的梯度匹配。

julia> ForwardDiff.derivative(s, 1.0) # Forward Mode AD 
i=1 
i=3 
i=5 
i=7 
i=9 
i=11 
0.540302303791887 

julia> Zygote.gradient(s, 1.0) # Reverse Mode AD 
i=1 
i=3 
i=5 
i=7 
i=9 
i=11 
(0.5403023037918872,) 
 
julia> cos(1.0) 
0.540302305868139 

可微編程中的深度學習

Zygote 是用來計算深度學習模型梯度的靈活后端，它的效率非常高。下面展示了一個典型的例子，研究人員使用 LSTM 學習莎翁的作品。下面的代碼展示了 Zygote 很多優秀的屬性，只要調用幾個 Julia 語言的一些便捷特性就能完成。

首先，定義的模型并沒有特殊的數據類型，因此可以直接啟用自動微分（AD）；模型的定義只需要確定前向傳播的計算流就行了，反向傳播只有在定義 BLAS 運算和基本數組操作等基本構建塊時才需要手動定義。

Zygote 還可以用于包裝模型損失的計算，明確指明計算邊界應該對于模型梯度的計算是可微的，但是其它部分的代碼，包括我們定義的 LSTM 層級，在寫的時候都不需要考慮自動微分過程。最后模型可以在 CPU、GPU 和谷歌的 TPU 上運行，它們不需要或只需要一點修改。

# Load data and alphabet 
alphabet, Xs, Ys = load_data("shakespeare_input.txt") 
# Define simple LSTM-based model to map from alphabet back on to alphabet, 
# predicting the next letter in a corpus of Shakespeare text. 
model = Chain( 
 LSTM(length(alphabet), 128), 
 LSTM(128, 128), 
 Dense(128, length(alphabet)), 
 softmax, 
) 
 
opt = ADAM(0.01) 
 
# Run through our entire dataset a few times 
for epoch_idx in 1:10, 
  (x_batch, y_batch) in zip(Xs, Ys)  
   
 # Calculate gradients upon the model for this batch of data, 
 # summing crossentropy loss across each time index in this batch 
 grads = |\Zygoteplain|.gradient(model) do model 
 return sum(crossentropy.(model.(x_batch), y_batch))  
 end 
 
# Update the model, then reset its internal LSTM states 
model = update!(opt, model, grads) 
Flux.reset!(model) 
end 

如上所示整個模型及訓練過程看起來非常簡潔，有點像 Keras 或 PyTorch 的編程風格，基本上只需要搭建模型、算梯度、更新參數等幾個重要步驟。即使用戶不太了解 Julia 的語法規則，但還是很容易讀懂的，它非常重要的部分就是通過 Zygote 求梯度。

Zygote 提供了一個成本極其低的自動微分接口。通過執行 source-to-source 的轉換，AD 轉換基本上沒有運行時開銷，因此它要比反向傳播的實際計算成本更有優勢。此外，Zygote 在 TPU pod 上訓練 ResNet 與 TensorFlow 有相同的性能水平。

表 1：在不同深度的 LSTM 網絡中，Zygote 每一個運算（OP）開銷的估計值。

 為了度量上面的結果，研究者對 LSTM 網絡的反向傳播做了一個基準測試，并研究運行時間與批量大小的關系，從而進一步估計自動微分系統中每一個操作的固定開銷。研究者最后在 i5 CPU 和 Julia 1.3 上完成了測試，結果如表 1 所示非常有競爭力。因為像 PyTorch 那樣的成熟深度學習框架，每個 op 的實際開銷都至少是 1µs。

這些消失的開銷提高了 AD 系統的效率和使用門檻，因此它能以非常細粒度的岔村集成到編程語言中，且還無需擔心性能問題。一般而言，自動微分的開銷越低，AD 系統的最小可行的內核就越小。因此我們只需要考慮反向傳播效率，就能設計出又小又高效的自動微分內核。

其他案例

除了用 Julia 實現 LSTM 模型的自動微分以外，論文還提供了使用可微分編程在其他機器學習領域中的應用示例。

投石機問題——強化學習

圖 3：使用神經網絡的替代解決反問題（inverse problem）

基于模型的強化學習對無模型強化學習具有優勢，因為高性能的智能體必須能夠對其環境動態進行擬合。但是，深度強化學習模型不能夠和真實環境結合，限制了它的發展。過去的研究通常使用機器學習框架，可以成功實現現實的物理引擎，但是這種工程的投入非常大，并且和現有的引擎相比有限制，因此在生物學或氣象學等領域中的應用較少。

Zygote 可以被用來解決控制問題，能夠將模型的反向傳播過程控制在一個對梯度的請求中。研究人員使用了投石機問題作為實驗案例。他們優化一個可以應對多個目標的神經網絡，神經網絡的輸入是目標的距離和當前的風速，而網絡的輸出則是投石機的設置（如射程和發射角度）。輸入來自一個模擬器，可以作為一個常微分方差，并計算落地點到目標的距離。研究人員然后對比實際結果，并進行反向傳播，調整網絡權重。研究人員的數據集是隨機從一系列目標和風速中選出的。這個智能體可以在筆記本電腦的 CPU 上訓練，并可以在常數時間內解決反問題，較直接優化投石機系統的強化學習方法快 100 倍。

計算機視覺

在本案例中，研究人員使用了一個原型生成器，基于期望最終生成的圖像，展示優化點光源位置的過程。研究人員定義了一個損失函數，將點光源作為輸入，產生圖像，并和參考圖像對比。按照慣例，梯度可以被提取，并用于更新點光源的位置。

julia> guess = PointLight(Vec3(1.0), 20000.0, Vec3(1.0, 2.0, -7.0)) 
 
julia> function loss_function(light) 
 rendered_color = raytrace(origin, direction, scene, light, eye_pos) 
 rendered_img = process_image(rendered_color, screen_size.w, 
 screen_size.h) 
 return mean((rendered_img .- reference_img).^2) 
 end 
 
julia> gs = gradient(x -> loss_function(x, image), guess) 

 圖 4、5、6：生成的圖片和點光源的關系。4（左）是最開始生成的圖像，5（中）是迭代 100 次的結果，6（右）目標圖像。