Expander OMP Sparse Inference
Implementation & benchmark of arXiv:2607.01799 — Expander Sparse Autoencoders: Parameter-Efficient Dictionaries for Mechanistic Interpretability
Source paper: Expander Sparse Autoencoders: Parameter-Efficient Dictionaries for Mechanistic Interpretability arXiv:2607.01799 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Use the expander decoder as a structured sparse-coding dictionary and replace dense OMP correlation steps with edge-wise gather-and-reduce operations. This is useful when codes must be inferred iteratively, including interpretable feature extraction, sparse retrieval, or an inference-time latent selector that cannot rely entirely on an amortized encoder.
Formulas
Mathematical statement
Orthogonal Matching Pursuit maintains a residual r^(t) = h - W x^(t) and selects the feature with the largest dictionary correlation. For a dense dictionary, the correlation vector is c^(t) = W transpose r^(t), costing O(mn). With W = A elementwise-multiplied by Theta and exactly d active coordinates per column, each correlation is c_j^(t) = sum over i with A_ij = 1 of w_ij r_i^(t), so all correlations require O(dn) multiply-adds and gathers. OMP selects j_t = argmax over unselected j of |c_j^(t)|, enlarges the support S, and refits coefficients by least squares on the selected columns. The residual is then recomputed as r^(t+1) = h - W_S x_S. Expansion of neighborhoods for sets up to 2k reduces support collisions, while column normalization makes correlation magnitudes comparable.
Implementation notes
(1) Integration point: implement an inference-only sparse encoder for a frozen dictionary, or use it as a teacher that generates sparse targets for a learned amortized encoder. Store the dictionary as edge arrays row[e], col[e], and weight[e] for e from 1 to dn. (2) Pseudocode: set residual R = H, support S_b = empty, and coefficient matrix X = zero for each batch item b. For t from 0 to k minus 1, initialize correlations C = zero; for each edge e, gather R[:, row[e]], multiply by weight[e], and reduce into C[:, col[e]]. Set selected columns to negative infinity, choose j_t as the maximum absolute correlation, and append it to the support. Form W_S from the selected columns, solve x_S = (W_S transpose W_S + lambda I)^(-1) W_S transpose h, and update R = H - W_S x_S. (3) The mathematical contribution is the O(dn) correlation formula and OMP support/refit recursion. Estimate speed from actual fused gather-reduce kernels; lambda, batching, and stopping thresholds are engineering choices. (4) First experiment: use 10,000 Pythia-70M residual vectors and compare dense batched OMP, expander OMP, and the amortized TopK encoder. Measure latency, peak memory, reconstruction error versus iteration count, and support overlap with dense OMP. Success is a substantial correlation-time reduction with comparable MSE at k selected features; report the break-even point where sparse-kernel overhead dominates.
Verification
Stage 1 · Toy mechanism gate: Failed ✗
Stage 2 · Mechanism transferred to benchmark: Not tested
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 9/10
Built dense and edge-wise OMP using fixed-degree sparse dictionaries stored as edge arrays. The core correlation identity was verified to floating-point precision (maximum error 4.4e-16), and both OMP implementations produced identical supports, coefficients, and reconstruction MSE (9.58e-5). However, the transparent unfused edge gather-reduce was slower than optimized dense BLAS by 1.8x–163x in the tested sizes, so the intended speedup was not observed; a fused kernel is needed for a fair systems-level test.
- Agent confidence
- 9/10
- Baseline
- Dense NumPy BLAS correlation; reconstruction MSE 9.5779e-5. Correlation times for n=128,256,512,1024 were 0.00701, 0.00191, 0.00393, and 0.000396 seconds.
- Idea
- Edge-wise correlation with d=6 and m=96, nominal arithmetic ratio d/m=0.0625. Correlation and OMP outputs matched exactly numerically, but edge/dense wall-time ratios were 1.82x, 13.56x, 13.07x, and 162.84x for n=128,256,1024 respectively.
- Note
- The reported timing sequence is noisy for small BLAS workloads; all measured edge implementations were slower, so no speed win was claimed.
Limitations: Only synthetic CPU experiments were run. The edge reduction is a Python loop rather than a fused CUDA or compiled scatter-reduce kernel; no Pythia residuals, amortized TopK encoder, GPU benchmark, peak-memory measurement, or large-scale end-to-end inference test was performed.
How to run: python3 expander_omp.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Artifacts
- expander_omp.py 4.4 KB View Raw
- report.md 1.6 KB View
- results.json 1.2 KB View Raw