Concept Heterogeneity-aware Representation Steering

TL;DR

This paper introduces CHaRS, a novel method for representation steering in large language models that accounts for concept heterogeneity by modeling representations as Gaussian mixtures and using optimal transport for input-dependent control.

Contribution

The paper proposes Concept Heterogeneity-aware Representation Steering (CHaRS), a new approach that models representation heterogeneity with Gaussian mixtures and applies optimal transport for more effective steering.

Findings

CHaRS outperforms global steering in various experimental settings.

Modeling representation heterogeneity improves control precision.

Input-dependent steering maps enhance behavioral control.

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 unimodal Gaussian distributions with identical covariance, 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.