Multimodal Model Diffing for Feature Discovery and Control

Overview§

MMDiff is a framework for interpreting and controlling multimodal LLMs by training sparse autoencoders (SAEs) on the multimodal model and comparing features against a base language model's SAE. The key idea is to identify features that are introduced or modified by multimodal training, then use these features for causal interventions.

Methodology§

Multimodal SAE Training§

A linear SAE with ReLU activation is trained on the residual stream activations of an MLLM for a given layer. The reconstruction loss is:

$$ \mathcal{L}_{\text{SAE}} = \| \mathbf{x} - \text{Dec}(\text{Enc}(\mathbf{x})) \|_2^2 + \lambda \| f(\mathbf{x}) \|_1 $$

where $\mathbf{x} \in \mathbb{R}^{d}$ is the activation, $f(\mathbf{x}) = \text{ReLU}(W_e \mathbf{x} + b_e)$ is the sparse code, and the decoder reconstructs $\hat{\mathbf{x}} = W_d f(\mathbf{x}) + b_d$. The L1 penalty promotes sparsity.

Feature Isolation via Model Diffing§

To isolate multimodal-specific features, MMDiff trains an SAE on the base LM (before multimodal adaptation) and another on the MLLM (after adaptation). The two SAEs are aligned using the decoder weight similarity, and features with low cosine similarity or features that only exist in the MLLM SAE are flagged as multimodal-altered features. Diffing can be formalized as selecting features $i$ where:

$$ \frac{| \cos(\mathbf{d}_i^{\text{base}}, \mathbf{d}_i^{\text{mllm}}) |}{ \max_k | \cos(\mathbf{d}_i^{\text{base}}, \mathbf{d}_k^{\text{mllm}})|} < \tau $$

Task-Specific Detection and Control§

For a given task, MMDiff computes per-token firing rates on positive and negative samples. Features with high contrastive firing ratio $r_i^+/r_i^-$ are selected as causal. Control is implemented by either removing the feature direction from the activation (ablation) or adding a scaled direction (steering). For steering, the modified activation is $\mathbf{x}' = \mathbf{x} + \alpha \mathbf{d}_i$.

Experiments and Results§

MMDiff was evaluated on LLaVA-MORE, PaliGemma 2, and InternVL3.5. Multimodal SAEs were trained on their residual streams. On visual-spatial understanding, removing discovered features dropped spatial accuracy by 12%; on OCR, accuracy dropped 17%. For multimodal safety, removing safety-related features reduced attack success rate by 24%, with no impact on general VQA. Steering improved spatial and OCR accuracy by +3.6% and +1.8% over a standard single-layer steering baseline.

Code Illustration§

The following PyTorch snippet sketches the multimodal SAE training and diffing logic:

import torch
import torch.nn as nn
import torch.nn.functional as F

class SAE(nn.Module):
    def __init__(self, d_model, d_sae):
        super().__init__()
        self.encoder = nn.Linear(d_model, d_sae, bias=False)
        self.decoder = nn.Linear(d_sae, d_model, bias=False)
        self.b_dec = nn.Parameter(torch.zeros(d_model))

    def forward(self, x):
        f = F.relu(self.encoder(x - self.b_dec))
        x_hat = self.decoder(f) + self.b_dec
        return x_hat, f

# Train SAEs on base and MLLM activations
base_sae = SAE(d_model, d_sae)
mllm_sae = SAE(d_model, d_sae)

# after training, diff feature directions
cos = F.cosine_similarity(base_sae.decoder.weight, mllm_sae.decoder.weight, dim=-1)
multimodal_features = (cos < tau).nonzero().squeeze()

# steer at inference
with torch.no_grad():
    activation = mllm_model.layer(x)
    _, f = mllm_sae(activation)
    intervention = activation + alpha * mllm_sae.decoder.weight[feature_idx]

This demonstrates the core pipeline: train multimodal SAEs, diff against base SAEs, and intervene on discovered feature directions.