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

Riepilogo Tecnico: ATHENA

Problematica

Nel campo dell'apprendimento per imitazione robotica, i modelli Vision-Language-Action (VLA) hanno dimostrato un potenziale significativo quando addestrati su dimostrazioni su larga scala. Tuttavia, le prestazioni di questi modelli dipendono criticamente dalla qualità dei dati piuttosto che solo dalla scala. L'aumento ingenuo dei dati di dimostrazione per il fine-tuning spesso comporta costi proibitivi pur offrendo guadagni di prestazione limitati o addirittura una degradazione delle stesse.

La sfida centrale affrontata da questo lavoro è la cura efficiente di sottoinsiemi di dimostrazioni di alta qualità per modelli VLA multitask con miliardi di parametri. I metodi esistenti di valutazione dei dati affrontano due barriere primarie quando applicati a questa scala:

1. Scalabilità Computazionale: I metodi tradizionali delle funzioni di influenza richiedono il calcolo del gradiente per ogni campione ($O(DP)$) e l'inversione della matrice Hessiana densa ($O(NP^2 + P^3)$), dove $D$ è il numero di parametri, $P$ è la dimensione di proiezione e $N$ è il numero di campioni. Questi costi sono intrattabili per modelli con miliardi di parametri.
2. Squilibrio Multitask: L'estensione ingenua dell'attribuzione di influenza single-task ai contesti multitask spesso risulta in una selezione avida che favorisce i task con segnali di gradiente più forti, portando a una cura dei dati sbilanciata tra i più di 50 task addestrati congiuntamente.

Metodologia: Il Framework ATHENA

Gli autori propongono ATHENA (Accelerated Multi-Task Heterogeneous Influence Functions), un framework progettato per scalare il calcolo delle funzioni di influenza a modelli VLA con miliardi di parametri, garantendo al contempo una cura multitask bilanciata.

1. Stima dell'Influenza Accelerata

ATHENA affronta i colli di bottiglia computazionali attraverso due innovazioni tecniche chiave:

• Featurizzazione del Gradiente Compressa tramite Kronecker: Invece di materializzare i gradienti completi dei parametri ($O(DP)$), ATHENA sfrutta la struttura di Kronecker dei gradienti degli strati lineari. Proiettando separatamente i fattori di attivazione ($x$) ed errore retropropagato ($\delta$) prima di formare il prodotto esterno, il framework riduce il costo di proiezione da $O(DP)$a$O(\sqrt{DP})$ per strato. Ciò evita il collo di bottiglia della memoria associato allo stoccaggio dei gradienti completi per modelli con miliardi di parametri.
• Approssimazione Truncata Casuale Rank-$r$ (RTA): Per evitare il costo $O(NP^2 + P^3)$ dell'inversione della matrice Hessiana densa proiettata ($G^\top G + \lambda I$), ATHENA approssima la matrice del gradiente compressa $G$ utilizzando le sue componenti spettrali casuali dominanti ($G \approx U_r \Sigma_r V_r^\top$). Ciò riduce il costo dominante dipendente dai dati da $O(NP^2)$ a $O(NPr)$, dove $r \ll P$.

2. Interazione di Influenza Multitask (MII)

Per prevenire una cura sbilanciata tra task eterogenei, ATHENA introduce un meccanismo di Interazione di Influenza Multitask. Esso decompone l'influenza di una dimostrazione $i$ (appartenente al task $c(i)$) in due componenti:

• Influenza Locale ($\hat{\Psi}^{c(i)}_{\pi\text{-inf}}$): L'effetto della dimostrazione sulle proprie rollout di valutazione del task.
• Influenza Cross-Task ($\hat{\Psi}^{\text{all-}c(i)}_{\pi\text{-inf}}$): L'effetto della dimostrazione sulle rollout di altri task.

Il punteggio finale di cura ($f^{MII}_i$) è formulato come il prodotto di utilità normalizzate derivate da queste due componenti. Ciò assicura che il sottoinsieme selezionato preservi gli esempi localmente critici tenendo conto delle interazioni cross-task, bilanciando così il ROI su tutti i 50 task.

Risultati Chiave

Gli autori hanno valutato ATHENA sul benchmark di simulazione RoboTwin 2.0 (50 task, 9,34 ore di dati) e su sei task di robot reali ALOHA (6,90 ore di dati).

• Efficienza Computazionale: ATHENA ottiene un incremento di velocità di 313,4× nel calcolo dell'influenza rispetto a un baseline non ottimizzato, riducendo il tempo da circa 8.054 ore-GPU a circa 25,7 ore-GPU per lo scenario con 50 task e 560k step.
• Prestazioni in Simulazione (RoboTwin 2.0): Utilizzando solo il 50% delle dimostrazioni, ATHENA ha eguagliato le prestazioni del fine-tuning congiunto su dati completi sotto valutazione pulita (tasso di successo del 43,36% vs 43,42%) e le ha superate sotto valutazione randomizzata (17,30% vs 15,44%). Ciò ha portato a un miglioramento cumulativo di 45,0 punti attraverso 50 task rispetto all'addestramento con dati completi.
• Prestazioni su Robot Reali: Su sei task di robot reali, ATHENA ha raggiunto un tasso di successo medio del 68,0% utilizzando solo il 66,7% dei dati. Questo ha superato il baseline del fine-tuning congiunto su dati completi (60,0%) e ha superato significativamente i baseline del fine-tuning single-task (46,7%).
• Trasferibilità: Si è dimostrato che i sottoinsiemi curati da ATHENA sul modello $\pi_0$ migliorano sostanzialmente le prestazioni del modello $\pi_{0.5}$, indicando che le dimostrazioni ad alta influenza generalizzano tra diverse varianti di modelli VLA.

Significato e Rivendicazioni

L'articolo sostiene che ATHENA fornisca una soluzione scalabile e fondata per la cura dei dati nel fine-tuning multitask di modelli VLA con miliardi di parametri. La sua importanza risiede in:

1. Scalabilità: È il primo framework di funzioni di influenza capace di operare su modelli con miliardi di parametri sfruttando le strutture di Kronecker e la troncatura casuale, superando i precedenti costi computazionali proibitivi.
2. Bilanciamento Multitask: Introduce una nuova regola di attribuzione che bilancia la cura dei dati tra diversi task, evitando le trappole della selezione avida dei metodi precedenti.
3. Alto ROI: Il framework dimostra che una cura dei dati di alta qualità può produrre prestazioni uguali o superiori all'addestramento con dati completi, utilizzando significativamente meno dimostrazioni, riducendo così i costi computazionali e di archiviazione associati al fine-tuning di policy robotiche su larga scala.

Gli autori segnalano i limiti riguardanti l'estensione di questi metodi alla cura dei dati di pre-training e l'alto costo della raccolta di rollouts di robot reali per la stima dell'influenza, identificandoli come aree per il lavoro futuro.