TruKAN: Towards More Efficient Kolmogorov-Arnold Networks Using Truncated Power Functions

Technische Samenvatting van "TruKAN: Naar Efficiëntere Kolmogorov-Arnold-Netwerken met Gebruik van Afgekante Vermogensfuncties"

Probleemstelling

Kolmogorov-Arnold-Netwerken (KANs) zijn opgekomen als een veelbelovend alternatief voor Multi-Layer Perceptrons (MLPs), waarbij ze superieure interpreteerbaarheid en het vermogen bieden om univariate randfuncties te leren. De standaard KANs vertrouwen echter op B-spline basisfuncties voor hun leerbare activaties. Het evalueren van deze B-splines vereist het recursieve de Boor–Cox-algoritme, wat computationeel duur is en vaak ervoor zorgt dat KANs tot wel 10 keer langzamer trainen dan equivalente MLPs. Bovendien kan de afhankelijkheid van complexe spline-algoritmen en vaste roosterstructuren leiden tot schaalbaarheidsproblemen in hoogdimensionale problemen en diepe architecturen, gekenmerkt door ruisende gradiënten en instabiele optimalisatie. Hoewel recente varianten zoals SineKAN en FourierKAN proberen deze problemen aan te pakken, introduceren zij vaak nieuwe afruilverschillen met betrekking tot interpreteerbaarheid of blijven zij beperkt tot synthetische benchmarks en eenvoudige datasets zoals MNIST.

Methodologie

De auteurs stellen TruKAN voor, een nieuwe architectuur die de canonieke KAN-topologie en leerbare activaties behoudt, maar de B-spline basis vervangt door een familie van afgekante vermogensfuncties (truncated power functions) afgeleid uit $k$-orde spline-theorie.

Kerncomponenten van de architectuur:

1. Afgekante Vermogensbasis: In plaats van B-splines gebruikt TruKAN afgekante vermogensfuncties van de vorm $(x - t_j)^k_+$. Theoretisch is de ruimte gespannen door deze functies identiek aan die van B-splines, maar hun computationele evaluatie is verschillend. Terwijl B-splines een compacte ondersteuning hebben (wat een $O(k^2)$ complexiteit per interval vereist via recursie), hebben afgekante vermogensfuncties een globale ondersteuning. De auteurs argumenteren echter dat in de praktijk hoge-resolutie roosters niet altijd noodzakelijk zijn, waardoor TruKAN efficiënt kan werken met een klein aantal knopen (knots).
2. Hybride Laagstructuur: Elke TruKAN-laag combineert twee componenten:
   • Een afgekante vermogenscomponent om lokale, stuksgewijze niet-lineariteiten te vangen.
   • Een polynomiale component (globale trend) om de hulp-SiLU-activatie die in standaard KANs wordt gebruikt, te vervangen. Deze polynomiale term zorgt voor een gladde, differentieerbare globale mapping, wat het gebrek aan een expliciete globale niet-lineariteit zoals SiLU compenseert.
3. Knopenconfiguraties: Het artikel introduceert twee varianten met betrekking tot de plaatsing van knopen:
   • Gedeelde Knopen (Shared Knots): Een enkele set knoopposities wordt gedeeld over alle outputkanalen.
   • Individuele Knopen (Individual Knots): Elk outputkanaal heeft zijn eigen set leerbare knoopposities.
   • Zowel vaste (vooraf gedefinieerde intervallen) als leerbare (geoptimaliseerd tijdens de training met beperkingen voor ordening en positiviteit) knoopstrategieën worden verkend.
4. Trainingsframework:
   • Optimalisatie: Om de instabiliteit van het trainen van diepe netwerken met afgekante vermogensfuncties aan te pakken (die last kunnen hebben van gradiëntexplosie door exponentiële groei/afname), gebruiken de auteurs een hybride optimalisatiestrategie met AdamW ingekapseld door LookAhead. Dit vervangt de standaard L-BFGS die in eerdere KAN-studies werd gebruikt, om beter aan te sluiten bij stochastische mini-batch training.
   • Normalisatie: De studie onderzoekt de impact van Layer Normalization (en Batch Normalization) binnen de classifier-heads, aangezien traditionele regularisatietechnieken zoals dropout gemengde resultaten hebben laten zien in KANs.
   • Backbone: De modellen worden geïntegreerd in een EfficientNet-V2 framework, waarbij de standaard lineaire classifier-head wordt vervangen door MLP, KAN, SineKAN of TruKAN lagen.

Belangrijkste Bijdragen

• Architecturale Innovatie: De vervanging van B-splines door afgekante vermogensfuncties, samen met een polynomiale globale term, creëert een model dat de expressiviteit van KAN behoudt terwijl het de evaluatie van de basisfunctie vereenvoudigt.
• Verbeterde Interpreteerbaarheid: Door gebruik te maken van afgekante vermogensfuncties wordt de bijdrage van elke knoop expliciet. Een niet-nul coëfficiënt bij een knoop $t_j$ geeft direct een intentionele $k$-de orde kromming aan dat specifieke punt aan, wat meer transparantie biedt dan de overlappende ondersteuning van B-splines.
• Schaalbaarheid en Efficiëntie: Het artikel demonstreert dat TruKAN effectief kan worden geschaald naar diepe architecturen en complexe visuele taken, waarmee het verder gaat dan de beperkingen van synthetische data in eerdere KAN-studies.
• Uitgebreide Benchmarking: De auteurs bieden een rigoureuze evaluatie over vier computer vision datasets (CIFAR-10, CIFAR-100, Oxford-IIIT Pets, STL-10) met behoud van eerlijke vergelijkingen door het aantal parameters gelijk te houden.

Experimentele Resultaten

De evaluatie werd uitgevoerd op vier datasets met EfficientNet-V2 backbones met variërende classifier-groottes (kleine en grote architecturen).

• Nauwkeurigheid en F1-score:
  • Kleine Architecturen: TruKAN-varianten, met name die met Layer Normalization en Gedeelde Knopen (TruKAN-SN), presteerden consequent beter dan of gelijk aan de baseline KAN en SineKAN modellen. Op CIFAR-10 bereikte TruKAN-IN de hoogste nauwkeurigheid (0.8927), waarmee het de beste MLP-baseline overtrof. Op STL-10 bereikte TruKAN-SN een nauwkeurigheid van 0.9341, waarbij het slechts een verwaarloosbaar verschil had met de genormaliseerde MLP.
  • Grote Architecturen: TruKAN-modellen vertoonden een superieure schaalbaarheid. Op CIFAR-10 bereikte TruKAN-IN een nauwkeurigheid van 0.8932, waarmee het de MLP-N baseline overtrof ondanks het gebruik van minder parameters. TruKAN-varianten domineerden consistent de KAN en SineKAN baselines over alle datasets, en behaalden vaak resultaten die vergelijkbaar met of beter waren dan de sterkste MLPs.
  • Impact van Normalisatie: Layer normalization bleek cruciaal voor de prestaties van TruKAN; het verbeterde de generalisatie aanzienlijk, vooral voor configuraties met individuele knopen die anders last hadden van overfitting.
• Computationele Efficiëntie:
  • Trainingstijd: Gedeelde-knoop TruKAN-modellen (TruKAN-FS) bereikten trainingsstap-tijden van ~0.0024s, wat een verbetering van 3-4x betekent ten opzijde van standaard KAN-modellen (0.0087s - 0.0099s).
  • Geheugengebruik: TruKAN-FS verbruikte slechts 118.04 MB aan piek GPU-geheugen, aanzienlijk minder dan KAN-PBT (570.84 MB) en KAN-PBF (691.79 MB). Echter, individuele knoop-varianten (TruKAN-FI) veroorzaakten hoge geheugenkosten (~2.6 GB) door de uitgebreide opslag van knopen.
  • Parameters: TruKAN-modellen behaalden over het algemeen vergelijkbare of betere prestaties met minder trainbare parameters dan hun KAN en SineKAN tegenhangers.

Betekenis en Claims

Het artikel claimt dat TruKAN er succesvol in slaagt de afruil tussen computationele efficiëntie en het naleven van KAN-principes aan te pakken. Door de recursieve B-spline algoritme te vervangen door afgekante vermogensfuncties en een polynomiale globale term, bereikt TruKAN:

1. Superieure Efficiëntie: Significante reducties in trainingstijd en geheugengebruik vergeleken met standaard KANs, waardoor ze bruikbaar zijn voor complexe visuele taken.
2. Behoud van Expressiviteit: Het model behoudt de theoretische benaderingscapaciteiten van KANs terwijl de nauwkeurigheid op real-world datasets wordt verbeterd.
3. Verbeterde Interpreteerbaarheid: De expliciete aard van afgekante vermogensfuncties maakt een duidelijkere inspectie van geleerde randfuncties en knoopbijdragen mogelijk vergeleken met B-splines of globale spectrale methoden zoals SineKAN.

De auteurs concluderen dat TruKAN een praktisch pad voorwaarts biedt voor KANs, waarbij zij voordelen aantonen die verder gaan dan de beperkte settings van eerdere studies. Zij erkennen dat hoewel de methode de stabiliteit verbetert, er nog steeds uitdagingen zijn met betrekking tot de integratie van andere regularisatietechnieken (zoals dropout) en het potentieel van mixed-precision training om de efficiëntie verder te verhogen. Toekomstig werk wordt voorgesteld om deze gebieden te verkennen en TruKAN toe te passen op reinforcement learning en tijdreeksvoorspelling.