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

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

Karpathy 同時(shí)也回復(fù)到：「我們正向前移動(dòng)了一點(diǎn)，與原來(lái)對(duì)程序有完整的定義不同，我們現(xiàn)在只是寫一個(gè)大致的架構(gòu)，這樣的架構(gòu)會(huì)通過(guò)權(quán)重參數(shù)化。如果我們有一個(gè)好的評(píng)估標(biāo)準(zhǔn)，那么最優(yōu)化算法就能幫我們找到更好的解。」

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

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

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

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

• GitHub項(xiàng)目地址：https://github.com/MikeInnes/zygote-paper

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

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

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

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

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

科學(xué)計(jì)算和機(jī)器學(xué)習(xí)的聯(lián)姻：可微編程

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

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

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

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

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

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

這個(gè)公式可以使用 Julia 語(yǔ)言進(jìn)行定義，為了運(yùn)行，請(qǐng)安裝 Julia 語(yǔ)言 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 語(yǔ)言可以更好的表示正弦函數(shù)，但是為了展示計(jì)算過(guò)程，這里使用了一個(gè)循環(huán)、一個(gè)條件語(yǔ)句、「isodd」和「factorial」函數(shù)，這些都是 Julia 的原生實(shí)現(xiàn)。在微分計(jì)算中，自動(dòng)微分直接工作。如下為 x = 1.0 開(kāi)始，i 的變化情況，計(jì)算得到的梯度，以及何時(shí)這個(gè)梯度和 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 

可微編程中的深度學(xué)習(xí)

Zygote 是用來(lái)計(jì)算深度學(xué)習(xí)模型梯度的靈活后端，它的效率非常高。下面展示了一個(gè)典型的例子，研究人員使用 LSTM 學(xué)習(xí)莎翁的作品。下面的代碼展示了 Zygote 很多優(yōu)秀的屬性，只要調(diào)用幾個(gè) Julia 語(yǔ)言的一些便捷特性就能完成。

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

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

# 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 

如上所示整個(gè)模型及訓(xùn)練過(guò)程看起來(lái)非常簡(jiǎn)潔，有點(diǎn)像 Keras 或 PyTorch 的編程風(fēng)格，基本上只需要搭建模型、算梯度、更新參數(shù)等幾個(gè)重要步驟。即使用戶不太了解 Julia 的語(yǔ)法規(guī)則，但還是很容易讀懂的，它非常重要的部分就是通過(guò) Zygote 求梯度。

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

表 1：在不同深度的 LSTM 網(wǎng)絡(luò)中，Zygote 每一個(gè)運(yùn)算（OP）開(kāi)銷的估計(jì)值。

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

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

其他案例

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

投石機(jī)問(wèn)題——強(qiáng)化學(xué)習(xí)

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

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

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

計(jì)算機(jī)視覺(jué)

在本案例中，研究人員使用了一個(gè)原型生成器，基于期望最終生成的圖像，展示優(yōu)化點(diǎn)光源位置的過(guò)程。研究人員定義了一個(gè)損失函數(shù)，將點(diǎn)光源作為輸入，產(chǎn)生圖像，并和參考圖像對(duì)比。按照慣例，梯度可以被提取，并用于更新點(diǎ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：生成的圖片和點(diǎn)光源的關(guān)系。4（左）是最開(kāi)始生成的圖像，5（中）是迭代 100 次的結(jié)果，6（右）目標(biāo)圖像。