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

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

Karpathy 同時也回復(fù)到：「我們正向前移動了一點，與原來對程序有完整的定義不同，我們現(xiàn)在只是寫一個大致的架構(gòu)，這樣的架構(gòu)會通過權(quán)重參數(shù)化。如果我們有一個好的評估標準，那么最優(yōu)化算法就能幫我們找到更好的解?！?/span>

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

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

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

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

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

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

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

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

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

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

科學計算和機器學習的聯(lián)姻：可微編程

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

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

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

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

用 Zygote 對正弦函數(shù)進行微分求導(dǎo)

研究人員以正弦函數(shù)為例，解釋可微編程的方法。如圖為正弦函數(shù) 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 語言可以更好的表示正弦函數(shù)，但是為了展示計算過程，這里使用了一個循環(huán)、一個條件語句、「isodd」和「factorial」函數(shù)，這些都是 Julia 的原生實現(xiàn)。在微分計算中，自動微分直接工作。如下為 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 很多優(yōu)秀的屬性，只要調(diào)用幾個 Julia 語言的一些便捷特性就能完成。

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

Zygote 還可以用于包裝模型損失的計算，明確指明計算邊界應(yīng)該對于模型梯度的計算是可微的，但是其它部分的代碼，包括我們定義的 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 的編程風格，基本上只需要搭建模型、算梯度、更新參數(shù)等幾個重要步驟。即使用戶不太了解 Julia 的語法規(guī)則，但還是很容易讀懂的，它非常重要的部分就是通過 Zygote 求梯度。

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

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

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

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

其他案例

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

投石機問題——強化學習

圖 3：使用神經(jīng)網(wǎng)絡(luò)的替代解決反問題（inverse problem）

基于模型的強化學習對無模型強化學習具有優(yōu)勢，因為高性能的智能體必須能夠?qū)ζ洵h(huán)境動態(tài)進行擬合。但是，深度強化學習模型不能夠和真實環(huán)境結(jié)合，限制了它的發(fā)展。過去的研究通常使用機器學習框架，可以成功實現(xiàn)現(xiàn)實的物理引擎，但是這種工程的投入非常大，并且和現(xiàn)有的引擎相比有限制，因此在生物學或氣象學等領(lǐng)域中的應(yīng)用較少。

Zygote 可以被用來解決控制問題，能夠?qū)⒛Ｐ偷姆聪騻鞑ミ^程控制在一個對梯度的請求中。研究人員使用了投石機問題作為實驗案例。他們優(yōu)化一個可以應(yīng)對多個目標的神經(jīng)網(wǎng)絡(luò)，神經(jīng)網(wǎng)絡(luò)的輸入是目標的距離和當前的風速，而網(wǎng)絡(luò)的輸出則是投石機的設(shè)置（如射程和發(fā)射角度）。輸入來自一個模擬器，可以作為一個常微分方差，并計算落地點到目標的距離。研究人員然后對比實際結(jié)果，并進行反向傳播，調(diào)整網(wǎng)絡(luò)權(quán)重。研究人員的數(shù)據(jù)集是隨機從一系列目標和風速中選出的。這個智能體可以在筆記本電腦的 CPU 上訓練，并可以在常數(shù)時間內(nèi)解決反問題，較直接優(yōu)化投石機系統(tǒng)的強化學習方法快 100 倍。

計算機視覺

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

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：生成的圖片和點光源的關(guān)系。4（左）是最開始生成的圖像，5（中）是迭代 100 次的結(jié)果，6（右）目標圖像。