機(jī)器之心報(bào)道

編輯：小舟、陳陳

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

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

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

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

最近，來(lái)自新加坡國(guó)立大學(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í)。具體來(lái)說(shuō)，該研究確定了三個(gè)關(guān)鍵挑戰(zhàn)：

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

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

(C3) 權(quán)重初始化。由于具有可學(xué)習(xí)的激活函數(shù)，KAN 中的權(quán)重初始化特別具有挑戰(zhàn)性，這對(duì)于實(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 的兼容性。通過(guò)在 CUDA 中實(shí)現(xiàn)這一點(diǎn)，該研究實(shí)現(xiàn)了更快的計(jì)算。

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

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

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

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

該研究在一系列視覺(jué)任務(wù)中實(shí)證驗(yàn)證了 KAT，包括圖像識(shí)別、目標(biāo)檢測(cè)和語(yǔ)義分割。結(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)確率，超過(guò)相同大小的 ViT 模型 3.1%。當(dāng)使用 ViT 的預(yù)訓(xùn)練權(quán)重進(jìn)行初始化時(shí)，準(zhǔn)確率進(jìn)一步提高到 82.7%。

不過(guò)，也有網(wǎng)友質(zhì)疑道：「自從有論文比較了具有相同參數(shù)大小的 MLP 模型和 KAN 模型的性能后，我就對(duì) KAN 持懷疑態(tài)度?？山忉屝运坪跏俏ㄒ坏玫骄薮筇嵘臇|西?！?/p>

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

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

Kolmogorov–Arnold Transformer （KAT）

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

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

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

KAT 整體架構(gòu)

正如其名稱(chēng)所暗示的那樣，KAT 用 KAN 層取代了視覺(jué) transformer 中的 MLP 層。

有理基函數(shù)

作者使用有理函數(shù)作為 KAN 層的基函數(shù)，而不是 B 樣條函數(shù)，即每個(gè)邊上的函數(shù)