AEGIS: Anchor-Enforced Gradient Isolation for Knowledge-Preserving Vision-Language-Action Fine-Tuning

1. 문제 정의: 교차 모달 그라디언트 비대칭성 (Cross-Modal Gradient Asymmetry)

기존의 사전 훈련된 비전 - 언어 모델 (VLM) 을 로봇 제어용 비전 - 언어 - 액션 (VLA) 모델로 미세 조정 (Fine-tuning) 할 때 발생하는 근본적인 문제를 지적합니다.

• 기하학적 불일치: VLM 은 방대한 양의 텍스트 - 이미지 쌍으로 교차 엔트로피 (Cross-Entropy, CE) 손실 함수를 통해 훈련되어 고차원의 의미론적 다양성 (Semantic Manifold) 을 형성합니다. 반면, 로봇 제어 작업은 흐름 매칭 (Flow-Matching) 기반의 연속적인 액션 예측을 위해 평균 제곱 오차 (MSE) 손실 함수를 사용하며, 이는 저랭크 (Low-rank) 이고 고강도의 그라디언트를 생성합니다.
• 파괴적 덮어쓰기: MSE 그라디언트가 VLM 백본으로 역전파될 때, 표준 옵티마이저는 이 좁고 강력한 물리적 그라디언트를 네트워크 전체에 무차별적으로 적용합니다. 이로 인해 VLM 이 가진 정교한 시각 - 언어 표현 (Visual-Question-Answering, VQA 능력) 이 파괴적으로 덮어쓰여 재학습 (Forgetting) 이 발생합니다.
• 기존 해결책의 한계:
  • Stop-Gradient: VLM 백본으로의 그라디언트 흐름을 차단하여 VQA 는 보존하지만, 연속적인 액션 감독 신호를 잃고 이산적 토큰 (Discrete Tokens) 에만 의존하게 되어 표현력이 제한됩니다.
  • LoRA (Low-Rank Adaptation): 파라미터 업데이트의 용량만 제한할 뿐, 그라디언트의 방향 (기하학) 을 제어하지 못합니다. 따라서 저랭크 서브스페이스 내에서도 파괴적인 업데이트가 발생하여 VQA 성능이 서서히 저하됩니다.
  • 공훈련 (Co-training): VQA 데이터를 섞어 학습시키는 방식은 연산 비용을 2 배로 증가시키고 원본 데이터에 대한 접근이 필요합니다.

2. 방법론: AEGIS (Anchor-Enforced Gradient Isolation System)

AEGIS 는 재학습 버퍼 (Replay Buffer) 나 추가 데이터 없이, 그라디언트 방향을 기하학적으로 제어하여 지식 보존을 달성하는 프레임워크입니다.

핵심 구성 요소

1. 정적 Wasserstein 앵커 (Static Wasserstein Anchor):

   • 미세 조정 전, 사전 훈련된 VLM 에 VQA 데이터를 입력하여 모든 트랜스포머 레이어의 활성화 통계량 (평균 $\mu$, 분산 $\sigma^2$) 을 계산합니다.
   • 이는 사전 훈련된 의미론적 다양성 (Manifold) 을 정의하는 Wasserstein-2 공간의 기준점 역할을 합니다.
   • 패딩 토큰을 제외하고 이미지와 텍스트 토큰 전체를 사용하여 완전한 의미론적 지표를 확보합니다.

2. Wasserstein-2 수송 페널티 (Wasserstein-2 Transport Penalty):

   • 각 학습 단계에서 현재 모델의 활성화 통계량과 사전 정의된 앵커 간의 Wasserstein-2 거리 ($W_2$) 를 계산합니다.
   • Bures Metric을 사용하여 평균 이동 (Mean Shift) 과 표준편차 불일치 (Std Mismatch) 를 분리하여 계산합니다.
   • 이 페널티는 실제 손실 함수에 추가되지 않으며, 앵커로 복원하는 방향의 그라디언트 ($g_{ot}$) 를 생성하는 데만 사용됩니다.

3. 레이어별 직교 그라디언트 투영 (Layer-wise Orthogonal Gradient Projection):

   • 이중 역전파 (Dual-Backward): 작업 손실 (MSE) 과 앵커 페널티 ($W_2$) 에 대해 순차적으로 역전파를 수행하여 각각의 그라디언트 ($g_{task}, g_{ot}$) 를 추출합니다.
   • Gram-Schmidt 직교화: 각 트랜스포머 레이어별로 두 그라디언트를 분석합니다.
     • $g_{task}$가 $g_{ot}$와 음의 상관관계 (파괴적 간섭) 를 보일 때, $g_{task}$에서 $g_{ot}$ 방향의 성분을 제거합니다.
     • 수식: $g_{final} = g_{task} - \alpha g_{ot}$ (여기서 $\alpha$는 투영 계수).
   • 레이어 단위 granularity: 전체 모델 단위나 텐서 단위 투영의 한계 (영합 문제, 어텐션 헤드 동기화 파괴) 를 피하기 위해 레이어 단위로 투영을 수행하여 트랜스포머 블록 내부 구조를 보존합니다.

3. 주요 기여 (Contributions)

1. 근본 원인 규명: VLA 미세 조정 중 발생하는 VQA 성능 저하의 원인이 '교차 모달 그라디언트 비대칭성' (고차원 CE 그라디언트와 저랭크 MSE 그라디언트의 기하학적 불일치) 임을 증명했습니다.
2. AEGIS 프레임워크 제안: 재학습 버퍼나 공훈련 데이터 없이, 사전 훈련된 활성화 다양성 위에 Wasserstein-2 앵커를 구축하고 레이어별 직교 투영을 수행하여 파괴적 그라디언트 간섭만 제거하는 새로운 방법을 제시했습니다.
3. 기하학적 통찰: LoRA 와 같은 기존 방법론이 실패하는 이유는 파라미터 용량 제한이 아니라 그라디언트 방향 제어의 부재임을 입증했습니다.

4. 실험 결과 (Results)

실험은 PaliGemma2-3B 모델과 LIBERO 로봇 조작 데이터셋을 사용하여 수행되었습니다.

• VQA 성능 보존:
  • Naive Fine-tuning: 1,500 스텝 후 VQA 손실이 급격히 증가하여 심각한 재학습 발생.
  • LoRA: 저하 속도는 늦추지만 여전히 VQA 손실이 지속적으로 증가.
  • AEGIS: 사전 훈련된 VQA 손실 수준을 거의 완벽하게 유지 (Naive 대비 2 차수, LoRA 대비 1 차수 개선). OK-VQA OOD 평가에서도 사전 훈련 모델과 유사한 정확도 (60.23% vs 60.15%) 를 기록했습니다.
• 작업 학습 수렴:
  • AEGIS 는 VQA 를 보존하면서도 Flow-Matching 액션 학습 수렴 속도는 Naive 파인튜닝과 유사하게 유지했습니다. 직교 투영이 작업 학습에 미치는 부정적 영향은 미미했습니다.
• 그라디언트 역학 분석:
  • 투영 효율성: 평균적으로 그라디언트 에너지의 0.62% 미만만 제거되었습니다. 이는 파괴적 간섭이 전체 그라디언트 중 매우 얇은 부분임을 의미하며, AEGIS 가 매우 정밀한 '외과적' 개입을 수행함을 보여줍니다.
  • 간섭 빈도: 약 51% 의 레이어에서 파괴적 간섭이 발생했으나, AEGIS 는 이를 정확히 식별하여 제거했습니다.

5. 의의 및 결론

• 지식 보존의 새로운 패러다임: AEGIS 는 데이터 재사용이나 모델 구조 변경 없이 그라디언트 기하학 (Gradient Geometry) 을 직접 조작하여 지식 보존 문제를 해결했습니다.
• 실용성: 재학습 버퍼나 추가 데이터 없이도 VLA 모델을 훈련할 수 있어, 대규모 로봇 학습 환경에서 연산 비용과 데이터 접근성 문제를 해결합니다.
• LoRA 와의 호환성: AEGIS 는 파라미터 용량을 제한하지 않으므로, LoRA 와 결합하여 용량과 방향을 동시에 제어하는 하이브리드 접근법으로도 확장 가능합니다.

결론적으로, AEGIS 는 VLM 기반 로봇 제어 모델이 새로운 물리적 작업을 학습하면서도 기존 언어 - 시각 능력을 잃지 않도록 하는 기하학적 그라디언트 분리 기술을 제시하여, VLA 모델의 실용적 배포에 중요한 기여를 합니다.