μVLA: Memory Token Tăng Tỉ Lệ Thành Công Từ 42% lên 84%

Tưởng tượng bạn đang lắp ghép một mảnh ghép hình, nhưng mỗi lần bạn nhìn vào bức tranh lại chỉ thấy một góc nhỏ, và góc đó thay đổi liên tục. Đó chính là bài toán partial observability — quan sát một phần — mà robot phải đối mặt mỗi ngày. Camera không nhìn thấy toàn cảnh, vật thể bị che khuất, tay robot che góc nhìn. Và VLA model truyền thống không có bộ nhớ nào để lưu lại "tôi vừa thấy gì ở bước trước."

Nhóm nghiên cứu từ AIRI (Artificial Intelligence Research Institute, Moscow) đã công bố μVLA (arXiv 2606.12497) — một giải pháp thanh lịch: thêm một tập nhỏ recurrent memory token vào backbone transformer, cho model "nhớ" được lịch sử quan sát qua các timestep. Kết quả: success rate trên MIKASA-Robo nhảy từ 42% lên 84% mà không cần thay đổi thêm gì trong kiến trúc, không cần auxiliary loss, không cần dataset mới.

Bài Toán: VLA Không Có Trí Nhớ

Các VLA model hiện tại như OpenVLA-OFT xử lý từng frame theo kiểu Markovian — chỉ nhìn vào observation hiện tại để ra quyết định. Điều này hoạt động tốt khi camera thấy toàn bộ scene, nhưng thất bại khi:

• Vật thể bị che khuất: Hộp đã bị tay robot che mất, nhưng robot vẫn cần biết vị trí của nó
• Target ra ngoài góc nhìn: Target di chuyển, camera không xoay theo
• Task yêu cầu nhớ giai đoạn trước: "Nhớ màu của cục khối bên trái rồi đặt vào đúng ô" — khi đặt, robot cần nhớ màu từ frame đầu episode

Benchmark MIKASA-Robo được thiết kế đặc biệt để thử thách điểm yếu này: 32 task manipulation với partial observability có chủ đích. Baseline OpenVLA-OFT không có memory chỉ đạt 0.42 average success rate trên 5 task huấn luyện — và chỉ 0.07 trên held-out tasks.

Ý Tưởng Cốt Lõi: Memory Token Trong Self-Attention

Giải pháp của μVLA đơn giản đến mức ngạc nhiên. Thay vì thiết kế lại toàn bộ kiến trúc, nhóm tác giả làm một việc nhỏ: thêm m token có thể học được vào sequence đầu vào của attention layer, và những token này được mang theo qua các timestep trong cùng một episode.

Cấu trúc input sequence với memory:

[BOS] [VIS] [PROPRIO] [M_t] [TEXT] [ACT] [STOP]

Trong đó M_t ∈ ℝ^{m×d} là memory state tại timestep t (m token, mỗi token có chiều d bằng hidden dim của backbone). Tại mỗi bước, attention layer nhận input là toàn bộ sequence này, và sau khi đi qua attention:

• Memory token output trở thành M_{t+1} — được đưa vào bước tiếp theo
• Observation và action token được xử lý bình thường để ra quyết định

Một forward pass duy nhất đồng thời đọc memory từ t-1 và ghi recurrent state cho t. Không cần pass riêng, không cần module phụ.

Memory-Action Guard: Chặn "Gian Lận"

Một insight quan trọng của paper: nếu để memory token attend đến action token từ bước trước, model có thể học cách "copy" action mà không thực sự encode thông tin scene — còn gọi là action leakage.

Giải pháp là memory-action guard: một attention mask đặc biệt ngăn memory token attend đến action token. Memory chỉ có thể attend đến: observation, proprioception, language, và memory state trước. Cách này buộc M_{t+1} phải encode thông tin scene có ý nghĩa, không phải copy action.

Kiến Trúc Chi Tiết

μVLA được xây dựng trên OpenVLA-OFT — phiên bản fine-tuned tối ưu tốc độ của OpenVLA, vốn dùng LLaMA2 7B làm backbone LLM và DINO + SigLIP làm vision encoder. Các tham số định nghĩa một μVLA variant:

Hai update rule:

TBPTT (Truncated Backpropagation Through Time): Backprop gradient qua K bước liên tiếp. Memory nhận gradient trực tiếp, học được rõ ràng hơn. Đây là method chính của paper.

EMA (Exponential Moving Average): Cập nhật memory theo công thức M_in[t+1] = α × M_out[t] + (1-α) × M_in[t]. Không cần backprop qua time, nhẹ hơn về VRAM nhưng gradient yếu hơn.

Khởi tạo: M_0 (đầu mỗi episode) là một learnable parameter chia sẻ — không phải zero, không phải random noise. Model học được "trạng thái khởi đầu tốt" cho memory.

Cài Đặt

μVLA yêu cầu một fork tùy chỉnh của transformers để support memory-augmented backbone. Nhóm tác giả đã pin version này trong pyproject.toml.

Yêu cầu phần cứng:

• Inference: 1 GPU với ~16 GB VRAM
• Training: 1-8 GPU với 40-80 GB VRAM (bfloat16); single 80GB GPU dùng --batch_size 4 --use_gradient_checkpointing True

git clone https://github.com/CognitiveAISystems/muVLA.git
cd muVLA
uv sync --python 3.10

Sau khi cài xong, làm theo tài liệu benchmark cụ thể:

• MIKASA-Robo: đọc SETUP.md rồi MIKASA.md (sections 2, 3, 4)
• LIBERO: đọc SETUP.md rồi LIBERO.md (sections 1, 2, 3, 4)

Pre-trained Checkpoints từ HuggingFace

Nhóm tác giả cung cấp 4 checkpoint đã merge LoRA weights trên Hugging Face (mu-vla org):

from huggingface_hub import snapshot_download

# Checkpoint tốt nhất cho memory tasks: m=64, K=2, TBPTT
snapshot_download(
    "mu-vla/mu-vla-openvla-oft-mikasa-robo-5-tasks-m64-k2-tbptt",
    local_dir="./checkpoints/mu-vla-k2"
)

Danh sách đầy đủ:

Training: Cách Huấn Luyện μVLA

μVLA được train end-to-end trên task-specific data với LoRA rank 32 trên top OpenVLA-OFT. Không cần pretraining riêng cho memory module — memory token được khởi tạo và học từ đầu cùng với fine-tuning.

# Enable memory (K=2, m=64 - cấu hình tốt nhất)
uv run python train.py \
  --base_model openvla/openvla-7b \
  --dataset mikasa_robo \
  --memory_tokens 64 \
  --tbptt_length 2 \
  --use_lora true \
  --lr_schedule cosine \
  --batch_size 4

# Baseline không có memory (ablation)
uv run python train.py \
  --base_model openvla/openvla-7b \
  --dataset mikasa_robo \
  --use_memory false \
  --use_lora true

Điểm quan trọng về training:

• Cosine LR schedule quan trọng hơn constant schedule — paper khuyến nghị bắt buộc
• Round-robin episodic dataloader cần thiết để maintain temporal order within episodes
• Memory token không cần warm-up riêng, nhưng epoch đầu unstable hơn — recommend ít nhất 3 seeds

Kết Quả Thực Nghiệm

MIKASA-Robo: Benchmark Partial Observability

MIKASA-Robo gồm 32 task manipulation trong 12 nhóm, được thiết kế đặc biệt để đo khả năng nhớ. Task điển hình là RememberColor — robot phải nhớ màu của vật thể từ frame đầu rồi execute task dựa trên màu đó khi vật đã bị che.

Tại sao K=2 win, không phải K=8?

Đây là insight thú vị nhất của paper. Với K lớn hơn, gradient chảy qua nhiều bước hơn nhưng lại gây ra vấn đề credit assignment — model khó biết gradient từ bước nào quan trọng. K=2 tìm được điểm ngọt: đủ để học temporal dependencies gần, không quá nhiều nhiễu từ xa.

RememberColor5 (task đại diện): μVLA K=2 đạt 0.93 vs baseline 0.35-0.40.

LIBERO: No Regression trên Fully Observable

Trên LIBERO — benchmark fully observable tiêu chuẩn gồm 4 suites — μVLA m=64, K=8 đạt 96.2% average success. Không có regression so với baseline. Điều này chứng minh memory token không "can thiệp" vào khả năng xử lý khi không cần nhớ gì thêm — model biết tự "ignore" memory khi nó không hữu ích.

Kiểm Tra Nhân Quả: Memory Có Thực Sự Được Dùng Không?

Nhóm tác giả làm hai thí nghiệm causal analysis để xác nhận memory không chỉ là "decoration":

Noise Injection Test: Inject random noise vào M_t tại inference time (không thay đổi gì khác). RememberColor5 tụt từ 0.94 → 0.09 với K=2. Memory bị corrupt → robot quên cue → fail hoàn toàn. Đây là bằng chứng causal rõ ràng.

Freeze-First Test: Freeze memory sau frame đầu tiên (giữ M_1 cố định cho toàn episode). Task dựa vào cue từ frame đầu vẫn hoạt động tốt; nhưng task phụ thuộc vào dynamics trong episode (ví dụ track vị trí vật di chuyển) bị collapse. Memory encode đúng loại thông tin mà từng loại task cần.

Plot cosine distance cho thấy memory thay đổi nhiều ở các bước quan trọng (khi task phase chuyển, khi vật thể biến mất) và ổn định khi scene không đổi — hành vi có ý nghĩa, không phải noise.

Inference: Một Chi Tiết Quan Trọng Dễ Bỏ Sót

Receding-horizon control là bắt buộc cho cue-recall tasks: query model tại MỖI environment step, chỉ execute action đầu tiên của chunk. Nếu dùng action chunking thông thường (chunk ≥ 2), performance trên cue-recall tasks collapse ngay.

Lý do: memory cần được cập nhật liên tục từ mỗi frame. Nếu execute nhiều action trước khi query lại model, thông tin cue bị stale trong memory.

# ĐÚNG: query mỗi bước
for step in range(episode_length):
    action_chunk = model.predict(obs, memory_state)
    memory_state = action_chunk.memory_output  # cập nhật memory
    env.step(action_chunk[0])  # chỉ execute action đầu tiên

# SAI: execute cả chunk
action_chunk = model.predict(obs, memory_state)
for action in action_chunk:
    env.step(action)  # memory stale trong 7 bước này!

So Sánh Với Các Hướng Tiếp Cận Khác

μVLA không phải cách duy nhất để thêm memory vào VLA. Bảng so sánh nhanh:

μVLA nằm ở điểm ngọt giữa đơn giản và hiệu quả. Không cần retrieval module, không cần auxiliary loss, không cần thêm dataset — chỉ cần thêm token và train lại trên task data hiện có.

Khi Nào Nên Dùng μVLA?

Phù hợp:

• Task có partial observability rõ ràng — camera góc hẹp, vật bị che
• Manipulation yêu cầu nhớ trạng thái qua nhiều bước (RememberColor-style)
• Bạn đang dùng OpenVLA-OFT và muốn cải thiện nhanh

Không phù hợp:

• Task fully observable (LIBERO-style): memory gain gần như 0, thêm vô nghĩa
• Inference cần action chunking lớn (chunk ≥ 2): mâu thuẫn với yêu cầu receding-horizon
• Budget training rất chặt: TBPTT K=2 tốn thêm VRAM và compute so với baseline

Góc Nhìn Kỹ Thuật

Điều làm μVLA nổi bật không phải chỉ là con số 84%. Điều đáng chú ý là sự kiểm soát khoa học chặt chẽ: từng biến được thử nghiệm riêng (m, K, update rule), có ablation rõ ràng, và có thí nghiệm causal để chứng minh mechanism thật sự hoạt động. So với nhiều paper VLA chỉ cho thấy "scaling works", μVLA cho bạn biết tại sao memory hoạt động và khi nào nó không cần thiết.

Paper cũng thành thật về giới hạn: K=2 mất ~75% recall khi episode dài hơn training horizon. Memory không extrapolate tốt sang unseen memory structures. Đây là những ràng buộc thực tế quan trọng cho engineer khi deploy.

Bài Viết Liên Quan

• OpenVLA-OFT: Fine-tuning VLA Nhanh Cho Robot Manipulation
• MemoryVLA++: Dual Memory Và World Model Cho Robotic Manipulation
• FM-VLA: Force Memory Token Cho Contact-Rich Manipulation