MAxBench: A Multinomial Concept Recovery Benchmark

AI summary

Overview

Research area: Mechanistic interpretability for large language models — specifically, how semantic concepts are geometrically represented in model activation spaces and how those representations can be used to steer model behavior.

Technical level: Intermediate to Advanced. Readers should be comfortable with basic linear algebra (subspaces, bases, singular value decomposition) and standard interpretability techniques like activation steering and sparse autoencoders.

Scope: This paper introduces MAxBench, a benchmark for comparing how well different methods recover multidimensional ("multinomial") concept representations from language model activations, evaluated across 10 localization methods, 6 concepts, 4 models, and 5 geometry types.

What This Paper Is About

Interpretability research has shown that binary concepts like refusal can be captured as a single direction in a language model's activation space and used to steer the model. However, many interesting concepts — Animals, Countries, Vehicles — are not binary: they contain many subcategories with many instances each, and the right geometry to represent them is unclear. This paper introduces MAxBench, a benchmark that evaluates how well different concept-recovery methods capture these multinomial concepts, using a geometry-agnostic sampling procedure to compare them fairly.

Key Contributions

1. A sampling-based steering framework for evaluating any concept representation regardless of its geometry — linear, affine, or non-linear — by drawing points from the recovered object and injecting them into the model's hidden state.

2. MAxBench, a benchmark with matching evaluation metrics that separates exclusivity (does the representation capture only the target concept, like precision) from completeness (does it cover the full range of concept values, like recall).

3. A large empirical comparison of 10 concept-localization methods spanning 5 geometry types (rank-one directions, linear subspaces, affine subspaces, collections of directions, and non-linear manifolds) across 6 concepts and 4 instruction-tuned models.

4. Affine generalizations of existing steering methods (DiffMean, Linear Probe, ReFT) that recover affine subspaces rather than directions through the origin, plus a new Schatten Probe that encourages low-rank classifier weights.

Main Findings

• Multidimensional beats one-dimensional. Subspace-based methods consistently outperform rank-one methods like DiffMean-r1 and ReFT-r1, confirming that multinomial concepts need more than a single direction.

• Affine beats linear — mainly because of the offset. Affine subspaces (with a non-zero offset set to the concept mean) steer more reliably and better cover the concept than their linear counterparts through the origin. The bulk of this advantage comes from where the representation is centered, not from which bases are chosen.

• Isotropic sampling is surprisingly strong on Gemma-3-1B. Sampling points isotropically around the concept centroid — without constraining to the recovered subspace — outperforms several linear and rank-one methods on that model, hinting at a "prototype-like" structure. This advantage does not hold for larger models, where basis directions matter more.

• Manifold methods are competitive when applicable. Spline-based non-linear manifolds match the best methods on ordered concepts like Days of the week, but cannot currently be applied to unordered concepts (Animals, Countries).

• Prompting remains hard to beat. No steering method consistently outperforms prompting, which dominates on the three larger models (Gemma-3-27B, Llama-3.1-8B, Qwen-3.5-4B). This replicates prior findings on binary concepts and suggests steering is best framed as a lightweight, interpretable alternative rather than a superior one.

• Rank has diminishing returns. MAxBench scores for affine methods are fairly flat across ranks 1 through 16, meaning low-dimensional subspaces suffice if the offset is chosen well.

Methodology in Plain English

The pipeline has three stages.

1. Recover a concept geometry. For each concept (Animals, Countries, Vehicles, Plants, Days, Years), the authors collect sentences that end with a concept instance (e.g., "The country that contains the Great Pyramid of Giza is ____"). They extract hidden activations from a middle-to-late layer of four instruction-tuned models. Then they apply ten different localization methods that each assume a different geometry: a single direction (DiffMean-r1, ReFT-r1), an affine subspace (PCA, Factor Analysis, Linear Probe, a new Schatten Probe, DiffMean extended to class-pair differences), a sparse dictionary of directions (SAE), a Gaussian mixture of subspaces (MFA), or a spline manifold (for ordered concepts).

2. Sample and intervene. For affine subspaces, they sample coefficients from a Gaussian, multiply by the basis matrix, and add the offset. For manifolds, they sample a position along the spline. Each sampled point is injected into the model's residual stream via a convex interpolation h' = (1-α)h + αs, with α decayed exponentially over generation steps to keep output fluent. Neutral AlpacaEval prompts serve as inputs.

3. Score the outputs. An LLM judge (Qwen-3.6-27B) rates each generation for concept presence (0-2), fluency (0-2), and instruction-following (0-2), plus a new diversity score computed as the normalized Shannon entropy of the distinct concept instances and semantic categories generated. The MAxBench score is the harmonic mean of concept, fluency, instruction, and diversity.

To test generalizability, they use 10 dataset splits — some holding out entire semantic categories, others holding out individual entities — so methods must perform well on unseen instances.

Why This Matters

Impact on research. Interpretability has mostly validated its tools on binary or rank-one concepts. This paper shows that conclusions from that setting do not automatically carry over: the geometry, the appropriate rank, and the best method all change when concepts have internal category structure. It also provides a shared benchmark and metric suite (especially the diversity score, which cleanly separates recall from precision) so future work can be compared fairly.

Real-world applications:

• Safer model steering — Better multinomial concept representations could let practitioners steer models toward or away from whole semantic families (e.g., avoiding all animal-related content in a children's application) rather than single keywords.
• Bias auditing — Probing how a model organizes concepts like Countries or Plants can reveal whether certain subcategories are stereotyped, underrepresented, or collapsed together.
• Controllable generation — Fine-grained topic control for creative writing, recommendation, or domain adaptation without retraining.
• Model diagnostics — Diversity scores provide a signal for whether a model's internal representations of a concept have "collapsed" onto a narrow subset.

Industry relevance. Steering is attractive because it is lighter than fine-tuning and more interpretable than prompting. This paper's finding that prompting still wins on large models is directly relevant to anyone choosing an intervention strategy in production: affinity-tuned steering is not yet a drop-in replacement, but affine subspaces close part of the gap cheaply.

Future Directions

• Extend manifold methods to unordered concepts. Spline-fitting currently only works for ordered values like Days and Years; Animals and Countries need different non-linear constructs.

• Investigate non-zero offsets more systematically. Since offsets drive most of the affine advantage, exploring alternative offset estimators (beyond the mean of positive activations) could yield further gains.

• Layer-wise evaluation. The paper fixes one layer per model largely due to SAE availability; how the optimal geometry varies across depth is unexplored.

• Reconcile the prototype-vs-direction debate. The strong isotropic baseline on Gemma-3-1B suggests concepts may be organized around centroids rather than directions — worth testing across more models and with better prototype-based methods.

• Close the gap with prompting. No steering method yet matches prompting on large models; understanding why might require richer, potentially non-linear, concept representations.

Target Audience

This paper is most valuable to mechanistic interpretability researchers, particularly those working on activation steering, concept localization, and representation geometry. It will also interest practitioners who use steering vectors or SAEs for model control and want to know which geometry to try when their concept is not binary. Readers without background in linear algebra or interpretability will find the paper dense, but the high-level findings — affine beats linear, offsets matter, prompting still leads — are accessible to anyone familiar with the broader steering literature.

Authors’ abstract

Fine-grained control of language model behaviors (e.g., steering) is among the more actionable outcomes of interpretability research. For binary concepts such as refusal, a single direction in activation space often suffices for steering. However, many concepts are not binary: Animals and Countries contain many subcategories, each with multiple instances. For these concepts, the search space over possible representation geometries is far larger than for binary concepts; it is thus not clear what geometries are most appropriate, nor what methods are most effective at recovering them. In this work, we introduce MAxBench, a geometry-agnostic evaluation framework for multinomial concept representations based on sampling from the recovered concept representation. We use MAxBench to compare 10 localization methods (covering 5 geometry types) across 6 concepts and 4 models. Using this framework, we find that (i) affine subspaces steer more reliably and have greater recall than rank-one or linear subspaces; (ii) much of this advantage is due to better non-zero offsets rather than the choice of bases; (iii) manifold steering is competitive with the best methods when applicable; and (iv) no method consistently outperforms prompting, in alignment with prior findings on binary concepts. These findings underscore the importance of expanding the scope of interpretability research and meta-evaluation to concepts with more varied structure.