Source-linked AI summary
Multi-Dimensional Featurizers Yield Fine-Grained Steering and Data Shaping in Robot Foundation Models
Aiden Swann, Mozes Jacobs, Mathilde Papillon, Siddharth Boppana, David Chanin, Dron Hazra, Usha Bhalla, Curt Tigges, Mac Schwager, Owen Lewis, Jack Merullo, Ekdeep Singh Lubana, Atticus Geiger, Thomas Fel, Matthew Kowal
TL;DR
Understanding the computations behind VLA actions is important for safe deployment, but precise behavioral steering remains challenging. The paper applies block-sparse featurizers to discover multidimensional task-relevant features, steer model behavior, and attribute behavior to training data. It reports fine-grained steering and identifies a small set of fine-tuning tokens associated with failure recovery.
Problem
VLA actions are poorly understood, and precise steering remains challenging, with prior work often reporting inconsistent, off-target, or weak effects.
Source: §1 · p. 2Method
The authors train block-sparse featurizers and baseline sparse autoencoders on unpooled VLA token activations, then use the recovered features for steering and feature-based data attribution.
Source: §1 · pp. 2–3Results
BSFs recover task-relevant features and outperform matched SAEs on recovering and intervening on task-relevant structure; fine-grained steering validates features, and fewer than 1.4% of fine-tuning tokens are associated with failure recovery.
Source: §1 · pp. 2–3Takeaways & Limitations
Targeted interventions establish behavioral relevance for recovered features, while feature-guided data masking selectively delays failure recovery and preserves task performance in the evaluation.
Source: §5 · p. 10The analysis does not fully explain how features emerge or participate in policy computation; some interventions are only evaluated open loop, and BSFs are not uniformly superior to SAEs on quantitative metrics.
Source: §5 · pp. 10–11
Abstract
from arXiv · showhide
Can better feature discovery improve control over robot behavior through internal interventions and training data? Recent work in vision and language models reveals multidimensional features, motivating block-sparse featurizers (BSFs), which represent each feature in a low-dimensional subspace. We investigate whether recovering this structure provides an interpretable interface for steering vision-language-action models and shaping the data from which they learn. We show that BSFs recover several multi-dimensional task-relevant features underlying the model's computation (e.g., perception, language, proprioceptive state, episode phase, and control variables). We causally validate these features through fine-grained steering: targeted changes in visual perception, semantic representation, and spatial localization produce predictable changes in downstream action prediction. Finally, we use feature-based data attribution to isolate where a specific behavior (failure recovery) arises in the training dataset, identifying fewer than $1.4\%$ of training tokens whose masking is sufficient to remove the targeted behavior. Overall, our results show how matching featurizers to representation geometry can provide more effective tools for understanding and changing robot behavior.
1 Introduction
Understanding the internal representations behind VLA actions matters for safe deployment, but precise behavioral steering remains challenging. This work applies block-sparse featurizers to discover task-relevant features, steer behavior, and trace behavior to training data.
1 Introduction
- Understanding VLA internal representations is presented as important for safe deployment as learned robotics models may see widespread deployment.A VLA maps visual observations, language instructions, and robot state to physical-world actions.Source: §1 · p. 2
- Prior VLA interpretability work reports inconsistent, off-target, or weak steering, and often uses pooled visual activations or small activation corpora.This leaves open whether larger-scale training on individual token activations improves feature discovery and steering.Source: §1 · p. 2
- BSFs recover task-relevant features and outperform matched SAEs in capturing computational geometry, while the study also strengthens SAE baselines with larger, unpooled activation corpora.The recovered features span perception, language, proprioception, task state, and episode progress.Source: §1 · p. 3
- Across layers and computational regimes, feature interventions steer perception, semantics, target-object selection, geometric localization, and downstream robot behavior.These experiments causally validate recovered features through fine-grained steering.Source: §1 · p. 3
- For failure recovery, masking fewer than 1.4% of fine-tuning tokens qualitatively transforms robot behavior.Feature-based attribution links the behavior to episodes in the fine-tuning data.Source: §1 · p. 3
2 From Sparse Directions to Sparse Subspaces
Sparse autoencoders represent features as individual directions, while block-sparse featurizers group coordinates into learned subspaces that can capture structured variation. This changes both what a feature can represent and how it is selected and evaluated.
2 From Sparse Directions to Sparse Subspaces
- An SAE must tile a curved variable such as the helix with scalar features, splitting one variable across features whose relationships must later be recovered.Figure 2 contrasts many SAE features tiling the helix with one three-coordinate BSF block reconstructing it.Source: §2 · p. 3, Figure 2 · p. 3
- A BSF replaces each scalar feature with a vector code and matrix decoder, allowing a feature to contribute within a learned subspace.The observed contribution can occupy a curve or surface within that subspace; block size bounds, but does not fix, its intrinsic dimension.Source: §2 · pp. 3–4
- A block’s norm indicates feature strength, while its within-block position indicates where the activation lies on the represented structure, invariant to orthogonal basis rotations.The paper displays block codes in a principal-component basis.Source: §2 · p. 4
- BSFs select whole coordinate groups together, using the block norm to measure feature presence and setting unselected blocks to zero.The encoder retains the k blocks with largest ℓ2 norm; comparisons distinguish active-feature budgets from active-coordinate budgets.Source: §2 · p. 4
- The standard objective minimizes expected squared reconstruction error, while the functional measure weights reconstruction errors by their effect on action likelihood.Errors in activation directions with no first-order effect on action likelihood need not be reconstructed.Source: §2 · pp. 4–5
- The featurizers are trained on separate image, text, and state tokens from approximately 52M activations across 45,344 episodes, giving SAE and BSF comparisons a shared corpus.The authors state that using the same corpus and pipeline makes featurizer structure the difference between the compared methods.Source: §2 · p. 5
3 Results
BSFs and SAEs are compared on reconstruction, action fidelity, and robotics-grounded measures, followed by steering experiments and examples of features tied to visual, language, and robot-state variables. The results show that BSFs capture task-relevant structure and support targeted interventions, while feature blocks organize variation within concepts and control variables.
3.1 Quantitative Comparison to SAEs
- BSFs produced larger per-feature ablation effects and higher task-variable decodability than SAEs, particularly in middle and later layers; Fisher-weighted BSFs matched or exceeded vanilla BSFs on both measures.Task-variable targets included end-effector position and displacement, episode progress, gripper state, and gripper-closure measures.Source: §3.1 · p. 6
3.2 Causal Validation of Features via Steering
- Geometric steering used a single BSF block to achieve 99.7% held-out recovery, while BSFs also provided more compact semantic and perceptual interventions without consistently improving transfer.Semantic steering with standard BSFs required fewer features while slightly improving held-out recovery; perceptual BSF variants used roughly 16–20 blocks versus 36 SAE latents, but SAE transfer was stronger.Source: §3.2 · pp. 6–7
- In closed-loop redirection, BSFs selected the redirected object in every trial and completed 93.3% of trials with 2.3 features on average, versus 88.3% selection and 75.0% completion with 2.5 SAE features.The evaluation used three object pairs, both steering directions, and randomized object and basket positions.Source: §3.2 · p. 7
3.3 Feature phenomenology
- Across feature examples, blocks organize visual variation within object categories, including marker appearance across manipulation stages and patches from different robot-arm parts.The marker and robot examples are Fisher-weighted BSFs from MolmoACT2 layer 7.Source: §3.3 · pp. 7–8, Figure 5 · p. 7
- A container feature in π0.5 activates across image and instruction tokens, placing related container words in distinct regions within one block.The passage contrasts this organization with three SAE features that tile much of the displayed text-token geometry.Source: §3.3 · p. 8, Figure 6 · p. 8
- Other control-relevant blocks organize gripper opening and side, track end-effector Cartesian coordinates despite joint-angle inputs, or encode camera position relative to the robot.The end-effector feature primarily activates on unused image-token positions near the robot base, a phenomenon the passage identifies as unique to π0.5 given its bidirectional attention mask.Source: §3.3 · p. 8, Figure 7 · p. 9
4 Application: Behavior Shaping via Data Attribution
The authors use the three-dimensional BSF feature F1515 to locate failed-grasp recovery behavior in training data, then test whether masking that data selectively changes recovery. Concept-based masking delays reopening while preserving task performance in the reported evaluation.
Identifying failure recovery features
- F1515 is a three-dimensional layer-nine BSF feature that activates during failed-grasp recovery in DROID demonstrations and real-robot rollouts.Its nonlinear, conical geometry is tiled by multiple SAE features but captured by one BSF block.Source: §4 · p. 9
From features to training data
- F1515 thresholding identifies recovery behavior in approximately 6.5% of DROID episodes, and the selected windows guide loss masking while observations remain unchanged.The mask removes supervision for action chunks overlapping selected windows, removing 1.38% of tokens; comparison methods use matched frame budgets and training recipes.Source: §4 · pp. 9–10
Targeted changes in recovery behavior
- 14.25 seconds versus 1.89 seconds: concept-based masking increases mean reopening time more than sevenfold, exceeding all comparison methods, while task success remains perfect with similar rollout durations.The unperturbed pick-and-place evaluation reports durations of 11.79 versus 12.13 seconds.Source: Table 2 · p. 9, §4 · p. 10
5 Conclusion
The conclusion presents BSFs as a way to recover multidimensional VLA features relevant to steering and training-data attribution. It also limits the claims: feature mechanisms remain incompletely explained, some steering is open loop, and BSFs do not outperform SAEs on every metric.
Limitations
- The analysis does not fully explain how features emerge or participate in policy computation, and some interventions are evaluated only open loop.Robust closed-loop visual and geometric steering would require tracking moving image regions and adapting to changing observations and robot state.Source: §5 · pp. 10–11
Conclusion
- BSFs recover interpretable multidimensional VLA features, and targeted interventions establish their behavioral relevance while feature-guided masking delays recovery and preserves task performance in the evaluation.The conclusion frames these as tools for steering behavior and tracing it to training data.Source: §5 · p. 10
Limitations
- BSFs are not uniformly superior: SAEs perform better on some quantitative metrics, and baseline strength highlights the importance of unpooled activations and training scale alongside feature geometry.The conclusion identifies closed-loop steering and understanding feature geometry in relation to proprioceptive inputs as future work.Source: §5 · pp. 10–11
A Behavior Shaping: Implementation Details
The implementation uses F1515 activation to identify recovery windows, compares its selections with several alternative labelers, and masks corresponding action-chunk losses during fine-tuning. The baselines are budget-matched to the F1515 windows.
Concept Based Attribution
- F1515 is evaluated on 81 failed-grasp recovery and 25 successful-grasp timesteps; it fires on 56 recovery timesteps and none of the successful grasps.A frame is active when block norm E ≥150 and active-token count C ≥50; consecutive active frames form events, with single-frame gaps bridged.Source: A Behavior Shaping: Implementation… · p. 16
Baselines
- The comparison labelers are difference of means, kNN, TRAK influence, and random windows, each scored on the same evaluated frames as the concept method.TRAK scores frame-level influence on recovery queries relative to successful-grasp queries; the other methods use their respective recovery-versus-success signals or random placement.Source: A Behavior Shaping: Implementation… · pp. 16–17
Selections
- Selected windows overlap little across methods: raw-frame Jaccard similarity with concept windows is 0.20 for difference of means, 0.09 for kNN, 0.02 for TRAK, and 0.00 for random.Figure 10 shows the highest-scoring frames for scored labelers and randomly drawn windows for the random baseline; random row order has no meaning.Source: Figure 10 · p. 16, A Behavior Shaping: Implementation… · p. 17
Training
- Policies are fine-tuned from π0.5-base on 98% of DROID, with selected frames masking the loss for any overlapping 16-step action chunk while leaving observations unchanged.Training runs for 25k steps; masked action-token fractions range from 1.27% to 1.44% across methods, with 1.38% for concept masking.Source: A Behavior Shaping: Implementation… · p. 17
B.1 Open-Loop Steering Protocol
The protocol tests whether compact feature sets can reproduce activation-patching effects on held-out steering inputs. It selects features greedily, applies edits at task-relevant tokens, and measures action-endpoint displacement relative to patching.
Evaluation setup
- Steering is evaluated on held-out inputs with fresh noise seeds, using distinct evaluation conditions for semantic, visual-token, and geometric edits.Semantic steering uses a mirrored layout, visual steering uses the other object’s instruction, and geometry steering uses a new scene under leave-one-pair-out evaluation.Source: §B.1 · p. 17
Finding features
- Features are ranked by activation differences between prompted and redirected conditions, then added and pruned until the set retains at least 75% of the activation-patching effect.The greedy procedure yields small sets but does not guarantee the smallest possible set.Source: §B.1 · p. 17
Applying features
- Edits target the early noun token, mid-layer image tokens, or late action-readout token, replacing selected subspace contributions or adding a fixed vector according to the steering type.Visual edits replace activations at 63 L7 image tokens in a fixed 9 × 7 table-region window; geometric edits add a vector at the L23 action-readout token.Source: §B.1 · p. 17
Measuring recovery
- Recovery compares the steered action-chunk endpoint displacement with the displacement under activation patching, relative to the unmodified policy and averaged over six held-out noise seeds.A value of 100% denotes a displacement as far as patching in the same direction.Source: §B.1 · p. 17
B.2 Closed-loop Steering
Closed-loop steering fixes feature sets before testing them on new scenes and tasks, with redirection measured across object pairs and directions.
B.2 Closed-loop Steering
- Closed-loop feature sets are discovered on both mirrored layouts, then tested on new scenes and tasks.This contrasts with the open-loop setup, which evaluates candidate semantic features on mirrored scenes and visual or geometric edits under specified alternative conditions.Source: Figure 11 · p. 18
- Each object pair is tested in both directions, with 10 trials per direction.Feature sets are calibrated on both mirrored layouts, fixed before evaluation, and applied throughout each episode to new scenes and tasks.Source: §B.2 · p. 18
- Table 3 reports direction-wise redirection trials out of 10, including grasp and successful basket-placement counts for A→B interventions.A→B means the instruction targets A while steering toward B.Source: Table 3 · p. 18
C Featurizer Training and Evaluation
The study collects unpooled token activations across several VLA models and trains matched BSF and SAE featurizers, alongside Fisher-weighted variants for MolmoAct2.
C Featurizer Training and Evaluation
- Approximately 50M activations are collected for each model-layer combination, keeping image, text, and state tokens separate.MolmoAct2 and π0.5 activations come from DROID; MolmoBot activations come from its simulation training distribution.Source: C Featurizer Training and Evaluation · p. 18
- Featurizers are trained to reconstruct clean activations from lightly corrupted inputs, with isotropic Gaussian noise added before encoding.The stated per-coordinate noise standard deviation is 0.1∥x∥2/√d, or 1% of the token’s energy.Source: C Featurizer Training and Evaluation · p. 18
- The main BSFs use b=3 and k=21, while matched SAEs use b=1 and k=63, giving both featurizers the same total coordinate count.They train for 10 epochs using Adam at 10^-4 on one H200; BatchTopK selects blocks during training and a fixed threshold is used at inference.Source: C Featurizer Training and Evaluation · pp. 18–19
- Fisher-weighted featurizers combine policy-action-sensitive functional error with an isotropic term that retains one fifth of the plain MSE weight.For MolmoAct2, the Fisher estimate uses gradients of teacher-forced on-policy action-token log-likelihoods from a 2,000-episode DROID subset.Source: C Featurizer Training and Evaluation · p. 19