ATHENA: Accelerated Multi-Task Heterogeneous Influence Functions for Robot Data Curation

Technische Samenvatting: ATHENA

Probleemstelling

In het domein van robot-imitatieleer hebben Vision-Language-Action (VLA) modellen aanzienlijk potentieel getoond wanneer ze worden getraind op grootschalige demonstraties. Echter, de prestaties van deze modellen zijn kritisch afhankelijk van de datakwaliteit in plaats van enkel de schaal. Het naïef opschalen van demonstratiedata voor fine-tuning brengt vaak verbijsterende kosten met zich mee, terwijl het slechts beperkte prestatiewinsten oplevert of zelfs tot degradatie leidt.

De centrale uitdaging die dit werk aanpakt, is de efficiënte curatie van hoogwaardige demonstratiesubsets voor miljarden-parameter, multitask VLA-modellen. Bestaande methoden voor datavaluatie ondervinden twee primaire barrières bij toepassing op deze schaal:

1. Computationele Schaalbaarheid: Traditionele influence function-methoden vereisen per-sample gradiëntberekening ($O(DP)$) en dichte Hessian-inversie ($O(NP^2 + P^3)$), waarbij $D$ het aantal parameters is, $P$ de projectiedimensie en $N$ het aantal samples. Deze kosten zijn onhandelbaar voor miljarden-parameter modellen.
2. Multitask Imbalans: Het naïef uitbreiden van single-task invloedsattributie naar multitask-omgevingen resulteert vaak in een hebzuchtige selectie die de voorkeur geeft aan taken met sterkere gradiëntsignalen, wat leidt tot een ongebalanceerde datacuratie over de meer dan 50 gezamenlijk getrainde taken.

Methodologie: Het ATHENA-framework

De auteurs stellen ATHENA voor (Accelerated Multi-Task Heterogeneous Influence Functions), een framework ontworpen om de berekening van influence functions te schalen naar miljarden-parameter VLA-modellen, terwijl een gebalanceerde multitask-curatie wordt gewaarborgd.

1. Versnelde Invloedschatting

ATHENA pakt computationele knelpunten aan via twee belangrijke technische innovaties:

• Kronecker Compressed Gradient Featurization: In plaats van volledige parametergradiënten te materialiseren ($O(DP)$), maakt ATHENA gebruik van de Kronecker-structuur van lineaire-laag gradiënten. Door de activatie ($x$) en de teruggepropageerde fout ($\delta$) factoren afzonderlijk te projecteren voordat de buitenproductvorming plaatsvindt, reduceert het framework de projectiekosten van $O(DP)$ naar $O(\sqrt{DP})$ per laag. Dit vermijdt de geheugenbottleneck van het opslaan van volledige gradiënten voor miljarden-parameter modellen.
• Rank-$r$ Random Truncated Approximation (RTA): Om de $O(NP^2 + P^3)$ kosten van het inverteren van de dichte geprojecteerde Hessian-matrix ($G^\top G + \lambda I$) te vermijden, benadert ATHENA de gecomprimeerde gradiëntmatrix $G$ met behulp van de leidende gerandomiseerde spectrale componenten ($G \approx U_r \Sigma_r V_r^\top$). Dit reduceert de dominante datagebonden kosten van $O(NP^2)$ naar $O(NPr)$, waarbij $r \ll P$.

2. Multitask Influence Interaction (MII)

Om ongebalanceerde curatie over heterogene taken te voorkomen, introduceert ATHENA een Multitask Influence Interaction mechanisme. Het deelt de invloed van een demonstratie $i$ (behorend tot taak $c(i)$) op in twee componenten:

• Lokale Invloed ($\hat{\Psi}^{c(i)}_{\pi\text{-inf}}$): Het effect van de demonstratie op de evaluatie-rollouts van de eigen taak.
• Cross-Task Invloed ($\hat{\Psi}^{\text{all-}c(i)}_{\pi\text{-inf}}$): Het effect van de demonstratie op rollouts van andere taken.

De uiteindelijke curatiescore ($f^{MII}_i$) wordt geformuleerd als het product van genormaliseerde nutswaarden afgeleid van deze twee componenten. Dit zorgt ervoor dat de geselecteerde subset kritieke lokale voorbeelden behoudt terwijl er rekening wordt gehouden met cross-task interacties, waardoor de ROI over alle 50 taken wordt gebalanceerd.

Belangrijkste Resultaten

De auteurs evalueerden ATHENA op de RoboTwin 2.0 simulatie-benchmark (50 taken, 9,34 uur aan data) en zes real-robot ALOHA-taken (6,90 uur aan data).

• Computationele Efficiëntie: ATHENA bereikt een 313,4× versnelling in invloedsberekening vergeleken met een niet-geoptimaliseerde baseline, waardoor de tijd wordt teruggebracht van ~8.054 GPU-uren naar ~25,7 GPU-uren voor het scenario met 50 taken en 560k stappen.
• Simulatieprestaties (RoboTwin 2.0): Gebruikmakend van slechts 50% van de demonstraties, matchte ATHENA de prestaties van full-data joint fine-tuning onder clean evaluatie (43,36% vs. 43,42% succespercentage) en overtrof het de baseline onder randomized evaluatie (17,30% vs. 15,44%). Dit resulteerde in een cumulatieve verbetering van 45,0 punten over 50 taken vergeleken met full-data training.
• Real-Robot Prestaties: Op zes real-robot taken behaalde ATHENA een gemiddeld succespercentage van 68,0% met slechts 66,7% van de data. Dit presteerde beter dan de full-data joint fine-tuning baseline (60,0%) en overtrof de single-task fine-tuning baselines (46,7%) aanzienlijk.
• Transfereerbaarheid: De door ATHENA gegenereerde gecureerde subsets op het $\pi_0$ model toonden aan de prestaties van het $\pi_{0.5}$ model substantieel te verbeteren, wat aangeeft dat hoog-invloedrijke demonstraties generaliseren over verschillende VLA-varianten.

Betekenis en Claims

Het artikel beweert dat ATHENA een principieel, schaalbaar oplossingsmodel biedt voor datacuratie in miljarden-parameter multitask VLA fine-tuning. De betekenis ligt in:

1. Schaalbaarheid: Het is het eerste influence function framework dat in staat is te opereren op miljarden-parameter modellen door gebruik te maken van Kronecker-structuren en random truncatie, waarmee de voorheen onhaalbare computationele kosten worden overwonnen.
2. Multitask Balans: Het introduceert een nieuwe attributieregel die datacuratie balanceert over diverse taken, waardoor de hebzuchtige selectie-valkuilen van eerdere methoden worden vermeden.
3. Hoge ROI: Het framework demonstreert dat hoogwaardige datacuratie prestaties kan opleveren die gelijk aan of beter zijn dan full-data training, terwijl er aanzienlijk minder demonstraties nodig zijn, wat de computationele en opslagkosten geassocieerd met het fine-tunen van grootschalige robot-policies verlaagt.

De auteurs merken beperkingen op met betrekking tot de uitbreiding van deze methoden naar pretraining data-curatie en de hoge kosten van het verzamelen van real-robot rollouts voor invloedschatting, en identificeren dit als gebieden voor toekomstig onderzoek.