Hyperparameter Trajectory Inference with Conditional Lagrangian Optimal Transport

Deze paper introduceert Hyperparameter Trajectory Inference, een methode die conditionele Lagrangiaanse optimale transport gebruikt om een surrogaatmodel te bouwen dat de veranderingen in de output van neurale netwerken bij verschillende hyperparameters voorspelt, waardoor kostbare hertraining wordt voorkomen.

De kern

De Kern: Een "Tijdmachine" voor AI-beslissingen

Stel je voor dat je een zelflerende robot (een neurale net) hebt die taken uitvoert, zoals het besturen van een auto of het geven van medisch advies. Deze robot is niet statisch; hij kan zich aanpassen. Maar om hem aan te passen, moet je een knop draaien. Laten we deze knop de "hyperparameter" noemen.

• Het probleem: Vaak moet je deze knop alvast instellen voordat de robot aan het werk gaat.
  • Voorbeeld: In een spelletje wil je misschien dat de robot agressief is (veel punten scoren) of voorzichtig is (niet crashen). Als je de knop op "agressief" zet, is hij snel maar onzeker. Zet je hem op "voorzichtig", dan is hij veilig maar traag.
  • De pijn: Stel dat de robot eenmaal is ingezet en de situatie verandert. Je wilt plotseling dat hij meer voorzichtig is. In het verleden moest je dan de hele robot opnieuw trainen (zoals een student die zijn hele jaar opnieuw moet leren omdat hij de toets niet haalde). Dit kost veel tijd, geld en energie.

De Oplossing: HTI (Hyperparameter Trajectory Inference)

De auteurs van dit paper hebben een slimme oplossing bedacht: HTI.

In plaats van de robot elke keer opnieuw te trainen, leren ze een "tussenpersoon" (een surrogaatmodel) die de hele geschiedenis van de robot kent. Deze tussenpersoon weet precies hoe de robot zich gedraagt als je de knop een beetje draait, een heel stuk draait, of ergens tussenin.

De Analogie: De Reisplanner
Stel je voor dat je drie foto's hebt van een reiziger:

1. Foto A: De reiziger staat in Amsterdam (Start).
2. Foto B: De reiziger staat in Utrecht (Midden).
3. Foto C: De reiziger staat in Rotterdam (Einde).

Normaal gesproken zou je denken: "Oké, hij is van A naar B en dan naar C." Maar wat als je wilt weten waar hij was op het exacte moment dat hij halverwege tussen A en B liep? Of wat als je wilt weten hoe hij zou lopen als hij een andere route had gekozen (een andere instelling)?

De oude methoden tekenden vaak een rechte lijn tussen de foto's. Maar mensen lopen niet altijd in rechte lijnen; ze lopen over paden, vermijden modder en volgen de drukte.

HTI doet iets beters:
Het kijkt naar de foto's en leert niet alleen waar de reiziger was, maar ook hoe hij liep. Het leert de "wetten van de beweging" van die specifieke reiziger.

• Als de reiziger graag door drukke straten loopt (een "dichtheids-bias"), leert HTI dat.
• Als de reiziger de kortste, meest energiezuinige route neemt (het "minimale actie-principe"), leert HTI dat ook.

Met deze kennis kan HTI nu een voorspelling doen: "Als je de knop nu op 'voorzichtig' zet, zie ik de robot hier lopen." Je hoeft de robot niet opnieuw te trainen; je vraagt gewoon aan de tussenpersoon: "Hoe zou de robot zich gedragen bij instelling X?" en hij geeft je het antwoord direct.

Hoe werkt het technisch? (De "Lagrange" Magie)

De auteurs gebruiken een wiskundig concept genaamd Optimal Transport (Optimale Vervoer).

• Vroeger: Men dacht dat het verplaatsen van data (van instelling A naar B) als het verplaatsen van zandkorrels in een rechte lijn was.
• Nu (HTI): Ze gebruiken een Lagrangiaan. Dit is een fancy woord voor een formule die de "energie" van de beweging beschrijft.

Ze leren twee dingen tegelijk:

1. De Potentiële Energie (U): Waar zit de "drukte"? Waar lopen de meeste robots? HTI zorgt ervoor dat de voorspelling door deze drukke, veilige gebieden loopt, in plaats van door lege, gevaarlijke gebieden.
2. De Kinetic Energie (K): Hoe beweegt de robot? Loopt hij rechtuit of maakt hij bochten? HTI leert de onderliggende "kaart" van de wereld van de robot, zodat hij weet dat je niet zomaar door een muur kunt lopen.

Waarom is dit geweldig? (Voorbeelden uit het paper)

1. Kankerbehandeling:

   • Situatie: Een arts wil een behandeling die de tumor verkleint, maar ook het immuunsysteem (NK-cellen) beschermt. De balans hangt af van de patiënt.
   • HTI: In plaats van voor elke patiënt een nieuwe AI te trainen, trainen ze één model dat alle mogelijke balanspunten kent. De arts kan tijdens het gesprek zeggen: "Deze patiënt is ouder, wees iets voorzichtiger," en de AI past zich direct aan zonder wachttijd.

2. Weersvoorspelling (Kwantiel-regressie):

   • Situatie: Je wilt weten hoe warm het morgen wordt. Soms wil je een voorspelling met een groot veiligheidsmarge (het kan heel koud worden), soms met een klein marge.
   • HTI: In plaats van tien verschillende modellen te trainen voor tien verschillende veiligheidsmarges, leert HTI één model dat de hele "spectrum" van onzekerheid begrijpt. Je kunt elk willekeurig veiligheidsniveau kiezen en het model geeft direct het juiste antwoord.

3. Robots die leren (Reinforcement Learning):

   • Situatie: Een robotarm die een object pakt. Soms wil je dat hij snel is (maar misschien onnauwkeurig), soms dat hij heel precies is (maar langzaam).
   • HTI: Je kunt de robot tijdens het werk "live" aansturen. "Oké, nu snel, nu precies," en de robot schakelt direct om zonder opnieuw te hoeven leren.

Samenvatting in één zin

Dit paper introduceert een slimme methode om een AI-model te leren dat alle mogelijke versies van zichzelf begrijpt, zodat je de instellingen (knoppen) van de AI op elk gewenst moment kunt veranderen zonder dat je hem opnieuw hoeft te trainen, net als het hebben van een perfecte reisplanner die elke mogelijke route voor je heeft uitgestippeld.

Technische samenvatting

Probleemstelling

Neurale netwerken (NN) vertonen vaak kritieke gedragsafwegingen die tijdens het ontwerpstadium worden ingesteld via hyperparameters (bijvoorbeeld beloningstaken in versterkende leer of kwantiel-doelen in regressie). Eenmaal ingezet, kunnen gebruikersvoorkeuren of omstandigheden veranderen, waardoor de oorspronkelijke instellingen suboptimaal worden. Het opnieuw trainen van het netwerk voor elke nieuwe instelling is echter vaak te duur of onhaalbaar.

De auteurs introduceren de taak Hyperparameter Trajectory Inference (HTI). Het doel van HTI is om, op basis van waargenomen data, te leren hoe de conditionele uitvoerverdeling van een NN ($p_{\theta_\lambda}(y|x)$) verandert als functie van de hyperparameter $\lambda$. Hierdoor kan een surrogaatmodel worden geconstrueerd dat uitvoeringen schat voor niet-geobserveerde hyperparameter-instellingen, zonder dat het netwerk opnieuw getraind hoeft te worden.

De uitdaging ligt in het feit dat de dynamiek van hyperparameters ($\lambda \mapsto p_{\theta_\lambda}(y|x)$) vaak niet-lineair en complex is, mede door de optimisatielandschappen van deep learning. Simpele interpolatiemethoden (zoals Conditional Flow Matching) leveren vaak onrealistische of "onhaalbare" trajecten op. Bovendien moet deze methode rekening houden met conditionele variabelen ($x$), wat de complexiteit vergroot ten opzichte van standaard Trajectory Inference (TI).

Methodologie: Conditional Lagrangian Optimal Transport (CLOT)

De auteurs stellen een methode voor die is gebaseerd op Conditional Lagrangian Optimal Transport (CLOT). In plaats van lineaire paden te veronderstellen, modelleren ze de overgang tussen distributies als een fysiek systeem dat wordt gestuurd door een Lagrangiaan.

De kerncomponenten zijn:

1. Lagrangiaan en Actie:
   De kostenfunctie voor het transport van massa wordt gedefinieerd via het principe van minste actie. Voor een pad $q_t$ is de actie $S(q|x) = \int_0^1 L(q_t, \dot{q}_t|x) dt$, waarbij de Lagrangiaan $L$ bestaat uit kinetische en potentiële energie:
   $$L(q_t, \dot{q}_t|x) = K(q_t, \dot{q}_t|x) - U(q_t|x) = \frac{1}{2}\dot{q}_t^T G(q_t|x) \dot{q}_t - U(q_t|x)$$

   • $K$ (Kinetische Energie): Wordt bepaald door een metriek $G(q|x)$ die de geometrie van het onderliggende manifold definieert.
   • $U$ (Potentiële Energie): Wordt gebruikt om inductieve biases in te bouwen.

2. Inductieve Biases:

   • Dichte Traversering (via $U$): De potentiële energie $\hat{U}(q|x)$ wordt zo ontworpen dat deze hoog is in gebieden met lage data-dichtheid en laag in dichte gebieden (geschat via een Nadaraya-Watson schatter). Dit dwingt de geodeet (het optimale pad) om door dichte regio's van de data-manifold te gaan, wat de realiteit van het traject verbetert.
   • Principe van Minste Actie (via $K$): De kinetische term zorgt ervoor dat het systeem de meest efficiënte route kiest binnen de geleerde geometrie.

3. Neurale Implementatie:

   • Metriek Parametrisatie ($G_{\theta_G}$): In tegenstelling tot eerdere werken die vaak vaste of 2D-specifieke metrieken gebruikten, parametriseren de auteurs $G_{\theta_G}$ via een eigendecompositie ($G = R E R^T$). Een neurale netwerk leert zowel de rotatiematrices ($R$) als de positieve eigenwaarden ($E$). Dit voorkomt degeneratie (waar de metriek naar nul gaat) en maakt de methode schaalbaar naar hogere dimensies.
   • Gecombineerd Leren: De methode gebruikt een min-max procedure om zowel de metriek $G_{\theta_G}$ als de Kantorovich-potentialen (voor de transportkaarten) en de geodeet-paden (geschat met splines) gezamenlijk te leren.
   • Sampling: Na training kunnen samples voor een nieuwe hyperparameter $\lambda_{target}$ worden gegenereerd door te starten bij een bekende randdistributie, de geleerde transportkaart toe te passen, en het pad langs de geodeet te evalueren op het gewenste tijdstip.

Belangrijkste Bijdragen

1. Introductie van HTI: Een nieuwe taakdefinitie om neurale netwerken aan te passen aan veranderende hyperparameters tijdens de inferentie via surrogaatmodellen.
2. CLOT-methode voor CTI: Een algemene methode voor Conditional Trajectory Inference die complexe, niet-Euclidische dynamiek leert uit spaarzame waarnemingen. Dit omvat het leren van een data-afhankelijke potentiaalterm ($U$) en een expressieve, conditionele metriek ($G$).
3. Uitbreiding naar Hogere Dimensies: Een nieuwe parametrisatie van de Riemannse metriek die degeneratie voorkomt en toepasbaar is op complexe, hoogdimensionale ruimtes.
4. Empirische Validatie: Uitgebreide experimenten die aantonen dat de methode superieur is aan bestaande alternatieven (zoals Directe regressie, Flow Matching, en eerdere OT-methoden) in diverse domeinen.

Resultaten

De auteurs testen hun methode op vier verschillende scenario's:

1. Synthetisch Semicirkel-experiment: Toont aan dat de volledige methode (met zowel $K$ als $U$) de ware niet-Euclidische paden het beste reconstrueert, terwijl methoden zonder inductieve biases falen in het volgen van de kromming of het vermijden van lege ruimtes.
2. Versterkend Leren (Cancer Therapy & Reacher):
   • In een kankerbehandelingssimulatie (DTR-Bench) leert het model een surrogaatbeleid dat de afweging tussen tumorreductie en het behoud van immuuncellen (NK-cellen) dynamisch kan aanpassen via de hyperparameter $\lambda_{NK}$.
   • Het surrogaatmodel bereikt een hogere gemiddelde beloning dan baselines en vereist slechts ~15 minuten training, vergeleken met ~3,5 uur voor het opnieuw trainen van een PPO-agent per instelling.
   • Ook presteert de methode goed bij niet-lineaire beloningsfuncties.
3. Kwantielregressie (Time-series Forecasting): Op het ETTm2-dataset leert de methode de dynamiek tussen extreme kwantielen ($\tau=0.01$ en $\tau=0.99$) en schat ze nauwkeurig voor tussenliggende waarden. Dit resulteert in de laagste Mean Squared Error (MSE) vergeleken met andere surrogaatmodellen.
4. Generatieve Modellen (Dropout): De methode interpolatie succesvol over de dropout-hyperparameter in diffusion modellen, wat resulteert in de laagste Wasserstein-afstand (WD) tot de grondwaarheid.

Betekenis en Impact

Dit werk biedt een krachtig kader voor het overwinnen van de rigiditeit van neurale netwerken in dynamische omgevingen. Door HTI te koppelen aan Conditional Lagrangian Optimal Transport, kunnen gebruikers:

• Inferentie-tijd aanpassing: Het gedrag van een NN aanpassen aan nieuwe voorkeuren of omstandigheden zonder kostbare hertraining.
• Efficiëntie: De reductie in rekentijd is aanzienlijk (van uren naar minuten in de geteste scenario's).
• Flexibiliteit: Het model kan worden gebruikt voor complexe afwegingen in medische toepassingen, robotica en forecasting, waar de "optimale" instelling context-afhankelijk is.

De paper benadrukt dat hoewel de huidige methode beperkt is tot één continue hyperparameter, het een fundament legt voor toekomstig werk gericht op multi-dimensionale hyperparameter-ruimtes en chaotische dynamieken.