Multimodal Model Diffing for Feature Discovery and Control

MMDiff is a multimodal model-diffing framework that trains multimodal sparse autoencoders (SAEs) and turns them into feature-level interfaces for discovering and controlling multimodal behavior in Multimodal Large Language Models (MLLMs). The framework supports three uses: (i) feature isolation, by diffing a base-LM SAE against its multimodal-adapted counterpart to identify features altered by multimodal training; (ii) task-specific feature detection, via per-token contrastive firing analysis that isolates causal features; and (iii) feature-level control, by causally removing or steering the discovered feature directions.

The pipeline has three stages. First, MMDiff fine-tunes SAEs on multimodal activations from an MLLM to obtain a feature dictionary aligned with the vision–language space. It initializes from base-LM SAE suites matched to each backbone: LLaMA-Scope Top-K SAEs for LLaVA-MORE (LLaMA-3.1-8B), Gemma-Scope JumpReLU SAEs for PaliGemma 2 (Gemma-2-2B), and Qwen-Scope Top-K SAEs for InternVL3.5-2B (Qwen3-1.7B). The SAEs are trained on cached activations from 50k VQAv2 image–question pairs, with text-only training achieving the lowest fraction of variance unexplained (FVU) and remaining most aligned with the base-LM dictionary.

Second, MMDiff isolates the subset of features that prefer visual tokens and undergo substantial geometric rotation between the base LM and the MLLM. This is done using two signals: geometric reorientation (decoder cosine similarity) and modality preference (visual energy). Features are defined as adapted if they satisfy both criteria: visual energy above a threshold and cosine similarity in the bottom 25%. This yields adapted sets of approximately 5%, 20%, and 13% of all features for MMDiff-Llama, MMDiff-Gemma, and MMDiff-Qwen, respectively.

Third, within this adapted set, MMDiff applies contrastive token firing between a target distribution and a baseline distribution, followed by selective filtering for lexical invariance, to obtain the task-specific feature set. The contrast picks up features that fire more often under the target shift, and the filter keeps only those that continue firing under neutral prompts sharing no target-specific lexicon. Features are retained as task-specific candidates if they have an odds ratio of at least 3 and a firing frequency gap of at least 0.05. This yields approximately 1,400 features for MMDiff-Gemma and 711 features for MMDiff-Llama for spatial tasks, 1,061 candidate unsafe features for safety, and 1,070 OCR-selective features.

MMDiff is evaluated on three MLLM families (LLaVA-MORE, PaliGemma 2, and InternVL3.5-2B) and three domains: visual-spatial understanding, multimodal safety, and OCR. For causal removal, MMDiff orthogonally projects out the feature direction at every transformer layer and only at text-token positions during inference. The results show that feature-level causal removal selectively degrades target behaviors by an average of 12% on spatial tasks and 17% on OCR, and reduces attack success rate by 24% on multimodal safety attacks, with no impact on VQA performance.

For steering, MMDiff-CAA combines multi-layer backbone contrastive activation steering at detected top feature layers with injection of a discovered feature's decoder direction at its feature-associated layer. This improves spatial and OCR accuracy by +3.6% and +1.8% on average over a standard single-layer steering baseline. The steering decomposition shows that single-layer CAA gives +8.96, multi-layer CAA adds +1.82, and injecting the feature's decoder direction adds a further +1.81, giving a total of +12.59.

The paper also includes ablations and analyses. It shows that model diffing matters: dropping the adapted-feature filter gives larger VSR drops but degrades general VQA by 24–26%, while the full pipeline is both causally effective and selective (−12.3% VSR, −0.1% VQA). A standard SAE trained directly on MLLM activations from random initialization does not recover the causally effective set. Cross-stage ablation on PaliGemma 2 shows that instruction tuning amplifies the causal contribution of spatial features by roughly 3× on average, indicating that these spatial behaviors are introduced during multimodal training rather than inherited from the pretrained variant. Attribution patching identifies task-specific heads, with top heads largely consistent across methods and clustering near each feature's home layer. Auto-interpretation with GPT-4o-mini provides qualitative validation of discovered features.

The paper concludes that multimodal SAEs can serve not only as interpretability tools, but as mechanisms for auditing, steering, and controlling MLLM behavior toward safer and more capable generations.

Improvements to AI systems:

1. Causal feature-level safety intervention: Integrate MMDiff’s adapted-feature detection and orthogonal projection removal directly into MLLM inference pipelines. This enables real-time suppression of unsafe feature directions (e.g., harmful response generation) without retraining, reducing attack success rates by 24% while preserving general VQA accuracy.

2. Task-specific capability steering: Use MMDiff’s contrastive firing + lexical invariance filter to identify and inject decoder directions for spatial reasoning or OCR tasks. The improved system can dynamically steer an MLLM toward higher accuracy on spatial benchmarks (+3.6%) and OCR (+1.8%) by activating only the relevant feature layers, without degrading unrelated capabilities.

3. Multimodal adaptation auditing: Deploy MMDiff’s geometric reorientation + visual energy criteria to automatically audit any new MLLM checkpoint. The system can flag which base-LM features were altered by multimodal training, enabling engineers to trace behavioral shifts (e.g., new spatial biases) back to specific feature directions and decide whether to keep, remove, or modify them.

4. Zero-shot feature discovery for novel tasks: Leverage MMDiff’s per-token contrastive firing to isolate causal features for any new task distribution (e.g., medical imaging, code generation with images). The improved system can, without manual annotation, extract 700–1,400 task-specific features per domain and immediately apply causal removal or steering, enabling rapid adaptation to unseen multimodal tasks.

5. Interpretable failure diagnosis: Use MMDiff’s auto-interpretation (GPT-4o-mini) on discovered features to generate natural-language explanations of why an MLLM fails on specific inputs. The system can then propose targeted feature-level corrections (e.g., “remove feature X associated with ignoring spatial prepositions”) and validate them via causal ablation, creating a closed-loop debugging tool.

6. Efficient SAE training for new backbones: Adopt MMDiff’s initialization from base-LM SAE suites and text-only fine-tuning on 50k VQAv2 pairs. The improved system can produce multimodal SAEs for new MLLMs with lower FVU and higher alignment, reducing training cost and enabling feature-level control for models that lack existing interpretability tooling.

What the improved AI system can do:

• Audit and control any MLLM’s multimodal behavior at the feature level, isolating and modifying specific capabilities (spatial, OCR, safety) with surgical precision.

• Steer generation toward desired task performance on-the-fly, using multi-layer CAA plus decoder-direction injection, without full model retraining.

• Detect and neutralize adversarial or unsafe multimodal inputs by removing causally relevant feature directions, while maintaining benign task performance.

• Explain its own failures and successes via auto-interpreted features, allowing developers to understand and fix model behavior in natural language.

• Transfer these capabilities to new MLLMs quickly, using a standardized SAE fine-tuning pipeline that preserves base-LM alignment and minimizes variance.

Multimodal Large Language Models (MLLMs) exhibit strong visual understanding, yet the internal features that cause these behaviors remain difficult to identify, audit, or control. While applicable to post-hoc inspection, hidden states that are decomposed into interpretable feature directions using sparse autoencoders (SAEs) neither readily isolate which features are changed by multimodal training, nor are they directly useful for targeted control. We introduce MMDiff, a multimodal model-diffing framework that trains multimodal SAEs and turns them into feature-level interfaces for discovering and controlling multimodal behavior. MMDiff supports three uses: (i) feature isolation, by diffing a base-LM SAE against its multimodal-adapted counterpart to identify features altered by multimodal training; (ii) task-specific feature detection, via per-token contrastive firing analysis that isolates causal features; and (iii) feature-level control, by causally removing or steering the discovered feature directions. We train multimodal SAEs for three MLLM families, LLaVA-MORE, PaliGemma 2, and InternVL3.5, and evaluate on visual-spatial understanding, multimodal safety, and OCR. MMDiff discovers sparse, causally specific features whose removal selectively degrades target behaviors by an average of 12% on spatial tasks and 17% on OCR, and reduces attack success rate by 24% on multimodal safety attacks, with no impact on VQA performance. Steering these features improves spatial and OCR accuracy by +3.6% and +1.8% on average over a standard single-layer steering baseline. These results show that multimodal SAEs can serve not only as interpretability tools, but as mechanisms for auditing, steering, and controlling MLLMs behavior toward safer and more capable generations.