重溫被Mamba帶火的SSM：HiPPO的一些遺留問題

©PaperWeekly 原創 · 作者 | 蘇劍林

單位 | 科學空間

研究方向 | NLP、神經網路

書接上文，在上一篇文章《重溫被Mamba帶火的SSM：線性系統和HiPPO矩陣》中，我們詳細討論了 HiPPO 逼近框架其 HiPPO 矩陣的推導，其原理是透過正交函式基來動態地逼近一個即時更新的函式，其投影係數的動力學正好是一個線性系統，而如果以正交多項式為基，那麼線性系統的核心矩陣我們可以解析地求解出來，該矩陣就稱為 HiPPO 矩陣。

當然，上一篇文章側重於 HiPPO 矩陣的推導，並沒有對它的性質做進一步分析，此外諸如“如何離散化以應用於實際資料”、“除了多項式基外其他基是否也可以解析求解”等問題也沒有詳細討論到。接下來我們將補充探討相關問題。

離散格式

假設讀者已經閱讀並理解上一篇文章的內容，那麼這裡我們就不再進行過多的鋪墊。在上一篇文章中，我們推匯出了兩類線性 ODE 系統，分別是：

其中是與時間無關的常數矩陣，HiPPO 矩陣主要指矩陣。在這一節中，我們討論這兩個 ODE 的離散化。

輸入轉換

在實際場景中，輸入的資料點是離散的序列，比如流式輸入的音訊訊號、文字向量等，我們希望用如上的 ODE 系統來即時記憶這些離散點。為此，我們先定義

其中就是離散化的步長。該定義也就是說在區間內，是一個常數函式，其值等於。很明顯這樣定義出來的無損原本序列的資訊，因此記憶就相當於記憶序列。

從變換到，可以使得輸入訊號重新變回連續區間上的函式，方便後面進行積分等運算，此外在離散化的區間內保持為常數，也能夠簡化離散化後的格式。

LegT版本

我們先以 LegT 型 ODE（1）為例，將它兩端積分

其中。根據的定義，它在區間內恆為，於是的積分可以直接算出來：

接下來的結果，就取決於我們如何近似的積分了。假如我們認為在區間內近似恆等於，那麼就得到前向尤拉格式

我們認為在區間內近似恆等於，那麼就得到後向尤拉格式

前後向尤拉都具有相同的理論精度，但後向通常會有更好的數值穩定性。如果要更準確一些，那麼認為在區間內近似恆等於，那麼得到雙線性形式：

這也等價於先用前向尤拉走半步，再用後向尤拉走半步。更一般地，我們還可以認為在區間內近似恆等於，其中，這就不進一步展開了。事實上，我們也可以完全不做近似，因為結合式（1）以及在區間中是常數，我們完全可以用“常數變易法” [1] 來精確求解出來，結果是

這裡的矩陣指數按照級數來定義，可以參考《恆等式 det(exp(A)) = exp(Tr(A)) 賞析》[2]。

LegS版本

現在輪到 LegS 型 ODE 了，它的思路跟 LegT 型基本一致，結果也大同小異。首先將式（2）兩端積分得到

根據定義，第二項積分的在恆為，所以它相當於的積分，可以直接積分出來得，當然直接換為一階近似也無妨，因為本身到的變換有很大自由度，這點誤差無所謂。至於第一項積分，我們直接採用精度更高的中點近似，得到

事實上，式（2）也可以精確求解，只需要留意到它等價於

這意味著只需要做變數代換，那麼 LegS 型 ODE 就可以轉化為 LegT 型 ODE：

利用式（9）得到（由於變數代換，時間間隔由變成）

然而，上式雖然是精確解，但不如同為精確解的式（9）好用，因為式（9）的指數矩陣部分是，跟時間無關，所以一次性計算完就可以了。但上式中在矩陣指數里邊，意味著在迭代過程中需要反覆計算矩陣指數，對計算並不友好，所以 LegS 型 ODE 我們一般只會用式（11）來離散化。

優良性質

接下來，LegS 是我們的重點關注物件。重點關注 LegS 的原因並不難猜，因為從推導的假設來看，它是目前求解出來的唯一一個能夠記憶整個歷史的 ODE 系統，這對於很多場景如多輪對話來說至關重要。此外，它還有其他的一些比較良好且實用的性質。

尺度不變

比如，LegS 的離散化格式（11）是步長無關的，我們只需要將代入裡邊，並記，就可以發現

步長被自動地消去了，從而自然地減少了一個需要調的超引數，這對於煉丹人士顯然是一個好訊息。注意步長無關是 LegS 型 ODE 的一個固有性質，它跟具體的離散化方式並無直接關係，比如精確解（14）同樣是步長無關的：

其背後的原因，在於 LegS 型 ODE 滿足“時間尺度不變性（Timescale equivariance）”——如果我們設代入 LegS 型 ODE，將得到

這意味著，當我們將換成時，LegS 的 ODE 形式並沒有變化，而對應的解則是換成了。這個性質的直接後果就是：當我們選擇更大的步長時，遞迴格式不需要發生變化，因為結果的步長也會自動放大，這就是 LegS 型 ODE 離散化與步長無關的本質原因。

長尾衰減

LegS 型 ODE 的另一個優良性質是，它關於歷史訊號的記憶是多項式衰減（Polynomial decay）的，這比常規 RNN 的指數衰減更緩慢，從而理論上能記憶更長的歷史，更不容易梯度消失。為了理解這一點，我們可以從精確解（16）出發，從式（16）可以看到，每遞迴一步，歷史資訊的衰減效應可以用矩陣指數來描述，那麼從第步遞迴到第步，總的衰減效應是

回顧 HiPPO-LegS 中的形式：

從定義可以看出，是一個下三角陣，其對角線元素為。我們知道，三角陣的對角線元素正好是它的特徵值（參考 Triangular matrix [3]），由此可以看到一個大小的矩陣，有個不同的特徵值，這說明矩陣是可對角化的，即存在可逆矩陣，使得，其中，於是我們有

可見，最終的衰減函式是的次函式的線性組合，所以 LegS 型 ODE 關於歷史記憶至多是多項式衰減的，比指數衰減更加長尾，因此理論上有更好的記憶力。

計算高效

最後，我們指出 HiPPO-LegS 的矩陣是計算高效（Computational efficiency）的。具體來說，直接按照矩陣乘法的樸素實現的話，一個的矩陣乘以的列向量，需要做次乘法，但 LegS 的矩陣與向量相乘則可以降低到次，更進一步地，我們還可以證明離散化後的（11）也可以在完成。

為了理解這一點，我們首先將 HiPPO-LegS 的 A 矩陣等價地改寫成

對於向量，我們有

這包含三種運算，第一項的是向量與做逐位相乘運算，第二項的則是向量與做逐位相乘，然後就是運算，最後乘以就是再逐位相乘向量，每一步都可以在內完成，因此總的複雜度是的。

我們再來看（11），它包含兩步“矩陣-向量”乘法運算，一是，是任意實數，剛才我們已經證明了是計算高效的，自然也是；二是，接下來我們將證明它也是計算高效的。這隻需要留意到求等價於解方程，利用上面給出的表示式，我們可以得到

記，那麼，代入上式得

整理得

這是一個標量的遞迴式，可以完全序列地計算，也可以利用 Prefix Sum 的相關演算法平行計算（參考這裡），計算複雜度為或者，總之相比都會更加高效。

傅立葉基

最後，我們以傅立葉基的一個推導收尾。在上一篇文章中，我們以傅立葉級數來引出了線性系統，但只推導了鄰近視窗形式的結果，而後面的勒讓德多項式基我們則推導了鄰近視窗和完整區間兩個版本（即 LegT 和 LegS）。那麼傅立葉基究竟能不能推導一個跟 LegS 相當的版本呢？其中會面臨什麼困難呢？下面我們對此進行探討。

同樣地，相關鋪墊我們不再重複，按照上一節的記號，傅立葉基的係數為

跟 LegS 一樣，為了記憶整個區間的訊號，我們需要一個的對映，為此選取最簡單的，代入後兩邊求導得到

分部積分得到

上一篇文章我們提到，HiPPO 選取勒讓德多項式為基的重要原因之一是可以分解為的線性組合，而傅立葉基的則不能做到這一點。但事實上，如果允許誤差的話，這個斷論是不成立的，因為我們同樣可以將分解為傅立葉級數：

這裡的求和有無限項，如果要截斷為有限項的話，就會產生誤差，但我們可以先不糾結這一點，直接往上代入得到

這樣一來

所以可以寫出

實際使用的時候，我們只需要截斷，就可以得到一個的矩陣。截斷帶來的誤差其實是無所謂的，因為我們在推導 HiPPO-LegT 的時候同樣引入了有限級數近似，那會我們同樣也沒考慮誤差，或者反過來講，對於特定的任務，我們會選擇適當的規模（即 N 的大小），而這個“適當”的含義之一，就是截斷帶來誤差對於該任務是可以忽略的。

對大多數人來說，傅立葉基的這個推導可能還更容易理解一些，因為勒讓德多項式對很多讀者來說都比較陌生，尤其是 LegT、LegS 推導過程中用到的幾個恆等式，而對於傅立葉級數大多數讀者應該或多或少都有所瞭解。

不過，從結果上來看，傅立葉基的這個結果可能不如 LegS 實用，一來它引入了複數，這增加了實現的複雜度，二來它推匯出的 A 矩陣不像 LegS 那樣是個相對較淡的下三角陣，因此理論分析起來也更為複雜。所以，大家權當它是一道深化對 HiPPO 的理解的練習題就好。