기술 요약: DiMaS – 시각-언어-행동 모델 조종을 위한 분포 매칭(Distribution Matching)

문제 정의

플로우 매칭(flow-matching) 액션 전문가(action experts)를 활용하는 비전-언어-행동(VLA) 모델은 로봇 조작 분야에서 강력한 정책으로 부상했습니다. 그러나 하나의 핵심적인 역량인 **미세한 행동 제어(fine-grained behavioral control)**는 여전히 충분히 탐구되지 않은 상태입니다. 표현 스티어링(representation steering)은 대규모 언어 모델(LLM) 및 비전-언어 모델(VLM)에서 행동적 특징이 일반적으로 선형 방향으로 인코딩되는 잘 확립된 해석 가능성 도구이지만, 본 논문은 고전적인 선형 스티어링 방식이 VLA 환경에서는 한계가 있음을 보여줍니다.

핵심 과제는 VLA가 이산적인 토큰이 아닌 연속적인 액션 궤적(continuous action trajectories)을 생성한다는 점입니다. 저자들은 이러한 시각-운동(visuomotor) 설정에서의 행동적 특징이 언어 모델과 같은 방식으로 선형적으로 조종(steerable)될 수 없다고 가설을 세웠습니다. 고정된 선형 방향을 따라 가산적 이동(additive shifts)을 가하는 기존 방식은 작업 성공률을 해치거나 원하는 변화를 신뢰성 있게 유도하지 못하면서 특정 행동 속성(예: 움직임 속도 또는 수직 변위)을 효과적으로 제어하는 데 실패합니다.

방법론: DiMaS

저자들은 플로우 매칭 VLA에 특화된 전략인 **DiMaS (Distribution Matching Steering)**를 제안합니다. 고정된 선형 방향을 따라 표현을 이동시키는 대신, DiMaS는 표현 간의 전송(transport)을 학습하는 매핑을 학습합니다.

핵심 구성 요소

1. 표현 소싱(Representation Sourcing): 이 방법은 VLA의 액션 전문가($f_A$)로부터 특정 레이어($l$)와 플로우 매칭 디노이징 단계($m$)에서의 잔차 스트림(residual stream) 표현($h$)을 추출합니다. 액션 전문가를 선택한 이유는 액션 전문가의 표현이 행동적 특징에 대해 더 높은 선형 분리성(linear separability)을 보이기 때문입니다.
2. 특징 추출(Feature Extraction): 예측된 액션으로부터 연속적인 스칼라 행동 특징(예: 말단 장치(end-effector)의 속도 또는 수직 변위)을 도출합니다.
3. 분포 분할(Distribution Partitioning): 표현을 특징이 없는 소스 분포($\mathcal{D}^-$)(낮은 값)와 특징이 존재하는 타겟 분포($\mathcal{D}^+$)(높은 값)로 분할합니다. 이는 잘 분리된 모집단을 보장하기 위해 경험적 분위수(empirical quantiles)를 기준으로 임계값을 설정하여 수행됩니다.
4. 최적 운송 맵(Optimal Transport Map): 스티어링 개입은 $\mathcal{D}^-$에서 $\mathcal{D}^+$로 표현을 이동시키는 전송 맵 $\mathcal{T}$로 정의됩니다. 이 맵은 단일 선형 방향을 가정하는 대신 전체 기하학적 구조를 존중하도록 저차원 싱크혼(low-rank Sinkhorn) 알고리즘을 사용하여 칸토로비치 최적 운송 목적 함수(특히 $W_2$ 거리)를 최소화함으로써 학습됩니다.
5. 추론 시 개입(Inference-Time Intervention):
   • 게이팅(Gating): 이진 분류기(선형 프로브)가 표현이 소스 분포에 속하는지(특징 부재)를 결정합니다. 스티어링은 특징이 부재할 때만 적용됩니다.
   • 보간(Interpolation): 작업 성공을 유지하기 위해, 이 방법은 표현을 타겟 분포로 완전히 전송하지 않습니다. 대신, 노브(knob) $\alpha \in [0, 1]$를 사용하여 원래의 표현과 전송된 표현 사이를 보간합니다.
   • 업데이트 규칙: $h' = (1-\alpha)h + \alpha(\mathcal{T} \circ P(h))$, 여기서 $P$는 소스 분포로의 최근접 이웃 투영(nearest-neighbor projection)입니다.

주요 기여

1. DiMaS 메커즘: 플로우 매칭 VLA에 특별히 적응된 분포 매칭(최적 운송) 기반의 새로운 스티어링 전략을 제공합니다. 이는 서로 다른 규모의 모델(SmolVLA 및 $\pi0.5$)에 걸쳐 검증되었습니다.
2. 일반화 분석: 행동 제어가 점점 더 이질적인 작업들 사이에서 어떻게 전이되는지를 특성화합니다. 일부 작업에서 학습된 스티어링이 홀드아웃(held-out) 작업 및 서로 다른 작업 세트에도 일반화될 수 있음을 입증하지만, 그 효과는 작업 다양성에 따라 달라짐을 보여줍니다.
3. 구조적 통찰: 액션 전문가의 표현 구조 분석을 통해 저자들은 왜 고전적인 선형 스티어링이 VLA에서 실패하는지 설명합니다. 그들은 행동적 특징이 **선형적으로 디코딩 가능(linearly decodable)**하지만(분류기에 의해 분리됨), 선형적으로 조종 가능(linearly steerable)하지는 않다(고정된 이동이 소스 분포를 타겟 분포로 변환할 수 없음)는 것을 보여줍니다. 이는 DiMaS의 분포 매칭 설계를 정당화합니다.

실험 결과

이 방법은 LIBERO 벤치마크(Object, Spatial, Goal, 10 suites)에서 두 가지 최첨단 VLA인 SmolVLA와 $\pi0.5$를 대상으로 평가되었습니다.

• 제어 효능(Control Efficacy): DiMaS는 증가 및 감소 양방향 모두에서 타겟 특징(움직임 속도 및 수직 변위)을 성공적으로 조절했습니다. 반면, 선형 스티어링 베이스라인(평균 차이 및 회귀 기반)과 프롬프트 주입 베이스라인은 일관성이 없었으며, 특징을 변화시키지 못하거나 잘못된 방향으로 변화시키는 경우가 많았습니다.
• 작업 성공률(Task Success): DiMapS는 속도를 조절할 때 작업 성공률을 대체로 보존했습니다. 그러나 수직 변위를 조절할 때는 성공률 하락이 발생했는데, 저자들은 이를 해당 특징이 작업 완료와 본질적으로 연결되어 있기 때문(예: 충분한 수직 리프트가 부족하면 물체 배치가 불가능함)이라고 설명합니다.
• 일반화(Generalization):
  • 작업 다양성: 여러 작업을 공동으로 학습하는 것(Setting 2)이 단일 작업에서 학습하는 것(Setting 1)보다 더 안정적인 효과를 냈습니다.
  • OOD 전이(Out-of-Distribution Transfer): DiMaS는 동일 세트 내의 홀드아웃 작업에 적용했을 때 효과적이었으며, 서로 다른 세트로 전이했을 때도 어느 정도 효과를 유지했습니다. 효과는 가장 다양한 세트(LIBERO-Goal)에서 가장 약하게 나타났는데, 이는 높은 행동적 이질성이 학습된 전송 맵을 어렵게 만든다는 것을 시사합니다.
  • 장기 호라이즌 작업(Long-Horizon Tasks): 이 방법은 더 긴 호라이즌 작업(LIBERO-10)에서도 효과를 유지했으며, 성공률에 대한 관리 가능한 비용을 지불하며 속도를 크게 줄였습니다.
• 보간 트레이드오프(Interpolation Trade-off): 보간 파라미터 $\alpha$를 통해 스티어링 강도와 성공률 사이의 조절 가능한 트레이드오프를 제공합니다. 중간값(예: $\alpha=0.5$)이 대부분의 성공적인 에피소드를 보존하면서 속도를 실질적으로 줄이는 유리한 균형을 제공하는 것으로 나타났습니다.

의의 및 주장

본 논문은 DiMapS가 VLA에서 **사후 행동 제어(post-hoc behavioral control)**를 위한 견고한 메커니즘을 제공한다고 주장하며, 사용자가 로봇이 '무엇을' 하는지뿐만 아니라 '어떻게' 수행하는지를 형성할 수 있는 격차를 해소합니다.

• 이론적 통찰: 이 연구는 VLA 영역에서 중요한 구분을 제공합니다: 행동적 특징은 선형적으로 디코딩 가능하지만 선형적으로 조종 가능하지는 않다는 것입니다. 이 발견은 LLM의 선형 스티어링 패러다임을 시각-운동 제어에 직접 적용하는 것에 도전하며 최적 운송의 사용을 정당화합니다.
• 실용적 유용성: DiMaS는 모델을 재학습하지 않고도 로봇의 행동을 형성할 수 있는 실용적인 인터페이스를 제공함으로써 작업 성공을 유지하면서 특정 행동을 제어합니다.
• 한계 및 겸손: 저자들은 수직 변위를 조종하는 것이 물리적 제약으로 인해 본질적으로 작업 성공에 위험할 수 있음을 인정합니다. 또한 현재 방법이 모든 타임스텝을 균등하게 다루며, 청킹된 출력(chunked outputs) 내의 표현과 액션 사이의 연관성은 근사치라는 점을 언급합니다.

요약하자면, DiMaS는 분포 매칭이 플로우 매칭 VLA를 조종하기 위해 필수적이고 효과적인 방법임을 입증하며, 로봇 조작을 위한 표현 수준 제어의 일반화된 프레임워크를 제시합니다.