Kolmogorov-Arnold Transformer

Kolmogorov-Arnold Transformer

September 16, 2024

Auteurs: Xingyi Yang, Xinchao Wang

cs.AI

Samenvatting

Transformers vormen de hoeksteen van moderne diepgaande leermodellen. Traditioneel vertrouwen deze modellen op multi-layer perceptron (MLP) lagen om informatie tussen kanalen te mengen. In dit artikel introduceren we de Kolmogorov-Arnold Transformer (KAT), een nieuw architectuur dat MLP-lagen vervangt door Kolmogorov-Arnold Network (KAN) lagen om de expressiviteit en prestaties van het model te verbeteren. Het integreren van KANs in transformers is echter geen gemakkelijke taak, vooral bij opschaling. Specifiek identificeren we drie belangrijke uitdagingen: (C1) Basale functie. De standaard B-spline functie die in KANs wordt gebruikt, is niet geoptimaliseerd voor parallel rekenen op moderne hardware, wat resulteert in langzamere inferentiesnelheden. (C2) Parameter- en Berekeningsefficiëntie. KAN vereist een unieke functie voor elk invoer-uitvoerpaar, wat de berekening extreem groot maakt. (C3) Gewichtsinitialisatie. Het initialiseren van gewichten in KANs is bijzonder uitdagend vanwege hun aanpasbare activatiefuncties, die cruciaal zijn voor het bereiken van convergentie in diepe neurale netwerken. Om de eerder genoemde uitdagingen te overwinnen, stellen we drie belangrijke oplossingen voor: (S1) Rationele basis. We vervangen B-spline functies door rationale functies om de compatibiliteit met moderne GPU's te verbeteren. Door dit te implementeren in CUDA bereiken we snellere berekeningen. (S2) Groeps-KAN. We delen de activatiegewichten via een groep neuronen om de rekenbelasting te verminderen zonder prestaties op te offeren. (S3) Variatiebehoudende initialisatie. We initialiseren de activatiegewichten zorgvuldig om ervoor te zorgen dat de activatievariatie over lagen behouden blijft. Met deze ontwerpen schaalt KAT effectief en presteert het beter dan traditionele op MLP gebaseerde transformers.

English

Transformers stand as the cornerstone of mordern deep learning. Traditionally, these models rely on multi-layer perceptron (MLP) layers to mix the information between channels. In this paper, we introduce the Kolmogorov-Arnold Transformer (KAT), a novel architecture that replaces MLP layers with Kolmogorov-Arnold Network (KAN) layers to enhance the expressiveness and performance of the model. Integrating KANs into transformers, however, is no easy feat, especially when scaled up. Specifically, we identify three key challenges: (C1) Base function. The standard B-spline function used in KANs is not optimized for parallel computing on modern hardware, resulting in slower inference speeds. (C2) Parameter and Computation Inefficiency. KAN requires a unique function for each input-output pair, making the computation extremely large. (C3) Weight initialization. The initialization of weights in KANs is particularly challenging due to their learnable activation functions, which are critical for achieving convergence in deep neural networks. To overcome the aforementioned challenges, we propose three key solutions: (S1) Rational basis. We replace B-spline functions with rational functions to improve compatibility with modern GPUs. By implementing this in CUDA, we achieve faster computations. (S2) Group KAN. We share the activation weights through a group of neurons, to reduce the computational load without sacrificing performance. (S3) Variance-preserving initialization. We carefully initialize the activation weights to make sure that the activation variance is maintained across layers. With these designs, KAT scales effectively and readily outperforms traditional MLP-based transformers.