KAN結(jié)合Transformer，真有團(tuán)隊(duì)搞出了解決擴(kuò)展缺陷的KAT

作者：機(jī)器之心 2024-09-19 13:04:03

來自 MIT 等機(jī)構(gòu)的研究者提出了一種非常有潛力的替代方法 ——KAN。

Transformer 是現(xiàn)代深度學(xué)習(xí)的基石。傳統(tǒng)上，Transformer 依賴多層感知器 (MLP) 層來混合通道之間的信息。

前段時(shí)間，來自 MIT 等機(jī)構(gòu)的研究者提出了一種非常有潛力的替代方法 ——KAN。該方法在準(zhǔn)確性和可解釋性方面表現(xiàn)優(yōu)于 MLP。而且，它能以非常少的參數(shù)量勝過以更大參數(shù)量運(yùn)行的 MLP。

KAN的發(fā)布，引起了AI社區(qū)大量的關(guān)注與討論，同時(shí)也伴隨很大的爭議。

而此類研究，又有了新的進(jìn)展。

最近，來自新加坡國立大學(xué)的研究者提出了 Kolmogorov–Arnold Transformer（KAT），用 Kolmogorov-Arnold Network（KAN）層取代 MLP 層，以增強(qiáng)模型的表達(dá)能力和性能。

論文標(biāo)題：Kolmogorov–Arnold Transformer
論文地址：https://arxiv.org/pdf/2409.10594
項(xiàng)目地址：https://github.com/Adamdad/kat

KAN 原論文第一作者 Ziming Liu 也轉(zhuǎn)發(fā)點(diǎn)贊了這項(xiàng)新研究。

將 KAN 集成到 Transformer 中并不是一件容易的事，尤其是在擴(kuò)展時(shí)。具體來說，該研究確定了三個(gè)關(guān)鍵挑戰(zhàn)：

(C1) 基函數(shù)。KAN 中使用的標(biāo)準(zhǔn) B 樣條（B-spline）函數(shù)并未針對現(xiàn)代硬件上的并行計(jì)算進(jìn)行優(yōu)化，導(dǎo)致推理速度較慢。

(C2) 參數(shù)和計(jì)算效率低下。KAN 需要每個(gè)輸入輸出對都有特定的函數(shù)，這使得計(jì)算量非常大。

(C3) 權(quán)重初始化。由于具有可學(xué)習(xí)的激活函數(shù)，KAN 中的權(quán)重初始化特別具有挑戰(zhàn)性，這對于實(shí)現(xiàn)深度神經(jīng)網(wǎng)絡(luò)的收斂至關(guān)重要。

為了克服上述挑戰(zhàn)，研究團(tuán)隊(duì)提出了三個(gè)關(guān)鍵解決方案：

(S1) 有理基礎(chǔ)。該研究用有理函數(shù)替換 B 樣條函數(shù)，以提高與現(xiàn)代 GPU 的兼容性。通過在 CUDA 中實(shí)現(xiàn)這一點(diǎn)，該研究實(shí)現(xiàn)了更快的計(jì)算。

(S2) Group KAN。通過一組神經(jīng)元共享激活權(quán)重，以在不影響性能的情況下減少計(jì)算負(fù)載。

(S3) Variance-preserving 初始化。該研究仔細(xì)初始化激活權(quán)重，以確保跨層保持激活方差。

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

實(shí)驗(yàn)結(jié)果表明：GR-KAN 計(jì)算效率高、易于實(shí)現(xiàn)，并且可以無縫集成到視覺 transformer（ViT）中，取代 MLP 層以實(shí)現(xiàn)卓越的性能。此外，該研究的設(shè)計(jì)允許 KAT 從 ViT 模型加載預(yù)訓(xùn)練權(quán)重并繼續(xù)訓(xùn)練以獲得更好的結(jié)果。

該研究在一系列視覺任務(wù)中實(shí)證驗(yàn)證了 KAT，包括圖像識(shí)別、目標(biāo)檢測和語義分割。結(jié)果表明，KAT 的性能優(yōu)于傳統(tǒng)的基于 MLP 的 transformer，在計(jì)算量相當(dāng)?shù)那闆r下實(shí)現(xiàn)了增強(qiáng)的性能。

如圖 1 所示，KAT-B 在 ImageNet-1K 上實(shí)現(xiàn)了 82.3% 的準(zhǔn)確率，超過相同大小的 ViT 模型 3.1%。當(dāng)使用 ViT 的預(yù)訓(xùn)練權(quán)重進(jìn)行初始化時(shí)，準(zhǔn)確率進(jìn)一步提高到 82.7%。

不過，也有網(wǎng)友質(zhì)疑道：「自從有論文比較了具有相同參數(shù)大小的 MLP 模型和 KAN 模型的性能后，我就對 KAN 持懷疑態(tài)度。可解釋性似乎是唯一得到巨大提升的東西。」

對此，論文作者回應(yīng)道：「的確，原始 KAN 在可解釋性上做得很好，但不保證性能和效率。我們所做的就是修復(fù)這些 bug 并進(jìn)行擴(kuò)展。」

還有網(wǎng)友表示，這篇論文和其他人的想法一樣，就是用 KAN 取代了 MLP，并質(zhì)疑為什么作者在嘗試一些已經(jīng)很成熟和類似的東西，難道是在炒作 KAN？對此，論文作者 Xingyi Yang 解釋道，事實(shí)確實(shí)如此，但不是炒作，根據(jù)實(shí)驗(yàn)，簡單地進(jìn)行這種替換是行不通的，他們在努力將這個(gè)簡單的想法變成可能的事情。

Kolmogorov–Arnold Transformer （KAT）

作者表示，標(biāo)準(zhǔn)的 KAN 面臨三大挑戰(zhàn)，限制了其在大型深度神經(jīng)網(wǎng)絡(luò)中的應(yīng)用。

它們分別是基函數(shù)的選擇、冗余參數(shù)及其計(jì)算、初始化問題。這些設(shè)計(jì)選擇使得原始版本的 KAN 是資源密集型的，難以應(yīng)用于大規(guī)模模型。

本文對這些缺陷設(shè)計(jì)加以改進(jìn)，以更好地適應(yīng)現(xiàn)代 Transformer，從而允許用 KAN 替換 MLP 層。

KAT 整體架構(gòu)

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

具體來說，對于 2D 圖像，作者首先將其平面化成 1D 序列，在此基礎(chǔ)上應(yīng)用 patch 嵌入和位置編碼，然后通過一系列 KAT 層進(jìn)行傳遞。對于層，可以執(zhí)行如下操作：

其中，表示層的輸出特征序列。

如圖所示，作者用兩層 KAN 替換兩層 MLP，同時(shí)保持注意力層不變。然而，簡單的替換不足以在大模型中實(shí)現(xiàn)可擴(kuò)展性。

最重要的是，在這里，作者引入了一種特殊的 Group-Rational KAN。作者使用有理函數(shù)作為 KAN 的基函數(shù)，并在一組邊之間共享參數(shù)。此外，作者還指定了權(quán)重初始化方案以確保穩(wěn)定的訓(xùn)練。這些改進(jìn)使得 KAT 更具可擴(kuò)展性并提高了性能。

有理基函數(shù)

作者使用有理函數(shù)作為 KAN 層的基函數(shù)，而不是 B 樣條函數(shù)，即每個(gè)邊上的函數(shù) ?? (??) 參數(shù)化為 ??、?? 階多項(xiàng)式 ?? (??)、??(??) 上的有理數(shù)。

標(biāo)準(zhǔn)形式轉(zhuǎn)化為：

至于為什么采用有理函數(shù)，作者表示從效率角度來看，多項(xiàng)式求值涉及簡單的運(yùn)算，非常適合并行計(jì)算。這使得有理函數(shù)對于大規(guī)模模型具有計(jì)算效率。

其次，從理論角度來看，有理函數(shù)可以比多項(xiàng)式更高效、更準(zhǔn)確地逼近更廣泛的函數(shù)。由于 B 樣條本質(zhì)上是局部多項(xiàng)式的和，因此有理函數(shù)在復(fù)雜行為建模方面比 B 樣條具有理論優(yōu)勢。

第三，從實(shí)踐角度來看，有理激活函數(shù)已經(jīng)成功用作神經(jīng)網(wǎng)絡(luò)中的激活函數(shù)。

Group KAN

作者表示，他們不必為每個(gè)輸入 - 輸出對學(xué)習(xí)一個(gè)獨(dú)特的基函數(shù)，而是可以在一組邊內(nèi)共享它們的參數(shù)。這減少了參數(shù)數(shù)量和計(jì)算量。這種參數(shù)共享和分組計(jì)算的方式一直是神經(jīng)網(wǎng)絡(luò)設(shè)計(jì)中的關(guān)鍵技術(shù)

圖 2 說明了原始 KAN、Group KAN 和標(biāo)準(zhǔn) MLP 之間的區(qū)別。Group KAN 通過在一組邊之間共享這些函數(shù)來減少參數(shù)數(shù)量。

除了節(jié)省參數(shù)數(shù)量外，這種分組還減少了計(jì)算需求。不同模型間參數(shù)數(shù)量和計(jì)算量的對比如下所示：

Variance-preserving 初始化

作者旨在初始化 Group-Rational KAN 中的 ??_??、??_?? 和 ?? 的值，其核心是防止整個(gè)層中的激活參數(shù)呈量級增長或減少，從而保持穩(wěn)定性。

實(shí)驗(yàn)

實(shí)驗(yàn)中，作者修改了原始 ViT 架構(gòu)，用 GR-KAN 層替換其 MLP 層。

圖像識(shí)別

實(shí)驗(yàn)結(jié)果表明，KAT 模型在 IN-1k（ ImageNet-1K ）數(shù)據(jù)集上的表現(xiàn)始終優(yōu)于其他模型。首先，GR-KAN 在 Transformer 架構(gòu)中的表現(xiàn)優(yōu)于傳統(tǒng)的基于 MLP 的混合器的性能。例如，KAT-S 模型的準(zhǔn)確率達(dá)到 81.2%，比 DeiT-S 模型高出 2.4%。

其次，原始 KAN 層面臨可擴(kuò)展性問題。ViT-T/S + KAN 的準(zhǔn)確率僅為 63% 左右，即使計(jì)算成本高得多。ViT-L + KAN 無法收斂，導(dǎo)致 NAN 錯(cuò)誤。本文解決了這些擴(kuò)展挑戰(zhàn)，從而使 KAT 模型能夠成功擴(kuò)展。

目標(biāo)檢測和實(shí)例分割

表 6 比較了不同骨干模型的性能。KAT 的表現(xiàn)始終優(yōu)于其他模型，尤其是在物體檢測方面，與 ViTDet 相比，其在 S 規(guī)模的模型上實(shí)現(xiàn)了 3.0 AP^box 增益，在 L 規(guī)模的模型上實(shí)現(xiàn)了 1.4 AP^box 增益。這種改進(jìn)在較小的模型中最為明顯，計(jì)算成本僅增加了 1 GFLOP。這表明 KAT 以最小的開銷提供了更好的準(zhǔn)確率。

語義分割

表 7 總結(jié)了分割結(jié)果。總體而言，KAT 比基于 ViT 的普通架構(gòu)表現(xiàn)出了競爭優(yōu)勢，比 DeiT-S 提高了 2.4%，比 DeiT-B 提高了 0.2%。這種性能提升伴隨著計(jì)算成本的輕微增加，反映在更高的 FLOP 上。與檢測結(jié)果類似，KAT 在較小的模型中顯示出更顯著的收益。然而，與具有分層架構(gòu)的模型（如 ConvNeXt）相比，它仍然有所不足，這些模型受益于更高效的架構(gòu)設(shè)計(jì)。