Concept Heterogeneity-aware Representation Steering
Abstract
Representation steering offers a lightweight mechanism for controlling the behavior of large language models (LLMs) by intervening on internal activations at inference time. Most existing methods rely on a single global steering direction, typically obtained via difference-in-means over contrastive datasets. This approach implicitly assumes that the target concept is homogeneously represented across the embedding space. In practice, however, LLM representations can be highly non-homogeneous, exhibiting clustered, context-dependent structure, which renders global steering directions brittle. In this work, we view representation steering through the lens of optimal transport (OT), noting that standard difference-in-means steering implicitly corresponds to the OT map between two identical distributions with differing first moments, yielding a global translation. To relax this restrictive assumption, we theoretically model source and target representations as Gaussian mixture models and formulate steering as a discrete OT problem between semantic latent clusters. From the resulting transport plan, we derive an explicit, input-dependent steering map via barycentric projection, producing a smooth, kernel-weighted combination of cluster-level shifts. We term this method Concept Heterogeneity-aware Representation Steering (CHaRS). Through numerous experimental settings, we show that CHaRS yields more effective behavioral control than global steering. The code is publicly available at https://github.com/lazizcodes/CHaRS.
Lay Summary
Most common representation steering methods apply a single global direction to control LLM behavior, implicitly assuming that the distributions of the source and target concepts differ only by their first moments. However, this assumption breaks down in practice due to the clustered, context-dependent structure of LLM representations. Our paper presents CHaRS, which reframes steering as a discrete optimal transport problem between semantic latent clusters, modelling source and target representations as Gaussian mixtures and deriving an explicit, input-dependent steering map via barycentric projection that produces a smooth, kernel-weighted combination of cluster-level shifts. CHaRS yields more effective and robust behavioral control than global steering across numerous experimental settings, with a theoretically grounded framework that contains the standard difference-in-means steering as a special case.