過去十年來，深度學(xué)習(xí)領(lǐng)域發(fā)展迅速，其一大主要推動(dòng)力便是并行化。通過 GPU 和 TPU 等專用硬件加速器，深度學(xué)習(xí)中廣泛使用的矩陣乘法可以得到快速評(píng)估，從而可以快速執(zhí)行試錯(cuò)型的深度學(xué)習(xí)研究。

盡管并行化已經(jīng)在深度學(xué)習(xí)研究中得到了廣泛的使用，但循環(huán)神經(jīng)網(wǎng)絡(luò)（RNN）和神經(jīng)常微分方程（NeuralODE）等序列模型卻尚未能完全受益于此，因?yàn)樗鼈儽旧硇枰獙?duì)序列長(zhǎng)度執(zhí)行序列式的評(píng)估。

序列評(píng)估已經(jīng)變成了訓(xùn)練序列式深度學(xué)習(xí)模型的瓶頸。這一瓶頸可能會(huì)使人們關(guān)注的研究方向偏離序列模型。

舉個(gè)例子，注意力機(jī)制和 transformer 在近些年中超過 RNN 成為了語言建模的主導(dǎo)技術(shù)，部分原因就是它們能以并行的方式訓(xùn)練。連續(xù)歸一化流（CNF）過去常使用的模型是 NeuralODE，現(xiàn)在卻轉(zhuǎn)向了訓(xùn)練過程不涉及到模擬 ODE 的新方向。

近期有一些嘗試復(fù)興序列 RNN 的研究工作，但它們的重心都是線性循環(huán)層 —— 可使用前綴掃描（prefix scan）來進(jìn)行并行化地評(píng)估，非線性循環(huán)層在其序列長(zhǎng)度上依然無法并行化。

近日，英國(guó) Machine Discovery 公司和牛津大學(xué)的一篇論文提出了一種新算法，可將 RNN 和 NeuralODE 等非線性序列模型的評(píng)估和訓(xùn)練工作并行化，并且他們宣稱這一算法還不會(huì)在「合理的數(shù)值精度」內(nèi)改變模型的輸出。

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

那么他們是怎么做到這一點(diǎn)的呢？

據(jù)介紹，他們引入了一種用于求解非線性微分方程的通用框架，其做法是將這些方程重新表述為二次收斂的定點(diǎn)迭代問題，這相當(dāng)于牛頓求根法。定點(diǎn)迭代涉及到可并行運(yùn)算和一個(gè)可并行地評(píng)估的逆線性算子，即使是對(duì)于 RNN 和 ODE 這樣的序列模型也可以。

由于是二次收斂，所以定點(diǎn)迭代的數(shù)量可以相當(dāng)小，尤其是當(dāng)初始起點(diǎn)接近收斂的解時(shí)。在訓(xùn)練序列模型方面，這是一個(gè)相當(dāng)吸引人的功能。由于模型參數(shù)通常是漸進(jìn)式更新的，所以之前訓(xùn)練步驟的結(jié)果可以被用作初始起點(diǎn)。

最重要的是，研究者表示，新提出的算法無需序列模型具備某種特定結(jié)構(gòu)，這樣一來，用戶不必改變模型的架構(gòu)也能收獲并行化的好處。 

DEER 框架：將非線性微分方程視為定點(diǎn)迭代

DEER 框架具有二次收斂性，并且與牛頓法存在關(guān)聯(lián)。這一框架可以應(yīng)用于一維微分方程（即 ODE），也可用于更高維的微分方程（即偏微分方程 / PDE）。該框架還可以應(yīng)用于離散差分方程以達(dá)到相同的收斂速度，這一特性可以應(yīng)用于 RNN。

使用該框架，用戶可以設(shè)計(jì)一種用于評(píng)估 RNN 和 ODE 的并行算法，并且不會(huì)對(duì)結(jié)果產(chǎn)生明顯的影響。

DEER 框架

令我們感興趣的輸出信號(hào)為 y (r)，其由 n 個(gè)在 d 維空間的信號(hào)構(gòu)成，其坐標(biāo)表示為 r。輸出信號(hào) y (r) 可能依賴于輸入信號(hào) x (r)，其關(guān)系是某個(gè)非線性的延遲微分方程（DE）：

其中 L [?] 是 DE 的線性算子，f 是非線性函數(shù)，其依賴于 P 個(gè)不同位置的 y 值、外部輸入 x 和參數(shù) θ 的。這是一個(gè)通用形式，足以表示各種連續(xù)微分方程，比如 ODE（當(dāng) L [?] = d/dt 且 r = t）、偏微分方程（PDE）、甚至用于 RNN 的離散差分方程。

現(xiàn)在，在左側(cè)和右側(cè)添加一項(xiàng)，其中 Gp (r) 是一個(gè)依賴于位置 r 的 n×n 矩陣。G_p 的值會(huì)在后面決定。現(xiàn)在 1 式就變成了：

2 式的左側(cè)是一個(gè)關(guān)于 y 的線性方程，在大多數(shù)情況下，其求解難度都低于求解非線性方程。在 3 式中，研究者引入了一個(gè)新符號(hào)，用以表示在給定邊界條件下求解 2 式左側(cè)的線性算子的線性算子。

3 式可被看作是一個(gè)定點(diǎn)迭代問題，即給定一個(gè)初始猜測(cè)，可以迭代地計(jì)算等式右側(cè)，直到其收斂。為了分析這種接近真實(shí)解的收斂性，這里將第 i 輪迭代時(shí)的 y 值記為，其中是滿足 3 式的真實(shí)解。將 y^(i) 代入 3 式可以得到 y^(i+1)，然后泰勒展開至一階，得：

其中 J_pf 是 f 在其第 p 個(gè)參數(shù)上的雅可比矩陣。根據(jù)上式，通過選擇

可讓 δy^(i+1) 的一階項(xiàng)為 0。

這表明，根據(jù)上式選擇矩陣 G_p，能以最快的速度收斂到解附近。這還表明，3 式和 5 式中的迭代相當(dāng)于在巴拿赫空間（Banach space）中實(shí)現(xiàn)牛頓法，因此能提供二次收斂性。

3 式中的迭代過程涉及到評(píng)估函數(shù) f、其雅可比矩陣和矩陣乘法，這些運(yùn)算可以使用現(xiàn)代加速器（如 GPU 和 TPU）來并行化處理。如果能以并行方式求解線性方程，那么整個(gè)迭代過程都可利用并行計(jì)算。在深度學(xué)習(xí)背景中，將非線性微分方程視為定點(diǎn)迭代問題來求解還有另一個(gè)優(yōu)勢(shì)，即可以將前一步驟的解（如果能放入內(nèi)存）用作下一訓(xùn)練步驟的起始猜測(cè)。如果起始猜測(cè)更好，則能減少尋找非線性微分方程的解所需的迭代步驟。

實(shí)際實(shí)現(xiàn)

為了將 3 式的 DEER 框架用于具體問題，需要遵循一些步驟。

第一步是將問題改寫成 1 式，定義變量 y、線性算子 L [?] 和非線性函數(shù) f (?)。

第二步是實(shí)現(xiàn)研究者所說的位移器函數(shù)（shifter function）。這個(gè)位移器函數(shù)是以 y (r) 的整體離散值為輸入，返回經(jīng)過位移的 y 值的列表，即 y (r ? s_p)，其中 p = {1, ..., P}。這個(gè)位移器函數(shù)可能需要一些附加信息，比如起始或邊界條件。這個(gè)位移器函數(shù)的輸出將會(huì)是非線性函數(shù)的輸入。

下一步（通常也是最難的一步）是根據(jù)矩陣列表 G_p 和在某些點(diǎn)離散的向量值 h 實(shí)現(xiàn)逆算子。這個(gè)逆算子可能也需要有關(guān)邊界條件的信息。

只要能提供算法 1 中的需求，就可以將 DEER 框架應(yīng)用于任意微分或差分方程。

并行化常微分方程（ODE）

ODE 的形式通常是 dy/dt = f (y (t), x (t), θ)，其中初始條件 y (0) 是已給定的。上面的 ODE 形式如果用 1 式表示，則有 r = t、L = d/dt、P = 1 和 s_1 = 0。這意味著 ODE 中的算子相當(dāng)于在給定初始條件 y (0) 時(shí)求解下面的線性方程。

假設(shè) G (t) 和 z (t) 是 t = t_i 和 t = t_{i+1} 之間的常量，分別為 G_i 和 z_i，則可以將 y_{i+1}=y_(t_i+1) 和 y_i = y (t_i) 之間的關(guān)系寫成：

其中 ?_i = t_{i+1} ? t_i，I 是單位矩陣，exp (?) 是矩陣指數(shù)。9 式可以使用并行前綴掃貓算法進(jìn)行評(píng)估。具體來說，首先可以為每個(gè)離散時(shí)間點(diǎn) t_i 定義一對(duì)變量，初始值 c_0=(I|y_0) 以及一個(gè)關(guān)聯(lián)算子

給定上面的初始值 c_0 和關(guān)聯(lián)算子，可以并行方式運(yùn)行關(guān)聯(lián)掃描以獲取上述算子的累積值。解 y_i 可從這個(gè)并行掃描算子的結(jié)果的第二個(gè)元素獲取。

并行化 RNN

循環(huán)神經(jīng)網(wǎng)絡(luò)（RNN）可以看作是一種離散版的 ODE。令索引 x 處的輸入信號(hào)為 x_i，前一狀態(tài)為 y_{i-1}，則當(dāng)前狀態(tài)可以寫成 y_i = f (y_{i-1}, x_i , θ)。

這個(gè)形式可以捕獲常見的 RNN 單元，比如 LSTM 和 GRU。而如果用 1 式來寫這個(gè)形式，則有 r = i、L [y] = y、P = 1 和 s_1 = 1。這意味著給定起始狀態(tài) y_0，可以通過求解下式來計(jì)算逆線性算子：

求解上式就相當(dāng)于求解前一小節(jié)的 9 式。這意味著也可以使用并行前綴掃描和 11 式中定義的關(guān)聯(lián)算子來將其并行化。

實(shí)驗(yàn)

圖 2 給出了新提出的方法在 V100 GPU 上所實(shí)現(xiàn)的速度提升。

這張圖表明，當(dāng)維度小、序列長(zhǎng)度長(zhǎng)時(shí)，取得的速度提升最大。但是，隨著維度增多，速度提升會(huì)下降。對(duì)前向 + 梯度計(jì)算的提速甚至超過僅前向計(jì)算的提速。

圖 3 比較了使用序列方法和 DEER 方法評(píng)估的 GRU 的輸出。

從圖 3 可以看出，使用 DEER 方法評(píng)估的 GRU 的輸出幾乎與使用序列方法獲得的輸出一樣。圖 3 (b) 中的小誤差源于單精度浮點(diǎn)的數(shù)值精度限制。

圖 4 (a, b). 給出了使用 DEER 方法和 RK45 方法時(shí)訓(xùn)練期間的損失變化情況。從圖中可以看到，相比于使用普通的 ODE 求解器，當(dāng)使用新提出的 DEER 方法時(shí)，訓(xùn)練速度可以提升 11 倍，并且這兩種方法的驗(yàn)證損失差別不大。

圖 4 (c, d) 比較了使用 DEER 方法和常用的序列方法時(shí)，GRU 網(wǎng)絡(luò)訓(xùn)練期間的驗(yàn)證準(zhǔn)確度。從圖中可以看到，使用 DEER 方法時(shí)的驗(yàn)證準(zhǔn)確度圖表與使用序列方法時(shí)的很相近。