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

기술 요약: ATHENA

문제 정의

로봇 모방 학습의 영역에서 시각-언어-행동(Vision-Language-Action, VLA) 모델은 대규모 데모니스트레이션 학습을 통해 상당한 잠재력을 입증해 왔습니다. 그러나 이러한 모델의 성능은 단순히 데이터의 규모보다는 데이터의 품질에 결정적으로 의존합니다. 파인튜닝을 위해 데모니스트레이션 데이터를 무작정 확장하는 것은 막대한 비용을 초 {비용을 발생시키면서도 성능 향상은 제한적이거나 오히려 성능 저하를 초래할 수 있습니다.

본 연구가 다루는 핵심 과제는 수십억 개의 파라미터를 가진 멀티태스크 VLA 모델을 위한 고품질 데모니스트레이션 서브셋을 효율적으로 큐레이션하는 것입니다. 기존의 데이터 가치 평가 방법들은 이 정도 규모에 적용될 때 두 가지 주요 장벽에 직면합니다:

1. 계산적 확장성: 전통적인 영향력 함수(influence function) 방법은 샘플당 그래디언트 계산($O(DP)$)과 조밀한 헤시안 역행렬 계산($O(NP^2 + P^3)$)을 요구하며, 여기서$D$는 파라미터 수,$P$는 투영 차원,$N$은 샘플 수입니다. 이러한 비용은 수십억 개의 파라미터를 가진 모델에는 감당할 수 없는 수준입니다.
2. 멀티태스크 불균형: 단일 태스크 영향력 귀속(influence attribution)을 멀티태스크 설정으로 단순히 확장하면, 그래디언트 신호가 강한 태스크를 선호하는 탐욕적 선택(greedy selection)이 발생하여, 공동으로 학습되는 50개 이상의 태스크 간 데이터 큐레이션 불균형을 초래합니다.

방법론: ATHENA 프레임워크

저자들은 ATHENA(Accelerated Multi-Task Heterogeneous Influence Functions)를 제안합니다. 이 프레임워크는 멀티태스크 큐레이션의 균형을 보장하면서 수십억 파라미터 VLA 모델에 대한 영향력 함수 계산을 확장할 수 있도록 설계되었습니다.

1. 가속화된 영향력 추정

ATHENA는 두 가지 핵심 기술 혁신을 통해 계산 병목 현상을 해결합니다:

• 크로네커 압축 그래디언트 특징 추출(Kronecker Compressed Gradient Featurization): 전체 파라미터 그래디언트($O(DP)$)를 직접 생성하는 대신, ATHENA는 선형 레이어 그래디언트의 크로네커 구조를 활용합니다. 외적(outer product)을 형성하기 전 활성화($x$)와 역전파 오차($\delta$) 요소를 각각 분리하여 투영함으로써, 프레임워크는 레이어당 투영 비용을 $O(DP)$에서$O(\sqrt{DP})$로 줄입니다. 이를 통해 수십억 파라미터 모델의 전체 그래디언트 저장에 따른 메모리 병목 현상을 방지합니다.
• Rank-$r$ 무작위 절단 근사(Rank-$r$ Random Truncated Approximation, RTA): 조밀한 투영 헤시안 행렬($G^\top G + \lambda I$)의 역행렬을 구하는 $O(NP^2 + P^3)$ 비용을 피하기 위해, ATHENA는 압축된 그래디언트 행렬 $G$를 주요 무작위 스펙트럼 성분($G \approx U_r \Sigma_r V_r^\top$)을 사용하여 근사합니다. 이는 지배적인 데이터 의존적 비용을 $O(NP^2)$에서 $r \ll P$인 $O(NPr)$로 감소시킵니다.

2. 멀티태스크 영향력 상호작용(Multitask Influence Interaction, MII)

이질적인 태스크 간의 불균형한 큐레이션을 방지하기 위해, ATHENA는 멀티태스크 영향력 상호작용 메커니즘을 도입합니다. 이는 데모니스트레이션 $i$(태스크 $c(i)$에 속함)의 영향력을 두 가지 구성 요소로 분해합니다:

• 로컬 영향력($\hat{\Psi}^{c(i)}_{\pi\text{-inf}}$): 해당 데모니스트레이션이 자신의 태스크에 대한 평가 롤아웃(evaluation rollouts)에 미치는 영향.
• 교차 태스크 영향력($\hat{\Psi}^{\text{all-}c(i)}_{\pi\text{-inf}}$): 해당 데모니스트레이션이 다른 태스크들의 롤아웃에 미치는 영향.

최종 큐레이션 점수($f^{MII}_i$)는 이 두 구성 요소에서 유도된 정규화된 효용(utility)의 곱으로 공식화됩니다. 이를 통해 선택된 서브셋이 로컬하게 중요한 예시를 보존하는 동시에 교차 태스크 상호작용을 고려하도록 하여, 50개 모든 태스크에 대한 투자 대비 수익(ROI)의 균형을 맞춥니다.

주요 결과

저자들은 RoboTwin 2.0 시뮬레이션 벤치마크(50개 태스크, 9.34시간의 데이터)와 6개의 실제 로봇 ALOHA 태스크(6.90시간의 데이터)에서 ATHANA를 평가했습니다.

• 계산 효율성: ATHENA는 최적화되지 않은 베이스라인에 비해 영향력 계산에서 313.4배의 속도 향상을 달성하였으며, 50개 태스크, 560k-스텝 시나리오에서 약 8,054 GPU-시간을 약 25.7 GPU-시간으로 단축했습니다.
• 시뮬레이션 성능 (RoboTwin 2.0): 데모니스트레이션의 **50%**만을 사용했을 때, ATHENA는 클린 평가(clean evaluation) 조건에서 전체 데이터 공동 파인튜닝의 성능과 일치하였고(성공률 43.36% vs 43.42%), 무작위 평가(randomized evaluation) 조건에서는 이를 능가하였습니다(17.30% vs 15.44%). 이는 전체 데이터 학습 대비 50개 태스크 전체에서 누적 45.0포인트의 개선을 가져왔습니다.
• 실제 로봇 성능: 6개의 실제 로봇 태스크에서 ATHENA는 단 **66.7%**의 데이터만을 사용하여 **68.0%**의 평균 성공률을 달성했습니다. 이는 전체 데이터 공동 파인튜닝 베이스라인(60.0%)보다 우수하며, 단일 태스크 파인튜닝 베이스라인(46.7%)을 크게 상회하는 수치입니다.
• 전이 가능성(Transferability): $\pi_0$ 모델에서 생성된 큐레이션된 서브셋이 $\pi_{0.5}$ 모델의 성능을 실질적으로 향상시킨다는 것을 보여주었으며, 이는 영향력이 높은 데모니스트레이션이 다양한 VLA 모델 변체 간에 일반화될 수 있음을 나타냅니다.

의의 및 주장

본 논문은 ATHENA가 수십억 파라미터 멀티태스크 VLA 파인튜닝을 위한 데이터 큐레이션에 있어 원칙적이고 확장 가능한 솔루션을 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

1. 확장성: 크로네커 구조와 무작위 절단을 활용함으로써, 이전의 막대한 계산 비용을 극복하고 수십억 파라미터 모델에서 작동할 수 있는 최초의 영향력 함수 프레임워크입니다.
2. 멀티태스크 균형: 다양한 태스크 간의 데이터 큐레이션 균형을 맞추는 새로운 귀속 규칙을 도입하여, 기존 방법들의 탐욕적 선택 문제를 피했습니다.
3. 높은 ROI: 고품질 데이터 큐레이션이 훨씬 적은 양의 데모니스트레이션을 사용하면서도 전체 데이터 학습과 대등하거나 더 나은 성능을 낼 수 있음을 입증함으로써, 대규모 로봇 정책 파인튜닝과 관련된 계산 및 저장 비용을 줄일 수 있습니다.

저자들은 이러한 방법들을 프리트레이닝(pretraining) 데이터 큐레이션으로 확장하는 문제와 영향력 추정을 위한 실제 로봇 롤아웃 수집의 높은 비용에 대한 한계를 언급하며, 이를 향후 연구 과제로 식별하였습니다.