Conditional Optimal Bridge for Riemannian Activation Steering
2026-07-12 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors introduce a new method called COBRAS to control large language models during use without retraining them. Unlike older methods that use rough guesses, their approach is based on solid math from something called the Schrödinger Bridge, making the control more precise and adaptive to each input. They tested COBRAS on various models and goals like helpfulness and truthfulness, finding it works better and avoids common problems seen with previous methods. Their work provides a clearer theory and better practical results for steering model behavior.
activation steeringlarge language modelsSchrödinger Bridgeentropic optimal transportlog-density ratioquery-adaptiveSinkhorn potentialsprobability flow ODEmodel alignmentout-of-distribution
Authors
Seyed Arshan Dalili, Ajay Narayanan Sridhar, Vijaykrishnan Narayanan, Mehrdad Mahdavi
Abstract
Activation steering offers a lightweight alternative to fine-tuning for controlling large language models at inference time. While many existing methods implicitly optimize a log-density-ratio objective between desired and undesired activation distributions, they do so heuristically rather than deriving it from a principled optimization problem. Moreover, these methods produce query-independent steering directions that can degrade performance on both in-distribution and out-of-distribution (OOD) inputs. We introduce \textsc{Cobras} (Conditional Optimal Bridge for Riemannian Activation Steering), which addresses both limitations by casting activation steering as a Schrödinger Bridge on the residual-stream hypersphere. This formulation yields, to our knowledge, the first principled derivation of the log-density-ratio steering objective from a well-posed optimization problem. Solving the bridge via entropic optimal transport and extracting the probability flow ODE recovers the widely used density-ratio gradient as a special case when the Sinkhorn potentials are uniform. Crucially, the Schrödinger potentials are evaluated at the current activation, making the resulting steering direction inherently query-adaptive. Empirically, across four models and three alignment axes (helpfulness, truthfulness, and detoxification), \textsc{Cobras} consistently outperforms prior activation steering baselines while avoiding the OOD degradation commonly observed in existing methods. The code can be found at https://github.com/arshandalili/cobras.