Scaling Diffusion Policy in Transformer to 1 Billion Parameters for Robotic Manipulation

ListenWhen Adding Layers Makes Your Robot Worse10:06

Show notes

We dig into **ScaleDP**, a Sept 2024 paper that diagnoses why Diffusion Policy in a transformer refuses to scale — and fixes it with two minimal moves borrowed from image DiTs. Swapping cross-attention conditioning for AdaLN and dropping the causal mask lets a diffusion action head scale cleanly from 10 million to 1 billion parameters, with a 21.6% average success-rate gain over DP-T on MetaWorld and big real-robot improvements. It's the policy-only precedent for the conditioning tricks now showing up all over visuomotor learning. **Sources** - Scaling Diffusion Policy in Transformer to 1 Billion Parameters for Robotic Manipulation

Subscribe to the channel: https://podcast.lumalabs.link/27a25f51ea81d41a8253f46c6133a35e/feed.xml

TL;DR

ScaleDP diagnoses why vanilla Diffusion Policy in transformer architecture (DP-T) fails to scale — cross-attention observation fusion produces large gradient magnitudes at depth, so adding layers hurts — and fixes it with two minimal architectural moves: replace cross-attention conditioning with AdaLN (regressing per-block scale/shift from timestep + observation embeddings, DiT-style) and remove causal masks in self-attention so each action embedding attends to both past and future actions in the predicted chunk. With these changes, Diffusion Policy scales cleanly from 10M to 1B parameters, with a 21.6% average success-rate improvement over DP-T across 50 MetaWorld tasks and 36.25%/75% real-robot gains on 4 single-arm and 3 bimanual tasks.

Key claims

Increasing DP-T depth degrades performance in MetaWorld (8 layers 80.1% → 12 layers 78.4% → 14 layers 74.6%), and the standard deviation of gradient magnitudes across layers rises with depth — evidence of an optimization-side failure at the fusion module [§III-A, Fig. 1].
Replacing cross-attention fusion with AdaLN blocks (regressing dimension-wise scale γ and shift β from the sum of timestep and observation embeddings, per Eqn. 2) normalizes the parameter distribution and keeps gradient magnitudes low even at 32 layers / 1B params [§III-B].
Removing masks in self-attention (non-causal attention over the predicted action chunk) is load-bearing on bimanual real-robot tasks: on Bimanual Stack Cube, ScaleDP-L jumps from 20% (masked) to 90% (unmasked) at fixed parameter count [Table IV].
ScaleDP scales cleanly from ScaleDP-Ti (10M) → ScaleDP-H (1B) with monotone gains in MetaWorld success rate at fixed 20-demo training budget, and larger models benefit more from additional demonstrations while small models plateau [Fig. 4b/4c/4d].
On 4 real single-arm Franka tasks, ScaleDP-H reaches 87.50% average vs DP-T’s 51.25% and ACT’s 66.25% [Table II]; ScaleDP-L→H improves by 13.75% and 8.33% average across the two real-robot setups [§IV-C].
ScaleDP-L (457M) exhibits visual generalization DP-T lacks — colored cube/mug swaps, shape-different novel objects (cup for cube, tape for mug), reduced-lighting scenes — with no training-time data augmentation, attributed to shape-based object recognition emerging from scale [§IV-D].

Method

The paper first isolates the DP-T scaling failure to the cross-attention observation fusion module by measuring per-layer gradient-magnitude standard deviation — this grows with depth and correlates with the observed success-rate drop. The fix is DiT-style AdaLN: instead of concatenating timestep and observation embeddings into a sequence attended to by the action tokens via cross-attention, sum the timestep and observation embedding vectors, project them through per-block MLPs into (γ, β) parameters, and apply LayerNorm(x) · (1 + γ) + β in place of vanilla LayerNorm. This changes the conditioning distribution rather than the input sequence and gives the model direct multiplicative control over the noise-action embedding distribution per block.

The second change removes the causal mask from self-attention over the action chunk. DP-T predicts ~10 actions per chunk but only executes the first, so the paper argues the causal constraint (each action only attends to earlier actions) hides useful right-context — non-causal attention lets each predicted action reason about the full trajectory it belongs to.

Five ViT-style configs sweep 10M → 1B (ScaleDP-Ti/S/B/L/H) at joint depth × width × heads: 8L×256d×4h → 32L×1280d×16h. Trained on MetaWorld with 20 demos per task and on real-robot demonstrations (100 trajectories per task; 40 for laptop-close), controlling 6D EE pose.

Results

MetaWorld (50 tasks, 4 difficulty tiers, 20 demos each): ScaleDP-Ti beats DP-T across all four difficulty levels at matched parameter count [Fig. 4a]; success rate grows monotonically from Ti → H with fixed demos, with the largest model benefiting most from added data [Fig. 4b–d].
Real single-arm Franka (4 tasks): DP-T 51.25%, ACT 66.25%, ScaleDP-S 58.75, ScaleDP-B 62.50, ScaleDP-L 73.75, ScaleDP-H 87.50 average [Table II]. Notably on “place tennis” DP-T reaches 5%, ScaleDP-B/L reach ≥50%.
Non-causal ablation (3 bimanual tasks): ScaleDP-S 33.33% → higher with unmask, ScaleDP-B 50.00% → higher, ScaleDP-L 66.66% → 90% on Bimanual Stack Cube (was 20%) — larger models gain more from removing the mask [Table IV].
Visual generalization: ScaleDP-L handles color swaps, novel-shape object substitutions (cup for blue cube, tape for gold mug), and reduced lighting without training-time augmentation; DP-T fails on all three. Distractor generalization is more fragile — adding a mouse to the scene breaks ScaleDP-L on some tasks [§IV-D].
Convergence: larger models converge to higher success rates and lower training losses on both MetaWorld Assembly and real-robot tasks [Fig. 4e/4f].

Why it’s interesting

ScaleDP is an early (Sept 2024, pre-VLA-boom) datapoint that the architecture-first stability story the wiki tracks under Training stability at scale — the DiT AdaLN pattern that fixes attention-logit and gradient-magnitude blowups without recourse to µP — transfers cleanly from image DiTs to visuomotor policies, and specifically identifies the observation-fusion module (not the attention core) as the load-bearing gradient-magnitude source in this regime. Complements the wiki’s VLA Models thread, which is dominated by post-2025 recipes bolting flow-matching or diffusion action heads onto VLM backbones (π*0.6: a VLA That Learns From Experience (RECAP), Embodied-R1.5: Evolving Physical Intelligence via Embodied Foundation Models): ScaleDP is the policy-only precedent showing that even without a language backbone, transformer diffusion action heads scale to 1B with the right conditioning mechanism, and it precedes the Behavior Prompting Policy: Demonstrations as Prompts for Manipulation BPP line’s own attention-pooled prompt / decoupled decoder design.