Multinomial concept representations improve language model steering accuracy
MAxBench: A Multinomial Concept Recovery Benchmark
Machine LearningArtificial IntelligenceComputation and Language
Summary
It's easier to guide a language AI's behavior when dealing with simple yes-or-no ideas, but many ideas have several parts, like different types of animals or countries. The authors created MAxBench, a tool that tests how well various methods find these complex ideas inside AI models. They found that using shapes called affine subspaces works better to steer the AI than simpler methods, and that some newer approaches also do well. However, none of the tested methods were better than just giving the AI carefully written instructions, which matches earlier findings.
What this means in practice
- •For machine learning engineers: Use affine subspace methods to more reliably guide language model outputs for concepts with multiple categories.
- •For natural language processing teams: Evaluate and select concept recovery methods for nuanced behavior control in large language models using MAxBench.
Authors
Divya Appapogu, Freya Behrens, Yonatan Belinkov, Aaron Mueller
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.