code window

顯示包含「變壓器」標籤的文章。顯示所有文章
顯示包含「變壓器」標籤的文章。顯示所有文章

2024年5月13日星期一

Transformer 語言模型內部運作機制的技術

這篇論文是一篇全面性的綜述,詳細介紹了目前用來解釋 Transformer 語言模型內部運作機制的技術。本文聚焦在生成式的 decoder-only 架構。作者在結論部分概述了這些模型的已知內部機制,揭示了當前流行方法和這個領域的積極研究方向之間的聯繫。 

論文出處: https://arxiv.org/pdf/2405.00208

數學符號說明

在本文中,我們採用以下數學符號:

- \( n \):序列長度

- \( V \):詞彙表

- \( t = \langle t_1, t_2, \ldots, t_n \rangle \):輸入的 token 序列

- \( x = \langle x_1, x_2, \ldots, x_n \rangle \):輸入的 token 嵌入序列

- \( d \):模型維度

- \( d_h \):注意力頭維度

- \( d_{FFN} \):前饋網路維度

- \( H \):注意力頭數量

- \( L \):層數

- \( x_{l,i} \in \mathbb{R}^d \):在位置 \( i \),第 \( l \) 層的殘差串流狀態

- \( x_{\text{mid},l,i} \in \mathbb{R}^d \):在位置 \( i \),第 \( l \) 層,經過注意力區塊後的殘差串流狀態

- \( f_c(x) \in \mathbb{R}^d \):在最後一個位置,組件 \( c \) 的輸出表示

- \( f_l(x) = x_{l,n} \in \mathbb{R}^d \):在最後一個位置,第 \( l \) 層的殘差串流狀態

- \( A_{l,h} \in \mathbb{R}^{n \times n} \):在第 \( l \) 層第 \( h \) 個注意力頭的注意力矩陣

- \( W_{l,h}^Q, W_{l,h}^K, W_{l,h}^V \in \mathbb{R}^{d \times d_h} \):在第 \( l \) 層第 \( h \) 個注意力頭的 queries、keys 和 values 權重矩陣

- \( W_{l,h}^O \in \mathbb{R}^{d_h \times d} \):在第 \( l \) 層第 \( h \) 個注意力頭的輸出權重矩陣

- \( W_{l}^{\text{in}} \in \mathbb{R}^{d \times d_{FFN}}, W_{l}^{\text{out}} \in \mathbb{R}^{d_{FFN} \times d} \):在第 \( l \) 層前饋網路的輸入和輸出權重矩陣

- \( W_E \in \mathbb{R}^{d \times |V|} \) 和 \( W_U \in \mathbb{R}^{|V| \times d} \):嵌入(Embedding)和去嵌入(Unembedding)矩陣


一個僅解碼器的模型 \( f \) 有 \( L \) 層,對一個序列的嵌入 \( x = \langle x_1, x_2, \ldots, x_n \rangle \) 進行操作,這些嵌入表示 tokens \( t = \langle t_1, t_2, \ldots, t_n \rangle \)。每個嵌入 \( x \in \mathbb{R}^d \) 是嵌入矩陣 \( W_E \in \mathbb{R}^{|V| \times d} \) 的一個列向量,其中 \( V \) 是模型詞彙表。中間層的表示,例如在位置 \( i \) 和層 \( l \),表示為 \( x_{l,i} \)。\( X \in \mathbb{R}^{n \times d} \) 代表將序列 \( x \) 表示為一個矩陣,嵌入堆疊為列。同樣地,對於中間表示,\( X_{l, \leq i} \) 是層 \( l \) 的表示矩陣,最多到位置 \( i \)。

遵循最近關於Transformers可解釋性的文獻,本文採用殘差串流的觀點來呈現架構。在這個觀點下,每個輸入嵌入都會透過注意力區塊和前饋網路區塊的向量相加來更新,產生殘差串流狀態(或中間表示)。最後一層的殘差串流狀態會透過去嵌入矩陣 \( W_U \in \mathbb{R}^{d \times |V|} \) 投影到詞彙空間,並透過softmax函數正規化以獲得詞彙上的機率分布,從中取樣出新的token。



在Transformer層中,LayerNorm是一種常見的操作,用於穩定深度神經網路的訓練過程。給定一個表示 \( z \),LayerNorm的計算為 \((z-\mu(z))/\sigma(z) \odot \gamma + \beta\),其中 \( \mu \) 和 \( \sigma \) 分別計算平均值和標準差,\( \gamma \in \mathbb{R}^d \) 和 \( \beta \in \mathbb{R}^d \) 是學習到的逐元素轉換和偏差。LayerNorm可以從幾何角度解釋,將減去平均值的操作視為將輸入表示投影到由法向量 \([1, 1, \ldots, 1] \in \mathbb{R}^d\) 定義的超平面上,然後將結果表示映射到 \(\sqrt{d}\) 範數的超球面。

注意力區塊由多個注意力頭組成。在解碼步驟 \( i \),每個注意力頭從先前位置(≤ \( i \))的殘差串流讀取,決定要關注哪些位置,從那些位置收集資訊,最後將其寫入當前的殘差串流。

注意力頭計算如下:


\[ \text{Attn}_{l,h}(X_{l-1, \leq i}) = \sum_{j \leq i} a_{l,h,i,j} x_{l-1,j} W_{l,h}^V W_{l,h}^O = \sum_{j \leq i} a_{l,h,i,j} x_{l-1,j} W_{l,h}^{OV} \]


可學習的權重矩陣 \( W_{l,h}^V \in \mathbb{R}^{d \times d_h} \) 和 \( W_{l,h}^O \in \mathbb{R}^{d_h \times d} \) 合併為 OV 矩陣 \( W_{l,h}^V W_{l,h}^O = W_{l,h}^{OV} \in \mathbb{R}^{d \times d} \),也稱為 OV (output-value) 電路。對於每個查詢(query)\( i \),給定當前的鍵(key)(\( \leq i \)),注意力權重計算為:


\[ a_{l,h,i} = \text{softmax}\left( \frac{x_{l-1,i} W_{l,h}^Q (X_{l-1, \leq i} W_{l,h}^K)^T}{\sqrt{d_k}} \right) \]


其中 \( W_{l,h}^Q \in \mathbb{R}^{d \times d_h} \) 和 \( W_{h}^Q (W_{h}^K)^T = W_{h}^{QK} \in \mathbb{R}^{d \times d} \) 組合為 QK (query-key) 電路 \[ W_{h}^Q (W_{h}^K)^T = W_{h}^{QK} \in \mathbb{R}^{d \times d} \]

QK 和 OV 電路可以視為負責分別從殘差串流讀取和寫入的單元。注意力區塊的輸出是各個注意力頭的總和,隨後加回到殘差串流:


\[ \text{Attn}_l(X_{l-1, \leq i}) = \sum_{h=1}^H \text{Attn}_{l,h}(X_{l-1, \leq i}) \]

\[ x_{\text{mid},l,i} = x_{l-1,i} + \text{Attn}_l(X_{l-1, \leq i}) \]


前饋網路區塊由兩個可學習的權重矩陣組成:\( W_{l}^{\text{in}} \in \mathbb{R}^{d \times d_{FFN}} \) 和 \( W_{l}^{\text{out}} \in \mathbb{R}^{d_{FFN} \times d} \)。\( W_{l}^{\text{in}} \) 從殘差串流狀態 \( x_{\text{mid},l,i} \) 讀取,其結果通過一個逐元素的非線性激活函數 \( g \),產生神經元激活。這些激活再由 \( W_{l}^{\text{out}} \) 轉換以產生輸出 \( \text{FFN}(x_{\text{mid},i}) \),然後加回到殘差串流:


\[ \text{FFN}_l(x_{\text{mid},l,i}) = g(x_{\text{mid},l,i} W_{l}^{\text{in}}) W_{l}^{\text{out}} \]

\[ x_{l,i} = x_{\text{mid},l,i} + \text{FFN}_l(x_{\text{mid},l,i}) \]


前饋網路的計算可以等同於鍵值記憶體檢索,\( W_{l}^{\text{in}} \) 中的列向量充當輸入序列上的模式檢測器(鍵),每個神經元激活加權了 \( W_{l}^{\text{out}} \) 的列向量(值)。元素式非線性在前饋網路內部創建了一個特權基底,鼓勵特徵與基底方向對齊。

預測層包含一個去嵌入矩陣 \( W_U \in \mathbb{R}^{d \times |V|} \),有時還有一個偏差。最後一個殘差串流狀態透過這個線性映射轉換,將表示轉換為下一個 token 的 logits 分布,再通過 softmax 函數轉換為概率分布。由於模型組件透過相加與殘差串流交互作用,未正規化的分數(logits)是透過組件輸出的線性投影獲得的。基於線性變換的性質,可以重新排列傳統的前向傳遞公式,使每個模型組件直接貢獻於預測的 token 的 logits:


\[ f(x) = x_L^n W_U = \left( \sum_{l=1}^L \sum_{h=1}^H \text{Attn}_{l,h}(X_{l-1}^{\leq n}) + \sum_{l=1}^L \text{FFN}_l(x_{\text{mid},l}^n) + x^n \right) W_U \]

\[ = \sum_{l=1}^L \sum_{h=1}^H \text{Attn}_{l,h}(X_{l-1}^{\leq n}) W_U + \sum_{l=1}^L \text{FFN}_l(x_{\text{mid},l}^n) W_U + x^n W_U \]


注意力頭 logits 更新 前饋網路 logits 更新

這個分解在定位負責預測的組件時扮演重要角色,因為它允許測量每個組件對預測 token 的 logits 的直接貢獻。

殘差網路就像淺層網路的集成,每個子網路定義了計算圖中的一條路徑。將前向傳遞分解為:


\[ f(x) = xW_U + xW_1^{OV} W_U + xW_1^{OV} W_2^{OV} W_U + xW_2^{OV} W_U \]


直接路徑 完整 OV 電路

虛擬注意力頭 (V-composition)

連結輸入嵌入與去嵌入矩陣的路徑稱為直接路徑。穿過單個 OV 矩陣的路徑稱為完整 OV 電路。涉及兩個注意力頭的路徑稱為虛擬注意力頭,執行 V-composition,因為兩個頭的順序寫入和讀取被視為 OV 矩陣的組合。

理解語言模型的內部運作機制意味著定位前向傳遞中負責特定預測的元素(輸入元素、表示和模型組件)。本文介紹了兩種不同類型的方法,允許定位模型行為:輸入歸因和模型組件歸因。

輸入歸因方法通常用於通過估計輸入元素(在語言模型的情況下是 tokens)對定義模型預測的貢獻來定位模型行為。對於像語言模型這樣的神經網路模型,梯度資訊經常被用作歸因目的的自然度量。基於梯度的歸因在這個上下文中涉及 Transformer 在點 \( x \) 處的一階泰勒展開,表示為 \( \nabla f(x) \cdot x + b \)。結果梯度 \( \nabla f_w(x) \in \mathbb{R}^{n \times d} \) 直觀地捕捉了模型對輸入中的每個元素在預測 token \( w \) 時的敏感度。雖然歸因分數是針對輸入 token 嵌入的每個維度計算的,但它們通常在 token 層面匯總,以獲得對個別 token 影響的更直觀概述。這通常是透過取梯度向量相對於第 \( i \) 個輸入嵌入的 \( L_p \) 範數來完成的:


\[ A_{\text{Grad},f_w}^{(x \rightarrow t_i)} = \| \nabla_{x_i} f_w(x) \|_p \]


透過梯度向量與輸入嵌入 \( \nabla_{x_i} f_w(x) \cdot x_i \) 的點積,稱為梯度 × 輸入方法,可以將這種敏感度轉化為重要性估計。然而,這些方法已知會出現梯度飽和和碎裂問題。這一事實促使引入了積分梯度和 SmoothGrad 等方法來過濾嘈雜的梯度資訊。例如,積分梯度近似基線輸入 \( \tilde{x} \) 和輸入 \( x \) 之間直線路徑上的梯度積分:


\[ \int_0^1 \nabla_{x_i} f_w(\tilde{x} + \alpha(x - \tilde{x})) d\alpha \]


隨後提出了適應文本輸入離散性的改編。最後,基於層次相關傳播 (Layer-wise Relevance Propagation, LRP) 的方法已廣泛應用於研究基於 Transformer 的語言模型。這些方法對梯度傳播使用自定義規則,以分解每一層的組件貢獻,確保它們的總和在整個網路中保持不變。

另一個流行的方法族是透過添加噪聲或消融輸入元素並測量對模型預測的結果影響來估計輸入重要性。例如,可以移除位置 \( i \) 處的輸入 token,結果機率差異 \( f_w(x) - f_w(x_{-x_i}) \) 可以用作其重要性的估計。如果給予 \( w \) 的 logit 或機率沒有改變,我們可以得出第 \( i \) 個 token 沒有影響的結論。

雖然原始模型內部資料(如注意力權重)通常被認為提供了不忠實的模型行為解釋,但最近的方法提出了注意力權重的替代方案來測量中間 token 級別的歸因。其中一些替代方案包括使用值加權向量和輸出值加權向量的範數,或使用向量距離來估計貢獻。这些方法的共同策略涉及使用注意力展開等技術聚合反映上下文混合模式的中間每層歸因,得到輸入歸因分數。

一個重要的限制是,歸因的輸出 token 屬於一個大的詞彙空間,在下一個詞預測中經常有語義上等價的 token 競爭機率質量。在這種情況下,歸因分數很可能誤代了驅動模型預測的幾個重疊因素,如語法正確性和語義適當性。最近的工作通過提出這些方法的對比公式來解決這個問題,為模型預測 token \( w \) 而不是替代 token \( o \) 提供反事實解釋。

另一個輸入歸因的維度涉及識別影響推理時特定模型預測的有影響力的訓練樣本。這些方法通常被稱為訓練數據歸因 (TDA) 或實例歸因方法,並被用來識別數據中的人工因素和語言模型預測中偏差的來源。最近的方法提出通過訓練運行模擬來執行 TDA。雖然已建立的 TDA 方法的適用性受到質疑,特別是由於其低效率,但這個領域最近的工作產生了更有效的方法,可以大規模應用於大型生成模型。

早期關於 Transformer 語言模型組件重要性的研究強調了模型能力的高度稀疏性。例如,即使刪除模型中相當一部分的注意力頭,也可能不會使其下游性能惡化。這些結果激發了一條新的研究路線,研究語言模型中的各種組件如何貢獻於其廣泛的能力。

讓我們稱組件 \( c \)(注意力頭或前饋網路)在特定層上對最後一個 token 位置的輸出表示為 \( f_c(x) \)。等式 (10) 中提出的分解允許我們測量每個模型組件對輸出 token \( w \in V \) 的直接 logit 歸因 (DLA):


\[ A_{DLA,f_w}^{(x \rightarrow c)} = f_c(x) W_U[:,w] \]


其中 \( W_U[:,w] \) 是 \( W_U \) 的第 \( w \) 列,即 token \( w \) 的去嵌入向量。實際上,組件 \( c \) 的 DLA 表示 \( c \) 對預測 token 的 logit 的貢獻,使用第 2.2 節中描述的模型組件的線性特性。



我們可以將模型的計算視為一個因果模型,並使用因果工具來闡明每個模型組件 \( c \in C \) 在不同位置對預測的貢獻。因果模型可以看作是一個有向無環圖 (DAG),其中節點是模型計算,邊是激活。



我們可以通過改變前向傳遞中由模型組件計算的某些節點值 \( f_c(x) \) 到另一個值 \( \tilde{h} \) 來干預模型,這被稱為激活修補。我們可以使用 do 算子表示這個干預:\( f(x \mid \text{do}(f_c(x) = \tilde{h})) \)。然後我們測量修補後預測的變化:


\[ A_{\text{Patch}, f}^{(x \rightarrow c)} = \text{diff}(f(x), f(x \mid \text{do}(f_c(x) = \tilde{h}))) \]


其中 \(\text{diff}(\cdot, \cdot)\) 函數的常見選擇包括 KL 散度和 logit/機率差異。修補後的激活 (\(\tilde{h}\)) 可以來自各種來源。一種常見的方法是創建一個具有分布 \(P_{\text{patch}}\) 的反事實數據集,其中一些關於任務的輸入信號被反轉。這種方法會導致兩種不同類型的消融:

- 重採樣干預,其中修補後的激活是從 \(P_{\text{patch}}\) 的單個樣本中獲得的,即 \(\tilde{h} = f_c(\tilde{x}), \tilde{x} \sim P_{\text{patch}}\)。

- 平均干預,其中多個 \(P_{\text{patch}}\) 樣本的平均激活用於修補,即 \(\tilde{h} = \mathbb{E}_{\tilde{x} \sim P_{\text{patch}}}[f_c(\tilde{x})]\)。

另外,修補激活的其他來源包括:

- 零干預,其中激活被替換為空向量,即 \(\tilde{h} = 0\)。

- 噪聲干預,其中新的激活是通過在受擾動的輸入上運行模型獲得的,例如 \(\tilde{h} = f_c(x + \epsilon), \epsilon \sim N(0, \sigma^2)\)。

在設計因果干預實驗時需要考慮的一個重要因素是設置的生態效度,因為零消融和噪聲消融可能使模型偏離自然激活分布,最終破壞組件分析的有效性。

機制可解釋性 (Mechanistic Interpretability) 子領域專注於將神經網路逆向工程為人類可理解的算法。MI 的最新研究旨在揭示電路的存在,電路是一組共同交互作用以解決任務的模型組件(子圖)。激活修補、logit 歸因和注意力模式分析是電路發現的常用技術。

邊緣修補和路徑修補利用每個模型組件輸入是其殘差串流中先前組件輸出之和這一事實,並考慮直接連接成對模型組件節點的邊。路徑修補將邊緣修補方法推廣到多條邊,允許更細粒度的分析。例如,使用等式中描述的淺層網路的前向傳遞分解,我們可以將圖中的單層 Transformer 視為由以下組件組成:


\[ f(x) = \text{Attn}(X^{\leq n}) W_u + \text{FFN}(\text{Attn}(X^{\leq n}) + x_n)W_u + x_n W_u \]



這裡的各個路徑包括:

- 直接從Attn到logits。

- 通過FFN再到logits的間接路徑。

其中每個發送節點Attn\( _L(X_{L-1}^{\leq n}) \)的副本都與單個路徑相關。在這個例子中,分別修補每個發送節點副本允許我們估計Attn\( _L(X_{L-1}^{\leq n}) \)對輸出logits\( f(x) \)的直接和間接效應。一般來說,我們可以將路徑修補應用於網路中的任何路徑,並測量頭部之間的組合、FFN或這些組件對logits的影響。

基於因果干預的電路分析存在幾個缺點:

1. 它需要為要評估的任務設計輸入模板,以及反事實數據集(即定義\( P_{\text{patch}} \))方面的大量努力。

2. 在獲得組件重要性估計後,需要人工檢查和領域知識來分離重要的子圖。

3. 已經表明,干預會在下游組件的行為中產生二階效應,在某些情況下甚至會引發類似於自我修復的補償行為。這種現象會使得難以得出關於每個組件作用的結論。

為了克服這些限制,Conmy等人提出了一種自動電路發現(ACDC)算法,通過迭代移除計算圖中的邊來自動識別電路的過程。然而,這個過程需要大量的前向傳遞(每個修補元素一次),在研究大型模型時變得不切實際。

修補的一個有效替代方法是基於梯度的方法,這些方法已經擴展到超越輸入歸因,以計算中間模型組件的重要性。例如,給定token預測\( w \),為了計算中間層\( l \)的歸因,表示為\( f_l(x) \),計算梯度\( \nabla f_w(f_l(x)) \)。Sarti等人將等式中的對比梯度歸因公式擴展到使用單次前向和後向傳遞來定位對正確延續的預測貢獻大於錯誤延續的組件。Nanda等人提出了邊緣歸因修補(EAP),包括修補前後預測差異的線性近似,以估計計算圖中每個邊的重要性。這種方法的主要優點是它只需要兩次前向傳遞和一次後向傳遞就可以獲得圖中每個邊的歸因分數。

另一個研究方向是在較低層次的神經網路中尋找可解釋的高層次因果抽象概念。這些方法涉及大量的計算搜索,並假設高層次變數與單元或神經元組對齊。為了克服這些限制,Geiger等人提出了分佈式對齊搜索(Distributed Alignment Search, DAS),它在通過梯度下降找到的低層次表示空間的非基底對齊子空間上執行分佈式交換干預(Distributed Intervention Interchange, DII)。DAS干預在使用語法評估尋找具有因果影響的特徵方面被證明是有效的,同時在分離實體的個別屬性的因果效應方面也有顯著效果。

最後,在Transformer語言模型的內部行為方面,本文總結了以下主要發現:

注意力區塊:

  • 位置頭:有些頭主要關注相對於正在處理的token的特定位置,如token本身、前一個token或下一個位置。
  • 子詞連接頭: 專門關注屬於與當前處理的token相同詞的前一個子詞token。
  • 語法頭: 一些頭關注具有與被處理token相關語法角色的token,明顯多於隨機基準。
  • 重複token:** 關注上下文中同一個token的先前出現。
  • 複製頭: OV矩陣表現出複製行為。
  • 歸納頭: 由兩個在不同層組合的頭組成,讓語言模型完成模式。一個早期的前一個token頭將第一個token A的資訊複製到B的殘差流,然後一個下游的歸納頭關注token B,增加B的可能性。
  • 複製抑制頭: 如果它們出現在上下文中並且當前殘差流正在自信地預測它,則減少它們關注的token的logit分數。

前饋網路區塊:

  • 神經元的輸入行為:有些神經元僅在特定位置範圍上激發;技能神經元,其激活與輸入提示的任務相關;概念特定神經元,其反應可用於預測上下文中概念的存在。
  • 神經元的輸出行為:有些神經元促進與特定語義和句法概念相關的token的預測;一小部分後層神經元負責做出在語言上可接受的預測;抑制不可能延續的神經元。
  • 多語義神經元: 早期層的大多數神經元專門用於n-gram集合,充當n-gram偵測器,大多數神經元在大量n-gram上激發。

殘差流:

  • 殘差流可以被視為Transformer中的主要通訊通道。直接路徑主要對應於bi-gram統計,而網路中的最新偏差根據詞頻轉移預測,促進高頻token。
  • 一些組件執行記憶體管理,以移除存儲在殘差流中的資訊。例如,有負特徵值的OV矩陣關注當前位置的注意力頭,以及輸入和輸出權重具有較大負餘弦相似性的前饋網路神經元。
  • 在殘差流中發現了離群值維度。這些維度展現出相對於其他維度的大幅度,與各向異性表示的生成相關聯。消融離群值維度已被證明會顯著降低下游性能,表明它們編碼任務特定的知識。

多組件行為:

  • 歸納機制: 是兩個組件(注意力頭)組合在一起以完成模式的一個明顯例子。最近的證據表明,多個注意力頭協同工作,在給定上下文樣本時創建描述任務的「函數」或「任務」向量。
  • Variengien 和 Winsor 研究了涉及回答可以在上下文中找到答案的請求的上下文檢索任務
  • 作者確定了一種在子任務和模型之間通用的高層次機制。具體而言,中間層處理請求,然後由後層的注意力頭執行從上下文中檢索實體的步驟。
  • 在 GPT-2 Small 中發現了用於間接賓語識別 (IOI) 任務的電路,主要包括:
    • 重複信號:重複token頭和涉及前一個token頭的歸納機制表明S(John)的重複性。這個資訊被最後位置的S-抑制頭讀取,它們在殘差流中寫入一個token信號,表明S被重複,以及S1 token的位置信號。
    • 名稱複製: 後層的名稱移動頭將它們在上下文中關注的名稱的資訊複製到最後的殘差流。然而,先前層S-抑制頭的信號修改了名稱移動頭的query,使得重複的名稱(在S1和S2中)受到較少關注,有利於複製間接賓語(IO),從而推動其預測。


這篇論文全面而深入地介紹了目前用於解釋 Transformer 語言模型內部運作的技術,並總結了通過這些方法得到的關於模型內部機制的見解。作者強調,雖然在可解釋性研究方面取得了顯著進展,但將這些見解應用於調試和改進未來模型的安全性和可靠性,為開發人員和用戶提供更好的工具來與之交互並理解影響其預測的因素,仍然是一個巨大的挑戰。未來可解釋性研究的發展將面臨從在模型組件空間運作的方法和分析轉向人類可解釋空間(即從模型組件到特徵和自然語言解釋)的挑戰性任務,同時仍然忠實地反映模型行為。

此外,跨學科研究將在擴大可解釋性分析的範圍方面發揮關鍵作用,以考慮從人的角度看模型解釋的感知和交互維度。最終,作者認為,確保對先進語言模型的內部機制的開放和便利訪問,將仍然是這一領域未來進展的基本先決條件。

這篇綜述論文對於理解 Transformer 語言模型的內部運作機制,以及目前在這個領域的最新研究進展,提供了全面而詳盡的概覽。透過系統性地介紹各種可解釋性技術,並深入探討它們揭示的模型內部行為,本文為相關研究人員提供了寶貴的參考。

不過,正如作者所指出,將這些見解應用到實際中仍面臨諸多挑戰。未來的可解釋性研究需要在忠實反映模型行為的同時,努力向更加貼近人類直觀理解的方向發展。這需要不同學科領域的通力合作。此外,開放模型內部機制的訪問,或許是這一領域能夠取得突破性進展的關鍵。


2024年5月4日星期六

Hello, KAN!

多層感知器(MLPs),也被稱為全連接前饋神經網絡,是當今深度學習模型的基礎組成部分。



MLPs在機器學習中扮演著至關重要的角色,因為它們是用於近似非線性函數的默認模型,這得益於通用近似定理所保證的表達能力。然而,MLPs真的是我們能構建的最佳非線性回歸器嗎?儘管MLPs被廣泛使用,但它們仍然存在一些顯著缺陷。例如,在Transformer模型中,MLPs幾乎消耗了所有非嵌入參數,並且在沒有使用後續分析工具的情況下,通常不如注意力層那樣易於解釋。

為了解決這些問題,本研究提出了一種MLPs的替代方案,稱為Kolmogorov-Arnold Networks(KANs)。與受通用近似定理啟發的MLPs不同,KANs的靈感來自於Kolmogorov-Arnold定理。KANs與MLPs一樣具有全連接結構,但兩者在激活函數的放置方式上有所不同。MLPs在節點("神經元")上放置固定的激活函數,而KANs則在邊("權重")上放置可學習的激活函數。因此,KANs完全沒有線性權重矩陣,取而代之的是每個權重參數都被一個以樣條線參數化的可學習一維函數所替換。KANs的節點只是簡單地對傳入信號進行求和,而不應用任何非線性變換。有人可能會擔心,由於KANs將MLPs的每個權重參數變成了樣條函數,因此會變得計算成本很高。然而,令人欣喜的是,KANs通常允許使用比MLPs小得多的計算圖。例如,本研究表明,在解決偏微分方程(PDE)問題時,一個2層寬度為10的KAN比4層寬度為100的MLP的均方誤差(MSE)小100倍(10^(-7)對10^(-5)),同時參數效率提高了100倍(10^2對10^4)。

事實上,使用Kolmogorov-Arnold定理來構建神經網絡的可能性已經被研究過。然而,大多數工作都局限於使用原始的深度為2、寬度為(2n+1)的表示,沒有機會利用更現代的技術(例如反向傳播)來訓練網絡。本研究的貢獻在於將原始的Kolmogorov-Arnold表示推廣到任意寬度和深度,並將其置於當前的深度學習應用背景下,通過廣泛的實證實驗來突出其作為AI+科學基礎模型的潛力。

儘管KANs有優雅的數學解釋,但它們實際上只是樣條和MLPs的結合,利用了各自的優點,避免了各自的缺點。樣條在低維度下表現出色,易於局部調整,並且能夠在不同分辨率之間切換。然而,樣條受到嚴重的維度詛咒問題的影響,因為它們無法利用組合結構。另一方面,MLPs由於其特徵學習能力,受維度詛咒的影響較小,但在低維度下的準確性不如樣條,因為它們無法很好地優化一維函數。為了準確學習一個函數,一個模型不僅需要學習組合結構(外部自由度),還需要能夠很好地逼近一維函數(內部自由度)。KANs正是這樣一種模型,因為它們在外部類似於MLPs,在內部類似於樣條。因此,KANs不僅能夠學習特徵(由於其與MLPs在外部的相似性),還能將這些學習到的特徵優化到很高的精度(由於其與樣條在內部的相似性)。例如,給定一個高維函數


\[ f(x_1, ..., x_N) = \exp\left(\frac{1}{N}\sum_{i=1}^N \sin^2(x_i)\right) \]

(1.1)




由於維度詛咒,樣條在N很大時會失效;MLPs可能會學到廣義加性結構,但它們在逼近指數函數和正弦函數方面效率很低,例如使用ReLU激活函數。相比之下,KANs可以很好地學習組合結構和一維函數,因此大大優於MLPs(見圖3.1)。

在本文中,本研究將通過大量的數值實驗來展示KANs在精度和可解釋性方面如何顯著改善MLPs。本文的結構如圖2.1所示。在第二節中,本研究介紹了KAN架構及其數學基礎,引入了網絡簡化技術以使KANs易於解釋,並提出了網格擴展技術以使KANs逐步提高精度。在第三節中,本研究展示了KANs在數據擬合和求解偏微分方程方面優於MLPs的表現:當數據中存在組合結構時,KANs可以克服維度詛咒,實現比MLPs更好的縮放規律。在第四節中,本研究展示了KANs的可解釋性,並將其應用於科學發現。本研究使用來自數學(結理論)和物理學(Anderson局域化)的兩個實例,證明了KANs可以作為幫助科學家(重新)發現數學和物理定律的有用"合作者"。第五節總結了相關工作。在第六節中,本研究通過討論廣泛的影響和未來方向來總結全文。代碼可在https://github.com/KindXiaoming/pykan獲取,也可以通過pip install pykan進行安裝。

Paper: https://arxiv.org/abs/2404.19756

Code: https://github.com/KindXiaoming/pykan

Quick intro: https://kindxiaoming.github.io/pykan/intro.html

Documentation: https://kindxiaoming.github.io/pykan/


2. Kolmogorov-Arnold 網絡(KAN)


多層感知器(MLPs)受到通用近似定理的啟發。相反,該研究關注的是可以通過一種新型神經網絡實現的Kolmogorov-Arnold表示定理,稱為Kolmogorov-Arnold網絡(KAN)。該研究在2.1節回顧Kolmogorov-Arnold定理,以啟發在2.2節中Kolmogorov-Arnold網絡的設計。在2.3節中,該研究提供了KANs的表達能力和神經縮放法則的理論保證。在2.4節中,該研究提出了一種網格擴展技術,以使KANs逐漸更加精確。在2.5節中,該研究提出了簡化技術,使KANs易於解釋。


2.1 Kolmogorov-Arnold表示定理(Kolmogorov–Arnold representation theorem)

弗拉基米爾·阿諾德(Vladimir Arnold)和安德烈·柯爾莫果洛夫(Andrey Kolmogorov)證明,如果f是在有界域上的多變量連續函數,則f可以表示為一個有限組合的連續一維函數和二元加法運算。更具體地說,對於一個平滑的\(f:[0, 1]^n → R\),


\[ f(x) = f(x_1, ..., x_n) = \sum_{q=1}^{2n+1} \Phi_q \left( \sum_{p=1}^n \phi_{q,p}(x_p) \right), \]

(2.1)


其中 \(\phi_{q,p} : [0, 1] → R\) 和 \(\Phi_q : R → R\)。從某種意義上說,他們證明了加法是唯一真正的多變量函數,因為其他所有函數都可以使用一維函數和求和來表示。人們可能會天真地認為這對機器學習是個好消息:學習一個高維函數歸結為學習一定數量的一維函數。然而,這些一維函數可能是不光滑甚至是分形的,因此在實踐中可能無法學習。由於這種病態行為,Kolmogorov-Arnold表示定理基本上在機器學習中被判了死刑,被認為在理論上是健全的但實際上無用。



然而,該研究對Kolmogorov-Arnold定理在機器學習中的用途持更加樂觀的態度。首先,該研究無需堅持原始方程(2.1),該方程僅具有兩層非線性和隱藏層中的少量項目(2n+1):該研究將網絡推廣到任意寬度和深度。其次,科學和日常生活中的大多數功能通常都是平滑的,並且具有稀疏的組合結構,這可能有助於實現平滑的Kolmogorov-Arnold表示。這裡的理念接近物理學家的心態,他們通常更關心典型案例而非最壞情況。畢竟,我們的物理世界和機器學習任務必須具有結構,才能使物理學和機器學習在所有方面都有用或具有普遍性。


2.2 KAN架構

假設該研究有一個監督學習任務,包含輸入輸出對{xi, yi},其中該研究希望找到一個函數f,使得對所有數據點yi ≈ f(xi)。方程(2.1)意味著如果該研究能找到適當的一維函數\(\phi_{q,p}\)和\(\Phi_q\),那麼該研究就完成了。這啟發該研究設計了一個明確參數化方程(2.1)的神經網絡。由於所有需要學習的函數都是一維函數,該研究可以將每個一維函數參數化為一條B-樣條曲線,具有可學習的局部B-樣條基函數係數(見圖2.2右)。現在該研究有了一個KAN的原型,其計算圖由方程(2.1)確切指定,並在圖0.1(b)中顯示(輸入維度n = 2),顯示為具有在邊上放置激活函數而非節點的兩層神經網絡(節點上執行簡單的求和),並在中間層寬度為2n+1。

正如所提及的,這樣的網絡被認為太簡單,無法在實踐中以平滑的樣條任意精確地近似任何函數!因此,該研究將KAN推廣為更寬更深。立即如何使KAN更深並不清楚,因為Kolmogorov-Arnold表示對應於兩層KAN。據該研究所知,尚無與更深KAN對應的“一般化”定理版本。

當該研究注意到MLPs和KANs之間的類比時,突破發生了。在MLPs中,一旦定義了一層(由線性變換和非線性組成),就可以堆疊更多層來使網路更深。為了建立深層KANs,該研究首先需要回答:“什麼是KAN層?”

結果顯示,一個KAN層,其輸入維度為nin,輸出維度為nout,可以被定義為一個1D函數的矩陣:


\[ \Phi = \{\phi_{q,p}\}, p = 1, 2, ..., nin, q = 1, 2, ..., nout, \]

(2.2)


其中函數\(\phi_{q,p}\)具有可訓練的參數。在Kolmogorov-Arnold定理中,內部函數形成了一個KAN層,nin = n 和 nout = 2n+1,外部函數形成了一個KAN層,nin = 2n+1 和 nout = 1。因此,方程(2.1)中的Kolmogorov-Arnold表示只是兩個KAN層的組合。現在變得清楚了,擁有更深的Kolmogorov-Arnold表示意味著什麼:只需堆疊更多的KAN層!

讓該研究引入一些符號。這一段可能有點技術性,但讀者可以參考圖2.2(左)以獲得具體的例子和直觀的理解。一個KAN的形狀由一個整數數組表示:


\[ [n_0, n_1, ..., n_L], \]

(2.3)


其中ni是計算圖中第i層的節點數。該研究用(l, i)表示第l層的第i個神經元,並用xl,i表示(l, i)-神經元的激活值。在第l層和第l+1層之間,有nlnl+1個激活函數:連接(l, i)和(l+1, j)的激活函數用


\[ \phi_{l,j,i}, l = 0, ..., L-1, i = 1, ..., nl, j = 1, ..., nl+1 \]

(2.4)


表示。\( \phi_{l,j,i} \)的預激活簡單地是\(xl,i\);\( \phi_{l,j,i} \)的後激活由\( \tilde{x}_{l,j,i} \equiv \phi_{l,j,i}(x_{l,i}) \)表示。第\((l + 1, j)\)神經元的激活值僅是所有傳入後激活的總和:


\[ x_{l+1,j} = \sum_{i=1}^{nl} \tilde{x}_{l,j,i} = \sum_{i=1}^{nl} \phi_{l,j,i}(x_{l,i}), j = 1, ..., nl+1. \]

(2.5)


以矩陣形式表示,這可以寫為:


\[ x_{l+1} = \begin{pmatrix} \phi_{l,1,1}( \cdot ) & \phi_{l,1,2}( \cdot ) & \cdots & \phi_{l,1,n_l}( \cdot ) \\ \phi_{l,2,1}( \cdot ) & \phi_{l,2,2}( \cdot ) & \cdots & \phi_{l,2,n_l}( \cdot ) \\ \vdots & \vdots & \ddots & \vdots \\ \phi_{l,n_{l+1},1}( \cdot ) & \phi_{l,n_{l+1},2}( \cdot ) & \cdots & \phi_{l,n_{l+1},n_l}( \cdot ) \end{pmatrix} x_l, \]

(2.6)


其中 \(\Phi_l\) 是對應於第 \(l\) 層的函數矩陣。一個通用的KAN網絡是 \(L\) 層的組合:給定一個輸入向量 \(x_0 \in \mathbb{R}^{n_0}\),KAN的輸出為


\[ \text{KAN}(x) = (\Phi_{L-1} \circ \Phi_{L-2} \circ \ldots \circ \Phi_1 \circ \Phi_0)x. \]

(2.7)


該研究還可以重寫上述方程,使其更類似於方程(2.1),假設輸出維度 \(n_L = 1\),並定義 \(f(x) \equiv KAN(x)\):


\[ f(x) = \sum_{i_{L-1}=1}^{n_{L-1}} \phi_{L-1, i_{L}, i_{L-1}}\left(\sum_{i_{L-2}=1}^{n_{L-2}} \cdots \left( \sum_{i_2=1}^{n_2} \phi_{2, i_3, i_2}\left( \sum_{i_1=1}^{n_1} \phi_{1, i_2, i_1}\left( \sum_{i_0=1}^{n_0} \phi_{0, i_1, i_0}(x_{i_0}) \right) \right) \cdots \right) \right), \]

(2.8)


這個表達式雖然複雜,但我們對KAN層的抽象化和視覺化更為清晰直觀。原始的Kolmogorov-Arnold表示方程(2.1)對應於形狀為 \([n, 2n + 1, 1]\) 的2層KAN。值得注意的是,所有操作都是可微的,因此該研究可以使用反向傳播訓練KANs。作為對比,一個MLP可以寫成交錯的仿射變換 \(W\) 和非線性 \(\sigma\):


\[ MLP(x) = (W_{L-1} \circ \sigma \circ W_{L-2} \circ \sigma \circ \cdots \circ W_1 \circ \sigma \circ W_0)(x). \]

(2.9)


很明顯,MLPs將線性變換和非線性分別作為 \(W\) 和 \(\sigma\) 處理,而KANs則將它們結合在 \(\Phi\) 中。在圖0.1的(c)和(d)中,該研究可視化了三層MLP和三層KAN,以澄清它們的差異。


實施細節

儘管KAN層方程(2.5)看起來非常簡單,但使其優化良好並非微不足道。關鍵技巧包括:

1. 殘差激活函數。該研究包括一個基函數 \(b(x)\)(類似於殘差連接),使得激活函數 \(\phi(x)\) 是基函數 \(b(x)\) 和樣條函數的總和:


   \[ \phi(x) = w(b(x) + \text{spline}(x)). \]

(2.10)


該研究設定


\[ b(x) = \text{silu}(x) = \frac{x}{1 + e^{-x}} \]

(2.11)


在大多數情況下。樣條函數 \( \text{spline}(x) \) 則被參數化為B-樣條的線性組合:


\[ \text{spline}(x) = \sum_i c_i B_i(x) \]

(2.12)


其中 \( c_i \) 是可訓練的。原則上 \( w \) 是多餘的,因為它可以被吸收進 \( b(x) \) 和 \( \text{spline}(x) \)。然而,該研究仍然包括這個 \( w \) 因子,以更好地控制激活函數的總體幅度。


2. 初始化規模。每個激活函數被初始化為使 \( \text{spline}(x) \approx 0 \)。\( w \) 根據Xavier初始化法則初始化,這種方法被用來初始化MLPs中的線性層。

3. 樣條網格的更新。該研究根據其輸入激活值即時更新每個網格,以解決樣條在有界區域內定義但激活值在訓練過程中可能超出固定區域的問題。


參數計數

為簡單起見,假設該研究設置一個網絡:

  1. 深度為 \( L \),
  2. 層寬度相等,\( n_0 = n_1 = \ldots = n_L = N \),
  3. 每個樣條的階數通常為 \( k = 3 \),在 \( G \) 個區間(共 \( G+1 \) 個網格點)。

則總共有 \( O(N^2 L (G + k)) \approx O(N^2 LG) \) 參數。相比之下,具有深度 \( L \) 和寬度 \( N \) 的MLP只需要 \( O(N^2 L) \) 參數,這似乎比KAN更有效。幸運的是,KANs通常需要比MLPs小得多的 \( N \),這不僅節省了參數,而且還實現了更好的泛化(參見圖3.1和圖3.3)並促進了解釋性。該研究指出,對於1D問題,可以取 \( N = L = 1 \),而該研究實現中的KAN網絡僅是一個樣條近似。對於更高維度,該研究通過下面的定理描述了KANs的泛化行為。




2.3 KAN的近似能力和縮放定律

回想方程(2.1)中,寬度為(2n+1)的兩層表示可能不是平滑的。然而,更深的表示可能帶來更平滑激活的優勢。例如,對於四變量函數


\[ f(x_1, x_2, x_3, x_4) = \exp\left(\sin(x_1^2 + x_2^2) + \sin(x_3^2 + x_4^2)\right) \]

(2.13)


可以被一個[4, 2, 1, 1] KAN平滑地表示,這是一個三層KAN,但可能不允許兩層KAN具有平滑激活。為了便於進行近似分析,該研究仍然假設激活的平滑性,但允許表示可以任意寬和深,如方程(2.7)。為了強調該研究KAN對有限集合網格點的依賴,該研究在下面使用\(ΦG_l\)和\(ΦG_l,i,j\)替代方程(2.5)和(2.6)中使用的符號\(Φ_l\)和\(Φ_l,i,j\)。


定理2.1(近似理論,KAT):設\(x = (x_1, x_2, ..., x_n)\)。假設一個函數\(f(x)\)允許表示


\[ f = (\Phi_{L-1} \circ \Phi_{L-2} \circ \ldots \circ \Phi_1 \circ \Phi_0)(x), \]

(2.14)


如方程(2.7)所示,其中每一個\(Φ_l,i,j\)都是\((k+1)\)次連續可微的。那麼存在一個依賴於f及其表示的常數C,使得有關於網格大小G的以下近似界限:存在k階B-樣條函數\(ΦG_l,i,j\),對於任何\(0≤m≤k\)有界限


\[ \|f - (\Phi^G_{L-1} \circ \Phi^G_{L-2} \circ \ldots \circ \Phi^G_1 \circ \Phi^G_0)(x)\|_{C^m} \leq CG^{-k-1+m}. \]

(2.15)


這裡採用\(C^m\)範數來測量高達\(m\)階的導數的幅度:


\[ \|g\|_{C^m} = \max_{|\beta| \leq m} \sup_{x \in [0,1]^n} |D^\beta g(x)|. \]


證明:通過經典的1D B-樣條理論和事實,即\(Φ_l,i,j\)作為連續函數可以在有界域上均勻有界,知道存在有限網格B-樣條函數\(ΦG_l,i,j\),使得對於任何\(0≤m≤k\),


\[ \|(\Phi_{l,i,j} \circ \Phi_{l-1} \circ \ldots \circ \Phi_1 \circ \Phi_0)(x) - (\Phi^G_{l,i,j} \circ \Phi_{l-1} \circ \ldots \circ \Phi_1 \circ \Phi_0)(x)\|_{C^m} \leq CG^{-k-1+m}, \]


具有獨立於G的常數C。固定這些B-樣條逼近。因此,通過定義殘差\(R_l\):


\[ R_l = (\Phi^G_{L-1} \circ \ldots \circ \Phi^G_{l+1} \circ \Phi_l \circ \Phi_{l-1} \circ \ldots \circ \Phi_0)x - (\Phi^G_{L-1} \circ \ldots \circ \Phi^G_{l+1} \circ \Phi^G_l \circ \Phi_{l-1} \circ \ldots \circ \Phi_0)x \]


滿足


\[ \|R_l\|_{C^m} \leq CG^{-k-1+m}, \]


具有獨立於G的常數。最後注意到


\[ f - (\Phi^G_{L-1} \circ \Phi^G_{L-2} \circ \ldots \circ \Phi^G_1 \circ \Phi^G_0)(x) = R_{L-1} + R_{L-2} + \ldots + R_1 + R_0, \]


(2.15)成立。


根據定理2.1的假設,KANs在有限網格大小下可以用一個與維度無關的殘差率很好地逼近函數,因此戰勝了維度的詛咒!這自然而然地發生,因為只用樣條來逼近1D函數。特別是對於\(m = 0\),恢復在\(L∞\)範數中的精確度,這反過來提供了在有限域上的均方根誤差(RMSE)的一個界限,給出了\(k+1\)的縮放指數。當然,常數C依賴於表示;因此它將依賴於維度。該研究將在未來的工作中討論常數對維度的依賴。

該研究指出,儘管Kolmogorov-Arnold定理(2.1)對應於形狀為\([d, 2d+1, 1]\)的KAN表示,它的函數並不一定是平滑的。另一方面,如果該研究能夠識別一個平滑的表示(可能需要額外的層或使KAN比理論預測的更寬),那麼定理2.1表明該研究可以戰勝維度的詛咒(COD)。這不應該令人驚訝,因為可以固有地學習函數的結構,並使有限樣本KAN近似可解釋。



神經縮放定律:與其他理論的比較

神經縮放定律是指隨著模型參數的增加,測試損失減少的現象,即\(ℓ ∝ N^(-α)\),其中\(ℓ\)是測試RMSE,\(N\)是參數數量,\(α\)是縮放指數。更大的\(α\)承諾通過簡單地擴大模型來實現更多的改進。不同的理論已被提出來預測\(α\)。Sharma和Kaplan建議,\(α\)來自於在具有內在維度\(d\)的輸入流形上的資料擬合。如果模型函數類是k階分段多項式(對於ReLU,k=1),則標準近似理論暗示\(α = (k+1)/d\)來自近似理論。這個界限受到維度的詛咒的困擾,所以人們尋求其他不依賴於\(d\)的界限,通過利用組合結構。特別是,Michaud等人考慮了只涉及一元(例如,平方、正弦、指數)和二元(+ 和 ×)操作的計算圖,發現\(α = (k+1)/d* = (k+1)/2\),其中\(d* = 2\)是最大元性。Poggio等人利用組合稀疏的想法,證明了給定函數類\(Wm\)(其導數直到第\(m\)階是連續的),需要\(N = O(ε^(-2/m))\)數量的參數來實現誤差\(ε\),這等效於\(α = m/2\)。該研究的方法,假設平滑的Kolmogorov-Arnold表示存在,將高維函數分解為幾個1D函數,給出\(α = k+1\)(其中k是樣條的分段多項式階數)。該研究選擇\(k = 3\)的立方樣條,所以\(α = 4\)是與其他工作相比最大且最好的縮放指數。該研究將在第3.1節展示,這個界限α = 4實際上可以用KANs實現,而以前的工作報告說,即使在滿足較慢的界限(例如,\(α = 1\))和迅速達到平台的情況下,MLPs也存在問題。


KAT和UAT之間的比較

全連接神經網路的威力是由通用逼近定理(UAT)證明的,該定理指出,給定一個函數和誤差容限\(\epsilon > 0\),具有\(k > N(\epsilon)\)個神經元的雙層網路可以將函數逼近至\(\epsilon\)的誤差內。然而,UAT並不保證\(N(\epsilon)\)如何隨\(\epsilon\)縮放。事實上,它遭受維度詛咒(COD)的困擾,並且在某些情況下,\(N\)已被證明隨\(d\)呈指數增長。KAT和UAT之間的差異是KAN利用了函數的內在低維表示,而MLP沒有。將展示KAN與符號函數很好地對齊,而MLP則不然。


2.4 精確性:網格擴展

原則上一個樣條可以被任意精確地調整以符合目標函數,只要網格可以被足夠細化。這一優良特性被KANs繼承。相比之下,MLPs沒有“細化”的概念。確實,增加MLPs的寬度和深度可以提高性能(“神經縮放定律”)。但這些神經縮放定律增長緩慢,且實現起來成本高昂,因為需要獨立訓練不同大小的模型。與此相反,對於KANs,可以首先訓練一個參數較少的KAN,然後通過簡單地細化樣條網格擴展到一個參數更多的KAN,無需從頭開始重新訓練大模型。

接下來,該研究描述如何執行網格擴展(如圖2.2右側所示),基本上是在一個粗糙的樣條上擬合一個新的細糙樣條。假設想在有界區域\([a, b]\)上用B-樣條的\(k\)階來近似一個一維函數\(f\)。一個具有G1區間的粗糙網格在\(\{t_0 = a, t_1, ..., t_{G1} = b\}\)處有網格點,這被擴充為\(\{t_{-k}, ..., t_{-1}, t_0, ..., t_{G1}, t_{G1+1}, ..., t_{G1+k}\}\)。有 \(G_1 + k\) 個B-樣條基函數,其中第 \(i\) 個B-樣條 \(B_i(x)\) 僅在 \([t_{-k+i}, t_{i+1}]\) 上非零。然後,\(f\) 在粗糙網格上用這些B-樣條基函數的線性組合表示:


\[ f_{\text{coarse}}(x) = \sum_{i=0}^{G_1+k-1} c_i B_i(x) \]

(2.16)


給定更細的網格,具有 \(G_2\) 區間,\(f\) 在細網格上相應地表示為:


\[ f_{\text{fine}}(x) = \sum_{j=0}^{G_2+k-1} c'_j B'_j(x) \]


這些 \(c'_j\) 的參數可以從 \(c_i\) 初始化,通過最小化 \(f_{\text{fine}}(x)\) 到 \(f_{\text{coarse}}(x)\) 的距離來獲得(在某些 \(x\) 的分佈上):


\[ \{c'_j\} = \arg \min_{\{c'_j\}} \mathbb{E}_{x \sim p(x)} \left[ \left( \sum_{j=0}^{G_2+k-1} c'_j B'_j(x) - \sum_{i=0}^{G_1+k-1} c_i B_i(x) \right)^2 \right] \]

(2.16)


這可以通過最小平方算法實現。對KAN中的所有樣條進行獨立的網格擴展。

玩具範例:階梯狀的損失曲線。使用一個玩具範例 f(x, y) = exp(sin(πx) + y^2) 來展示網格擴展的效果。在圖 2.3(左上)中,我們展示了一個 [2, 5, 1] KAN 的訓練和測試 RMSE。網格點的數量開始為 3,每 200 個 LBFGS 步驟增加到更高的值,最終達到 1000 個網格點。很明顯,每次細化網格時,訓練損失下降得比之前更快(除了最細的 1000 點網格,優化可能由於不良的損失地形而停止工作)。然而,測試損失先下降然後上升,呈現 U 形,這是由於偏差-方差權衡(欠擬合 vs. 過擬合)。我們推測,當參數數量與數據點數量匹配時,在插值閾值處會達到最佳測試損失。由於我們的訓練樣本為 1000,而 [2, 5, 1] KAN 的總參數為 15G(G 是網格區間的數量),我們預期插值閾值為 G = 1000/15 ≈ 67,這與我們實驗觀察到的值 G ∼ 50 大致一致。

小型 KAN 泛化性能更好。這是我們可以達到的最佳測試性能嗎?請注意,合成任務可以由 [2, 1, 1] KAN 精確表示,所以我們訓練一個 [2, 1, 1] KAN 並在圖 2.3 右上方呈現訓練動態。有趣的是,它可以實現比 [2, 5, 1] KAN 更低的測試損失,具有更清晰的階梯結構,由於參數更少,插值閾值被延遲到更大的網格尺寸。這突顯了選擇 KAN 架構的微妙之處。如果我們不知道問題結構,我們如何確定最小的 KAN 形狀?在第 2.5 節中,我們將提出一種通過正則化和剪枝自動發現這種最小 KAN 架構的方法。

縮放規律:與理論比較。我們也對測試損失如何隨網格參數數量的增加而降低感興趣。在圖 2.3(左下)中,[2,1,1] KAN 的縮放大致為 test RMSE ∝ G^−3。然而,根據定理 2.1,我們預期 test RMSE ∝ G^−4。我們發現樣本之間的誤差並不均勻。這可能歸因於邊界效應 [18]。事實上,有一些樣本的誤差顯著大於其他樣本,使得整體縮放放緩。如果我們繪製平方損失的中位數(非平均值)的平方根,我們得到接近 G^−4 的縮放。儘管存在這種次優性(可能是由於優化),KAN 在數據擬合(圖 3.1)和 PDE 求解(圖 3.3)方面仍然具有比 MLP 更好的縮放規律。此外,訓練時間隨網格點數 G 的增加而呈現良好的縮放,如圖 2.3 右下方所示。

外部與內部自由度。KAN 強調的一個新概念是外部與內部自由度(參數)之間的區別。節點如何連接的計算圖表示外部自由度("dofs"),而激活函數內部的網格點是內部自由度。KAN 受益於同時具有外部 dofs 和內部 dofs。外部 dofs(MLP 也有但樣條沒有)負責學習多個變量的組合結構。內部 dofs(樣條也有但 MLP 沒有)負責學習單變量函數。

 2.5 解釋性:簡化KANs並使它們可互動

上一小節留下的一個問題是,不知道如何選擇最佳匹配數據集結構的KAN形狀。例如,如果知道資料集是通過符號公式 \(f(x, y) = \exp(\sin(\pi x) + y^2)\) 生成的,那麼知道一個 [2, 1, 1] KAN 能夠表達這個函數。但在實踐中不知道這些信息,所以最好有方法能夠自動確定這種形狀。這個想法是從一個足夠大的KAN開始,並用稀疏正則化訓練它,隨後進行剪枝。該研究將展示,這些剪枝後的KAN比未剪枝的更易於解釋。為了使KANs最大限度地可解釋,該研究在2.5.1節提出了一些簡化技術,並在2.5.2節示例了用戶如何與KANs互動以使它們更易於解釋。

2.5.1 簡化技術

1. 稀疏化:對於MLPs,使用L1正則化線性權重以傾向於稀疏化。KANs可以適應這一高級思想,但需要兩個修改:

  •  KANs中沒有線性“權重”。線性權重被可學習的激活函數替代,所以應定義這些激活函數的L1範數。
  •  該研究發現對於KANs的稀疏化,僅使用L1正則化是不夠的;相反,還需要額外的熵正則化。

定義激活函數 \( \phi \) 的L1範數為其在 \( N_p \) 個輸入上的平均幅度,即


\[ |\phi|_1 \equiv \frac{1}{N_p} \sum_{s=1}^{N_p} |\phi(x^{(s)})| \]

(2.17)


對於一個KAN層 \( \Phi \),具有 \( n_{\text{in}} \) 個輸入和 \( n_{\text{out}} \) 個輸出,定義 \( \Phi \) 的L1範數為所有激活函數L1範數的總和,即


\[ |\Phi|_1 \equiv \sum_{i=1}^{n_{\text{in}}} \sum_{j=1}^{n_{\text{out}}} |\phi_{i,j}|_1 \]

(2.18)


此外,定義 \( \Phi \) 的熵為


\[ S(\Phi) \equiv -\sum_{i=1}^{n_{\text{in}}} \sum_{j=1}^{n_{\text{out}}} \frac{|\phi_{i,j}|_1}{|\Phi|_1} \log \left( \frac{|\phi_{i,j}|_1}{|\Phi|_1} \right) \]

(2.19)


總訓練目標 \( \ell_{\text{total}} \) 是預測損失 \( \ell_{\text{pred}} \) 加上所有KAN層的L1和熵正則化:


\[ \ell_{\text{total}} = \ell_{\text{pred}} + \lambda \left( \mu_1 \sum_{l=0}^{L-1} |\Phi_l|_1 + \mu_2 \sum_{l=0}^{L-1} S(\Phi_l) \right) \]

(2.20)


其中,\( \mu_1 \) 和 \( \mu_2 \) 是相對大小,通常設置為 \( \mu_1 = \mu_2 = 1 \),而 \( \lambda \) 控制整體正則化大小。



2. 可視化。當我們將一個KAN可視化時,為了感受量級,我們將一個激活函數\(\phi_{l,i,j}\)的透明度設置為與\(\tanh(\\beta A_{l,i,j})\)成正比,其中\(\beta = 3\)。因此,量級較小的函數會顯得褪色,使我們能夠專注於重要的函數。

3. 修剪。在使用稀疏化懲罰進行訓練之後,我們可能還想將網路修剪為更小的子網路。我們在節點層面(而不是邊層面)對KAN進行稀疏化。對於每個節點(例如第\(l\)層的第\(i\)個神經元),我們定義其輸入和輸出分數為

$$I_{l,i} = \max_k(|\phi_{l-1,k,i}|1),\ O{l,i} = \max_j(|\phi_{l+1,j,i}|_1),$$

(2.21)

如果輸入和輸出分數都大於閾值超參數(\\theta = 10^{-2}\)(預設值),則認為一個節點是重要的。所有不重要的神經元都會被修剪掉。

4. 符號化。在我們懷疑某些激活函數實際上是符號函數(例如\(\cos\)或log)的情況下,我們提供了一個介面來將它們設置為指定的符號形式。fix_symbolic(l,i,j,f)可以將\((l,i,j)\)激活設置為\(f\)。然而,我們不能簡單地將激活函數設置為精確的符號公式,因為其輸入和輸出可能存在偏移和縮放。因此,我們從樣本中獲取預激活\(x\)和後激活\(y\),並擬合仿射參數\((a,b,c,d)\),使得\(y \approx cf(ax + b) + d\)。擬合通過\(a\),\(b\)的迭代網格搜索和線性回歸完成。

除了這些技術,我們還提供了額外的工具,允許用戶對KAN進行更細粒度的控制,列在附錄A中。

2.5.2 人類與KANs的互動實例

該研究提出了一些簡化KANs的技術,可以將這些簡化選項視為可以點擊的按鈕。使用者通過與這些按鈕互動,可以決定哪個按鈕是最有希望點擊的,使KANs更易於解釋。用以下示例展示了使用者如何與KAN互動以獲得最大限度可解釋的結果。

考慮回歸任務


\[ f(x, y) = \exp(\sin(\pi x) + y^2) \]

(2.22)


假設有資料點 \( (x_i, y_i, f_i) \),\( i = 1, 2, \ldots, N_p \),一位假想的使用者Alice有興趣弄清楚符號公式。Alice與KANs互動的步驟如下(圖2.4所示):

步驟1:訓練與稀疏化。從一個完全連接的[2, 5, 1] KAN開始,用稀疏化正則化訓練可以使其變得相當稀疏。隱藏層中的5個神經元中有4個看起來是無用的,因此希望將它们剪除。

步驟2:剪枝。自動剪枝觀察到除了最後一個以外,所有隐藏層神經元都被丢棄,只留下一個 [2, 1, 1] KAN。激活函數是已知的符號函数。

步驟3:設定符号函數。假设用户能從KAN圖中正確猜測这些符號公式,他們可以設置:

```

fix_symbolic(0,0,0,'sin')

fix_symbolic(0,1,0,'x^2')

fix_symbolic(1,0,0,'exp')

```

(2.23)


如果使用者沒有領域知識或不清楚這些激活函數可能對應哪些符號函數,提供了一個 `suggest_symbolic` 函數來建議符號候選者。

步驟4:進一步訓練。在將所有激活函數符號化之後,唯一剩下的參數是仿射參數。繼續訓練這些仿射參數,當看到損失下降到機器精度時,知道找到了正確的符號表達式。

步驟5:輸出符號公式。使用 Sympy 計算輸出節點的符號公式。使用者獲得的公式為 \(1.0e1.0y^2+1.0sin(3.14x)\),這是真正的答案(只顯示了\(π\)的兩位小數)。

為什麼不用符號迴歸(SR)?

在這個例子中使用符號迴歸似乎是合理的。然而,一般來說,符號迴歸法比較脆弱,且難以調試。它們要麼最終成功返回結果,要麼失敗,而不提供可解釋的中間結果。相較之下,KAN通過梯度下降在函數空間中進行連續搜索,因此其結果更加連續且穩定。此外,由於KAN的透明度,使用者可以比使用符號迴歸更好地控制KAN。該研究將KAN的可視化方式比喻為向使用者展示KAN的「大腦」,使用者可以對KAN進行「手術」(調試)。符號迴歸通常不提供這種控制層級。將在第4.4節展示此類實例。更普遍地,當目標函數不是符號表達時,符號迴歸將失敗,但KAN仍然可以提供有意義的輸出。例如,特殊函數(如貝塞爾函數)是不可能被符號迴歸學習的,除非事先提供,但KAN可以使用樣條進行數值近似(見圖4.1(d))。


3 KANs的精確性


在這一節中,該研究將展示KANs在各種任務(如回歸和偏微分方程求解)中表示函數的有效性超過了MLPs。在比較兩類模型時,公平地比較它們的精度(損失)和復雜度(參數數量)是合理的。該研究將展示KANs比MLPs展示出更有利的帕雷托前緣。此外,在3.5節中,該研究將顯示KANs能夠自然地進行持續學習而不會出現災難性遺忘。



3.1 玩具資料集


在第2.3節中,理論建議測試RMSE損失 \( \ell \) 隨模型參數 \( N \) 的增加而按 \( \ell \propto N^{-4} \) 縮放。然而,這依賴於存在柯爾莫哥洛夫-阿諾德表示定理(Kolmogorov-Arnold representation theorem),也被稱為超曲面定理(superposition theorem),是數學分析中的一個重要定理。它由俄羅斯數學家安德烈·柯爾莫哥洛夫(Andrey Kolmogorov)在1956-1957年提出,並由其學生弗拉基米爾·阿諾德(Vladimir Arnold)在1959年完成證明。
該定理的主要內容如下:

任意一個連續函數 \(f(x_1,\ldots,x_n)\),其中自變量 \((x_1,\ldots,x_n)\) 屬於 \(n\) 維單位立方體 \([0,1]^n\),都可以表示為:
\(f(x_1,\ldots,x_n) = \sum_{q=0}^{2n} \Phi_q \left(\sum_{p=1}^n \psi_{q,p}(x_p)\right)\)
其中 \(\Phi_q\) 是連續函數,\(\psi_{q,p}\) 是 \([0,1]\) 上的單變量函數。

這個定理說明,任何多維連續函數都可以由一些單變量函數的疊加來表示,而這些單變量函數與原函數的維度無關。這一結果在當時是非常令人驚訝的,因為它表明高維函數可以用低維函數來表示。

在神經網絡領域,這個定理啟發了新的網絡結構的設計,即論文中提到的KAN(Kolmogorov-Arnold network)。傳統的神經網絡使用固定的激活函數,而KAN使用可學習的、類似於定理中 \(\psi_{q,p}\) 的單變量函數作為激活函數,並將它們放在網絡的邊上而不是節點上。這種結構使得KAN具有更強的表達能力和可解釋性。

然而,柯爾莫哥洛夫-阿諾德表示定理在實踐中也存在一些問題。定理本身是一個存在性結果,並沒有給出如何構造這些單變量函數的方法。此外,定理中的單變量函數可能高度非光滑,這在實際應用中可能會帶來數值不穩定性。儘管如此,該定理仍然是啟發神經網絡設計的重要理論基礎,並為進一步的理論研究提供了方向。。為了作為一個合理性檢驗,構造了五個我們知道具有平滑KA表示的例子:

  1.  \( f(x) = J_0(20x) \),這是貝塞爾函數。由於它是一個一元函數,它可以被樣條表示,這是一個 [1, 1] KAN。
  2.  \( f(x, y) = \exp(\sin(\pi x) + y^2) \)。我們知道它可以被一個 [2, 1, 1] KAN準確表示。
  3.  \( f(x, y) = xy \)。我們從圖4.1知道它可以被一個 [2, 2, 1] KAN準確表示。
  4.  一個高維例子 \( f(x_1, \ldots, x_{100}) = \exp\left(\frac{1}{100} \sum_{i=1}^{100} \sin^2(\pi x_i)\right) \),它可以被一個 [100, 1, 1] KAN表示。
  5.  一個四維例子 \( f(x_1, x_2, x_3, x_4) = \exp\left(\frac{1}{2} (\sin(\pi (x_1^2 + x_2^2)) + \sin(\pi (x_3^2 + x_4^2)))\right) \),它可以被一個 [4, 4, 2, 1] KAN表示。

通過每200步增加網格點數量來訓練這些KANs,總共覆蓋 \( G = \{3, 5, 10, 20, 50, 100, 200, 500, 1000\} \)。該研究還訓練了具有不同深度和寬度的MLPs作為基線。無論是MLPs還是KANs都使用LBFGS訓練了總共1800步。將KANs和MLPs的測試RMSE作為參數數量的函數繪製在圖3.1中,顯示KANs比MLPs有更好的縮放曲線,特別是對於高維例子。為了比較,將從KAN理論預測的紅色虛線(\( \alpha = k+1 = 4 \))和從Sharma & Kaplan 預測的黑色虛線(\( \alpha = (k+1)/d = 4/d \))繪製在圖中。KANs幾乎可以達到更陡峭的紅線,而MLPs即使在達到較慢的黑線時也難以收斂並迅速達到平台期。另外還注意到,對於最後一個例子,2層的KAN [4, 9, 1]的表現比3層的KAN(形狀 [4, 2, 2, 1])差得多。這突顯了更深層KANs的更大表達能力,這對MLPs也是一樣:更深的MLPs比較淺的具有更多的表達能力。


 3.2 特殊函數


上述結果的一個限制是,該研究假設了「真實」KAN形狀的知識。在實踐中,不知道KA表示的存在。即使被告知存在這樣的KA表示,也不知道KAN形狀。多變量特殊函數是這種情況,因為如果多變量特殊函數(例如,貝塞爾函數 \( f(\nu, x) = J_\nu(x) \))能用KA表示表達,只涉及一元函數和加法,這在數學上將是驚人的。以下展示:
  1. 尋找(近似的)緊湊KA表示特殊函數是可能的,這揭示了從柯爾莫哥洛夫-阿諾德表示的角度看特殊函數的新數學屬性。
  2. KANs在表示特殊函數方面比MLPs更有效和精確

該研究收集了15個在數學和物理中常見的特殊函數,彙總在表2中。選擇了固定寬度為5或100的MLPs,並掃描了深度 \(\{2, 3, 4, 5, 6\}\)。運行了帶剪枝和不帶剪枝的KANs。


  • 未剪枝的KANs:固定了KAN的形狀,寬度設置為5,深度也在 \(\{2,3,4,5,6\}\) 中掃描。
  • 帶剪枝的KANs。使用第2.5.1節中的稀疏化(\( \lambda = 10^{-2} \) 或 \( 10^{-3} \))和剪枝技術來獲得一個從固定形狀KAN剪枝得到的更小的KAN。
每個KAN都初始化為有 \( G = 3 \) 的網格點,使用LBFGS訓練,每200步增加網格點數量,覆蓋 \( G = \{3, 5, 10, 20, 50, 100, 200\} \)。對於每個超參數組合,我們運行3個隨機種子。

對於每個資料集和每個模型家族(KANs或MLPs),在(參數數量,RMSE)平面上繪製了帕雷托前緣,顯示在圖3.2中。KANs的表現一致優於MLPs,即KANs在相同參數數量下可以達到比MLPs更低的訓練/測試損失。此外,報告了自動發現的特殊函數KANs的(令人驚訝的緊湊)形狀,列在表2中。一方面,解釋這些緊湊表達在數學上意味著什麼是有趣的(在附錄F的圖F.1和F.2中包括了KAN的插圖)。另一方面,這些緊湊的表達暗示了將高維查找表拆分為幾個一維查找表的可能性,這可以在推理時節省大量內存,並且(幾乎可以忽略不計)進行一些加法運算的開銷。

3.3 費曼數據集


第3.1節的設置是當我們清楚地知道「真實」KAN形狀時。第3.2節的設置是當我們明確不知道「真實」KAN形狀時。這部分調查位於中間的設置:鑑於數據集的結構,我們可能會手工構造KAN,但我們不確定它們是否是最優的。在這個範疇中,比較人工構造的KAN和通過剪枝(在第2.5.1節中的技術)自動發現的KAN是有趣的。





當有資料集的結構資訊時,可以手動構建KAN,但可能不確定這些結構是否完全準確或有效。例如,可能基於資料的某些特性猜測一個特定的KAN結構,然後嘗試用這個結構進行訓練和優化,但隨後的剪枝和調整可能顯示出一個更有效的結構。

這種探索對於理解KAN在實際應用中的彈性和潛力至關重要。透過對比人工構建的模型與經過剪枝精簡後的模型,可以更好地評估KAN設計選擇的有效性,並進一步優化模型的表現和效率。

費曼資料集。費曼資料集從費曼的教科書中收集了許多物理方程式。對於至少有兩個變數的問題感興趣,因為單變數問題對於KAN來說是微不足道的(它們簡化為一維樣條)。費曼資料集的一個樣本方程式是相對論速度加法公式 \( f(u, v) = \frac{u+v}{1+uv} \)。通過隨機抽取 \( u_i \in (-1, 1) \) 和 \( v_i \in (-1, 1) \),並計算 \( f_i = f(u_i, v_i) \),可以建立資料集。給定許多元組 \( (u_i, v_i, f_i) \),訓練一個神經網路,目的是從 \( u \) 和 \( v \) 預測 \( f \)。該研究感興趣的是:
(1) 神經網路在測試樣本上的表現如何;
(2) 可以從神經網路中學到多少關於問題結構的知識。

比較四種神經網路:
  1. 人工構造的KAN。給定一個符號公式,將其改寫為Kolmogorov-Arnold表示。例如,要將兩個數 \( x \) 和 \( y \) 相乘,可以使用恆等式 \( xy = \frac{(x+y)^2}{4} - \frac{(x-y)^2}{4} \),對應於一個 [2, 2, 1] KAN。構造的形狀列在表中的「人工構造KAN形狀」一欄。
  2. 不帶修剪的KAN。我們將KAN形狀固定為寬度5,深度在 {2, 3, 4, 5, 6} 中掃描。
  3. 帶修剪的KAN。使用節中的稀疏化(\( \lambda = 10^{-2} \) 或 \( 10^{-3} \))和修剪技術,從中的固定形狀KAN得到一個更小的KAN。
  4. 固定寬度為20的MLP,深度在 {2, 3, 4, 5, 6} 中掃描,激活函數從 {Tanh, ReLU, SiLU} 中選擇。

每個KAN初始化為 \( G = 3 \),用LBFGS訓練,每200步增加網格點數量以涵蓋 \( G = \{3, 5, 10, 20, 50, 100, 200\} \)。對於每個超參數組合,嘗試3個隨機種子。對於每個資料集(方程式)和每種方法,在表中報告了隨機種子和深度上最佳模型(最小KAN形狀,或最低測試損失)的結果。發現MLP和KAN在平均表現上相當。對於每個資料集和每個模型族(KAN或MLP),在參數數量和RMSE損失構成的平面上繪製帕累托前沿,如附錄D中的圖D.1所示。推測費曼資料集太簡單了,不足以讓KAN做出進一步改進,因為變數依賴關係通常是平滑或單調的,這與經常表現出振盪行為的特殊函數的複雜性形成對比。

自動發現的KAN比人工構造的KAN更小。在表中的兩列中報告了修剪後的KAN形狀;一列是可以達到合理損失(即測試RMSE小於 \(10^{-2}\))的最小修剪KAN形狀;另一列是達到最低測試損失的修剪KAN。為了完整起見,在附錄D(圖D.2和D.3)中可視化了所有54個修剪後的KAN。有趣的是,自動發現的KAN形狀(無論是最小的還是最佳的)通常都比我們的人工構造更小。這意味著Kolmogorov-Arnold表示可能比我們想像的更有效。同時,這可能會使解釋變得微妙,因為資訊被壓縮到一個比我們習慣的更小的空間。

以相對論速度合成公式 \( f(u, v) = \frac{u+v}{1+uv} \) 為例。構造相當深,因為假設 \( u \) 和 \( v \) 的乘法需要兩層,\( 1+uv \) 的倒數需要一層,\( (u+v) \) 和 \( \frac{1}{1+uv} \) 的乘法需要另外兩層,總共需要五層。然而,自動發現的KAN只有兩層深!事後看來,如果回憶相對論中的快度技巧,實際上是意料之中的:定義兩個"快度" \( a \equiv \text{arctanh} u \) 和 \( b \equiv \text{arctanh} v \)。在快度空間中,速度的相對論合成是簡單的加法,即 \( \frac{u+v}{1+uv} = \tanh(\text{arctanh} u + \text{arctanh} v) \),這可以用一個兩層的KAN實現。假裝我們不知道物理中快度的概念,我們可能會直接從KAN中發現這個概念,而無需通過試錯的符號操作。

KAN的可解釋性,有助於科學發現,是第四節的主題。

3.4 求解偏微分方程式

考慮具有零Dirichlet邊界條件的Poisson方程。對於 \(\Omega = [-1, 1]^2\),考慮Partial Differential Equation\(\begin{aligned}u_{xx} + u_{yy} &= f \text{ in } \Omega, \\u &= 0 \text{ on }\partial\Omega.\end{aligned}\)
(3.2)





考慮的資料 \(f\) 為 \(-\pi^2(1 + 4y^2) \sin(\pi x) \sin(\pi y^2) + 2\pi \sin(\pi x) \cos(\pi y^2)\),其中 \(u = \sin(\pi x) \sin(\pi y^2)\) 是真實解。使用物理信息神經網路(PINN)的框架來求解這個PDE,損失函數為\(\text{loss}\text{pde} = \alpha \text{loss}i + \text{loss}b := \alpha \frac{1}{n_i}\sum{i=1}^{n_i} |u{xx}(z_i) + u{yy}(z_i) - f(z_i)|^2 + \frac{1}{n_b}\sum_{i=1}^{n_b} u^2(z_i),\)



其中用 \(\text{loss}_i\) 表示內部損失,通過在區域內均勻採樣 \(n_i\) 個點 \(z_i = (x_i, y_i)\) 來離散化和計算,用 \(\text{loss}_b\) 表示邊界損失,通過在邊界上均勻採樣 \(n_b\) 個點來離散化和計算。 \(\alpha\) 是平衡兩項效果的超參數。

使用相同的超參數 \(n_i = 10000\), \(n_b = 800\) 和 \(\alpha = 0.01\) 來比較KAN架構和MLP架構。

測量 \(L^2\) 範數和能量(\(H^1\))範數中的誤差,發現KAN在使用更小的網路和更少的參數時,實現了更好的縮放規律和更小的誤差;見圖3.3。因此,推測KAN可能有潛力作為PDE模型簡化的良好神經網路表示。

3.5 持續學習

災難性遺忘是當前機器學習中的一個嚴重問題。當人類掌握一項任務並轉向另一項任務時,他們不會忘記如何執行第一項任務。不幸的是,神經網路並非如此。當一個神經網路在任務1上訓練,然後轉向在任務2上訓練時,網路很快就會忘記如何執行任務1。人工神經網路和人腦之間的一個關鍵區別是,人腦在空間上局部放置了功能截然不同的模塊。當學習一個新任務時,結構重組只發生在負責相關技能的局部區域,其他區域保持不變。大多數人工神經網路,包括MLP,都沒有這種局部性的概念,這可能是災難性遺忘的原因。

KAN具有局部可塑性,可以通過利用樣條的局部性來避免災難性遺忘。這個想法很簡單:由於樣條基函數是局部的,一個樣本只會影響少數附近的樣條係數,遠處的係數保持不變(這是理想的,因為遠處的區域可能已經存儲了想要保留的信息)。相比之下,由於MLP通常使用全局激活函數,例如ReLU/Tanh/SiLU等,任何局部變化都可能不受控制地傳播到遠處的區域,破壞存儲在那裡的信息。

用一個玩具示例來驗證這個直覺。1D回歸任務由5個高斯峰組成。圍繞每個峰值的數據按順序(而不是一次全部)呈現給KAN和MLP,如圖3.4頂行所示。每個訓練階段後的KAN和MLP預測分別顯示在中間行和底行。正如預期的那樣,KAN僅對當前階段存在資料的區域進行重塑,保持先前區域不變。相比之下,MLP在看到新的數據樣本後對整個區域進行重塑,導致災難性遺忘。

這裡只是在一個極其簡單的示例上展示了初步結果,演示人們如何可能利用KAN中的局部性(得益於樣條參數化)來減少災難性遺忘。然而,目前尚不清楚方法是否可以推廣到更現實的設置,將此留待未來的工作。還想研究我們的方法如何與持續學習中的最先進方法聯繫和結合。

4 KAN是可解釋的


在本節中,展示KAN是可解釋和可互動的,這要歸功於我們在2.5節中開發的技術。不僅想在合成任務上測試KAN的使用(第4.1和4.2節),而且想在真實的科學研究中測試。證明KAN可以(重新)發現紐結理論中高度非平凡的關係(第4.3節)和凝聚態物理學中的相變邊界(第4.4節)。由於其準確性(上一節)和可解釋性(本節),KAN可能成為AI+科學的基礎模型。

4.1 有監督的玩具資料集

首先檢驗KAN揭示符號公式中組合結構的能力。下面列出了六個示例,KAN在圖4.1中可視化。KAN能夠揭示這些公式中存在的組合結構,以及學習正確的單變量函數。以下是一些示例:

(a) 乘法 \( f(x, y) = xy \)。一個 [2, 5, 1] KAN 被修剪為 [2, 2, 1] KAN。學到的激活函數是線性和二次方的。從計算圖中,看到它計算 \( xy \) 的方式是利用 \( 2xy = (x + y)^2 - (x^2 + y^2) \)。

(b) 正數的除法 \( f(x, y) = x/y \)。一個 [2, 5, 1] KAN 被修剪為 [2, 1, 1] KAN。學到的激活函數是對數函數和指數函數,KAN 通過利用恆等式 \( x/y = \exp(\log x - \log y) \) 來計算 \( x/y \)。

(c) 數值到類別。任務是將 [0, 1] 中的實數轉換為其第一位小數(作為 one-hot),例如,\( 0.0618 \to [1, 0, 0, 0, 0, \cdots] \),\( 0.314 \to [0, 0, 0, 1, 0, \cdots] \)。請注意,激活函數被學習為位於相應十進制數字附近的尖峰。



(d) 特殊函數 \( f(x, y) = \exp(J_0(20x) + y^2) \)。符號回歸的一個局限性是,如果特殊函數不作為先驗知識提供,它將永遠無法找到特殊函數的正確公式。KAN可以學習特殊函數——高度波動的Bessel函數 \( J_0(20x) \) 被KAN(數值地)學習。

(e) 相變 \( f(x_1, x_2, x_3) = \tanh(5(x_1^4 + x_2^4 + x_3^4 - 1)) \)。相變在物理學中備受關注,因此希望KAN能夠檢測相變並識別正確的序參量。使用 tanh 函數來模擬相變行為,序參量是 \( x_1 \), \( x_2 \), \( x_3 \) 的四次方項的組合。KAN訓練後出現了四次方依賴和 tanh 依賴。這是第4.4節中討論的局域化相變的一個示例。

(f) 更深的組合\(f(x_1, x_2, x_3, x_4) = \sqrt{(x_1 - x_2)^2 + (x_3 - x_4)^2}\)。為了計算這個,需要恆等函數、平方函數和平方根,這至少需要一個三層KAN。事實上,該研究發現[4, 3, 3, 1] KAN可以自動修剪為[4, 2, 1, 1] KAN,這正好對應於期望的計算圖。

來自費曼資料集和特殊函數資料集的更多示例在附錄D和F的圖D.2、D.3、F.1、F.2中可視化。

4.2 無監督玩具數據集

通常,科學發現被表述為有監督的學習問題,即給定輸入變量\(x_1, x_2, \cdots, x_d\)和輸出變量\(y\),想找到一個可解釋的函數\(f\)使得\(y \approx f(x_1, x_2, \cdots, x_d)\)。然而,另一種科學發現可以表述為無監督學習,即給定一組變量\((x_1, x_2, \cdots, x_d)\),想發現變量之間的結構關係。具體來說,想找到一個非零的\(f\)使得

$$f(x_1, x_2, \cdots, x_d) \approx 0.$$  
(4.1)



例如,考慮一組特徵\((x_1, x_2, x_3)\)滿足\(x_3 = \exp(\sin(\pi x_1) + x_2^2)\)。那麼一個有效的\(f\)是\(f(x_1, x_2, x_3) = \sin(\pi x_1) + x_2^2 - \log(x_3) = 0\),意味著\((x_1, x_2, x_3)\)的點形成了一個由\(f = 0\)指定的2D子流形,而不是填充整個3D空間。

如果可以設計出求解無監督問題的算法,它相對於有監督問題有相當大的優勢,因為它只需要特徵集\(S = (x_1, x_2, \cdots, x_d)\)。另一方面,有監督問題試圖根據其他特徵預測特徵的子集,即將\(S = S_\text{in} \cup S_\text{out}\)分成要學習的函數的輸入和輸出特徵。如果沒有領域專業知識來指導分割,那麼有\(2^d - 2\)種可能性使得\(|S_\text{in}| > 0\)且\(|S_\text{out}| > 0\)。通過使用無監督方法,可以避免這個指數級大的有監督問題空間。這種無監督學習方法對4.3節中的紐結數據集很有價值。Google Deepmind團隊手動選擇signature作為目標變量,否則他們將面臨上述組合問題。這提出了一個問題,是否可以直接處理無監督學習。在下面介紹我們的方法和一個玩具示例。

通過將無監督學習問題轉化為對所有\(d\)個特徵的有監督學習問題來解決無監督學習問題,而不需要選擇分割。基本思想是學習一個函數\(f(x_1, \ldots, x_d) = 0\)使得\(f\)不是0函數。為此,類似於對比學習,定義正樣本和負樣本:正樣本是真實數據的特徵向量。負樣本通過特徵損壞構造。為了確保每個拓撲不變量的整體特徵分佈保持不變,通過在整個訓練集上隨機排列每個特徵來進行特徵損壞。現在要訓練一個網路\(g\)使得\(g(x_\text{real}) = 1\)且\(g(x_\text{fake}) = 0\),這將問題轉化為有監督問題。但是,請記住,最初想要\(f(x_\text{real}) = 0\)且\(f(x_\text{fake}) \neq 0\)。可以通過令\(g = \sigma \circ f\)來實現這一點,其中\(\sigma(x) = \exp(-\frac{x^2}{2w^2})\)是一個高斯函數,寬度$w$很小,這可以方便地通過一個形狀為\([..., 1, 1]\)的KAN來實現,其最後一個激活設置為高斯函數\(\sigma\),所有前面的層形成\(f\)。除了上述修改外,其他一切與有監督訓練相同。

現在演示無監督範式對合成示例有效。考慮一個6D資料集,其中\((x_1, x_2, x_3)\)是相依變量,使得\(x_3 = \exp(\sin(x_1)+x_2^2)\); \((x_4, x_5)\)是相依變量,\(x_5 = x_4^3\); \(x_6\)與其他變量無關。在圖4.2中,展示了對於seed = 0,KAN揭示了\(x_1\)、\(x_2\)和\(x_3\)之間的函數依賴關係;對於另一個seed = 2024,KAN揭示了\(x_4\)和\(x_5\)之間的函數依賴關係。初步結果依賴於隨機性(不同的種子)來發現不同的關係;在未來,希望研究一種更系統、更可控的方式來發現一組完整的關係。即便如此,目前狀態下的工具可以為科學任務提供見解。在4.3節中介紹了在紐結資料集上的結果。

4.3 數學應用:紐結理論(Knot Theory)

紐結理論是低維拓撲學中的一個主題,它闡明了三維流形和四維流形的拓撲方面,並在生物學和拓撲量子計算等領域有多種應用。從數學上講,紐結\(K\)是\(S^1$到$S^3\)的嵌入。如果兩個紐結\(K\)和\(K'\)可以通過\(S^3\)的環境空間變形相互變形,則它們在拓撲上是等價的,在這種情況下,我們寫成\([K] = [K']\)。

一些紐結在拓撲上是平凡的,這意味著它們可以平滑地變形為標準圓。紐結具有多種變形不變特徵\(f\),稱為拓撲不變量,如果\(f(K) \neq f(K')\),它們可用於表明兩個紐結在拓撲上是不等價的,\([K] \neq [K']\)。在某些情況下,拓撲不變量在本質上是幾何的。例如,雙曲紐結\(K\)的補空間\(S^3 \backslash K\)承認規範的雙曲度量\(g\),使得\(\text{vol}_g(K)\)是一個稱為雙曲體積的拓撲不變量。其他拓撲不變量在本質上是代數的,例如Jones多項式。

鑑於紐結在數學中的基本性質及其應用的重要性,研究ML是否可以帶來新結果是很有趣的。例如,強化學習被用來建立某些紐結的帶性,這排除了許多潛在的反例來反駁光滑4d Poincaré猜想。

有監督學習和人類領域專家被用來得出一個關於代數和幾何紐結不變量的新定理。在這種情況下,梯度顯著性識別出有監督問題的關鍵不變量,這使領域專家提出了一個猜想,隨後被改進和證明。研究KAN是否可以在同一問題上取得良好的可解釋結果,該問題預測紐結的signature。他們從研究紐結理論數據集中得到的主要結果是:

(1) 他們使用網路歸因方法發現signature \(\sigma\)主要依賴於子午距\(\mu\)(實部\(\mu_r\),虛部\(\mu_i\))和經向距\(\lambda\)。 
(2) 人類科學家後來發現\(\sigma\)與斜率\(\equiv \text{Re}(\frac{\lambda}{\mu}) = \frac{\lambda\mu_r}{\mu_r^2+\mu_i^2}\)高度相關,並推導出\(|2\sigma - \text{slope}|\)的界。

下面展示,KAN不僅以更小的網路和更多的自動化方式重新發現了這些結果,而且還提出了一些有趣的新結果和見解。







為了研究(1),將17個紐結不變量視為輸入,將signature視為輸出。signature(偶數)被編碼為 one-hot 向量,網路用交叉熵損失訓練。發現一個極小的[17, 1, 14] KAN能夠達到81.6%的測試準確率(而Deepmind的4層寬度300的MLP達到78%的測試準確率)。[17, 1, 14] KAN(G = 3, k = 3)有\(\approx 200\)個參數,而MLP有\(\approx 3 \times 10^5\)個參數,如表4所示。值得注意的是,KAN可以同時比MLP更準確和更節省參數。

在可解釋性方面,根據每個激活的大小來縮放其透明度,因此無需特徵歸因就可以立即清楚哪些輸入變量很重要(見圖4.3左):signature主要依賴於\(\mu_r\),略微依賴於\(\mu_i\)和\(\lambda\),而對其他變量的依賴很小。然後,在三個重要變量上訓練一個[3, 1, 14] KAN,獲得78.2%的測試準確率。結果與文獻中的結果有一個微妙的區別:他們發現signature主要依賴於\(\mu_i\),而該研究發現signature主要依賴於\(\mu_r\)。這種差異可能是由於微妙的算法選擇,但促使該研究進行以下實驗:
(a)消融研究。表明\(\mu_r\)對準確性的貢獻大於\(\mu_i\)(見圖4.3):例如,僅\(\mu_r\)就可以達到65.0%的準確率,而僅\(\mu_i\)只能達到43.8%的準確率。 
(b)找到一個只涉及\(\mu_r\)和\(\lambda\)的符號公式(在表5中),但可以達到77.8%的測試準確率。

為了研究(2),即獲得\(\sigma\)的符號形式,我們將問題表述為回歸任務。使用2.5.1節中介紹的自動符號回歸,可以將訓練好的KAN轉換為符號公式。訓練形狀為[3, 1]、[3, 1, 1]、[3, 2, 1]的KAN,其對應的符號公式顯示在表5 B-D中。很明顯,通過使用更大的KAN,準確性和複雜性都會增加。因此,KAN提供的不僅僅是一個單一的符號公式,而是一個完整的公式帕累托前沿,在簡單性和準確性之間進行權衡。然而,KAN需要額外的歸納偏置來進一步簡化這些方程,以重新發現中的公式(表5 A)。該研究測試了兩種情況:
(1)在第一種情況下,假設真實公式具有多元Pade表示(兩個多元Taylor級數的商)。首先訓練[3, 2, 1],然後將其擬合到Pade表示。可以獲得表5中的公式E,它與Deepmind的公式有相似之處。 
(2)假設除法對KAN來說不是很直觀,因此訓練兩個KAN(一個用於分子,另一個用於分母)並手動將它們相除。令人驚訝的是,最終得到了公式F(在表5中),它只涉及\(\mu_r\)和\(\lambda\),儘管\(\mu_i\)也被提供但被KAN忽略了。

到目前為止,已經重新發現了主要結果。令人矚目的是,KAN使這一發現變得非常直觀和方便。與其使用特徵歸因方法(這是很好的方法),不如直接盯著KAN的可視化。此外,自動符號回歸也使符號公式的發現變得容易得多。

在下一部分中,我們提出了一種Deepmind論文中沒有包括的新範式"AI for Math",我們旨在使用KAN的無監督學習模式來發現紐結不變量中的更多關係(除了signature)。





無監督學習 正如在4.2節中提到的,無監督學習是更有前景的設置,因為它避免了手動劃分輸入和輸出變量,而手動劃分有組合多種可能性。在無監督學習模式下,將所有18個變量(包括signature)視為輸入,使它們處於同等地位。紐結數據是正樣本,通過隨機洗牌特徵來獲得負樣本。訓練一個[18, 1, 1] KAN來分類給定的特徵向量是否屬於正樣本(1)或負樣本(0)。手動將第二層激活設置為在零處中心為一的高斯函數,因此正樣本在(接近)零處有激活,隱含地給出了紐結不變量之間的關係 

$$\sum_{i=1}^{18} g_i(x_i) = 0$$

其中\(x_i\)表示一個特徵(不變量),\(g_i\)是相應的激活函數,可以從KAN圖中直接讀出。用\(\lambda = {10^{-2}, 10^{-3}}\)訓練KAN以偏好稀疏的輸入組合,種子為\(\text{seed} = {0, 1, \cdots, 99}\)。所有200個網路可以分為三個簇,具有代表性的KAN顯示在圖4.4中。這三組相依變量是:
  1. 第一組相依變量是signature、子午距的實部和經向距(加上另外兩個可以因(3)而移除的變量)。這就是上面研究的signature依賴關係,所以看到這個依賴關係在無監督模式下再次被發現是非常有趣的。
  2. 第二組變量涉及尖點體積\(V\)、子午平移\(\mu_r\)的實部和經向平移\(\lambda\)。它們的激活函數看起來都像對數函數(可以通過2.5.1節中暗示的符號功能驗證)。因此,關係是\(-\log V + \log \mu_r + \log \lambda = 0\),等價於\(V = \mu_r \lambda\),這在定義上是正確的。然而,在沒有任何先驗知識的情況下發現這個關係,這是令人欣慰的。
  3. 第三組變量包括短測地線\(g_r\)的實部和內射半徑。它們的激活看起來定性相同,但相差一個負號,因此推測這兩個變量存在線性相關。我們繪製2D散點圖,發現\(2r\)上限為\(g_r\),這也是一個眾所周知的關係。
有趣的是,KAN的無監督模式可以重新發現幾個已知的數學關係。好消息是KAN發現的結果可能是可靠的;壞消息是還沒有發現任何新東西。值得注意的是,選擇了一個淺層KAN以便於簡單可視化,但如果存在,更深層的KAN可能會發現更多關係。希望在未來的工作中研究如何用更深層的KAN發現更複雜的關係。

4.4 物理應用:Anderson局域化

Anderson局域化是量子系統中無序導致電子波函數局域化的基本現象,導致所有傳輸停止。在一維和二維中,縮放論證表明,對於無限小量的隨機無序,所有電子本征態都是指數局域化的。相比之下,在三維中,臨界能量形成了一個相邊界,將延展態與局域化態分開,稱為遷移率邊緣。對這些遷移率邊緣的理解對於解釋各種基本現象至關重要,如固體中的金屬-絕緣體轉變,以及光在光子器件中的局域化效應。因此,有必要開發展現遷移率邊緣的微觀模型,以便進行詳細研究。在低維中開發此類模型通常更實用,在低維中引入準周期性而不是隨機無序也可能導致分隔局域化相和延展相的遷移率邊緣。此外,分析遷移率邊緣的實驗實現可以幫助解決關於相互作用系統中局域化的爭論。事實上,最近的幾項研究集中在確定此類模型並推導它們的遷移率邊緣的精確解析表達式上。

在這裡,將KAN應用於從準周期緊束縛模型生成的數值資料,以提取它們的遷移率邊緣。我們特地檢查三類模型:馬賽克模型(MM)、廣義Aubry-André模型(GAAM)和修正的Aubry-André模型(MAAM)。對於MM,證明KAN能夠準確提取遷移率邊緣作為能量的1D函數。對於GAAM,發現從KAN獲得的公式與基本事實非常接近。對於更複雜的MAAM,再次展示了該框架符號可解釋性的另一個例子。用戶可以通過"協作"的方式簡化從KAN(和相應的符號公式)獲得的複雜表達式,在協作中,人類生成假設以獲得更好的匹配(例如,對某些激活函數的形式做出假設),然後KAN可以快速進行假設檢驗。

為了量化這些模型中態的局域化,通常使用反參與率(IPR)。第\(k\)個本征態\(\psi^{(k)}\)的IPR由下式給出

$$\text{IPR}_k = \frac{\sum_n |\psi_n^{(k)}|^4}{(\sum_n |\psi_n^{(k)}|^2)^2}$$ 
(4.2)

其中求和是對格點指標進行的。在這裡,使用相關的局域化度量——態的分形維數,由下式給出


$$D_k = -\frac{\log(\text{IPR}_k)}{\log(N)}$$ 
(4.3)


其中\(N\)是系統大小。\(D_k = 0(1)\)表示局域化(延展)態。

馬賽克模型(MM) 我們首先考慮由哈密頓量定義的一類緊束縛模型


$$H = t\sum_n (c_{n+1}^\dagger c_n + \text{H.c.}) + \sum_n V_n(\lambda, \phi)c_n^\dagger c_n,$$ (4.4)


其中$t$是最近鄰耦合,\(c_n(c_n^\dagger)\)是格點\(n\)處的湮滅(產生)算符,勢能\(V_n\)由下式給出


\[ V_n(\lambda, \phi) = \lambda \cos(2\pi nb + \phi) \cdot \mathbf{1}_{\{j = m\kappa\}} \]
(4.5)


為了引入準周期性,將\(b\)設置為無理數(選擇\(b\)為黃金比例\(\frac{1+\sqrt{5}}{2}\))。\(\kappa\)是一個整數,準周期勢以間隔\(\kappa\)出現。對於這個模型,能量(\(E\))譜通常包含由遷移率邊緣分隔的延展和局域化區域。有趣的是,這裡發現的一個獨特特徵是,即使對於任意強的準周期勢(即系統中總是存在延展態,與局域化態共存),遷移率邊緣也存在。

遷移率邊緣可以用\(g(\lambda, E) \equiv \lambda - |f_\kappa(E)| = 0\)來描述。\(g(\lambda, E) > 0\)和\(g(\lambda, E) < 0\)分別對應於局域化相和延展相。因此,學習遷移率邊緣取決於學習"序參量"\(g(\lambda, E)\)。誠然,這個問題可以通過許多其他理論方法來解決這類模型,將在下面演示,KAN框架已經準備好並方便地接受來自人類用戶的假設和歸納偏置。

假設一個假想的用戶Alice,她是凝聚態物理學的一名新博士生,並獲得一個[2, 1] KAN作為該任務的助手。首先,她理解這是一個分類任務,因此明智的做法是使用fix_symbolic功能將第二層中的激活函數設置為sigmoid。其次,她意識到學習整個2D函數\(g(\lambda, E)\)是不必要的,因為最終她只關心由\(g(\lambda, E) = 0\)確定的\(\lambda = \lambda(E)\)。這樣做,假設\(g(\lambda, E) = \lambda - h(E) = 0\)是合理的。Alice只需再次使用fix_symbolic功能將\(\lambda\)的激活函數設置為線性。現在Alice訓練KAN網路,方便地獲得遷移率邊緣,如圖4.5所示。Alice既可以獲得直觀的定性理解(底部),也可以獲得定量結果(中部),與基本事實(頂部)吻合良好。



廣義Andre-Aubry模型(GAAM)接下來,考慮由 Hamiltonian定義的一類緊束縛模型


$$H = t\sum_n (c_{n+1}^\dagger c_n + \text{H.c.}) + \sum_n V_n(\alpha, \lambda, \phi)c_n^\dagger c_n,$$ 
(4.6)

其中\(t\)是最近鄰耦合,\(c_n(c_n^\dagger)\)是格點\(n\)處的湮滅(產生)算符,勢能\(V_n\)由下式給出


$$V_n(\alpha, \lambda, \phi) = 2\lambda\frac{\cos(2\pi nb + \phi)}{1 - \alpha \cos(2\pi nb + \phi)},$$
 (4.7)

對於\(\alpha \in (-1, 1)\),它是光滑的。為了引入準周期性,我們再次將\(b\)設置為無理數(特別地,我們選擇\(b\)為黃金比例)。和之前一樣,我們想得到遷移率邊緣的表達式。對於這些模型,遷移率邊緣由閉合形式表達式給出,


$$\alpha E = 2(t - \lambda).$$ 
(4.8)

隨機採樣模型參數:\(\phi\)、\(\alpha\)和\(\lambda\)(設置能量尺度\(t = 1\)),並計算能量本征值以及相應本征態的分形維數,這構成了訓練資料集。 

在這裡,要學習的"序參量"是\(g(\alpha, E, \lambda, \phi) = \alpha E + 2(\lambda - 1)\),遷移率邊緣對應於\(g = 0\)。再次假設Alice想要確定遷移率邊緣,但只能訪問IPR或分形維數資料,因此她決定使用KAN來幫助她完成任務。Alice希望模型盡可能小,因此她可以從一個大模型開始,使用自動修剪來獲得一個小模型,或者她可以根據對給定問題複雜性的理解來猜測一個合理的小模型。無論哪種方式,假設她得到一個[4, 2, 1, 1] KAN。首先,她將最後一個激活設置為sigmoid,因為這是一個分類問題。



她用一些稀疏正則化訓練她的KAN達到98.7%的準確率,並在圖4.6(a)步驟1中可視化訓練後的KAN。她觀察到KAN完全沒有選擇$\phi$,這使她意識到遷移率邊緣與\(\phi\)無關(與公式(4.8)一致)。此外,她觀察到幾乎所有其他激活函數都是線性或二次方的,因此她打開自動符號捕捉,將函數庫限制為僅線性或二次方。之後,她立即得到一個已經是符號的網路(如圖4.6(a)步驟2所示),具有相當(甚至略好)的準確率98.9%。通過使用symbolic_formula功能,Alice方便地得到\(g\)的符號形式,如表6 GAAM-KAN auto(第三行)所示。也許她想劃掉一些小項並將係數捕捉為小整數,這使她接近真正的答案。

如果Alice使用符號回歸方法,這個假設的故事將完全不同。如果她幸運的話,SR可以返回完全正確的公式。然而,在絕大多數情況下,SR不會返回有用的結果,而且Alice不可能"調試"或與符號回歸的底層過程交互。此外,在運行SR之前,Alice可能會感到不舒服/缺乏經驗,無法將符號項的函數庫作為先驗知識提供給SR。相比之下,在KAN中,Alice不需要將任何先驗信息放入KAN。她可以先盯著訓練後的KAN獲得一些線索,只有在那時她才決定她想做出哪些假設(例如,"所有激活都是線性或二次方的"),並在KAN中實現她的假設。儘管KAN不太可能立即返回正確答案,但KAN總是會返回一些有用的東西,Alice可以與之合作來完善結果。

修正的Andre-Aubry模型(MAAM)考慮的最後一類模型由Hamiltonian定義


$$H = \sum_{n \neq n'} te^{-p|n-n'|} (c_n^\dagger c_{n'} + \text{H.c.}) + \sum_n V_n(\lambda, \phi)c_n^\dagger c_n,$$ 
(4.9)

其中\(t\)是空間中指數衰減耦合的強度,\(c_n(c_n^\dagger)\)是格點\(n\)處的湮滅(產生)算符,勢能\(V_n\)由下式給出


$$V_n(\lambda, \phi) = \lambda \cos(2\pi nb + \phi),$$ 
(4.10)


和之前一樣,為了引入準周期性,將\(b\)設置為無理數(黃金比例)。對於這些模型,遷移率邊緣由閉合形式表達式給出,


$$\lambda \cosh(p) = E + t = E + t_1 \exp(p)$$ 
(4.11)

其中定義\(t_1 \equiv te^{-p}\)為最近鄰跳躍強度,下面設\(t_1 = 1\)。

假設Alice想要確定MAAM的遷移率邊緣。這個任務更加複雜,需要更多人類智慧。和上一個例子一樣,Alice從一個[4, 2, 1, 1] KAN開始訓練,但得到的準確率在75%左右,低於可接受的水平。然後她選擇一個更大的[4, 3, 1, 1] KAN,成功獲得98.4%的準確率,這是可以接受的(圖4.6(b)步驟1)。Alice注意到KAN沒有選擇\(\phi\),這意味著遷移率邊緣與相位因子\(\phi\)無關(與公式(4.11)一致)。如果Alice打開自動符號回歸(使用包含exp、tanh等的大型函數庫),她會得到表6-MAAM-KAN auto中的複雜公式,準確率為97.1%。但是,如果Alice想找到一個更簡單的符號公式,她會想使用手動模式,自己完成符號捕捉。在此之前,她發現訓練後的[4, 3, 1, 1] KAN可以被修剪為[4, 2, 1, 1],同時保持97.7%的準確率(圖4.6(b))。Alice可能認為除了依賴\(p\)的激活函數外,所有激活函數都是線性或二次方的,並通過使用fix_symbolic手動將它們捕捉為線性或二次方。捕捉並重新訓練後,更新的KAN如圖4.6(c)步驟3所示,保持97.7%的準確率。從現在開始,Alice可能會根據她的先驗知識做出兩個不同的選擇。在一種情況下,Alice可能已經猜到\(p\)的依賴關係是\(\cosh\),因此她將\(p\)的激活設置為\(\cosh\)函數。她重新訓練KAN,獲得96.9%的準確率(圖4.6(c)步驟4A)。在另一種情況下,Alice不知道\(\cosh p\)的依賴關係,因此她追求簡單性,再次假設\(p\)的函數是二次方的。她重新訓練KAN,獲得95.4%的準確率(圖4.6(c)步驟4B)。如果她同時嘗試了這兩種情況,她會意識到在準確性方面\(\cosh\)更好,而在簡單性方面二次方更好。這些步驟對應的公式列在表6中。很明顯,Alice進行的手動操作越多,符號公式就越簡單(準確性略有犧牲)。KAN有一個"旋鈕",用戶可以調整以在簡單性和準確性之間進行權衡(有時簡單性甚至可以提高準確性,如GAAM案例中所示)。

5 相關工作

Kolmogorov-Arnold定理與神經網路。Kolmogorov-Arnold定理(KAT)與神經網路之間的聯繫在文獻中並不新鮮,但內函數的病態行為使KAT在實踐中看起來並不樂觀。大多數這些先前的工作都堅持使用原始的2層寬度為\((2n + 1)\)的網路,表達能力有限,其中許多工作甚至早於反向傳播。因此,大多數研究都建立在理論之上,只有相當有限或人工的玩具實驗。該研究的貢獻在於將網路推廣到任意寬度和深度,在當今的深度學習潮流中重新喚起並將其情境化,並強調其作為AI+科學基礎模型的潛在作用。

神經縮放定律(NSL)。NSL是測試損失相對於模型大小、資料、計算等呈冪律行為的現象。NSL的起源仍然是個謎,但有競爭理論,包括內在維度、任務量化、資源理論、隨機特徵、組合稀疏性和最大奇異性。本文通過表明,如果高維函數具有平滑的Kolmogorov-Arnold表示,則可以驚人地縮放為1D函數(這是人們可以希望的最佳界限),為這一空間做出了貢獻。論文為神經縮放定律帶來了新的樂觀情緒,因為它承諾了有史以來最快的縮放指數。在實驗中表明,這種快速神經縮放定律可以在合成數據集上實現,但未來的研究需要解決以下問題:KA表示是否存在於一般任務中?如果存在,訓練是否在實踐中找到這些表示?

機制可解釋性(MI)。MI是一個新興領域,旨在機械地理解神經網路的內部工作原理。MI研究大致可分為被動和主動MI研究。大多數MI研究是被動的,專注於理解用標準方法訓練的現有神經網路。主動MI研究試圖通過設計本質上可解釋的架構或開發訓練方法來明確鼓勵可解釋性來實現可解釋性。該研究的工作屬於第二類,其中模型和訓練方法在設計上是可解釋的。

可學習激活。在機器學習中,可學習激活在神經網路中的想法並不新鮮。可訓練激活函數以可微方式學習或以離散方式搜索。激活函數被參數化為多項式、樣條、sigmoid線性單元或神經網路。KAN使用B樣條來參數化其激活函數。該研究還介紹了關於可學習激活網路(LAN)的初步結果,其性質介於KAN和MLP之間,為了專注於主論文中的KAN,結果被推遲到附錄B。

符號回歸。有許多現成的基於遺傳算法的符號回歸方法(Eureka 、GPLearn 、PySR )、基於神經網路的方法(EQL 、OccamNet )、基於物理啟發的方法(AI Feynman )和基於強化學習的方法。KAN與基於神經網路的方法最相似,但與之前的工作不同,激活函數是在符號捕捉之前連續學習的,而不是手動固定的。

物理信息神經網路(PINN)和物理信息神經算子(PINO)。在3.4小節中,演示了KAN可以取代在求解PDE時使用MLP來施加PDE損失的範式。參考Deep Ritz方法、PINN用於PDE求解,以及Fourier神經算子、PINO 、DeepONet 用於學習解映射的算子學習方法。在上述所有網路中,都有可能用KAN取代MLP。

AI for Mathematics。正如我們在4.3小節中看到的,AI最近被應用於紐結理論中的幾個問題,包括檢測一個紐結是否是未紐結或帶紐結,以及預測紐結不變量並揭示它們之間的關係。

6 討論

在本節中,我們從數學基礎、算法和應用的角度討論KAN的局限性和未來方向。

數學方面:我們提出了KAN的初步數學分析(定理2.1),但對它們的數學理解仍然非常有限。Kolmogorov-Arnold表示定理在數學上已經得到了徹底的研究,但該定理對應於形狀為[n, 2n+1, 1]的KAN,這是KAN的一個非常受限的子類。在更深的KAN上的經驗成功是否意味著數學中的某些基本原理?一個有吸引力的廣義Kolmogorov-Arnold定理可以定義超出2層組合的"更深"Kolmogorov-Arnold表示,並可能將激活函數的平滑性與深度聯繫起來。假設存在一些函數,它們不能在原始(2層)Kolmogorov-Arnold表示中平滑表示,但可能在3層或更深層中平滑表示。我們能否使用這種"Kolmogorov-Arnold深度"的概念來表徵函數類?

算法方面:該研究討論以下內容:
  1. 準確性。架構設計和訓練中的多個選擇沒有得到充分研究,因此替代方案可能進一步提高準確性。例如,樣條激活函數可能被徑向基函數或其他局部核取代。可以使用自適應網格策略。
  2. 效率。KAN運行緩慢的一個主要原因是不同的激活函數無法利用批量計算(大量數據通過相同的函數)。實際上,可以在激活函數全部相同(MLP)和全部不同(KAN)之間進行插值,將激活函數分組為多個組("多頭"),其中一個組內的成員共享相同的激活函數。
  3. KAN和MLP的混合。與MLP相比,KAN有兩個主要區別:
    • 激活函數在邊上而不是在節點上,
    • 激活函數是可學習的而不是固定的。
  4. 哪種變化對解釋KAN的優勢更重要?在附錄B中介紹了初步結果,其中研究了一個模型,它具有(ii),即激活函數是可學習的(如KAN),但不具有(i),即激活函數在節點上(如MLP)。此外,還可以構建另一個模型,其激活函數是固定的(如MLP),但在邊上(如KAN)。
  5. 自適應性。得益於樣條基函數的內在局部性,可以在KAN的設計和訓練中引入自適應性,以提高準確性和效率:參見多層次訓練的想法,如中的多網格方法,或如中的多尺度方法的域相關基函數。
應用方面:提出了一些初步證據,表明KAN在與科學相關的任務中比MLP更有效,例如擬合物理方程和求解PDE。期望KAN在求解Navier-Stokes方程、密度泛函理論或任何其他可以表述為回歸或求解PDE的任務中也很有前景。還希望將KAN應用於與機器學習相關的任務,這需要將KAN集成到當前的架構中,例如將transformer中的MLP替換為KAN,人們可以出"kansformer"。

作為AI+科學的"語言模型"的KAN 大型語言模型之所以如此變革,是因為它們對任何能說自然語言的人都有用。科學的語言是函數。KAN由可解釋的函數組成,因此當人類用戶盯著KAN時,就像用函數的語言與之交流一樣。這一段旨在促進AI-科學家-協作範式,而不是特定工具KAN。就像人們使用不同的語言交流一樣,預計在未來,KAN將只是AI+科學的語言之一,儘管KAN將是最早實現AI和人類交流的語言之一。然而,在KAN的支持下,AI-科學家-協作範式從未如此簡單和方便,使我們重新思考要如何接近AI+科學的範式:想要AI科學家,還是想要幫助科學家的AI?(完全自動化)AI科學家的內在困難在於,很難使人類偏好量化,這將人類偏好編入AI目標。事實上,不同領域的科學家對哪些函數是簡單或可解釋的可能有不同的感受。因此,科學家擁有一個能說科學語言(函數)並能方便地與個別科學家的歸納偏置交互以適應特定科學領域的AI更為可取。

最終要點:我應該使用KAN還是MLP?

目前,KAN的最大瓶頸在於其緩慢的訓練。在給定相同數量的參數的情況下,KAN通常比MLP慢10倍。應該誠實地說,該研究有努力優化KAN的效率,因此KAN的緩慢訓練更像是一個工程問題,未來有待改進,而不是一個根本的限制。如果想要快速訓練模型,應該使用MLP。然而,在其他情況下,KAN應該與MLP相當或更好,這使得它們值得嘗試。圖6.1中的決策樹可以幫助決定何時使用KAN。簡而言之,如果你關心可解釋性和/或準確性,並且緩慢的訓練不是主要問題,建議嘗試KAN。



參考文獻(略)


附錄

A KAN功能

表7包含了用戶可能會發現有用的常見功能。

B 可學習激活網路(LAN)

B.1 架構

除了KAN,該研究還提出了另一種可學習激活網路(LAN),它們幾乎是MLP,但具有參數化為樣條的可學習激活函數。KAN對標準MLP進行了兩項主要更改:
  1. 激活函數變為可學習而不是固定的; 
  2. 激活函數放置在邊上而不是節點上。為了區分這兩個因素,該研究還提出了可學習激活網路(LAN),它只有可學習的激活,但仍然在節點上,如圖B.1所示。

對於寬度為\(N\)、深度為$L$和網格點數為\(G\)的LAN,參數數量為\(N^2L + NLG\),其中\(N^2L\)是權重矩陣的參數數量,\(NLG\)是樣條激活的參數數量,除MLP外,這引起的開銷很小,因為通常\(G \ll N\),所以\(NLG \ll N^2L\)。LAN與MLP相似,因此可以從預訓練的MLP初始化,並通過允許可學習的激活函數進行微調。一個示例是使用LAN來改進SIREN,在B.3節中介紹。

LAN和KAN的比較
    
    LAN的優點:
  1. LAN在概念上比KAN更簡單。它們更接近標準MLP(唯一的變化是激活函數變為可學習)。
  2. LAN比KAN更容易擴展。LAN/KAN在節點/邊上有可學習的激活函數。因此,LAN/KAN中的激活參數分別隨\(N\)/\(N^2\)縮放,其中\(N\)是模型寬度。
    LAN的缺點:
  1. LAN似乎不太可解釋(權重矩陣難以解釋,就像MLP中一樣);
  2. LAN似乎也不如KAN準確,但仍然比MLP更準確。與KAN一樣,如果LAN的激活函數由樣條參數化,LAN也承認網格擴展。




B.2 LAN可解釋性結果

我們在圖B.2中介紹了LAN的初步可解釋性結果。對於圖4.1中KAN完全可解釋的相同示例,由於權重矩陣的存在,LAN似乎可解釋性要差得多。首先,權重矩陣不如可學習的激活函數容易解釋。其次,權重矩陣帶來了太多的自由度,使得可學習的激活函數過於不受約束。我們使用LAN的初步結果似乎暗示,去除線性權重矩陣(通過在邊上具有可學習的激活,如KAN)對於可解釋性是必要的。

B.3 擬合圖像(LAN)

隱式神經表示將圖像視為2D函數\(f(x, y)\),其中像素值\(f\)是像素的兩個坐標\(x\)和\(y\)的函數。為了壓縮圖像,這樣的隱式神經表示(\(f\)是神經網路)可以在保持幾乎原始圖像質量的同時實現令人印象深刻的參數壓縮。SIREN提出使用具有周期性激活函數的MLP來擬合函數\(f\)。在LAN中自然考慮其他激活函數。然而,由於我們將LAN激活初始化為平滑的,而SIREN需要高頻特徵,LAN不能立即工作。請注意,LAN中的每個激活函數都是基函數和樣條函數的總和,即\(\phi(x) = b(x) + \text{spline}(x)\),我們將\(b(x)\)設置為正弦函數,與SIREN中的設置相同,但讓\(\text{spline}(x)\)是可訓練的。對於MLP和LAN,形狀都是[2,128,128,128,128,128,1]。我們使用Adam優化器對它們進行訓練,批量大小為4096,學習率為$10^{-3}$訓練5000步,學習率為$10^{-4}$訓練5000步。如圖B.3所示,LAN(橙色)可以實現比MLP(藍色)更高的PSNR,這是由於LAN微調激活函數的靈活性。我們表明,也可以從MLP初始化LAN,並進一步微調LAN(綠色)以獲得更好的PSNR。在我們的實驗中,我們選擇了\(G = 5\),因此額外的參數增加大約是原始參數的\(G/N = 5/128 \approx 4\%\)。


C 超參數的影響

我們在圖C.1中顯示了超參數對\(f(x, y) = \exp(\sin(\pi x) + y^2)\)案例的影響。為了獲得可解釋的圖,我們希望活躍激活函數的數量盡可能少(理想情況下為3)。

  1.  我們需要熵懲罰來減少活躍激活函數的數量。如果沒有熵懲罰,會有許多重複函數。
  2.  結果可能取決於隨機種子。對於一些不幸的種子,修剪後的網絡可能比需要的更大。
  3.  整體懲罰強度\(\lambda\)有效地控制稀疏性。
  4.  網格數\(G\)也對可解釋性有微妙影響。當\(G\)太小時,由於每個激活函數不是很有表現力,網絡傾向於使用集成策略,使得解釋更加困難。
  5.  分段多項式階數\(k\)對可解釋性只有微妙影響。然而,它的行為有點像隨機種子,在這個玩具示例中沒有顯示任何可見的模式。

D 費曼KAN

我們在費曼資料集(第3.3節)上包含了更多結果。圖D.1顯示了每個費曼數據集的KAN和MLP的帕累托前沿。圖D.3和D.2將每個費曼方程擬合任務的最小KAN(在測試RMSE < 10−2 的約束下)和最佳KAN(具有最低測試RMSE損失)可視化。


E 關於網格大小的備註

對於PDE和回歸任務,當我們在均勻網格上選擇訓練數據時,當網格大小更新到一個大水平時,我們會看到訓練損失突然增加(即性能突然下降),相當於空間一個方向上的不同訓練點。這可能是由於更高維度中B樣條的實現,需要進一步調查。

F 特殊函數的KAN

我們在特殊函數數據集(第3.2節)上包含了更多結果。圖F.2和F.1將每個特殊函數擬合任務的最小KAN(在測試RMSE < 0.01的約束下)和最佳KAN(具有最低測試RMSE損失)可視化。



幾點見解和思考:



1. Kolmogorov-Arnold定理在神經網路領域的應用前景廣闊。論文將KA定理與深度學習結合,提出了KAN這一新穎的神經網路架構。KAN不僅在各種任務上展現出優異的性能,其可解釋性也遠勝於傳統的MLP。這為深度學習模型的可解釋性研究提供了新的思路。

2. KAN在速度上的不足可能限制了其應用。論文坦承,目前KAN的訓練速度比MLP慢10倍左右。作者認為這更多是工程實現的問題,未來有優化的空間。KAN要真正發揮作用,還需要在效率上有所突破。

3. 無監督學習可能是AI輔助科學發現的重要範式。論文在紐結理論的實驗中展示了無監督學習的威力。相比於有監督學習需要專家標註數據,無監督學習可以直接從原始數據出發尋找隱藏的關係,更符合科學發現的實際過程。

4. AI與科學家的協同是大勢所趨。KAN作為一種會"說"科學語言(數學函數)的AI系統,為AI與科學家的溝通提供了渠道。未來,懂科學語言的AI助手或許能夠與科學家進行更加深入的交流,共同推進科學發現的進程。

5. 神經網路的數學基礎仍有待加強。儘管KAN取得了令人矚目的表現,但我們對其背後的數學原理還缺乏深入理解。Kolmogorov-Arnold定理能否推廣到更深層次?網路深度與函數平滑性是否存在內在聯繫?這些問題的解答,需要數學和深度學習領域的專家共同努力。

6. 對科學家而言,擁有正確的工具選擇十分重要。論文最後給出的決策樹,為科學家在KAN和MLP間做出選擇提供了參考。一方面,科學家要善於利用新工具(如KAN)來加速研究;另一方面,也要理性看待新工具的局限性,根據具體問題選擇最合適的方法。

這是一篇將數學理論與深度學習緊密結合,並在多個科學問題上進行實踐的高質量論文。它不僅展示了KAN的強大性能,更重要的是,它體現了一種AI助力科學發現的新範式。儘管KAN目前還不夠完善,但其背後的思想值得我們深入探索。期待未來有更多類似的工作,推動AI與科學的深度融合。












2024年5月1日星期三

語言模型的演算法進展與Scaling Laws

引言

近年來,語言模型(Language Models, LLMs)在自然語言處理領域取得了顯著的進展。LLM的預訓練成本高昂,但scaling laws使我們能夠從成本較低的訓練過程中準確預測更大規模訓練的性能。最近的研究大幅改善了我們對LLM scaling laws的實際理解,可以在模型參數數量和訓練tokens與參數之比兩方面進行外推。

LLM發展背景

先前的研究發現了語言模型的scaling laws,顯示隨著模型變大和預訓練資料增加,模型性能會提高。然而,這些研究主要關注計算最佳情況,且性能評估多以下一個token預測損失為主。在實務應用中,語言模型很少以最佳計算方式訓練,通常使用更多tokens訓練較小模型(過度訓練)以降低推理成本,且更關注下游任務實際表現。這引出疑問:scaling laws是否適用於這些更實用的情況?

演算法進展與規模擴展

2014至2024年間,演算法改進和運算能力提升顯著改善了語言模型性能。新演算法帶來的進步相當於每5至14個月運算能力翻倍。儘管近年來性能改進更多源自於大規模計算擴展,但Transformer架構和Chinchilla縮放規律等創新分別帶來了近兩年和8到16個月的算法進步。



過度訓練的LLM

在實務中,通常採用過度訓練模式,即以更多tokens訓練較小模型,以在達到有競爭力指標的同時降低推理成本。研究發現,過度訓練模型同樣遵循冪律scaling laws,token倍增器M值的改變只影響y截距,形成一組平行線。因此,我們可從scaling laws中預測更大過度訓練模型的性能。

下游任務性能預測 

利用scaling laws,我們可繪製LLM在下游評估任務上的平均top-1錯誤率,發現其隨著模型在C4資料集上的驗證損失呈指數下降。這揭示了下游任務性能的明確趨勢,據此提出了top-1錯誤的scaling laws,可用較少計算預測大模型在下游任務上的錯誤率。


實證評估與應用價值

通過訓練不同規模的LLM並使用scaling laws外推,實驗顯示我們可以較低成本地預測大規模訓練的驗證損失和下游任務錯誤率,相對誤差分別僅有0.7%和3.6%。這證實了scaling laws在實踐中預測大規模訓練性能的價值。





個人見解

LLM的快速發展得益於算法創新和計算規模擴張的共同推動,但後者的貢獻更為關鍵。Scaling laws作為一種低成本預測工具,在提高LLM研發效率方面有重要價值。隨著對scaling laws理解的加深,相信我們能更好地指導語言模型的發展方向與資源配置。不過,當前的scaling laws仍有局限,如對特定創新影響的估計能力不足、難以區分算法改進和計算擴展的貢獻等。未來還需要更多更好的數據,來建立更精準和全面的scaling laws模型。

此外,我認為LLM發展的關鍵是平衡效能提升與成本控制。過度訓練等策略雖以效率為代價換取效能,但對於實際部署至關重要。如何在scaling laws框架下找到最優的訓練配置,是一個值得深入探索的問題。總的來說,語言模型的進展令人振奮,但要實現更通用和強大的AI,仍需在算法和計算兩個維度協同發力,並以切實應用為導向,審慎權衡效能和成本。Scaling laws作為連接理論和實踐的紐帶,在其中扮演著不可或缺的角色。

2024年4月29日星期一

擴張注意力: 改良的注意力機制

 擴張注意力技術(Dilated Attention)可以將Transformer模型的上下文窗口擴展到10億個token。

擴張注意力是一種改良的注意力機制。它讓模型可以關注輸入序列中更遠的部分,且計算成本不會大幅增加。這個概念借鑑了卷積神經網路中的擴張卷積。 

在標準的Transformer模型中,自注意力機制讓輸入序列中的每個token都能關注其他所有token,導致計算複雜度為O(N^2),其中N是序列長度。

當嘗試擴展到非常長的序列(例如10億個token)時,由於計算和記憶體的限制,平方複雜度變得不可行。

擴張注意力引入了一個擴張因子,它改變了token相互關注的模式。其理念是讓token以越來越大的距離關注其他token,跳過中間的token。

擴張注意力的運作方式

1. 擴張因子: 這是一個超參數,決定了被關注的token之間的距離。例如,擴張因子為2意味著每個token關注序列中間隔一個的其他token。

2. 多層結構: 通過堆疊具有不同擴張因子的多個層,模型可以從各種距離收集信息,創建更全面的表示。

3. 稀疏連接: 這導致注意力矩陣中的稀疏連接,將複雜度從O(N^2)降低到更易處理的程度,例如O(N log N)甚至是線性複雜度,具體取決於擴張模式。

擴張注意力如何實現10億token的上下文

1. 記憶效率: 擴張注意力大大減少了記憶體需求,允許處理更長的序列。

2. 計算效率: 通過降低複雜度,它使得在合理的計算時間內處理更大的序列成為可能。

3. 保持長距離依賴: 儘管降低了複雜度,擴張注意力仍然允許捕捉文本中的長距離依賴關係,這對許多NLP任務至關重要。

與HuggingFace或其他框架整合

在流行的深度學習框架中實現擴張注意力通常需要定制注意力層。您可能需要修改現有的注意力機制以納入擴張因子、更新遮罩,並確保正確處理反向傳播。

總之,擴張注意力表示著將基於注意力的模型擴展到處理極長序列的重要一步。通過引入擴張因子並利用稀疏連接,它使得處理高達10億個token的上下文成為可能,在大幅降低計算要求的同時保持了注意力的本質。

擴張注意力與稀疏注意力的區別

擴張注意力和稀疏注意力都是降低Transformer中注意力機制計算複雜度的機制,但它們通過不同的方式實現這一目標。讓我們探討一下兩者的區別:

稀疏注意力選擇性地只關注序列中的一部分token,創建一個許多條目為零的稀疏矩陣。

稀疏模式可以是固定的或學習的,它定義了哪些token關注哪些其他token。常見的模式包括帶狀、滑動窗口或其他預定義的結構。(滑動窗口是Swinformer使用的變體)

根據稀疏模式的不同,稀疏注意力可能無法確定地捕捉長距離依賴。

1.相似性: 兩者都旨在降低計算複雜度,並允許模型處理更長的序列。

2.區別:

- 模式: 擴張使用規則的跳躍模式,而稀疏可以使用各種模式,包括學習的模式。

- 靈活性**: 稀疏注意力在定義注意力模式方面提供了更大的靈活性,而擴張注意力具有更剛性、規則的結構。

-覆蓋範圍: 由於規則的跳躍,擴張注意力可能會錯過一些更精細的局部細節,而稀疏注意力捕捉依賴關係的能力取決於所使用的特定模式。

- 實現複雜度: 稀疏注意力通常需要專門的算法和硬件來進行高效的稀疏矩陣乘法,而擴張注意力可以用標準的稠密矩陣運算來實現。

總之,儘管這兩種機制實現了相似的目標,但它們採用了不同的策略、權衡和最適合的應用。擴張注意力提供了一種實現長距離依賴的直接方式,但會損失一些粒度,而稀疏注意力可以允許更複雜和特定於任務的注意力模式,但通常代價是更大的實現複雜度。


為什麼語言Transformer需要從視覺Transformer引入擴張?

LONGNET中的擴張注意力

- 應用: 專為序列建模設計,將Transformer擴展到極長序列(如10億個token)。

- 概念: 隨著距離的增長,關注區域呈指數擴張,使模型能夠捕捉序列中的長距離依賴關係。

- 複雜度: 同時保持線性複雜度,使模型能夠處理極長的序列而不犧牲性能。

- 整合: 可作為Transformer模型中標準注意力的直接替代,可應用於一般語言任務。

視覺Transformer與語言Transformer在任務特性和序列長度上有所不同,這導致它們在注意力機制的設計上採取了不同的方法。以下是視覺Transformer的擴張注意力在語言Transformer中的應用動機:

1. 序列長度的差異:

   - 視覺Transformer通常處理較短的序列,如圖像分塊後的序列。

   - 語言Transformer需要處理更長的序列,如整篇文章或對話。

2. 捕捉長距離依賴的需求:

   - 在圖像任務中,相鄰片段間的局部關係通常足以捕捉關鍵信息。

   - 在語言任務中,捕捉詞語或句子之間的長距離語義依賴至關重要。

3. 計算效率的考量:

   - 視覺Transformer的序列較短,採用稠密注意力計算代價相對較小。

   - 語言Transformer面臨更長序列,需要更高效的注意力機制來控制計算成本。

4. 跨領域知識借鑑:

   - 擴張注意力最初在視覺任務中展現出優越性能,吸引了NLP研究者的興趣。

   - 將擴張注意力引入語言Transformer,有望在保持計算效率的同時提升模型性能。

儘管視覺和語言任務在序列特性上有所不同,但它們在注意力機制設計上有著共通的目標:捕捉關鍵的依賴關係,同時兼顧計算效率。語言Transformer借鑑並改進了視覺Transformer的擴張注意力,以更好地滿足語言任務的特定需求。這種跨領域的知識融合和創新,推動了Transformer架構在不同任務上的進一步發展和應用。

人工智慧技術的發展常常需要跨領域的思想碰撞和借鑑。視覺和語言雖然是不同的領域,但它們在底層的模型架構和算法設計上往往有許多共通之處。擴張注意力從視覺Transformer到語言Transformer的演進過程,正是這種跨領域融合創新的一個絕佳例證。

SambaNova SN40L: 利用Dataflow和專家組合(COE)來克服AI記憶牆的大模型

摘要 GPT-4等整體式大型語言模型(LLM)為現代生成AI應用鋪路。然而,大規模訓練、服務及維護整體式LLM仍然極其昂貴和充滿挑戰。現代AI加速器計算能力與記憶體比例的不成比例增長已經造成了記憶體壁障,需要新的方法來部署AI。最近的研究顯示,許多小型專家模型的組合,每個模型參數...