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 · N/A · 2026

Summary

Abstract Summary

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.

Key Points

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.

From Sparse Directions to Sparse Subspaces

The goal of a featurizer is to factorize model activations into a set of interpretable concepts ( Huben et al., 2024 ; Bricken et al., 2023 ; Gao et al., 2025 ; Bussmann et al., 2024 ; Rajamanoharan et al., 2024 ; Fel et al., 2025 ; Hindupur et al., 2025 ; Bussmann et al., 2025 ) . Moreover, each featurizer also implicitly encodes a hypothesis about how activations are generated Hindupur et al. (2025) .

Choosing a featurizer for a VLA therefore requires matching this geometry to the variables the policy represents and computes.

From Sparse Directions to Sparse Subspaces — Multi-Dimensional Featurizers Yield Fine-Grained Steering and Data Shaping in Robot Foundation Models — Figure 1: Block-sparse featurizers (BSFs) recover multi-dimensional robotics concepts. Examples from vision-language-action models encode end-effector position (top), fruit-and-vegetable semantics (middle), and proprioceptive state (bottom). Each point is a token’s 3D feature code, colored by PCA coordinates, task or semantic labels, and sparse-autoencoder (SAE) assignments across columns. Activations from MolmoAct2 L25, L21, L13 (top to bottom). Left overlays share the PCA colors. SAE features tile distinct regions within each block, showing how BSFs group related variation into a single feature.

Results

Given the featurizers covered in Section 2 , we now turn to the features they discover and the understanding and control they admit. We first perform a quantitative comparison of BSFs to SAEs ( Section 3.1 ) and causally validate the discovered features through open- and closed-loop steering experiments ( Section 3.2 ).

We then review a broad range of feature phenomenology, spanning visual, language, control, and proprioceptive features ( Section 3.3 ).

Results — Multi-Dimensional Featurizers Yield Fine-Grained Steering and Data Shaping in Robot Foundation Models — Figure 2: BSFs better capture low-dimensional manifolds. (b) One three-coordinate BSF block reconstructs the helix. (c) Many individual SAE features tile the helix.

Application: Behavior Shaping via Data Attribution

We next ask whether a discovered feature can identify training data responsible for a policy behavior, and whether filtering that data predictably changes the learned policy. We focus on F1515 (Figure 9 ), a three-dimensional BSF feature in layer nine of $\pi_{0.5}$ that activates during recovery from failed grasps .

We identified it using a small dataset of our own rollouts Figure 9 shows the feature’s geometry and activations during repeated grasp failures and corrections in a DROID demonstration (top) and a real-robot rollout (bottom).

Application: Behavior Shaping via Data Attribution — Multi-Dimensional Featurizers Yield Fine-Grained Steering and Data Shaping in Robot Foundation Models — Figure 3: Reconstruction, dimensionality, and behavioral relevance of SAE and BSF features in MolmoAct2. Left three panels (layer 13): Held-out reconstruction $R^{2}$ and absolute recorded-action log-likelihood change versus active blocks $k$ (lower change is better; mean $\pm$ SEM, 48 episodes), and within-block stable rank versus block size at $k=32$ (bands: interquartile range; gray: full dimension). Colors denote block size ($b=1$: SAE); solid/dashed lines denote Vanilla/Fisher training. Right two panels (across layers): Three-dimensional Vanilla and Fisher BSFs versus SAEs at approximately matched latent widths and active-coordinate budgets (21 blocks versus 63 latents). Ablation effects measure mean predicted end-effector displacement over the top seven BSF blocks or 21 SAE latents by activation mass. Decodability is the best single-feature ridge-probe $R^{2}$ on held-out episodes, averaged over ten robot-state and task-progress variables.

Sources and Demonstrations

Figures are reproduced from Swann et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.