LaMem-VLA: 로봇 조작을 위한 Vision-Language-Action 모델의 이중 잠재 메모리

일반적인 VLA 모델은 현재 관찰에 주로 의존해 장기적 과업 진행 상황을 반영하지 못하는 경향이 있다. LaMem-VLA는 과거 경험을 동일한 연속 임베딩 공간에서 저장하고 재구성해 추론 과정에 직접 결합함으로써 시간적 의존성을 모델 내부에서 처리할 수 있게 했다. 이 방식은 장기 의존성 처리에 필요한 컨텍스트를 확장하면서 입력 길이 폭발을 피하는 대안적 경로를 제공한다.

기존 VLA 모델은 대부분의 의사결정을 현재 관찰과 명령에만 의존하는 Markovian 가정을 따른다. 이로 인해 과거의 상태 전이, 완료된 단계, 과업의 진행 상황 등의 정보가 반영되지 않아 장기적이고 시계열 의존성이 강한 조작 과업에서 실패가 잦았다. 따라서 과거 경험을 효과적으로 보존하고 추론 과정에서 직접 사용하도록 구조화하는 것이 필요했다. LaMem-VLA의 핵심 원리는 과거 경험을 외부의 보조 컨텍스트가 아니라 모델의 연속 임베딩 공간 내의 '네이티브' 토큰으로 표현하는 것이다. 이를 위해 과거 관찰과 행동 히든 상태를 각각 short-term과 long-term 볼트에 저장하고, 현재의 멀티모달 인지 상태로 컨텍스트화된 쿼리를 만들어 관련 단위를 검색한다. 검색된 다수의 단위를 메모리 콘덴서가 고정 길이의 학습 가능한 슬롯으로 압축하면, 이 슬롯들을 현재 입력 시퀀스 앞에 병합해 self-attention 내에서 관찰·명령·메모리가 동시 상호작용하게 된다. 이 접근은 세 가지 측면에서 기존과 차별화된다. 첫째, 메모리가 모델의 내부 추론 루프 안으로 통합되어 행동 토큰 형성 과정에 직접 영향을 준다. 둘째, dual-scale 설계로 시각적 단서와 장기 의미 정보를 분리 보관해 각각의 검색·응축 전략을 적용할 수 있다. 셋째, 응축된 고정 길이 인터페이스는 검색 크기 변화에도 추론 컨텍스트의 길이를 일정하게 유지해 계산 비용 폭증을 막는다.

전체 접근은 네 단계 모듈로 구성되어 있다. 먼저 Vision-Language 백본이 현재의 RGB 관찰과 언어 명령을 임베딩 토큰으로 변환하고 학습 가능한 액션 쿼리를 추가해 조작 관련 잠재 표현을 생성한다. 그다음 메모리 큐레이터가 과거 경험을 short-term과 long-term 두 볼트로 분해해 저장한다. 검색 단계에서는 멀티모달 인지 상태에서 컨텍스트 기반 쿼리를 생성하고, 이 쿼리와 learnable query slots를 이용해 두 볼트에서 Top-K 단위를 선택한다. short-term 검색은 키-값 쌍의 키와 쿼리의 코사인 유사도로 랭킹하며 long-term은 mean-pooled 행동 히든 상태와의 유사도로 랭킹한다. Top-K 선택은 미분 불가 연산으로 처리되며 이후 응축 모듈이 선택된 단위를 입력으로 받는다. 메모리 콘덴서는 학습 가능한 메모리 슬롯(Ts, Tl)과 소형 transformer 기반 메모리 포머(ℱv, ℱc)를 사용해 검색된 증거를 고정 길이의 잠재 메모리 토큰으로 재구성한다. 재구성된 Ms와 Ml은 학습 가능한 소스 임베딩과 합쳐져 VLA 입력 시퀀스 맨 앞에 삽입된다. 최종적으로 VLM은 이 확장 시퀀스에 대해 self-attention을 수행해 메모리-그라운디드 액션 토큰을 생성하고, 이 토큰을 확산 기반 액션 전문가로 전달해 연속 액션 청크를 복원한다.

SimplerEnv-Bridge에서 LaMem-VLA는 평균 성공률 73.9%를 기록해 같은 벤치마크의 기존 베이스라인 CogACT보다 16.6 포인트 높은 성과를 보였다. 세부 과제별로 Put Spoon on Towel 83.3%, Put Carrot on Plate 75.0%, Stack Cube 41.7%, Put Eggplant in Basket 95.8%를 달성했다. 이 결과는 dual-scale 잠재 메모리의 주입이 에피소드 내 시각 단서와 장기 과업 진행을 결합해 성능을 향상시켰음을 시사한다. LIBERO 평가에서 LaMem-VLA는 전체 평균 97.6%의 성공률을 얻어 MemoryVLA 대비 1.1 포인트, CogACT 대비 4.4 포인트의 개선을 보였다. 각 스위트별로 Spatial 98.8%, Object 99.0%, Goal 97.2%, Long-10 95.8%, Long-90 97.0%를 기록해 장기 추론 과제에서 특히 안정적인 성능을 나타냈다. 추가적으로 장기 스위트에서의 이득은 short-term과 long-term 메모리의 상보적 기능 덕분이라는 증거와 일치했다. 제거 및 예산 민감도 실험에서 두 스케일의 메모리를 모두 제거하면 SimplerEnv에서 57.3%, LIBERO-90에서 92.1%로 성능이 크게 떨어졌다. short-term 또는 long-term 중 하나만 제거해도 성능이 각각 약 7~9 포인트 하락해 두 스트림의 보완성이 확인되었다. 검색 예산 K와 잠재 토큰 수 (Ls, Ll)에 대한 실험은 중간 수준의 검색(예: K=8)과 균형 잡힌 토큰 설정(Ls=8, Ll=4)이 성능과 비용의 균형점임을 보여주었다.

아키텍처 구성은 Prismatic 7B VLM 백본과 약 300M 파라미터의 diffusion action expert로 구성되며 입력은 단일 224×224 RGB 프레임과 언어 명령이다. 시각 인코더가 추출한 비주얼 토큰과 토크나이즈된 명령 토큰, 학습 가능한 액션 쿼리를 연결해 VLA 토큰 시퀀스를 구성한다. short-term vault는 키-값 쌍으로 시각 증거를 보관하고 long-term vault는 타임스텝별 행동 히든 상태를 누적한다. 메모리 업데이트는 용량 L을 초과하면 인접한 키의 코사인 유사도가 가장 높은 페어를 평균화해 중복을 통합하는 방식으로 동작한다. 검색은 컨텍스트 평균 풀된 쿼리 벡터와 저장된 키들 간의 코사인 유사도로 Top-K를 선택하며 이 Top-K 선택은 비미분 연산으로 처리된다. 검색으로 얻은 Zshort과 Zlong는 각각 소형 transformer 기반의 메모리 포머 ℱv와 ℱc에 의해 학습 가능한 슬롯 Ts와 Tl로 응축된다. 메모리 주입은 Ms, Ml에 학습 가능한 소스 바이어스(bs, bl)를 더해 VLA 입력 맨 앞에 붙이는 방식으로 이루어진다. 이렇게 확장된 시퀀스에서의 self-attention은 메모리 토큰, 관찰, 명령 및 액션 쿼리 간의 상호작용을 직접 허용한다. 액션 생성은 메모리-그라운디드 액션 토큰을 조건으로 하는 확산 정책으로 수행되며 학습은 노이즈 예측 MSE 손실을 사용하고 추론은 DDIM 10 스텝 샘플링으로 수행된다. 학습 및 하이퍼파라미터 측면에서 짧은·긴 메모리 볼트의 최대 용량 L은 16으로 설정되었고 검색 예산 K는 8로 설정되었다. 응축 후의 잠재 토큰 수는 Ls=8, Ll=4로 균형을 맞추었고, 전체 학습은 8개의 NVIDIA H800 GPU에서 PyTorch FSDP로 진행되어 전체 배치 크기 256, 학습률 2×10^-5를 사용했다. 쿼리 빌더 ℬ와 메모리 포머들은 마스킹 어텐션을 가진 2-레이어 transformer 블록으로 구현되어 있다.

논문에서 명시한 주된 한계는 실험이 모두 시뮬레이션 환경에서 수행되었다는 점이다. 따라서 실세계 로봇 플랫폼으로의 직접적인 성능 이전 가능성은 후속 실험을 통해 검증되어야 한다. 또한 응축된 잠재 토큰이 계산 효율을 제공하지만 토큰 수를 늘릴수록 self-attention 비용이 증가하므로 실제 로봇 배치 시 효율-정확도 균형을 추가로 고려해야 한다.