KAN結合Transformer，真有團隊搞出了解決擴充套件缺陷的KAT

机器之心發表於2024-09-19
ORM套件
Transformer 是現代深度學習的基石。傳統上，Transformer 依賴多層感知器 (MLP) 層來混合通道之間的資訊。

前段時間，來自 MIT 等機構的研究者提出了一種非常有潛力的替代方法 ——KAN。該方法在準確性和可解釋性方面表現優於 MLP。而且，它能以非常少的引數量勝過以更大引數量執行的 MLP。

KAN的釋出，引起了AI社群大量的關注與討論，同時也伴隨很大的爭議。

而此類研究，又有了新的進展。

最近，來自新加坡國立大學的研究者提出了 Kolmogorov–Arnold Transformer（KAT），用 Kolmogorov-Arnold Network（KAN）層取代 MLP 層，以增強模型的表達能力和效能。
圖片
  • 論文標題：Kolmogorov–Arnold Transformer
  • 論文地址：https://arxiv.org/pdf/2409.10594
  • 專案地址：https://github.com/Adamdad/kat

KAN 原論文第一作者 Ziming Liu 也轉發點讚了這項新研究。
圖片
將 KAN 整合到 Transformer 中並不是一件容易的事，尤其是在擴充套件時。具體來說，該研究確定了三個關鍵挑戰：

(C1) 基函式。KAN 中使用的標準 B 樣條（B-spline）函式並未針對現代硬體上的平行計算進行最佳化，導致推理速度較慢。
(C2) 引數和計算效率低下。KAN 需要每個輸入輸出對都有特定的函式，這使得計算量非常大。
(C3) 權重初始化。由於具有可學習的啟用函式，KAN 中的權重初始化特別具有挑戰性，這對於實現深度神經網路的收斂至關重要。

為了克服上述挑戰，研究團隊提出了三個關鍵解決方案：

(S1) 有理基礎。該研究用有理函式替換 B 樣條函式，以提高與現代 GPU 的相容性。透過在 CUDA 中實現這一點，該研究實現了更快的計算。
(S2) Group KAN。透過一組神經元共享啟用權重，以在不影響效能的情況下減少計算負載。
(S3) Variance-preserving 初始化。該研究仔細初始化啟用權重，以確保跨層保持啟用方差。

結合解決方案 S1-S3，該研究提出了一種新的 KAN 變體，稱為 Group-Rational KAN (GR-KAN)，以取代 Transformer 中的 MLP。

實驗結果表明：GR-KAN 計算效率高、易於實現，並且可以無縫整合到視覺 transformer（ViT）中，取代 MLP 層以實現卓越的效能。此外，該研究的設計允許 KAT 從 ViT 模型載入預訓練權重並繼續訓練以獲得更好的結果。

該研究在一系列視覺任務中實證驗證了 KAT，包括影像識別、目標檢測和語義分割。結果表明，KAT 的效能優於傳統的基於 MLP 的 transformer，在計算量相當的情況下實現了增強的效能。
圖片
如圖 1 所示，KAT-B 在 ImageNet-1K 上實現了 82.3% 的準確率，超過相同大小的 ViT 模型 3.1%。當使用 ViT 的預訓練權重進行初始化時，準確率進一步提高到 82.7%。

不過，也有網友質疑道：「自從有論文比較了具有相同引數大小的 MLP 模型和 KAN 模型的效能後，我就對 KAN 持懷疑態度。可解釋性似乎是唯一得到巨大提升的東西。」
圖片
對此，論文作者回應道：「的確，原始 KAN 在可解釋性上做得很好，但不保證效能和效率。我們所做的就是修復這些 bug 並進行擴充套件。」
圖片

還有網友表示，這篇論文和其他人的想法一樣，就是用 KAN 取代了 MLP，並質疑為什麼作者在嘗試一些已經很成熟和類似的東西，難道是在炒作 KAN？對此， 論文作者 Xingyi Yang 解釋道，事實確實如此，但不是炒作，根據實驗，簡單地進行這種替換是行不通的，他們在努力將這個簡單的想法變成可能的事情。
圖片
Kolmogorov–Arnold Transformer （KAT）

作者表示，標準的 KAN 面臨三大挑戰，限制了其在大型深度神經網路中的應用。

它們分別是基函式的選擇、冗餘引數及其計算、初始化問題。這些設計選擇使得原始版本的 KAN 是資源密集型的，難以應用於大規模模型。

本文對這些缺陷設計加以改進，以更好地適應現代 Transformer，從而允許用 KAN 替換 MLP 層。

KAT 整體架構

正如其名稱所暗示的那樣，KAT 用 KAN 層取代了視覺 transformer 中的 MLP 層。

具體來說，對於 2D 影像圖片，作者首先將其平面化成 1D 序列，在此基礎上應用 patch 嵌入和位置編碼，然後透過一系列 KAT 層進行傳遞。對於圖片層，可以執行如下操作：

圖片

其中，圖片表示圖片層的輸出特徵序列。

如圖所示，作者用兩層 KAN 替換兩層 MLP，同時保持注意力層不變。然而，簡單的替換不足以在大模型中實現可擴充套件性。

最重要的是，在這裡，作者引入了一種特殊的 Group-Rational KAN。作者使用有理函式作為 KAN 的基函式，並在一組邊之間共享引數。此外，作者還指定了權重初始化方案以確保穩定的訓練。這些改進使得 KAT 更具可擴充套件性並提高了效能。

有理基函式

作者使用有理函式作為 KAN 層的基函式，而不是 B 樣條函式，即每個邊上的函式 𝜙 (𝑥) 引數化為 𝑚、𝑛 階多項式 𝑃 (𝑥)、𝑄(𝑥) 上的有理數。
圖片
標準形式轉化為：
圖片
至於為什麼採用有理函式，作者表示從效率角度來看，多項式求值涉及簡單的運算，非常適合平行計算。這使得有理函式對於大規模模型具有計算效率。

其次，從理論角度來看，有理函式可以比多項式更高效、更準確地逼近更廣泛的函式。由於 B 樣條本質上是區域性多項式的和，因此有理函式在複雜行為建模方面比 B 樣條具有理論優勢。

第三，從實踐角度來看，有理啟用函式已經成功用作神經網路中的啟用函式。
圖片
Group KAN

作者表示，他們不必為每個輸入 - 輸出對學習一個獨特的基函式，而是可以在一組邊內共享它們的引數。這減少了引數數量和計算量。這種引數共享和分組計算的方式一直是神經網路設計中的關鍵技術

圖 2 說明了原始 KAN、Group KAN 和標準 MLP 之間的區別。Group KAN 透過在一組邊之間共享這些函式來減少引數數量。
圖片
除了節省引數數量外，這種分組還減少了計算需求。不同模型間引數數量和計算量的對比如下所示：
圖片
Variance-preserving 初始化

作者旨在初始化 Group-Rational KAN 中的 𝑎_𝑚、𝑏_𝑛 和 𝑤 的值，其核心是防止整個層中的啟用引數呈量級增長或減少，從而保持穩定性。
圖片
實驗

實驗中，作者修改了原始 ViT 架構，用 GR-KAN 層替換其 MLP 層。
圖片
影像識別

實驗結果表明，KAT 模型在 IN-1k（ ImageNet-1K ） 資料集上的表現始終優於其他模型。首先，GR-KAN 在 Transformer 架構中的表現優於傳統的基於 MLP 的混合器的效能。例如，KAT-S 模型的準確率達到 81.2%，比 DeiT-S 模型高出 2.4%。
圖片
其次，原始 KAN 層面臨可擴充套件性問題。ViT-T/S + KAN 的準確率僅為 63% 左右，即使計算成本高得多。ViT-L + KAN 無法收斂，導致 NAN 錯誤。本文解決了這些擴充套件挑戰，從而使 KAT 模型能夠成功擴充套件。

目標檢測和例項分割

表 6 比較了不同骨幹模型的效能。KAT 的表現始終優於其他模型，尤其是在物體檢測方面，與 ViTDet 相比，其在 S 規模的模型上實現了 3.0 AP^box 增益，在 L 規模的模型上實現了 1.4 AP^box 增益。這種改進在較小的模型中最為明顯，計算成本僅增加了 1 GFLOP。這表明 KAT 以最小的開銷提供了更好的準確率。
圖片
語義分割

表 7 總結了分割結果。總體而言，KAT 比基於 ViT 的普通架構表現出了競爭優勢，比 DeiT-S 提高了 2.4%，比 DeiT-B 提高了 0.2%。這種效能提升伴隨著計算成本的輕微增加，反映在更高的 FLOP 上。與檢測結果類似，KAT 在較小的模型中顯示出更顯著的收益。然而，與具有分層架構的模型（如 ConvNeXt）相比，它仍然有所不足，這些模型受益於更高效的架構設計。
圖片
作者介紹

Xingyi Yang 現在是新加坡國立大學（NUS）三年級博士生，導師是 Xinchao Wang 教授，這篇論文就是師徒兩人合作完成的。

Xingyi Yang 於 2021 年在加州大學聖地亞哥分校獲得碩士學位，並於 2019 年在東南大學獲得電腦科學學士學位。

Xinchao Wang 目前是新加坡國立大學電氣與計算機工程系（ECE）的助理教授，研究興趣包括人工智慧、計算機視覺、機器學習、醫學影像分析和多媒體。

相關文章