Interpreting and Steering Spatial Reasoning via Multimodal Sparse Autoencoders (SAEs)

Posted on October 5, 2026 | Mechanistic Interpretability | Vision-Language Models | Representation Engineering

Sparse Autoencoders Mechanistic Interpretability Vision-Language Models Spatial Reasoning Causal Steering PyTorch

Code & Repository: https://github.com/kernel-loophole/autoencoder →

Linear probing and direct linear interventions ($h_t' = h_t + \alpha \cdot \text{probe}$) in vision-language models can produce collateral semantic changes because learned linear directions are frequently entangled with non-target properties such as color, texture, object identity, shape, background, and language attributes.

This research codebase investigates whether Multimodal Sparse Autoencoders (SAEs) provide a more interpretable, disentangled feature basis for spatial concepts (such as containment, inside/outside, left/right, above/below, front/behind, depth ordering).

Project Status — Experimental Research Framework This repository is designed for hypothesis testing and controlled falsification. Empirical conclusions regarding SAE superiority or feature monosemanticity have not yet been established and are subject to ongoing experimental verification.

1. Motivation

Linear probing and direct linear interventions ($h_t' = h_t + \alpha \cdot \text{probe}$) in vision-language models (VLMs) can produce collateral semantic changes because learned linear directions are frequently entangled with non-target visual and linguistic properties, such as object color, shape, texture, identity, and background context.

$$h_t' = h_t + \alpha \cdot w_{\text{probe}}$$

This project investigates whether Sparse Autoencoders (SAEs) can extract a more interpretable, disentangled feature basis that enables precise causal steering of spatial reasoning while preserving unrelated semantic attributes.

2. Hypotheses

H1 (Representation): Spatial reasoning features in VLMs are localized in intermediate multimodal residual stream layers and vision-language projector representations rather than early visual encoders or final text generation tokens alone.

H2 (Sparse Feature Discovery): Multimodal SAEs can discover sparse, highly selective latent features corresponding to primitive spatial concepts (containment, adjacency, relative positioning, depth ordering).

H3 (Disentanglement): Individual SAE features exhibit significantly lower attribute leakage and higher purity than direct linear probe direction vectors when tested on counterfactual scene pairs.

H4 (Causal Steering): Intervening on sparse SAE features ($z_i' = z_i + \alpha$ or $z_i' = c$) steers spatial predictions with higher target specificity and lower collateral semantic mutation than linear probe steering.

3. Core Research Questions

RQ1 — Representation: Where inside a vision-language model are spatial relations represented?
Target activation sources: Vision encoder outputs, vision-language projector, multimodal token embeddings, intermediate transformer residual streams, selected attention/MLP layer outputs.

RQ2 — Feature Discovery: Can sparse autoencoders discover interpretable features corresponding to spatial concepts?
Concepts explored: inside, outside, left/right, above/below, front/behind, near/far, containment, overlap.
Experimental design: Allows 1-to-many, many-to-1, distributed, and polysemantic feature mappings without hardcoding 1-to-1 assumptions.

RQ3 — Disentanglement: Are SAE features less entangled with irrelevant attributes than ordinary linear probe directions?
Evaluated confounds: Object color, shape, texture, object identity, background, language templates, image position, object size.

RQ4 — Causal Steering: Can individual SAE features be causally manipulated to alter spatial reasoning while minimizing collateral semantic shift?
Intervention strategies: Compares add, clamp, suppress, and scale in the SAE latent space ($h \to z \to z' \to h'$).

4. Experimental Variables & Baselines

Experimental Variables

Independent Variables:
- Intervention method: Linear probe vector vs. SAE latent feature vs. Random feature vs. Matched control feature.
- Intervention strength ($\alpha \in [-10.0, 10.0]$).
- Target layer / representation source.
- Scene attribute variations (colors, shapes, background grid, counterfactual attribute swaps).
Dependent Variables:
- Target spatial relation prediction accuracy.
- Feature selectivity & purity index.
- Counterfactual invariance score.
- Specificity score & collateral semantic shift.
- KL divergence of output logit distributions.

Baselines

Unintervened Model Baseline: Clean forward pass without any intervention.
Linear Probe Intervention Baseline: Direct linear direction steering ($h' = h + \alpha \cdot w_{\text{probe}}$).
Random Feature Steering Control: Modifying arbitrary random SAE features with equal norm.
Matched-Control Feature Steering: Modifying correlated non-spatial SAE features (e.g., color-selective features).

5. System Architecture

The framework implements an end-to-end modular pipeline consisting of seven decoupled stages:

                    ┌─────────────────────────┐
                    │ Synthetic Scene Data    │
                    │ & Counterfactual Pairs  │
                    └────────────┬────────────┘
                                 │
                                 ▼
                    ┌─────────────────────────┐
                    │ Vision-Language Model   │
                    │      (VLM Adapter)      │
                    └────────────┬────────────┘
                                 │
                                 ▼
                    ┌─────────────────────────┐
                    │ Activation Extraction   │
                    │ & Layer Storage         │
                    └────────────┬────────────┘
                                 │
                                 ▼
                    ┌─────────────────────────┐
                    │ Sparse Autoencoder      │
                    │      (SAE Training)     │
                    └────────────┬────────────┘
                                 │
                                 ▼
                    ┌─────────────────────────┐
                    │ Feature Discovery &     │
                    │ Correlation Analysis    │
                    └────────────┬────────────┘
                                 │
                                 ▼
                    ┌─────────────────────────┐
                    │ Causal Steering         │
                    │ Intervention Engine     │
                    └────────────┬────────────┘
                                 │
                                 ▼
                    ┌─────────────────────────┐
                    │ Disentanglement &       │
                    │ Causal Evaluation       │
                    └─────────────────────────┘

6. Mathematical Formulations

Sparse Autoencoder (SAE) Formulation

Given extracted model residual activations $x \in \mathbb{R}^d$, the Sparse Autoencoder projects $x$ into an expanded sparse latent representation $z \in \mathbb{R}^M$:

$$\begin{aligned} z &= \text{ReLU}\left(W_{\text{enc}} (x - b_{\text{dec}}) + b_{\text{enc}}\right) \\ \hat{x} &= W_{\text{dec}} z + b_{\text{dec}} \end{aligned}$$

With an $L_1$-penalized reconstruction loss:

$$\mathcal{L}_{\text{SAE}}(x) = \|x - \hat{x}\|_2^2 + \lambda \|z\|_1$$

Or alternatively using Top-K SAEs:

$$z = \text{TopK}\left(\text{ReLU}\left(W_{\text{enc}} (x - b_{\text{dec}}) + b_{\text{enc}}\right), k\right)$$

Causal Steering Intervention

During forward pass inference, we intercept activation $h$ at the target layer, compute latent activations $z$, manipulate the target spatial feature index $j$, and reconstruct into the residual stream:

$$\begin{aligned} z'_j &= z_j + \alpha \quad \text{or} \quad z'_j = c \\ h' &= h + W_{\text{dec}} (z' - z) \end{aligned}$$

7. Directory Structure & Core Modules

.
├── README.md
├── LICENSE
├── pyproject.toml
├── .gitignore
├── .env.example
├── configs/
│   ├── model/           # VLM specifications (Qwen2-VL, LLaVA, Mock)
│   ├── data/            # Synthetic scene & counterfactual configs
│   ├── sae/             # SAE architectures (L1, Top-K, loss)
│   ├── extraction/      # Activation caching & layer storage configs
│   ├── intervention/    # Causal steering strategies & strengths
│   └── experiments/     # End-to-end evaluation runs
├── src/
│   └── spatial_sae/
│       ├── data/         # Synthetic scene rendering & counterfactual pairs
│       ├── models/       # VLM adapters & forward hook interceptors
│       ├── extraction/   # Layer activation collectors & chunked disk storage
│       ├── sae/          # Sparse Autoencoder PyTorch models & trainers
│       ├── probing/      # Linear probe baselines & probe direction steering
│       ├── intervention/ # SAE causal steering engine & runtime hooks
│       ├── evaluation/   # Spatial flip rate, attribute specificity & KL drift
│       ├── experiments/  # Compositional binding & relation transfer benchmarks
│       └── utils/        # Utilities, config loaders & logging
├── scripts/
│   ├── prepare_data.py
│   ├── extract_activations.py
│   ├── train_sae.py
│   ├── analyze_features.py
│   ├── run_probe.py
│   └── run_intervention.py
├── tests/
├── notebooks/
├── checkpoints/
└── outputs/

Core Modules Breakdown

Module	Description
`spatial_sae.data`	Manages synthetic scene rendering, counterfactual pair generation, schemas, and PyTorch datasets.
`spatial_sae.models`	VLM adapters (`VLMAdapter`, `HuggingFaceVLMAdapter`, `MockVLMAdapter`), representation definitions, and forward hook managers.
`spatial_sae.extraction`	Layer activation collection, disk storage shards (`ActivationStorage`), and extraction pipelines (`ActivationExtractor`).
`spatial_sae.sae`	PyTorch Sparse Autoencoder (`SparseAutoencoder`), $L_1$ / Top-K loss functions (`SAELoss`), trainer (`SAETrainer`), metrics, and feature discovery tools.
`spatial_sae.probing`	Linear probe baselines (`LinearProbe`) and direction steering comparison ($h' = h + \alpha \cdot w_{\text{probe}}$).
`spatial_sae.intervention`	Steering intervention engine (`SAEInterventionEngine`), PyTorch forward steering hooks, and strategies (`add`, `clamp`, `suppress`, `scale`).
`spatial_sae.evaluation`	Quantitative metrics for spatial accuracy, feature selectivity, counterfactual invariance, specificity score, and KL divergence.
`spatial_sae.experiments`	Standardized experiment execution scripts and compositional binding evaluation.

8. Quick Start (Synthetic Dry Run)

You can run a complete end-to-end synthetic pipeline (dataset generation, dry-run extraction, SAE training, feature discovery, and causal intervention) without downloading a VLM model:

Step 1: Generate Synthetic Dataset & Counterfactual Pairs

python scripts/prepare_data.py --config configs/data/default.yaml --num_pairs 100

Step 2: Extract Activations (Dry Run with Mock VLM)

python scripts/extract_activations.py --config configs/extraction/default.yaml --dry_run

Step 3: Train Sparse Autoencoder (SAE)

python scripts/train_sae.py --config configs/sae/default.yaml

Step 4: Discover Features & Analyze Correlations

python scripts/analyze_features.py --config configs/sae/default.yaml

Step 5: Run Linear Probe Baseline vs Causal Steering Interventions

# Linear Probe Baseline
python scripts/run_probe.py --config configs/experiments/default.yaml

# SAE Causal Steering & Compositional Binding Experiment
python scripts/run_intervention.py --config configs/intervention/default.yaml --compositional

9. Real VLM Model Experiments

To configure and run experiments on a real open-source Vision-Language Model (such as Qwen/Qwen2-VL-7B-Instruct or llava-hf/llava-1.5-7b-hf):

Step 1: Edit Model Config (`configs/model/default.yaml`)

model_name: "Qwen/Qwen2-VL-7B-Instruct"
adapter_type: "huggingface"
device: "cuda"
dtype: "bfloat16"
target_layer: "model.layers.16"

Step 2: Extract Activations from Real VLM

python scripts/extract_activations.py --config configs/extraction/default.yaml --model_config configs/model/default.yaml

Step 3: Train SAE & Run Interventions

python scripts/train_sae.py --config configs/sae/default.yaml
python scripts/run_intervention.py --config configs/intervention/default.yaml

10. Evaluation Metrics, Ablations & Failure Modes

Core Evaluation Metrics

Metric	Definition & Objective
Spatial Accuracy	Correctness of target spatial relation classification.
Spatial Selectivity & Purity	Correlation with spatial label vs. maximum correlation with non-spatial attribute confounds.
Counterfactual Invariance	Cosine similarity of feature activations between counterfactual scene pairs (e.g., A: red ball in blue box vs B: blue ball in red box).
Intended Effect Rate	Percentage of steered forward passes where model prediction matches the targeted spatial relation.
Specificity & Collateral Shift	Frequency of unintended mutations to object identity/color bindings; measured alongside KL divergence of output logit distributions.

Planned Ablations

SAE Expansion Factor: Varying latent dictionary capacity ($4\times, 8\times, 16\times$).
Sparsity Loss Formulation: Comparing $L_1$ penalty regularization vs. Top-K exact activation selection.
Intervention Locality: Steering at a single intermediate layer vs. multi-layer simultaneous interventions.

Expected Failure Modes

Polysemantic Features: Single SAE features responding to both a spatial relation AND specific object identities.
Distributed Spatial Concepts: Spatial concept represented across a linear combination of multiple SAE features rather than an isolated monosemantic unit.
Reconstruction Degradation: High SAE reconstruction error distorting general language generation capabilities and token coherence.

11. Reproducibility Protocol

Deterministic Seeding: Explicit global random seed management (set_seed()) synchronized across Python, NumPy, and PyTorch.
Configuration Tracking: Exact configuration files (configs/**/*.yaml) serialized and saved alongside all checkpoint weights and evaluation runs.
Full Metadata Recording: Comprehensive logging of activation extraction shapes, layer designations, trained SAE loss trajectories, and benchmark metrics.

12. Repository & Source Code

The full research codebase, experimental configs, and reproduction scripts are hosted on GitHub:
https://github.com/kernel-loophole/autoencoder →

← Back to all blog posts Home →