Technische Samenvatting: Tensor Network Kernel Machines: Een JAX-framework voor Machine Learning en Nietlineaire Systeemidentificatie

Probleemstelling

Het ontwikkelen van nietlineaire modellen die een balans vinden tussen een hoge expressieve kracht en computationele efficiëntie blijft een aanzienlijke uitdaging in machine learning en nietlineaire systeemidentificatie. Traditionele benaderingen, zoals kernelmethoden, Gaussische processen en diepe neurale netwerken, kampen vaak met schaalbaarheidsproblemen waarbij de modelcomplexiteit of de computationele kosten evenredig groeien met het aantal parameters of trainingsvoorbeelden. Specifiek leidt het representeren van nietlineaire mapping via expliciete feature maps in hoogdimensionele ruimtes tot een exponentiële explosie in het aantal parameters, wat directe optimalisatie voor hoogdimensionele problemen onhaalbaar maakt.

Methodologie

Het artikel introduceert Tensor Network Kernel Machines (TNKM), een framework dat schaalbaarheid aanpakt door expliciete nietlineaire feature-representaties te combineren met compacte, low-rank tensor-netwerk parameterisaties. De kern van de wiskundige formulering representeert een nietlineaire functie $f(x)$ als een lineair model in een hoogdimensionele feature-ruimte:
$$f(x) = \phi(x)^\top w$$
waarbij $\phi(x)$ een nietlineaire feature map is en $w$ de parametervector is. Om de vloek van dimensionaliteit die inherent is aan $w$ te overwinnen, hanteert de TNKM-benadering twee belangrijke strategieën:

1. Tensor-product Feature Maps: De feature map $\phi(x)$ wordt geconstrueerd als een tensorproduct van eendimensionale maps die opereren op individuele inputvariabelen, wat natuurlijk voortvloeit uit productkernels.
2. Low-Rank Tensor Decompositie: In plaats van de volledige parametervector $w$ op te slaan, worden de parameters gerepresenteerd met behulp van low-rank tensordecompositie, specifiek de Canonical Polyadic (CP) en Tensor-Train (TT) formaten. Dit reduceert de parametercomplexiteit van exponentieel naar lineair (of bijna lineair) ten opzichte van de inputdimensie.

De auteurs hebben tnkm ontwikkeld, een open-source Python-bibliotheek gebouwd op JAX, om dit framework te implementeren. De bibliotheek hanteert een modulaire architectuur die drie primaire componenten scheidt:

• Feature Maps: Ondersteunt polynomiale, Fourier en Volterra (geheugen-gebaseerde) features.
• Tensor-Network Parameterizations: Implementeert CP- en TT-decompositie via specifieke modelklassen (bijv. CPKM, TTKM).
• Optimalisatiealgoritmen: Biedt twee complementaire strategieën:
  • Alternating Least Squares (ALS): Maakt gebruik van de multilineaire structuur voor snelle, deterministische convergentie op least-squares problemen.
  • Gradient-Based Optimization: Maakt gebruik van JAX' automatische differentiatie en de Optax-bibliotheek (bijv. Adam, SGD) voor flexibiliteit bij willekeurige differentieerbare verliesfuncties.

Belangrijkste Bijdragen

Het artikel beschrijft vier hoofdbijdragen:

1. Verenigd JAX Framework: Een bibliotheek met minimale afhankelijkheden (tnkm) die verschillende feature maps, tensor-netwerk backbones (CP en TT) en optimalisatiestrategieën verenigt onder één enkele interface.
2. Schaalbare Trainingsimplementaties: Integratie van structuur-bewuste ALS-optimalisatie naast algemene gradiënt-gebaseerde methoden, wat toepassing mogelijk maakt voor zowel regressie als bredere leer-taken.
3. Modulaire Software Architectuur: Een ontwerp dat het systematisch wisselen van componenten (feature maps, modelklassen, optimizers) mogelijk maakt en hulpmiddelen biedt voor gegevenspre-processing, simulatie en evaluatie specifiek voor nietlineaire systeemidentificatie.
4. Uitgebreide Experimentele Evaluatie: Een rigoureuze benchmarking van het framework op nietlineaire regressie- en systeemidentificatietaken, waarbij een competitieve nauwkeurigheid met efficiënte trainingstijden wordt aangetoond.

Experimentele Resultaten

De auteurs hebben tnkm op twee fronten geëvalueerd:

1. Vergelijking van Optimalisatiemethoden (Airfoil Self-Noise Dataset):

• ALS vs. Adam: ALS convergeerde snel naar een validatie Mean Squared Error (MSE) van 0,102, wat de Kernel Ridge Regression (KRR) baseline (0,15) dicht benaderde, in slechts 1,16 seconden. In tegenstelling hiertoe had Adam aanzienlijk meer tijd nodig (3,25s voor mini-batch) en meer epochs om een hogere fout (0,145) te bereiken.
• Conclusie: ALS biedt superieure efficiëntie voor least-squares TNKM-problemen, terwijl gradiënt-gebaseerde methoden de nodige flexibiliteit bieden voor niet-least-squares doelstellingen.

2. Benchmarks voor Nietlineaire Systeemidentificatie:
Het framework werd getest op drie standaardbenchmarks: Coupled Electric Drives, Cascaded Tanks en Silverbox.

• Coupled Electric Drives: TNKM behaalde een RMSE die vergelijkbaar is met de best presterende methoden (GPNARX en dynoNet), maar met een trainingstijd van < 0,1 seconde, terwijl concurrenten seconden tot minuten nodig hadden.
• Cascaded Tanks: TNKM bood een competitieve nauwkeurigheid (RMSE 0,347) onder black-box benaderingen, en presteerde aanzienlijk beter dan methoden zoals SUBNET en PNARX in termen van trainingstijd (< 0,1s vs. 1-60 min).
• Silverbox: TNKM behaalde prestaties die vergelijkbaar zijn met leidende black-box methoden (GPNARX, SUBNET) met trainingstijden van 1–60 seconden. Hoewel het de nauwkeurigheid van physics-informed benaderingen (NL-LFR) niet evenaarde, demonstreerde het een effectieve stabiliteit bij voorspellingen over lange tijdshorizonten.

Betekenis en Claims

Het artikel claimt dat tnkm de reproduceerbare ontwikkeling en toepassing van tensor-netwerk-gebaseerde leermethoden faciliteert. Door het ontkoppelen van feature maps, tensor-parameterisaties en optimalisatie-algoritmen, stelt de bibliotheek gebruikers in staat om systematisch TNKM-configuraties te verkennen zonder de onderliggende workflow te wijzigen.

De auteurs positioneren TNKM als een tussenweg tussen klassieke kernelmethoden en diepe neurale netwerken. Net als kernelmethoden gebruikt TNKM expliciete feature-representaties, wat meer controle en de mogelijkheid biedt om voorkennis te integreren, maar vervangt het de volledige parametervector door een compact tensor-netwerk om schaalbaarheid te garanderen. In tegenstelling tot neurale netwerken leert het niet direct features uit de data, wat de auteurs erkennen als een beperking die handmatige feature-selectie vereist.

De betekenis van het werk ligt in het bieden van een praktische, uitbreidbare softwarematige basis die de kloof overbrugt tussen theoretische tensor-netwerkmethoden en real-world nietlineaire systeemidentificatie, waarbij wordt aangetoond dat deze modellen een hoge voorspellende nauwkeurigheid kunnen bereiken terwijl ze lage computationele kosten behouden.