IROS 2026 Poster

Shallow-π

Knowledge Distillation for Efficient Flow-based Vision-Language-Action Models

Boseong Jeon¹ · Yunho Choi³ · Taehan Kim²

¹42dot, software center of Hyundai Motor Group · ²Ludo Robotics, KRAFTON startup · ³Samsung Research · Research performed during Samsung Research.

Paper

Project

Run flow-based VLAs efficiently on edge devices by making the transformer shallow, then distilling the original model.

Transformer layers18 → 6

Faster inference>2×

Success-rate drop2%

Real robots testedGrippers + hands

Why Go Shallow?

Flow cost≈ N_steps× C_tokens× L_depth

N_stepsStep reduction mostly shortens the small diffusion policy head.

C_tokensGPU parallelism already helps; aggressive image-token pruning is risky for precision tasks.

L_depthSequential and hard to parallelize; it scales VLM and policy-head parameters.

CUDA time decreases strongly with transformer depth reduction — **Depth gives the bigger latency win.** Shallow layers cut repeated sequential transformer compute.

Comparison of shallow VLA scopes across openVLA, CogACT, Gr00t, SmoLVLA, and Shallow-pi styles — **Depth touches both parts.** Our student reduces VLM and action-head layers.

Previous Methods for Shallow VLA

Adjacent layer representation similarity and success drop across layer indices — **Choosing layers is fragile.** Similarity does not reliably predict success drop.

LIBERO 10 success rate decreases as more transformer layers are skipped — **Skipping quickly hurts success.** Aggressive manual skipping collapses performance.

Layer skippingmanual thresholds, task-dependent rules

Learned selectorsextra routing, training and inference complexity

Can we make it simple and automatic?Make a shallow student directly, then recover capability with knowledge distillation.

Shallow-π Is Surprisingly Simple

Shallow-pi knowledge distillation diagram

1. Make the student shallow Uniformly subsample layers in both the VLM backbone and action head.

2. Learn task loss Supervise the student with ground-truth flow velocity. \(L_{\mathrm{task}}=\mathbb{E}[\lVert v_\theta(\cdot)-u\rVert_2^2]\)

3. Mimic teacher denoising Match the teacher's predicted velocity. \(L_{\mathrm{kd}}=\mathbb{E}[\lVert v_\theta(\cdot)-v_\phi(\cdot)\rVert_2^2]\)

4. Distill attention Match middle-layer action-to-VL attention. \(L_{\mathrm{attn}}=\mathbb{E}[\mathrm{KL}(\mathrm{Attn}_{\phi}^{a\to vl}\Vert\mathrm{Attn}_{\theta}^{a\to vl})]\)

That's it. No learned layer selector. No manual inference threshold. No modified inference pipeline.

Ablation: each distillation term contributes, and middle-layer action attention is the strongest target.

Loss design
\(L_{\mathrm{task}}\)	\(L_{\mathrm{kd}}\)	\(L_{\mathrm{attn}}\)	L9	L6	L4
✓	✗	✗	95.4	93.0	92.3
✓	✓	✗	96.2	93.9	93.9
✓	✓	✓	96.8	94.6	94.2

Attention layer
	L9	L6	L4
Initial	94.4	93.9	91.0
Middle	96.8	94.6	94.2
Later	95.6	94.1	92.8

Attention target
	L9	L6	L4
Action only	96.8	94.6	93.9
All tokens	0.0	61.5	0.0

Just Benchmark (LIBERO)

Layer distillation moves models toward the upper-left: high success with lower CUDA time and fewer FLOPs.

Bubble size encodes FLOPs. Layer-distilled variants keep teacher-level success while cutting CUDA time.

Baseline Token compression Small backbone Layer distillation