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

Resumo Técnico: ATHENA

Declaração do Problema

No domínio do aprendizado por imitação robótica, modelos de Visão-Linguagem-Ação (VLA) têm demonstrado um potencial significativo quando treinados em demonstrações de larga escala. No entanto, o desempenho desses modelos é criticamente dependente da qualidade dos dados, e não apenas da escala. Escalar ingenuamente os dados de demonstração para o ajuste fino (fine-tuning) frequentemente acarreta custos proibitivos, gerando ganhos de desempenho limitados ou até mesmo degradação.

O desafio central abordado por este trabalho é a curadoria eficiente de subconjuntos de demonstrações de alta qualidade para modelos VLA multitarefa com bilhões de parâmetros. Os métodos existentes de avaliação de dados enfrentam duas barreiras primárias quando aplicados a esta escala:

1. Escalabilidade Computacional: Métodos tradicionais de função de influência exigem o cálculo do gradiente por amostra ($O(DP)$) e a inversão de Hessiana densa ($O(NP^2 + P^3)$), onde $D$ é o número de parâmetros, $P$ é a dimensão de projeção e $N$ é o número de amostras. Esses custos são intratáveis para modelos de bilhões de parâmetros.
2. Desequilíbrio Multitarefa: A extensão ingênua da atribuição de influência de tarefa única para configurações multitarefa frequentemente resulta em uma seleção gananciosa que favorece tarefas com sinais de gradiente mais fortes, levando a uma curadoria de dados desequilibrada entre as mais de 50 tarefas treinadas conjuntamente.

Metodologia: O Framework ATHENA

Os autores propõem o ATHENA (Accelerated Multi-Task Heterogeneous Influence Functions), um framework projetado para escalar o cálculo da função de influência para modelos VLA de bilhões de parâmetros, garantindo simultaneamente uma curadoria multitarefa equilibrada.

1. Estimativa de Influência Acelerada

O ATHENA aborda os gargalos computacionais através de duas inovações técnicas principais:

• Featurização de Gradiente Comprimida por Kronecker: Em vez de materializar gradientes de parâmetros completos ($O(DP)$), o ATHENA explora a estrutura de Kronecker dos gradientes das camadas lineares. Ao projetar os fatores de ativação ($x$) e o erro retropropagado ($\delta$) separadamente antes de formar o produto externo, o framework reduz o custo de projeção de $O(DP)$ para $O(\sqrt{DP})$ por camada. Isso evita o gargalo de memória do armazenamento de gradientes completos para modelos de bilhões de parâmetros.
• Aproximação de Truncamento Aleatório de Rank-$r$ (RTA): Para evitar o custo de $O(NP^2 + P^3)$ da inversão da matriz Hessiana projetada densa ($G^\top G + \lambda I$), o ATHENA aproxima a matriz de gradiente comprimida $G$ usando seus componentes espectrais aleatórios líderes ($G \approx U_r \Sigma_r V_r^\top$). Isso reduz o custo dominante dependente dos dados de $O(NP^2)$ para $O(NPr)$, onde $r \ll P$.

2. Interação de Influência Multitarefa (MII)

Para evitar a curadoria desequilibrada entre tarefas heterogêneas, o ATHENA introduz um mecanismo de Interação de Influência Multitarefa. Ele decompõe a influência de uma demonstração $i$ (pertencente à tarefa $c(i)$) em dois componentes:

• Influência Local ($\hat{\Psi}^{c(i)}_{\pi\text{-inf}}$): O efeito da demonstração em suas próprias execções (rollouts) de avaliação de tarefa.
• Influência entre Tarefas ($\hat{\Psi}^{\text{all-}c(i)}_{\pi\text{-inf}}$): O efeito da demonstração em execuções de outras tarefas.

O escore de curadoria final ($f^{MII}_i$) é formulado como o produto de utilidades normalizadas derivadas desses dois componentes. Isso garante que o subconjunto selecionado preserve exemplos localmente críticos enquanto considera as interações entre tarefas, equilibrando assim o ROI em todas as 50 tarefas.

Principais Resultados

Os autores avaliaram o ATHENA no benchmark de simulação RoboTwin 2.0 (50 tarefas, 9,34 horas de dados) e em seis tarefas de robô real ALOHA (6,90 horas de dados).

• Eficiência Computacional: O ATHENA alcança uma aceleração de 313,4× no cálculo de influência em comparação com um baseline não otimizado, reduzindo o tempo de ~8.054 horas de GPU para ~25,7 horas de GPU para o cenário de 50 tarefas e 560k passos.
• Desempenho de Simulação (RoboTwin 2.0): Usando apenas 50% das demonstrações, o ATHENA igualou o desempenho do ajuste fino conjunto de dados completos sob avaliação limpa (43,36% vs. 43,42% de taxa de sucesso) e o superou sob avaliação randomizada (17,30% vs. 15,44%). Isso resultou em uma melhoria cumulativa de 45,0 pontos em 50 tarefas comparado ao treinamento com dados completos.
• Desempenho em Robô Real: Nas seis tarefas de robô real, o ATHENA alcançou uma taxa média de sucesso de 68,0% usando apenas 66,7% dos dados. Isso superou o baseline de ajuste fino conjunto de dados completos (60,0%) e superou significativamente os baselines de ajuste fino de tarefa única (46,7%).
• Transferibilidade: Subconjuntos curados pelo ATHENA no modelo $\pi_0$ mostraram aumentar substancialmente o desempenho do modelo $\pi_{0.5}$, indicando que demonstrações de alta influência generalizam entre variantes de modelos VLA.

Significância e Alegações

O artigo afirma que o ATHENA fornece uma solução principista e escalável para a curadoria de dados no ajuste fino de modelos VLA multitarefa com bilhões de parâmetros. Sua significância reside em:

1. Escalabilidade: É o primeiro framework de função de influência capaz de operar em modelos de bilhões de parâmetros ao aproveitar estruturas de Kronecker e truncamento aleatório, superando custos computacionais anteriormente proibitivos.
2. Equilíbrio Multitarefa: Introduz uma regra de atribuição inovadora que equilibra a curadoria de dados entre diversas tarefas, evitando as armadilhas de seleção gananciosa de métodos anteriores.
3. Alto ROI: O framework demonstra que a curadoria de dados de alta qualidade pode gerar um desempenho igual ou superior ao treinamento com dados completos, utilizando significativamente menos demonstrações, reduzindo assim os custos computacionais e de armazenamento associados ao ajuste fino de políticas robóticas de larga escala.

Os autores observam limitações quanto à extensão destes métodos para a curadoria de dados de pré-treinamento e o alto custo de coleta de execuções de robôs reais para estimativa de influência, identificando estes como áreas para trabalhos futuros.