OxyGen: 멀티태스크 병렬 처리를 위한 시각-언어-행동 모델의 통합 KV 캐시 관리

로봇이 물건을 옮기면서 동시에 사람과 대화하거나 주변을 기억하는 멀티태스크 능력을 갖추려면 매우 빠른 추론 속도가 필요하다. 기존 시스템은 동일한 시각 정보를 중복 계산하여 속도가 느렸으나, 이 논문은 계산 자원을 통합 관리하여 로봇의 실시간 제어 성능을 유지하면서도 언어 처리량을 획기적으로 높였다.

Transformer 기반 모델은 추론 시 이전 단계의 계산 결과인 KV 캐시를 메모리에 저장해 재사용한다. 하지만 기존 로봇 제어 시스템은 행동 생성과 언어 생성을 별개의 독립된 과정으로 처리하여, 동일한 카메라 영상을 입력받음에도 불구하고 각각 KV 캐시를 따로 생성하는 비효율이 발생한다. 이는 제한된 온디바이스 자원에서 연산 낭비와 메모리 부족을 초래한다. OxyGen은 통합 KV 캐시 관리자를 통해 이 문제를 해결한다. 공통된 시각 정보를 처리하는 Backbone에서 나온 KV 캐시를 중앙에서 관리하며, 이를 행동 전문가와 언어 전문가가 동시에 참조하게 만든다. 즉, 한 번의 계산으로 두 가지 작업을 모두 수행할 수 있는 구조를 만든 것이다. 또한, 로봇 제어는 매 프레임 정해진 시간 내에 끝나야 하는 하드 데드라인이 있는 반면, 대화는 조금 늦어져도 되는 소프트 데드라인을 가진다. OxyGen은 이 두 작업의 시간적 특성 차이를 이용해, 행동 제어는 즉시 처리하고 언어 생성은 여러 프레임에 걸쳐 묶어서 처리하는 연속 배칭 기법을 적용하여 전체적인 시스템 효율을 극대화한다.

MoT VLA 추론을 Prefill 단계와 Generation 단계로 분리함. [시각 관측값 $o_t$를 입력으로] -> [Backbone $\Theta_{VLM}$ 연산을 수행해] -> [KV 캐시 $\mathcal{K}_t$를 얻고] -> [이는 특정 작업에 종속되지 않고 관측 정보를 인코딩한 범용적 데이터임]. 통합 KV 캐시 관리자 $\mathcal{M}$은 $\mathcal{K}_t$의 저장, 조회, 갱신, 삭제 인터페이스를 제공함. [저장된 $\mathcal{K}_t$를 입력으로] -> [확산 모델 기반의 Denoising 연산을 수행해] -> [행동 chunk $A_t$를 얻고] -> [이는 로봇의 관절 제어 명령으로 사용됨]. 교차 프레임 연속 배칭 알고리즘으로 언어 생성 요청을 배치 상태 $\hat{\sigma}$로 결합함. [개별 요청의 KV 캐시와 토큰 버퍼를 입력으로] -> [배치 차원으로 쌓고 연결하는 연산을 수행해] -> [통합 배치 $\hat{\mathcal{K}}$와 $\hat{\mathbf{y}}$를 얻고] -> [이는 GPU가 여러 대화 요청을 한 번의 실행으로 병렬 처리하게 함].

NVIDIA RTX 4090 GPU에서 $\pi_{0.5}$ 모델을 사용하여 LIBERO, DROID, ALOHA 벤치마크에서 평가함. OxyGen은 기존의 순차적 실행 방식 대비 최대 3.7배의 속도 향상을 기록했으며, 행동 제어 주기를 19.1Hz에서 70.5Hz로, 언어 처리량을 57.2 tok/s에서 212.9 tok/s로 동시에 끌어올림. Ablation Study 결과, 교차 작업 KV 공유만으로도 약 1.4배의 초기 속도 향상을 얻었으며, 교차 프레임 연속 배칭을 추가했을 때 긴 문장 생성 시에도 로봇 제어 주기가 60Hz 수준으로 일정하게 유지됨을 확인함. 에너지 효율성 측면에서도 OxyGen은 배칭 크기가 커짐에 따라 요청당 에너지 소비를 최대 78%까지 절감함. 이는 모델 가중치에 대한 메모리 접근 횟수가 줄어들었기 때문이며, 메모리 오버헤드는 기존 대비 약 15% 증가하는 수준에 그침.

OxyGen은 KV 캐시를 작업 간 및 프레임 간 공유 가능한 상태로 추상화함. 각 요청을 $(\mathcal{K}_t, \mathbf{y}_t, \delta_t)$ 형태의 재개 가능한 생성 상태로 정의하여, 제어 주기 사이사이에 언어 생성을 중단하고 재개할 수 있도록 설계함. 수학적으로 행동 생성은 $p_{\Theta_{Act}}(A_t | \mathcal{K}_t)$로, 언어 생성은 $p_{\Theta_{Lang}}(y_t | y_{t,1:j-1}, \mathcal{K}_t)$로 모델링됨. 두 전문가 모듈이 동일한 조건부 변수 $\mathcal{K}_t$를 공유함으로써, 연산 복잡도가 높은 시각 인코딩 과정을 단일화함. 비대칭적 데드라인 문제를 해결하기 위해, 행동 작업에는 엄격한 시간 제약을 부여하고 언어 작업에는 처리량 극대화 목적 함수를 적용함. 연속 배칭 시 GPU의 병렬 처리 능력을 활용하여, 단일 요청 처리 시 발생하는 낮은 하드웨어 활용도 문제를 해결하고 연산 비용을 여러 요청으로 분산시킴.