Probe-Free Low-Rank Activation Intervention (2025)

Abstract: Language models (LMs) can produce texts that appear accurate and coherent but contain untruthful or toxic content. Inference-time interventions that edit the hidden activations have shown promising results in steering the LMs towards desirable generations. Existing activation intervention methods often comprise an activation probe to detect undesirable generation, triggering the activation modification to steer subsequent generation. This paper proposes a probe-free intervention method FLORAIN for all attention heads in a specific activation layer. It eliminates the need to train classifiers for probing purposes. The intervention function is parametrized by a sample-wise nonlinear low-rank mapping, which is trained by minimizing the distance between the modified activations and their projection onto the manifold of desirable content. Under specific constructions of the manifold and projection distance, we show that the intervention strategy can be computed efficiently by solving a smooth optimization problem. The empirical results, benchmarked on multiple base models, demonstrate that FLORAIN consistently outperforms several baseline methods in enhancing model truthfulness and quality across generation and multiple-choice tasks. Our implementation can be found at https://github.com/nguyenngocbaocmt02/EFI.

## introduction transformer-based language models (lms) have revolutionized generative modeling for natural language processing (nlp) @xcite @xcite . the lm training pipeline in various nlp areas typically involves three steps: (i) provide the model with prompts that include demonstrations of the task, (ii) learn from the examples, and (iii) * equal contribution † equal contribution make predictions without additional training. in the generating process, there is evidence @xcite @xcite showing that lm output can contain undesirable contents, such as inaccurate answers, toxic answers, or answers without linguistic meaning. one of the predominant solutions to the above issues is activation intervention @xcite @xcite . compared with previous studies on model editing @xcite and supervised fine-tuning @xcite , activation intervention alters the model activations during inference time. therefore, it does not need to alter model weights using a subset of text samples, requiring fewer computational resources. however, new challenges must be addressed in the activation intervention task: (q1) how to improve the quality of activation intervention through appropriate modeling with theoretical underpinnings? (q2) how to enhance the computational efficiency of activation intervention by designing a low-cost intervention strategy? for question (q1), @xcite reveals latent knowledge inside the internal activations of a language model and finds a direction in activation space that satisfies logical consistency properties. following this observation, several works focus on finding a good intervention systematically by twostage methods @xcite @xcite . in the first stage, they train a classifier (the probe) on the activations of a network to identify a subset of attention heads that are most important for learning the specific task. in the second step, they propose intervention policies based on the probe and geometric transformation, e.g., linear operations @xcite , addictive offset bias vector @xcite , linear transformation using chance constrained programming @xcite , and householder transformation @xcite . the preliminary success of these activation intervention methods motivates us to improve the desirable generation of lms. specifically, we want to design an inference activation intervention method that does not require the 'detect and rectify' procedure, addressing (q1) and (q2) simultaneously. our core idea is to depict the region of the desirable answers and consider the related low-rank transformation mapping for intervention. based on these, the parameters in the intervention mapping are the solutions to a nonlinear low-rank optimization problem, which minimizes the distance between the intervention vector and its projection on the region. compared with the previous region modelings of the lm activation space @xcite , we give the concrete formulation of the region with an analytical projection operator under the suitable distance measure. ## nonlinear low-rank intervention mapping this section defines the nonlinear low-rank activation intervention maps. in our probe-free intervention framework, this mapping serves as the computationally efficient intervention function during inference time. for clarity, we consider a specific layer ℓ of the transformer network, and the index of the layer ℓ will be omitted in the following statement. this layer contains h attention heads; each head is of dimension d, so the output of this layer is a @xmath0. for an input x i , the activation in the layer ℓ is denoted by a @xmath1; a i2 ; • • • ; a ih ] ∈ r d . our inferencetime intervention will modify the activation vector head-wise for a ih for all @xmath2. the construction of the low-rank mapping involves two key considerations: (i) the intervention strategy should be generalized to different inputs, and (ii) the intervention strategy should not change a desirable answer too much. motivated by these requirements, we construct our efficient nonlinear low-rank intervention as follows: for a better understanding of the mapping constructed, we consider the degenerate case: if w is a zero matrix and ϕ is a linear function, then l(a) = b and f (a) = (i + br ⊤ )a + s. this represents the classical bilinear low-rank mapping, which does not depend on the input vector a. ## probe-free intervention framework we propose a probe-free intervention framework based on the low-rank mapping described above. the framework performs the following tasks: (i) models the desirable answer region, and (ii) computes the intervention parameters by solving the associated smooth optimization problem. we begin by presenting the data matrix and the framework setup. from the training data, we collect the activations of the desirable output and form a matrix g @xmath3, where |g(q)| denotes the cardinality of the set g(q). each column of g q represents the activation of one desirable answer. we construct the undesirable matrix b @xmath4. let m q denote a generic manifold representing the region of desirable answer data points. the nonlinear lowrank mapping f is trained by solving: we now describe one specific instance of our framework. the manifold m q for question q is chosen as an ellipsoid of the form theorem 4.1 gives the optimization problem under the mahalanobis distance with analytical projection expression. theorem 4.1 (ellipsoidal manifold and mahalanobis distance). suppose that the manifold, the distance measure, and the projection operator are chosen as in (3). problem (2) becomes the objective function for the optimization problem is straightforward to compute, as it transforms the distance between (i + l(a)r)a + s and its projection into the distance between (i + l(a)r)a + s and a fixed vector μq . readers can refer to the appendix b for the proof of theorem 4.1. we now discuss how to derive the parameters to specify the manifold. the main challenge here is the high dimensionality of the activation vectors (@xmath5, while for each question q, we observe fewer than ten samples in total foot_1 . next, we present a practical method for estimating μq , σq , and the radius ρ q . from the training data, we compute the question-specific mean vector for question q and the mean vector computed by all questions due to the noise incurred when computing μ+ q with small sample sizes, additional information must be incorporated to accurately construct the mean vector μq . we propose the following extrapolation scheme: for the covariance matrix, we assume that σ-1 q is constant across all q. this assumption is a fundamental component of linear discriminant analysis in machine learning @xcite . we begin by computing the empirical covariance matrix: to estimate radius ρ, the heuristic approach is to find the minimum ρ such that all samples remain within the desirable answer region, centered at μ+ q . for question q, we define: the degenerated bilinear low-rank optimization problem has many applications in the statistical estimation field (e.g., matrix sensing, matrix recovery, pca) @xcite . the literature shows that simple methods such as gradient descent can converge to the ground truth under mild assumptions with good initialization. in contrast, our objective function (w, r, b, s) → f(w, r, b, s) cannot be represented by the factorization formulation and contains an extra nonlinear activation function, making it challenging to derive similar theoretical guarantees. however, the first-order algorithm is still a good choice with lightweight computational complexity. we emphasize that the activation function ϕ(•) = tanh(•) introduces smoothness to the objective function, thereby stabilizing the optimization processes during gradient descent. furthermore, we propose to use preconditioned gradient descent, a powerful approach that has gained popularity in (bilinear) low-rank optimization. this method effectively balances the scale at each iteration, accelerating convergence while maintaining the same computational complexity. ## empirical results this section presents the empirical results of our proposed algorithm florain. section 6.1 clarifies the setting of our experiments on the truth-fulqa dataset, including details about datasets, tasks, metrics, baselines, and computational resources. section 6.2 showcases the superiority of our proposed methods to other baselines. we introduce the above contents in the toxic comments classification challenge dataset in appendix c. ## conclusion and future directions in this paper, we introduced florain, a novel probe-free low-rank intervention framework for activation intervention. the framework aims to minimize the distance between the post-intervention vector and its projection onto the desirable answer region, an ellipsoid that is carefully estimated using first-order and second-order statistical information. unlike existing intervention methods, our probefree approach does not require classifiers to identify the responsible heads or layers for intervention. instead, it formulates an optimization problem that incorporates region modeling and analytical projection. additionally, the intervention is performed in a single layer of the output transformer network, allowing for efficient parallel computation and mitigating distribution shift phenomena. a potential direction for future research is to extend the region modeling approach, such as by considering the subspace spanned by the desirable answer matrix as the desirable answer region.