💻 computer science

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

The paper proposes ATHENA, a scalable influence function framework that accelerates data curation for billion-parameter multitask Vision-Language-Action models by leveraging Kronecker structures and rank-r approximations to achieve significant speedups while matching full-data performance with substantially fewer demonstrations.

Original authors: Tao Xu, Jiaxin Wang, Runhao Zhang, Jiayi Guan, Xianchao Zeng, Weixi Song, Xinyu Zhou, Zhetao Chen, Guang Chen, Yong-Lu Li

Published 2026-06-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tao Xu, Jiaxin Wang, Runhao Zhang, Jiayi Guan, Xianchao Zeng, Weixi Song, Xinyu Zhou, Zhetao Chen, Guang Chen, Yong-Lu Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a super-intelligent robot how to do a hundred different jobs, like stacking bowls, picking up fruit, or stamping envelopes. You have a massive library of video demonstrations showing humans doing these tasks.

The problem? The library is huge, and not every video is helpful. Some videos show perfect moves, while others show clumsy mistakes or irrelevant actions. If you just feed the robot all the videos, it gets confused, wastes time, and might even learn bad habits. But if you try to pick the "best" videos by watching them all manually, it would take forever.

This is where ATHENA comes in. Think of ATHENA as a super-smart, ultra-fast librarian that helps you curate the perfect training library for your robot.

Here is how it works, broken down into simple concepts:

1. The Problem: Too Much Data, Too Slow

In the past, trying to figure out which specific video helps the robot learn the most was like trying to count every grain of sand on a beach to find the one that makes the best sandcastle.

The Scale: Modern robot brains (called VLA models) are massive, with billions of "neurons" (parameters).
The Bottleneck: Traditional methods required re-training the robot's brain over and over again, removing one video at a time to see if the robot got better or worse. With billions of parameters, this is computationally impossible—it would take thousands of years.
The Multitask Mess: If you have 50 different tasks, a simple "pick the best video" approach often picks videos that are great for one task (like stacking bowls) but useless or even harmful for the others. It's like hiring a chef who is amazing at making sushi but terrible at baking bread, and then trying to teach them both at the same time.

2. The Solution: ATHENA's Two Superpowers

ATHENA solves these problems with two main tricks:

Trick A: The "Shortcut" (Accelerated Calculation)

Instead of calculating the impact of every single video by doing the heavy math on the entire billion-parameter brain, ATHENA uses a clever mathematical shortcut.

The Analogy: Imagine trying to measure the weight of a giant ship. A normal method would require weighing every single plank of wood individually. ATHENA, however, realizes the ship is built in a specific, repeating pattern (like a stack of identical bricks). Instead of weighing every plank, it weighs a few representative bricks and uses a formula to instantly know the total weight.
The Result: This shortcut makes the calculation 313 times faster. What used to take weeks of computer time now takes just a few hours.

Trick B: The "Balanced Judge" (Multitask Interaction)

Once ATHENA can calculate the scores quickly, it needs to decide which videos to keep.

The Analogy: Imagine a talent show with 50 different categories (singing, dancing, juggling, etc.). A bad judge might pick only the best singers because they have the loudest voices, leaving the jugglers out. ATHENA acts like a fair producer. It asks two questions for every video:
1. "How much does this video help the specific task it belongs to?" (Local Influence)
2. "How much does this video help the other 49 tasks?" (Global Influence)
The Result: It creates a balanced library where the robot learns a little bit of everything, ensuring it doesn't become a master of one thing and a failure at the rest.

3. The Results: Less Data, Better Performance

The researchers tested ATHENA in two ways:

In the Simulation (The Video Game): They used a robot simulator with 50 different tasks and 2,500 hours of video data.
- The Claim: ATHENA was able to train the robot using only 50% of the data (half the videos) and still get the same (or better) results as training with 100% of the data.
- The Metaphor: It's like a student who reads only half the textbook but gets a higher grade than the student who read the whole thing, because they studied the right pages.
On Real Robots (The Physical World): They tested on six real-world tasks (like picking fruit or wiping a board) using real robot arms.
- The Claim: Using only 66.7% of the data, ATHENA helped the real robots succeed more often than if they had used all the data or if they had tried to train each task separately.
- The Metaphor: It's like a coach who cuts out the boring, repetitive drills from a training camp, leaving only the high-impact exercises, resulting in a team that performs better in the actual game.

Summary

ATHENA is a tool that helps robot developers stop wasting time and money on bad data. By using mathematical shortcuts to speed up the process and a "fair judge" system to balance different tasks, it allows robots to learn faster and better using a smaller, higher-quality set of demonstrations.

Key Takeaway: You don't need more data to make a robot smarter; you need better data, and ATHENA is the tool that finds it.

Technical Summary: ATHENA

Problem Statement

In the realm of robot imitation learning, Vision-Language-Action (VLA) models have demonstrated significant potential when trained on large-scale demonstrations. However, the performance of these models is critically dependent on data quality rather than just scale. Naively scaling demonstration data for fine-tuning often incurs prohibitive costs while yielding limited performance gains or even degradation.

The central challenge addressed by this work is the efficient curation of high-quality demonstration subsets for billion-parameter, multitask VLA models. Existing data valuation methods face two primary barriers when applied to this scale:

Computational Scalability: Traditional influence function methods require per-sample gradient computation ($O(DP)$) and dense Hessian inversion ( $O(NP^2 + P^3)$ ), where $D$ is the number of parameters, $P$ is the projection dimension, and $N$ is the number of samples. These costs are intractable for billion-parameter models.
Multitask Imbalance: Naively extending single-task influence attribution to multitask settings often results in greedy selection that favors tasks with stronger gradient signals, leading to imbalanced data curation across the 50+ jointly trained tasks.

Methodology: The ATHENA Framework

The authors propose ATHENA (Accelerated Multi-Task Heterogeneous Influence Functions), a framework designed to scale influence function computation to billion-parameter VLA models while ensuring balanced multitask curation.

1. Accelerated Influence Estimation

ATHENA addresses computational bottlenecks through two key technical innovations:

Kronecker Compressed Gradient Featurization: Instead of materializing full parameter gradients ($O(DP)$), ATHENA exploits the Kronecker structure of linear-layer gradients. By projecting the activation ( $x$ ) and backpropagated error ( $\delta$ ) factors separately before forming the outer product, the framework reduces the projection cost from $O(DP)$ to $O(\sqrt{DP})$ per layer. This avoids the memory bottleneck of storing full gradients for billion-parameter models.
Rank- $r$ Random Truncated Approximation (RTA): To avoid the $O(NP^2 + P^3)$ cost of inverting the dense projected Hessian matrix ( $G^\top G + \lambda I$ ), ATHENA approximates the compressed gradient matrix $G$ using its leading randomized spectral components ( $G \approx U_r \Sigma_r V_r^\top$ ). This reduces the dominant data-dependent cost from $O(NP^2)$ to $O(NPr)$, where $r \ll P$ .

2. Multitask Influence Interaction (MII)

To prevent imbalanced curation across heterogeneous tasks, ATHENA introduces a Multitask Influence Interaction mechanism. It decomposes the influence of a demonstration $i$ (belonging to task $c(i)$ ) into two components:

Local Influence ( $\hat{\Psi}^{c(i)}_{\pi\text{-inf}}$ ): The effect of the demonstration on its own task's evaluation rollouts.
Cross-Task Influence ( $\hat{\Psi}^{\text{all-}c(i)}_{\pi\text{-inf}}$ ): The effect of the demonstration on rollouts from other tasks.

The final curation score ( $f^{MII}_i$ ) is formulated as the product of normalized utilities derived from these two components. This ensures that the selected subset preserves locally critical examples while accounting for cross-task interactions, thereby balancing the ROI across all 50 tasks.

Key Results

The authors evaluated ATHENA on the RoboTwin 2.0 simulation benchmark (50 tasks, 9.34 hours of data) and six real-robot ALOHA tasks (6.90 hours of data).

Computational Efficiency: ATHENA achieves a 313.4× speedup in influence computation compared to an unoptimized baseline, reducing the time from ~8,054 GPU-hours to ~25.7 GPU-hours for the 50-task, 560k-step scenario.
Simulation Performance (RoboTwin 2.0): Using only 50% of the demonstrations, ATHENA matched the performance of full-data joint fine-tuning under clean evaluation (43.36% vs. 43.42% success rate) and exceeded it under randomized evaluation (17.30% vs. 15.44%). This resulted in a cumulative 45.0-point improvement across 50 tasks compared to full-data training.
Real-Robot Performance: On six real-robot tasks, ATHENA achieved a 68.0% average success rate using only 66.7% of the data. This outperformed the full-data joint fine-tuning baseline (60.0%) and significantly surpassed single-task fine-tuning baselines (46.7%).
Transferability: Curated subsets generated by ATHENA on the $\pi_0$ model were shown to substantially enhance the performance of the $\pi_{0.5}$ model, indicating that high-influence demonstrations generalize across VLA model variants.

Significance and Claims

The paper claims that ATHENA provides a principled, scalable solution for data curation in billion-parameter multitask VLA fine-tuning. Its significance lies in:

Scalability: It is the first influence function framework capable of operating on billion-parameter models by leveraging Kronecker structures and random truncation, overcoming previous prohibitive computational costs.
Multitask Balance: It introduces a novel attribution rule that balances data curation across diverse tasks, avoiding the greedy selection pitfalls of previous methods.
High ROI: The framework demonstrates that high-quality data curation can yield performance equal to or better than full-data training while using significantly fewer demonstrations, thereby reducing the computational and storage costs associated with fine-tuning large-scale robot policies.

The authors note limitations regarding the extension of these methods to pretraining data curation and the high cost of collecting real-robot rollouts for influence estimation, identifying these as areas for future work.