Constrain Alignment with Sparse Autoencoders
Authors: Qingyu Yin
Chak Tou Leong
Hongbo Zhang
Minjun Zhu
Hanqi Yan
Qiang Zhang
Yulan He
Wenjie Li
Jun Wang
Yue Zhang
Linyi Yang
Venue: ICML 2025 poster (2025)
Keywords: Alignment
SAE
LLM
Primary area: social_aspects->accountability_transparency_and_interpretability

Abstract: The alignment of large language models (LLMs) with human preferences remains a key challenge. While post-training techniques like Reinforcement Learning from Human Feedback (RLHF) and Direct Preference Optimization (DPO) have achieved notable success, they often experience computational inefficiencies and training instability. In this paper, we propose \textbf{F}eature-level constrained \textbf{P}reference \textbf{O}ptimization (FPO), a novel method designed to simplify the alignment process while ensuring stability. FPO leverages pre-trained Sparse Autoencoders (SAEs) and introduces feature-level constraints, allowing for efficient, sparsity-enforced alignment. Our approach enjoys efficiency by using sparse features activated in a well-trained sparse autoencoder and the quality of sequential KL divergence by using the feature-level offline reference. Experimental results on benchmark datasets demonstrate that FPO achieves an above 5\% absolute improvement in win rate with much lower computational cost compared to state-of-the-art baselines, making it a promising solution for efficient and controllable LLM alignments.

Constrain Alignment with Sparse Autoencoders
Qingyu Yin 12* Chak Tou Leong 3 Hongbo Zhang?
Minjun Zhu Hanqi Yan Qiang Zhang'
Yulan He* Wenjie Li
Jun Wang
Yue Zhang?
Linyi Yang
Abstract
forcement Learning from Human Feedback (RLHF) (Chris- tiano et al., 2017; Ouyang et al., 2022) has emerged as lead- ing technique, integrating human feedback to guide models towards producing valuable and useful outputs. Despite its success, RLHF involves complex mechanisms such as reward modeling and policy gradients, which introduce sig- nificant training complexity and computational cost (Zheng et al., 2023b; Rafailov al., 2024). To address these lim- itations, Direct Preference Optimization (DPO) (Rafailov et al., 2024) has been proposed as more efficient alter- native. Unlike reward-based methods such as Proximal Policy Optimization (PPO) (Schulman et al., 2017). DPO directly adjusts the model's output probabilities based on human preferences, reducing training complexity and com- putational cost. DPO-like approaches can offer more stable and faster alignment process by bypassing the chal- lenges associated with reward models and policy updates, making it a compelling solution for efficient LLM alignment since DPO uses a reference model to stabilize post-training.
The alignment of large language models (LLMs) with human preferences remains a key challenge. While post-training techniques like Reinforce- ment Learning from Human Feedback (RLHF) and Direct Preference Optimization (DPO) have achieved notable success, they often experience computational inefficiencies and training insta- bility. In this paper, we propose Feature-level constrained Preference Optimization (FPO), a novel method designed simplify the alignment process while ensuring stabiliity FPO leverages pre-trained Sparse Autoencoders (SAEs) and in- troduces feature-level constraints, allowing for efficient, sparsity-enforced alignment. Our ap- proach enjoys efficiency by using sparse features activated in a well-trained sparse autoencoder and the quality of sequential KL. divergence by using the feature-level offline reference. Experimental results on benchmark datasets demonstrate that FPO achieves an above 5% absolute improvement in win rate with much lower computational cost compared to state-of-the-art baselines, making it a promising solution for efficient and controllable LLM alignments. Code is available at Feature- Alignment.
Recent advancements in DPO focus on mainly two direc- tions: efficiency i.e., further simplifying the constraints of DPO, and controllability i.e., keeping the balance between alignment and generation diversity. In terms of simplicity, methods like SimPO (Meng et al., 2024) and Odds Ra- tio Preference Optimization (ORPO) (Hong et al., 2024) climinate the need for a reference model by using the av- erage log probability of sequences as an implicit normal- izer, thereby reducing memory usage and computational demands. However, DPO's performance is sensitive the strength of constraints from the reference policy (Liu et al., 2024), and these reference-free free alignment approaches (Hong et al., 2024; Meng et al., 2024) can compromise control, resulting in unstable training. In terms of controllability, Token-level Direct Preference Optimization (TDPO) (Zeng et al., 2024) introduces token-level rewards and sequen- tial Kullback-Leibler (KI.) divergence (Kullback & Leibler, 1951) to tackle issues related to linguistic coherence, diver- sity, and stability. However, it comes at the cost of increased computational complexiry, introducing an additional sequen- tial KL and depending on reference models, complicating the loss computation.
1. Introduction
Aligning large language models (LLMs) with human values and practical objectives is critical challenge in AI develop- ment (Wang et al., 2023). Post-training methods, including fine-tuning (Wei et al., 2022; Chung et al., 2024) and align- ment strategies (Tunstall et al., 2023), have played a signif- icant role in refining LLM behavior. Among these, Rein-
'Equal contribution 'Zhejiang University "Westlake Univer- sity ' 'Hnnkknn Polytechnic University KKin's 's College London 2 Uniiversity College London "Southern University of Science and Technology. Correspondence to: Linyi Yang and Yue Zhang <yan- glinyiued@@ gmail.com, zhangyue @ westlaak..dd..com
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).
A natural hypothesis arises: "Is there a method that can strike the right balance between efficiency and controlla-

---

Constrain Alignment with Sparse Autocncoders
molMules (z)
melan (z)
Phase 1: Offiine Caching
Phase 2: FPO Training
valale)
Dnpo flog
0
layer E
Improvements
MSE
9
layer C
layer
LainPo
log (Buula))
log mo(ai
Reference-Free Unstable
Nolpule) Tre(Bulo)
mv(ayl)) mee(bille)
Umoro
flog
Blog
STeOO Mt)
B yoll
Token-level
Lyro
log mo(s)a))))
laul
log mo(suli)) T Gin (at, Mar 31)
Reference- Needed
Stable
Offine-RRffreene
Figure 1:
Left. The DPO objective loss function and its two main improvement directions: SimPO and TDPO. SimPO focuses on simplifying the reference model, while TDPO concentrates on controlling the alignment process to enhance generation diversity. Right. The pipeline of FPO consists of sparse autoencoders and the feature-level MSE constraints.
bility?" In response, we propose FPO, Feature-level Con- strained Direct Preference Optimization (See Figure 1), in- troducing an efficient and controllable method for constrain- ing the model the feature level. Here feature refers to a salient piece of information for the model decision (Huben et al., 2024). Intuitively, adjusting the model us- ing feature-level preferences allows fine-grained adjustment that minimizes the side impact, by avoiding the negative influence of spurious features in course-grann grained control such as token level regularization (Zeng et al., 2024).
results that meet or even exceed the effectiveness of sequen- tial KL, at a significantly lower computational cost (17.6% reductions compared TDPO2 as shown in Figure 4 Left). Additionally, we introduce detailed ablation studies to show that method maintains stable performance over differ- ent temperatures and the selection of SAE layers.
Overall, we show that FPO enjoys the efficiency of SimPO by using the offline reference control, while also the con- straint quality of sequential KL by using the sparse feature- level constraints. To our knowledge, this is the first ap- proach that integrates sparse feature-level constraints into LLM alignment. By incorporating sparse autoencoders with token-level DPO, FPO makes practically meaningful and theoretically solid improvements over existing preference optimization methods along three dimensions: simplicity of implementation, efficiency, and generation diversity.
To achieve that, we derive the FPO objective by contrasting SimPO and DPO, showing the constraint term that SimPO misses. We then add such term by introducing the feature- level constraints as an alternative to the costly sequential KL (Zeng et al., 2024). We use Sparse Autoencoders (SAEs) (Huben et al., 2024), which generate representations where only a few features are active, enhancing computational efficiency (See Figure 2 Right). Furthermore, regularization in the coefficient space promotes sparsity, stability, and uniqueness in the model's representations. Since SAEs produce sparse representations, only a few dozen out of 16,000 features are active at any given time (Lieberum et al., 2024). Compared to SimPO, FPO is as efficient memory and time complexity, yet has improved controllability due to feature-level constraints; compared constraint-basse methods like TDPO, FPO matches the computational and memory efficiency of methods such SimPO, and has potentially improved performance as feature-level control can give stronger generalization than token-level control. A contrast between FPO, DPO, SimPO and TDPO is shown Figure 1.
2. Preliminary
Direct Preference Optimization (DPO). DPO, derived from Reinforcement Learning from Human Feedback (RLHF), provides direct way to align Language Mod- els (LLMs) with human preferences without explicitly using a reward model. In practice, an LLM is prompted with a sequence æ (c.g., a question) to generate a corresponding sequence y (e.g., an answer), where both and y consist of tokens. DPO maps the reward function r(x, y) to the opti- mal policy by minimizing the reverse KI. divergence from reference model. This results in the following equation for the reward:
r(x, () = Blogg [b(2))) ())(()(
+ Bog Z(x),
(1)
Our experiments demonstrate that FPO consistently out- performs tate-of-tthe-an methods based different sizes of backbone LLMs, achieving up to 5% absolute improve- ments in win rate (See Table 2) based on AlpacaEval-2 and Arena-Hard benchmarks, up to 0.5 scores on MT-Bench and competitive output diversity. By constraining the shifts of these features during the training process, we can achieve
where n0(- and )T((-1z) are policy (i.e, the LLM for post-training) and reference (i.e., the base LLM) models, respectively. 3 is the coefficient that governs the strength of the KL divergence penalty, Z(x) is the partition function. To align with human preferences, DPO uses the Bradley-Terry (BT) model for pairwise comparisons. By incorporating the
2

---

Constrain Alignment with Sparse Autoencoders
Method
Reference
Efficiency
Constraint
Value
SFT
Free
High
Weak
1.0
DPO
Offline
High
Weak
SimPO
Free
High
Weak
TDPO
Needed
Low
Strong/ Dense
FPO(Ours)
Offiine
High
Strong Sparse
0.0
10
20
30
40
50
Feature Rank Num
Figure 2: Left. Top-50 SAE feature activation value distribution in Gemma-2-24 We ranked the activated feature by its activation value. The vertical axis represents the activation values, while the horizontal axis shows the rank of the maximum activation values. This plot illustrates the sparsity of SAE-out of 16,000 features, fewer than 50 have significant activation values. Right. Comparison of existing alignment methods on (1) if they need to load a reference model when training the policy model. (2) Memory consumption. (3) Their ability to control the generation diversity.
reward function into the BT model and using the negative log-likelihood, DPO computes the loss:
where ŌTDPO, (I, Mws 3i) the KL divergence difference between the preferred and less preferred completions:
(To; Trat) - -B (o..). [log a (u(2,, Uw w))], ux ( Yuus y1) B(log
òIDRO, (I. y1) ((Droro, (x, ui merl( - DTDPO, (æ, yws Terll
(3)
To (B (x)
log (o(3n(a)
Tief (Tef(
and the sequential KL divergence between policy and refer- ence output with sequence length is defined as
Here, represents the dataset with human preference pairs. Yu and y are the preferred and less preferred completions, respectively. DPO provides direct way to align L.LMS with human preferences without the explicit use of a reward model, leveraging preference comparisons.
Droro (z, yi Tret (eot||To) = (Froffl " (
To further stabilize the gradient within the optimization, an improved loss function CIDFO, is given by replacing the regularization ÓTDPO, with:
Simple Preference Optimization (SimPO). SimPO sim- plifies DPO by removing the need for a reference model and aligning rewards directly with the length-normale log-likelihood the policy model's output. The SimPO loss function can be formulated as:
Oroeag (I, thes 9) = a (8 (a, yi me)]
(4)
-g (3 (z, Vwi mael To))
(simP((Ta) = -a WDDD (log (u(a,, yor vR))
where a is additional hyperparameter to balance between alignment and regularization, 3 is the coefficient that gov- erns the strength of the KL divergence, and sg denotes the stop-gradient operator. Unlike DPO, TDPO introduces token-level forward KL divergence, allowing for finer con- trol over model aligmment and diversity in generation, also introducing additional computational overhead.
3
B [3d]
u(x, Nw: Mi)
log To (B
log me (31lz) y.
yw
where is positive margin ensuring the reward for the preferred response exceeds that of the less preferred one by at least 7. However, while SimPO is computationally efficient, the lack of reference control (Roy et al., 2021) results in instability, as the reference model can stabilize training and improving performance (Liu al., 2024).
Sparse Autoencoders (SAE). SAEs provide method for recovering monosemantic, interpretable features, enhancing the steerability of language models, where individual neu- rons activate in semantically diverse contexts. SAEs aim to reconstruct internal representations with sparsely activated features, disentangling the representations into interpretable components. Given the latent representation of model h Rd its sparse activation Rmm is computed as:
Token-Level Direct Preference Optimization (TDPO). Leven-Level Direct Preference Optimization (TDPO) re- fines the DPO framework by operating at the token level, accounting for the sequential nature of text generation. The first version of TDPO loss function given by:
LYDPO, (ITe; meet) - B (a. )w [log (u(æ, Me 9)))
c Ce.e + b), = WT.
(5)
u(x, Mwn 31
where Wane E Roxd and Wase E mmxd are the learned weight matrices, b E Rmm is the bias vector, m the number
To (JWl()))
,9 log (xrlz)) EIeeo, (Iz, Vws Vt) (v(())
log Tra(Yu (x)
3

---

Constrain Alignment with Sparse Autocncoders
of latent features with m > d, and h is the reconstructed input, computing loss:
This similarity reveals that the reward difference in DPO can be interpreted combination of log probability differ- ence with adaptive margin "ief from the reference model, whereas SimPO calculates the average log probability differ- ence with a fixed margin. Based on the above observation, we can reframe the loss function of DPO and its two variants into a unified form:
LSAE (1) = 11h - 111 al|el],
(6)
where a controls the sparsity of the hidden representation. The norm on c enforces sparsity, ensuring only a small number of features are active at any given time (See Figure 2 Left for visualization of SAE's sparsity).
(Tro Werr P....w..))) log (u(z,, Uw: Ve. mi)]]
3. Feature-Level Direct Preference Optimization
u(I,, Mwe, (e)
3
log To (x))
3
log no (vnlzr) (8)
1y
M
-Tof-LNN KEro ywu, vt).
In the right table of Figure 2, we present comparison of FPO with other methods from three perspectives: reference model usage, efficiency, and constraint control, which is dis- tinguished from existing methods in the following aspects:
We summarize the specific implementations for DPO, SimPO and TDPO in the form of Equation (8) in Table 1. SimPO eliminates the reference model from the alignment training by using a fixed margin and omitting constraints, which reduces memory and computational costs. However, it has been criticized that completely removing reference models leads to instability (Liu et al., 2024). Our approach begins by applying the length normalization technique of SimPO to the original implicit reward difference of DPO:
Reference-free methods such as SimPO and ORPO are memory and computation efficient. However, they struggle with instability brought by the lack of refer- ence constraints.
Alignment methods with KI control on output log- its, like TDPO and KTO (Ethayarajh et al., 2024)', are powerful yet controllable, but their sequential KL based on output probabilities makes them costly.
B To( Manl
G
To (yu (z)
log
log
Ter (x)
The (x)
B
log Tver(
(ze)
log Tee((B
(9)
|Mull
Interpretability methods such as SAE are widely used for interpreting the inner representations of LLMs due to their sparse and monosemantic activations (Chen et al., 2017; Huben et al., 2024). However, this fea- ture has not yet been applied in areas outside of inter- pretability.
B
g log No (vo(()) LN
log To (o
Maell
y
Equation (9) suggests using average log probability differ- ence as the Log Probability Difference (LPD) term and introducing an adaptive margin with length normalization as the Margin. The length-normalized margin "hef-LN enhances the stability by using reference model to calculate an adap- tive margin for each preference pair. We consider an offline caching technique to minimize the computational overhead introduced by the reference model.
DPO with Reference-base Target Margin. To begin, we examine the loss functions of DPO and its enhanced variants, specifically SimPO and TDPO. By comparing Section and Equation (2), we notice that TDPO and DPO share an identical implicit reward difference term: 3 log Furly (x)
Essentially, TDPO can be viewed as an extension of DPO, where a KL constraint 6(z, Yw: 31) is incorporated into the sigmoid function (-) in addition to the implicit reward difference. Taking step further, we can isolate Teef from each implicit reward term:
(M(l]>)) - log Hur(yilz) (z)
Ta(M) (n))
Feature-level Constraints. Currently, the use of con- straints 6(x, yw) in alignment processes typically fol- lows KL divergence-based approach shown in Equation (3) and 4. However, this method has a significant issue: for most LLMs, which generally have a very large vocabulary, where we assume the vocabulary size is V. For each batch with an input length of T, the resulting output probabilities have a size of V T. This work adopts Gemma (Lieberum et al., 2024), an advanced open-sourced LLM series, which has massive vocabulary size of 265K. For an input length of 1024, this results in probabilities matrix containing approximately 262M elements, which is nearly 1/10 the size of its 2B version model. Therefore, computing the KL divergence incurs considerable computational overhead to DPO-enhancing methods such as TDPO.
3 log Te(ymnla) (x) Ter( 1)) - 8 log
Ta (2)
Ter (r)
B log To (ww (z)) log No (ux(2²) (log )w)() log Tee (ve((()())) (r))
We can that Equation (7) shares similar form with the reward difference calculation of SimPO in Section 2.
The loss function of KTO is similar that of TDPO in terms of its use of KI. divergence.

---

Constrain Alignment with Sparse Autocncoders
Table 1: Specific implementations of Log Probability Difference (LPD), Margin, and Constraint in Equation (8) for DPO, its variants SimPO and TDPO, and the proposed FPO.
Method
LPD
Margin
Constraint
Constraint Type
DPO
log mal (yo [x) 3 log No (B((())
"Yee
SimPO
8
log To (y (az)
eer log mo (In(ar)
(a constant)
TDPO,
3 log Ta (M (x) 9 log (ax) (z)
'Yre
oYDRo, (2,, Mw. 31))
KL Divergence
FPO
# Mho
log To (Uwlz)
Iual log To (uilz))
Thef-L
óero (z7, Vuu 31)
MSE
LLMs generate these sizable output probabilities by project- ing their internal representations onto vocabulary space. In contrast to this, SAE is found to be capable of projecting these representations onto a sparse feature space. Motivated by the efficient nature of sparsity, we leverage the sparse feature activations from SAE to approximate the function of KL divergence. Specifically, for the output representation (((+)) from layer E of the model at position i, we can ob- tain its sparse activation clt.() using an SAE as described in Equation (5). Since KI, divergence measures the difference between two probability distributions, we employ MSE as the loss to measure the discrepancy between the sparse acti- vation from the two models. To further improve efficiency, instead of calculating the sum of biken-wise discrepancy like TDPO, we first perform average pooling for the sparse activation across tokens and then calculate the MSE between pooled sparse activations, which gives us a more efficient sequential discrepancy:
means pre-calculating and caching the results related to the reference model needed for training and then reading them during the training loop. This approach allows us free up the reference model during alignment with only a small and acceptable I/O demand.
To explore an implementation for Equation (8) that enjoys the advantages of SimPO, such as length normalization, while ensuring stability, first, we pre-compute and store the margin "Yef-N using the length normalization for each pref- erence pair. Since it is scalar, it only occupies ((N) space to store it, where N is the number of preference pairs. Next, for the feature-level constraint, we pre-compute and store the sparse activation of each sample in the training dataset following the computation in Equation (10). Consequently, we only need to pre-compute and store one sparse activation d for each sample, which requires 0(2- - k) space. This results in a significantly smaller space requirement com- pared constraints used in TDPO, where the vocabulary size is V, for each batch with N preference pairs, requiring much larger space of O(2N- V). By combining all the above results, we arrive at the loss function for FPO:
Dero (z, yj Radl
Σ SE4
)
=
(10)
where pooled sparse activation ج = Et ctt,, Is top, (indices(e) (( U topy (indicess( ((.ce) and (-) returns the indices of the & largest elements. We focus on measuring the MSE between the largest activations to capture the discrepancy in dominant features, as these are likely to be the most influential. Echoing the strategy of TDPO, we replace DTDPO in ôTDPO, (z, Fa 3t) with Df as a plug-and-plaa -lay efficient approximation. This results in a feature-level constraint ofro (z, Bw. v1).
LePO (er)) = - -B ww wDnD log o (u(z; yws ni)}
B
u(I,, Me: U1)
log To (ro (z)
log To (o((] (11)
V Syro Nw. Bt).
4. Experimental Setup
Model and Training Settings. Our model selection is guided by two key principles: scalability and transparency. For scalability, we first select series of models spanning dif- ferent parameter sizes, including Gemma-2-28 and Gemma- 2-9B (Team et al., 2024) This ensures that we can evaluate our approach's 's performance as the model parameters scale and assess its robustness across diverse model architectures.
Building Offline Reference Margin and Constraint. We have justified the implementation of the key components in Equation (7), which is a SimPO-like reward difference with a reference-based adaptive margin and a feature-level constraint. At first glance. the reference model appears to be deeply involved in both the calculation of the margin and the constraint, making its complete elimination challenging.
For transparency. we exclusively select foundational models, which have not undergone supervised fine-tuning (SFT) or alignment processes. We begin by fine-tuning these models
Therefore, instead of directly removing the reference model, we propose a more appropriate approach: separating the computation of the reference model from the training pro- cess by computing output offline. Offline computation
We select Gemma-scope it provides pre-trained SAEs (Lieberum et al., 2024) for all layers.
5

---

Constrain Alignment with Sparse Autoencoders
Table 2: Left: Performance comparison of different methods for Gemma-2-2B and Gemma-2-9B across various benchmarks (AlpacaEval-2, Arena- Arena-Hard, and MT-Bench), compared to Supervised Fine- Tuning (SFT), DPO and variants. Length controlled Winning Rate: WR-L; Winning Rate: WR. Right: Comparison of FPO and other baseline methods in terms of the trade-off between Alignment Acc(accuracy) and Diversity iee, Diversity (Entropy) on the UltraFeddbcck dataset.
Gemma-2-2B
Gemma-2-91
Method
Alpa-2Eval-2
Arena-Hard MT-BBench
AlpacaEval-2
Arena-Hare MT-Bench
Method
Acc(%) f
1
FPO v.s.
WR-1(%) (%)
(%)
Score
(R)
(%)
WR (%)
Score
DPO
59.9
1.66
SFT
54.7
55.1
53.2
+0.5
51.2
52.4
53.4
+0.3
TDPO-1
63.2
1.65
DPO
51.7
50.8
51.6
+0.1
51.0
51.0
51.2
+0.1
TDPO-2
64.2
1.68
TDPO-1
51.5
54,4
51.4
+0.3
50.8
50.2
51.8
+0.1
SimPO
63.4
1.64
TDPO-2
50.9
54.0
50.6
+0.2
50.2
49.9
49.5
0.0
FPO
64.1
1.68
SimPO 51.1
52.2
51.4
+0.4
50.2
51.8
51.0
+0.2
using a unified conversational format provided by the Ha- los dataset, applying it to the Uttrachtt-200K (Ding et al.,
et al., 2023; Dubois et al., 2024), and Arena-Hard (Li et al., 2024; Chiang et al., 2024). These benchmarks are designed to test the models' conversational abilities across broad spectrum of tasks and have gained significant adoption in the research community. AlpacaEval 2 includes 805 questions derived from five different datasets, while MTT-Bench spans eight categories with a total of 80 questions. Arena-Hard, the most recent release, builds on MTT-Bench by introducing 500 complex technical mooblem-solving queries.
2023). Dataset for initial instruction tuning. This establishes baseline conversational capability and ensures that all our methods are compared quently, we employ the UltraFeedback (Cui et al., 2024). Dataset to align the SFT models using various methods. This approach maintains transparency and control through- out the process, as all data and methods are open-sourced across the experimental setup.
consistent SFT model. Subse-
We follow the standard protocols for each benchmark in evaluations, by computing the AScore as the margin be- tween FPO and other methods. The metries evaluated in- clude Length Controlled Winning Rate (WR-L) and Win- ning Rate (WR) for AlpacaEval-2 and Arena-Hard and score from 1-10 for MTT-Bench. For all methods, we use GPT-4 Turbo (Achiam et al., 2023) as the evaluator.
For the hyperparameters related alignment methods, such as and 3. we initially refer to the hyperparameter settings from the corresponding papers. If these settings are explic- itly provided, we directly adopt their configurations. For configurations that are not given, we perform a hyperparam- eter search to determine the optimal values. Regarding the training hyperparameters, we standardize the batch size to 32, the learning rate to 5 x 10-7, and warm-up period of 150 steps, after which the learning rate remains constant, set the epoch as 1. We employ the Adam (Kingma, 2014) and RMSProp optimizers (Graves, 2013) for Gemma- 2-2B and Gemma-2-98 respectively.
For analyzing the alignment and diversity trade-off of our method, following Zeng et al. (2024), in experiments, validate and compare FPO against several strong alignment baselines, including DPO (Rafailov et al., 2024), SimPO (Meng et al., 2024), TDPO1, and TDPO2 (Zeng et al., 2024).
5. Results and Discussions
Baseline Methods. Regarding our baseline comparison methods, we primarily compare three categories of ap- proaches. The first category consists of our foundational methods, including instruction fine-tuning (SFT) and DPO itself. Here, SFT refers to the model's 's performance after the first-stage fine-tuning, while DPO refers to the direct ap- plication of DPO for further alignment following SFT. The second category includes methods with explicit KL control and efficient reference-free methods. We select the TDPO series ie.,, TDPO-1, TDPO-2 and SimPO, as they currently represent the state-of-the-an in these two classes of methods IDPO-enhancing and -PO-simplified), respectively.
FPO Consistently Outperforms Strong Baselines on Three Benchmarks. We evaluate the performance dif- ferences between FPO and other methods across three key aspects: training accuracy, generation diversity, and perfor- mance on downstream tasks. In terms of downstream tasks, we assess the model's performance including the winning rate or score on the AlpacaFval2 Benchmark, Arena Hard, and MT Bench. As shown in Table 2, FPO achieves highly competitive results, with up to 5.08% improvement in winning rate compared other methods when testing on Gemma-2-22 2-22. Additionally, based on Gemma-2-9B, we ob- serve a consistent improvement in our method compared to baselines. However, the performance improvements on the 9B model introduced by FPO are limited compared to
Evaluation Benchmarks. We evaluate our models on three widely-used open-ended instruction- follooiinnn bench- marks: MTT-Bench (Zheng al., 2023a), AlpacaEval (Li
6

---

Constrain Alignment with Sparse Autocncoders
Chosen KL Divergence
Rejected KL Divergence
Divergence Margin
0.75
TDPD2
TDP02
x10
0.65
PPO Other Mathods
TDP01
TDP01
TPPO2
TDPO2
32 Divergence
0.60
SimPo
DPQ
X Divelgence
0.60
SmmP
6,0
SimPo
Winning Rate
DFO TOPOL
SmPP
FPO
FPO
Margin
0.55
0.45
0.45
4.5
0.50
Tw
0.40
0.15
0.15
0.0
0.8
1.2
1.0
2.4
0.0
0.6
1.2
1.8
2.4
0.0
0.6
3.2
1.8
2.4
0.35
0.2
0.4
0.6
0.8
1.0
Step
X100
Step
X10'
Step
K10
Temperature
Figure 3: Left 1. KI. Divergence on the preferred responses (chosen). Left 2. KI. Divergence on the dispreferred responses
(rejected). Right 1. KL Divergence margin i.e., [&akLL (I, yts Teell (eel|(TT) (Ta) 3 SegKL (I, Mui Teef (eel|(Te) Right 2. Win rates of FPO v.s. other methods above the improvements based on Gemma-2-2B different sampling temperatures.
the 2B model. We argue that this is because, with the same width of the SAE, smaller models, due to their lower com- plexity, achieve a more thorough decomposition features, filtering more noisy features, and leading to more accurate constraints.
ure 4, we consider this result to be highly competitive. The efficiency of FPOis reflected primarily in two aspects: (I) Offline Processing. FPO does not require an additional ref- erence model to be loaded during training, but only incurs minimal I/O overhead read pre-stored information at each step, specifically the one-dimensional tensors needed for training. This process be efficiently handled by the dataloader. (2) Sparsity. Due to the sparse activation, only need to process the activated values, reducing com- putational overhead. To validate its efficiency, we tested the memory consumption of different methods during train- ing. In terms of memory usage, FPO maintains nearly the same level of memory consumption as reference-free meth- ods like SimPO. Compared to methods that introduce more computation, such as TDPO, FPO achieves approximately 17% memory optimization.
5.1. The Trade-off between Controllability and Efficiency.
Accuracy vs. Diversity, We measure the training accu- racy on the UltraFeedback dataset, which is defined as the probability that the chosen answer's token-wiss wise probabil- ities exceed those of the rejected answer. Table 2 shows the model's generation diversity by measuring the entropy of the top 100 results on AlpacaEval2,, where the indi- cates higher values are preferable. We use bold to show the pes--pefformig result across all metrics, and underline to denote the second-best result. The results indicate that FPO achieved the second-highest training accuracy, only behind TDPO2, outperforms other baselines, and has the highest diversity. We also demonstrate that FPO exhibits entropy levels comparable to methods like TDPO-2, which excel in controlling output diversity, indicating the effectiveness of FPO.
It important to note that, compared to reference-free meth- ods like SimPO, FPO still requires pre-computation the reference model's log probabilities and SAE feature ac- tivations. However, this reduces the peak computational and memory demands, making the model easier to on smaller devices with lower costs. Considering that scaling computational resources is generally more challenging than extending runtime, we believe this represents a reason- able trade-off between performance and cost.
FPO Yields Better Controllability and Efficiency Trade- off. Using Gemma-2-2B as the base model, we first con- duct dialogue fine-tuning and proceed with the testing phase. For the calculation of KL divergence, we consistently apply TDPO's sequential KL divergence method. Specifically, we compute the KL divergence of the policy model relative to the reference model for both the preferred response (i.e., chosen) and the dispreferred response (i.e., rejected). The re- sults (See Table 2) indicate that, due to FPO's excellent KI. control and well-designed reward structure, it achieves per- formance comparable to other methods while maintaining lower computational costs.
Consistency between MSE Loss and KL Divergence. In TDPO and KTO, the use of KL divergence serves to con- strain the margin between the model's preferred response (chosen) and dispreferred response (rejected), thereby al- lowing for better control over the dispreferred responses. We also evaluated the margin between chosen and rejected responses under MSE Loss across 32 response sets (see Figure 4). The results indicate a high degree of consistency between the constraints enforced by MSE Loss and those enforced by KL. divergence (see Figure 3 and Figure 4). Through these constraints, the model reduces the deviation in the distribution of dispreferred responses.
Hardware Efficiency of FPO. Given the efficiency of FPO compared to TDPO2, as shown in the left one in Fig-
7

---

Constrain Alignment with Sparse Autoencoders
Table 3: Ablation Study on SAE layer selection, hyperparameters a and stop-gradient operator (Grad. sg. for short). We perform experiments on Gemma-2-2b, with the 25th layer's 's residual SAE used to evaluate the effects of varying a and applying a stop-gradient. We search for the best settings considering the trade-off between Alignment (accuracy) and Diversity (entropy).
Search Strategy: Layer Selection
Search Strategy: o Selection / Stop-Gradient
layer &
7
7
13
13
19
19
25
25
a Grad sg.
0.1
0.5
1
2
0.1
0.5
1
2
SAE type
Res
ML.P
Res
ML.P
Res
MLP
Res
MI.P
Yes
Yes
Yes
Yes
Ace (%)
57.2
57,4
59.1
61.3
59.7
62.4
63.6
63.4
64.1
63.7
63.4
61.9
64.0
63.6
62.7
62.1
H f
1.645
1.609
1.612
1.637
1.644
1.654
1.680
1.671
1.630
1.642
1.666
1.643
1.652
1.680
1.682
1.679
Chosen and Releciod HSE Loss / Margin
output have a more significant impact on the final result. If the encoder is inserted too early, the later layers do not receive gradients from the MSE loss, which negatively af- fects the model's performance. Regarding the choice of a, we find that although larger yields stronger constraint effects while also limits the model's alignment performance. Therefore, we select 0.5 as the optimal a. Our tests on the stop-gradient operator demonstrate its effectiveness, which is consistent with TDPO.
55
Gigabytes(B8g)
50.1
2:21
50
47.6
45.1
45.5
A5
80
0.54
35
30
0
Smi/
TOPO-1
ToPo-3
rro staha
OF0
YOPO-1 YoPo-2
Method
Method
Figure 4: Left. GPU memory consumption on a sin- gle H100 with all methods. We average the average GPU memory in 1,000 steps at the beginning of the train- ing. Right. Feature-level MSE Loss of all methods after the whole alignment process. Here margin is defined as ID (z, Mi; ( - BD (I, Mws meel| J. The close correspondence between the MSE Loss margin reduction and KI. divergence margin reduction supports the validity of our approach.
Varying Sampling Temperatures. To investigate the per- formance variation of FPO under different sampling tem- peratures, we designed a set of temperature comparison experiments based on the ArenaHard dataset. We config- ured five different softmax sampling temperatures: 0.2, 0.4, 0.6, 0.8, and 1.0. Then, for each of these temperature set- tings, we sampled responses from all tested methods across the first 100 questions of the ArenaHard dataset. We com- pared FPO's sampling results with those of other methods, using GPT-4 Turbo as the judge, and calculated a winning rate based on the win-loss results for each comparison. A winning rate greater than 50% indicates that FPO achieved better alignment. As shown in Figure 4. the results show that, across multiple temperature settings, FPO outperforms other methods in at least 3-4 temperature conditions.
5.2. Ablation Study
To validate the insertion position of the SAE encoder and the settings of other hyperparameters, we conduct an abla- tion study as shown in Table 3. We train Gemma-2-2B on UltraFeedback for one epoch to evaluate the performance of different configurations. In terms of metrics, we focus on accuracy and diversity (measured by entropy) to balance alignment and diversity. Regarding the insertion position of the SAE encoder, we test the following: (1) Inserting at different layers, including shallow, middle, and deep layers. (2) Inserting the encoder after the residual stream, i.c., im- mediately after the residual connection to extract features, versus inserting it after the output of the MLP layer. We did not test the insertion after the attention output, as SAE is designed to capture more polysemous features in the MLP layer and the final residual output. Prior work supports this design. (3) Varying the value of a, which affects the strength of the constraint. (4) The use of the stop-gradient operator. From Table 3, we show that inserting the encoder closer to the final output leads to better performance. We hypothesize that this is because the layers near the final
5.3. FPO Achieves Accurate Control Over Model Capabilities
A key advantage of FPO lies in its ability precisely control model capabilities, While efficiency and reduced memory usage are significant outcomes, the accurately in controlling model behavior stems from FPO's underlying mechanism.
To substantiate this, experiments were conducted to accu- rately regulate specific capabilities during alignment, while leaving others unaffected and maintaining strong overall performance. These experiments assessed four critical do- mains: (1) Instruction Following: Evaluated using IFEval, focusing on tasks such as JSON formatting, capitalizing, highlighting text, lowercasing, creating bullet lists, and us- ing quotations. (2) Multilingual Capability: Assessed with MultiAlpaca and WildChat datasets, covering French, Cana-
8

---

Constrain Alignment with Sparse Autocncoders
Table 4: FPO Accurate Control Experiment Highlights
Domain
Capability
Method/Setting
Metric
Value
FPO (8 = 0 on format features)
0.46 0.00
Instruction Following
JSON Format
FPO (B = 1 on all format features)
Accuracy
0.03 / 0.00
FPO (3 = 1 on JSON feature)
0.04 / 0.00
FPO (B = 0 on French feature)
80
Multilingual
French
Output Rate %
FPO (3 = 1 on French feature)
3
FPO (8 0)
30
Safety
FPO (3 1 on safety-related features)
Atlack Success Rate %
80
FPO (3 1 on harmful features)
5
FPO (B 0)
65
Sentiment
Positive
Positive Sentiment Ratio
FPO (3 1 Positive feature)
5
dian French, German, and Italian, with English questions from MKQA. (3) Safety: Measured using Jailbreak Bench and AdvBench. (4) Sentiment: Analyzed with the Twitter Financial News Sentiment dataset.
to harmful features decreased the attack success rate to 5% (indicating enhanced safety), compared to a 30% success rate when O. For Sentiment, setting 3 on positive sentiment features reduced positive sentiment expression 5% and increased negative sentiment to 95%. Conversely, targeting negative sentiment features with 3 resulted in 93% positive and 7% negative sentiment outputs.
The experimental setup involved the Gemma2-2B model with an SAE width of 16k. Features relevant to each do- main were identified by first selecting the top 50 features most aligned with each target domain based on single-token activations, followed by validating their global relevance through average activations across target datasets. The model trained with learning rate of 2e-5, batch size of 64, and maximum sequence length of 2048. using an Adam optimizer and cosine learning rate schedule with 10% warmup steps over 1 epoch. Hyperparameters such as a for TDPO2 and FPO were set to 0.5, 3 for SimPO to 2, and /3 for other methods to 0.1.
6. Conclusion
In conclusion, we proposed FPO, novel method for effi- cient and stable alignment of large language models using feature-level constraints. By leveraging sparse autoencoders and pre-computed offline references, FPO reduced the com- putational overhead traditionally associated with alignment methods like DPO and TDPO. Our experimental results demonstrate that FPO achieved significant improvements in alignment aceuracy and diversity while maintaining low resource consumption.
The results demonstrate FPO's capacity for targeted control. For example, in Instruction Following, by adjusting the 3 value for all format-related features to 1, FPO significantly reduced the model's propensity to use these formats (e.g., JSON format accuracy dropped to 0.03 with instruction and 0.00 without, compared to 0.46 and 0.00 respectively when = 0). Conversely, setting = for these features (ef- fectively removing the constraint) maintained or enhanced these abilities (e.g., JSON format accuracy of 0.46 with instruction). When targeting specific feature like JSON formatting by setting its 3 = 1, the JSON format accuracy dropped to 0.04 (with instruction), while other formatting abilities like capitalizing (0.73) or highlighting (0.55) re- mained high.
A Acknowledgement
This publication has been supported by the National Natural Science Foundation China (NSFC) Key Project under Grant Number 62336006.
Impact Statement
This paper proposes Feature-level Constrained Direct Pref- erence Optimization (FPO) to address the challenges of computational inefficiency and training instability in align- ing large language models (LLMMs) with human preferences. Its impact spans multiple aspects. In academic research, FPO offers novel perspective by being the first inte- grate sparse feature-level constraints into LLM alignment, inspiring further exploration in this area. Methodologically. it enriches the field by using Sparse Autoencoders (SAEs) to approximate KL divergence, balancing efficiency and controllability. On the industry front, FPO promotes the
Similar precise control was observed in other domains. For Multilingual Capability, setting / = for French-related features drastically reduced French output (3% output rate) while other languages like German (45%) and Italian (50%) were less affected compared to when / 0 for French features (French output rate of 80%). In Safety, applying to salety-related features increased the attack success rate to 80% (indicating reduced safety), whereas applying it

---

Constrain Alignment with Sparse Autocncoders
safe and reliable application of LLMs in critical sectors, such healthcare, finance, and education, while its use of open source models and datasets fosters the development of the open source ecosystem, encouraging collaboration and innovation among researchers.
Huben, R., Cunningham, H., Smith, L. R., Ewart, A., and Sharkey, L. Sparse autoencoders find highly inter- pretable features in language models. In The Twelfth International Conference on Learning Representations, 2024. URL https:/ wwwp//wpweweweeeeee//e net/ forum? id-F7 6bwRSLeK.
References
Kingma, D. Adam: method for stochastic optimization. arXiv preprint arXiv: 141.:9:41,.6680,, 2014.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L.,, Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.087. 087774 2023.
Kullback, S. and Leibler, R. A. On information and suf- ficiency. The annals of mathematical statistics, 22(1): 79-86, 1951.
Chen, Y., Ge, D., Wang, M., Wang, Z., Ye, Y., and Yin, H. Strong up-hardness for sparse optimization with con- cave penalty funetions. International Conference on Machine Learning, pp. 740-747. PMLR, 2017.
Li, T. Chiang, W., Frick, E., Dunlap. L., Zhu, B., Gonzalez, J. E., and Stoica, From live data to high-quality benchmarks: The arena-hard pipeline, April 2024. URL https: /t//mmmys.../ orgpaysgrr//bg////// 2024 04 aa--arena-haadd
Chiang, W.-L., Zheng. L., Sheng, Y., Angelopoulos, A. N., Li, T., Li, D., Zhang, H., Zhu, B., Jordan, M., Gonzalez, E.,, and Stoica, I. Chatbot arena: An open platform for evaluating Ilms by human preference, 2024.
Li, X., Zhang, T., Dubois, Y., Taori, R.. Guirajani, I., Guestrin, C., Liang, P., and Hashimoto, T. B. Alpacae- val: An automatic evaluator of instruction-followinn models. https: //gpttttb..//.gg ctt.uuubccootttttt-laaaaaaa alpaca_evar eval, 2023.
Christiano, P. F., Leike, J., Brown, T., Martic, M., Legg, S., and Amodei, D. Deep reinforcement learning from human preferences. Advances in neural information pro- cessing systems, 30, 2017.
Lieberum, T., Rajamanoharan, S., Conmy, A., Smith, L., Sonnerat, N., Varma, V., Kramár, J., Dragan, A., Shah, R., and Nanda, N. Gemma scope: Open sparse autoen- coders everywhere all at once on gemma arXiv preprint arXiv:2408.055: 14700,0551,,, 2024.
Chung. H. W., Hou, L., Longpre. S., Zoph, B., Tay, Y., Fe- dus, W., Li, Y., Wang. X., Dehghani, M., Brahma, S., et al. Scaling instruction-finetunee language models. Journal of Machine Learning Research, 25(70):1-53. 2024.
Liu, Y., Liu, P., and Cohan, A. Understanding reference policies in direct preference optimization, 2024. URL https:/ www.i..rg//aaaiaaaaaaaaaoa wttpo//gg//./////e/eee///e//////
Cui, G., Yuan, L.,, Ding, N., Yao, G., Zhu, W., Ni, Y., Xie, G., Liu, Z., and Sun, M. Ultrafeeback:: Boosting language models with high-quality feedback, 2024. URL https: //openreview f dd=pNko=3IVVWWW
Meng. Y., Xia, M., and Chen, D. Simpo: Simple preference optimization with a reference-free reward. arXiv preprint arXiv:2405. 2024.
Ding, N., Chen, Y., Xu, B., Qin, Y., Zheng, Z., Hu, S., Liu, Z., Sun, M., and Zhou, B. Enhancing chat language mod- els by scaling high-quality instructional conversations. arXiv preprint arXiv:2305.14233 2023.
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., al. Training language models follow instructions with human feedback. Advances in neural information processing systems, 11227730227744 2022.
Dubois, Y., Galambosi, B., Liang, P., and Hashimoto, B. Length-controlled alpacaeval: A simple way debias automatic evaluators. arXiv preprint arXiv:2404.04475. 2024.
Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model, 2024.
Ethayarajh, K., Xu, W. Muennighoff, N., Jurafsky, D., and Kiela, D. Kto: Model alignment as prospect theoretic optimization. arXiv preprint arX 2024.
Roy, J., Girgis, R., Romoff, J., Bacon, P.-L., and Pal, C. Di- rect behavior specification via constrained reinforcement learning. arXiv preprint 12. 2021.
Graves, A. Generating sequences with recurrent neural networks. arXiv preprint 100.:::3300000000 2013.
Hong. J., Lee, N., and Thorne, J. Reference- free monolithic preference optimization with odds ratio. arXiv preprint arXiv:2403.07661, 2024.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv: 177.0217 2017.
10

---

Constrain Alignment with Sparse Autocncoders
Team, G., Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhupatiraju, S., Hussenot, L.,, Mesnard, T., Shahriari, B., Ramé, A., Ferret, J., Liu, P., Tafti, P., Friesen, A., Casbon, M.. Ramos, S., Kumar, R., Lan, C. L., Jerome, S., Tsit- sulin, A., Vieillard, N., Stanczyk, P., Girgin, S., Momchev, N., Hoffman, M., Thakoor, S., Grill, J.-B., Neyshabur, B., Bachem, O., Walton, A., Severyn, A., Parrish, A., Ahmad, Hutchison, A., Abdagic, A., Carl. A., Shen, A., Brock, A., Coenen, A., Laforge, A., Paterson, A., Bastian, B., Piot, B., Wu, B., Royal, B., Chen, C., Kumar, C., Perry, C., Welty, C., Choquette-Chone
language models with human: survey. arXiv preprint arXiv:2307.12966, 2023.
Wei, J., Bosma, M., Zhao, V., Guu, K., Yu, A. W., Lester, B., Du, N., Dai, A. M., and Le, Q. V. Finetuned language models zero-shot learners. In International Confer- ence Learning Representations, 2022. URL https:
Zeng. Y., Liu, G., Ma, W., Yang, N., Zhang, H., and Wang, J. Token-level direct preference optimization. arXiv preprint arXiv:2404. 1999, 2024.
C. A., Sinopalnikov. D., Weinberger, D., Vijaykumar, D., Rogoziñska, D., Her- bison, D., Bandy, E., Wang. E., Noland, E., Moreira, E., Senter, E., Eltyshev, E., Visin, F., Rasskin, G., Wei, G. Cameron, G., Martins, G., Hashemi, H.. Klimezak- Plucińska, H., Batra, H., Dhand, H., Nardini, I., Mein, J.. Zhou, J., Svensson, J., Stanway, J., Chan, J., Zhou, P. Carrasqueira, J., Iljazi, J., Becker, J., Fernandez, J., van Amersfoort, J., Gordon, J.. Lipschultz, J., Newlan, J., yeong Ji, J., Mohamed, K., Badola, K., Black, K., Milli- can, K., McDonell, K., Nguyen. K., Sodhia, K., Greene, K. Sjoesund, I. L.. Usui, L.., Sifre, L., Heuermann, L.., Lago, L., McNealus, L., Soares, L. B., Kilpatrick, L., Dixon, L., Martins, L., Reid, M., Singh, M. Iverson, M., Görner, M., Velloso, M., Wirth, M., Davidow, M., Miller, M., Rahtz, M., Watson, M., Risdal, M., Kazemi, M., Moynihan, M., Zhang, M., Kahng, M., Park, M., Rahman, M., Khatwani, M., Dao, N., Bardoliwalla, N., Devanathan, N., Dumai, N., Chauhan, N., Wahltinez. O., Botarda, P., Barnes, P., Barham, P., Michel, P., Jin, P., Georgiev, P., Culliton, P., Kuppala, P., Comanescu, R., Merhej, R., Jana, R., Rokni, R. A., Agarwal, R., Mullins, R., Saadat, S., Carthy, S. M., Perrin, S., Amold, S. M. R., Krause, S., Dai, S., Garg, S., Sheth, S., Ronstrom, S., Chan, S., Jordan, T., Yu, T., Eccles, T., Hennigan, T., Kocisky, T., Doshi, T., Jain, V., Yadav, V., Meshram, V., Dharmadhikari, V., Barkley, W.. Wei, W., Ye, W., Han, W. Kwon, W. Xu, X., Shen, Z., Gong, Z., Wei, Z., Cotruta, V., Kirk, P., Rao, A., Giang, M., Peran, L., Warkentin, T., Collins, E., Barral, J., Ghahramani, 7., Hadsell, R., Sculley. D., Banks, J., Dragan, A., Petrov, S., Vinyals, O., Dean, J., Hassabis, D., Kavukcuoglu, K., Farabet, C., Buchatskaya, E., Borgeaud, S., Fiedel, N., Joulin, A., Kenealy, K., Dadashi, R., and Andreev, A. Gemma 2: Improving open language models a practical size, 2024. URL https::////att /arivv org/ abs 2abs/2400 2408 00118.
Zheng, L., Chiang, W.-L.,, Sheng, Y., Zhuang, S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing. E., et al. Judging Ilm-as-a-judge with mt-bench and chatbot arena. Ad- vances Neural Information Processing Systems, 36: 46595-46623, 2023a.
Zheng, R., Dou, S., Gao, S., Hua, Y., Shen, W., Wang, B., Liu, Y., Jin, S., Zhou, Y., Xiong, L., et al. Delve into ppo: Implementation matters for stable rlhf. In NeurIPS 2023 Workshop Instruction Tuning and Instruction Following, 2023b.
Tunstall, L., Beeching, E., Lambert, N., Rajani, N., Ra- sul, K., Belkada, Y., Huang, S., von Werra, L.,, Fourrier, C., Habib, N., et al. Zephyr: Direct distillation of Im alignment. arXiv preprint arXiv:2310.16944 2023.
Wang, Y., Zhong, W., Li, L., Mi, F., Zeng, X., Huang, W., Shang, I.,, Jiang, X., and Liu, Q. Aligning large

---

Constrain Alignment with Sparse Autocncoders
A. Training Settings
Model Name
Gemma-2-2b
Parameters
2B
Method
SFT
DPO
TDPO-1
TDPO-2
SimPO
0.5
0.5
0.1
0.1
0.1
0.1
0.5
learning rate
5 x 10 7
5 x 10 -7
5 x 10
5 x 10
x 10
7
x 10
optimizer
Adam
Adam
Adam
Adam
Adam
Adam
warmup steps activation checkpoint
150 True None
150
150
150
150
150
True
True
True
True
True
SAE width GPU(s)
None
None
None
None
16k
4 H100
Model Name
Gemma-2-9b
Parameters
9B
Method
SFT
DPO
TDPO-1
TDPO-2
SimPO
0.5
0.5
learning rate
5 x 10 RMSprop
5 10
5 10
5 10
5 10
x 10
optimizer
RMSprop 150
RMSprop 150
RMSprop
RMSprop
RMSprop
warmup steps
150 True None
150
150
150
activation checkpoint SAE width GPU(s)
True None
True None
True None
True
True
None
16k
H100
Table 5: Hyperparameters for Gemma-2-22 and Gemma-2-9h
B. Bounding KL Divergence with MSE of Sparse Activation
L.L J.L Theorem B.1. Let ne and Teet be two models with final layer outputs h22 .hee Rdd position Let co Cref E R be their respective sparse activation generated by SAE. Under certain conditions, minimizing the MSE between these sparse activation values leads to reduction in the upper bound of the KL divergence between their token probability distributions.
Proof. We begin by establishing key definitions and conditions:
Definition B.2 (Sparse Activations).
LLL ReLU(W h.b + b)
(12)
Definition B.3 (Token Logits and Probabilities).
ph = softmax(2") W oue hi4
(13)
Definition B.4 (KL Divergence).
V DKL (plor (1pb) = (io (i) p (2)
(14)
[Accurate Reconstruction] The SAE reconstructs hidden representations accurately, i.e., for some small >>> 0:
WWT.c ch n <c
(15)
[Bounded Operator Norm]
|||] S M for K W. WT and some M out
(16)
12

---

Constrain Alignment with Sparse Autocncoders
[Small Logit Differences] The difference in logits A=' = 2/ zyot is small enough for the quadratic approximation of the
KL divergence to hold. A small Az" generally exists since (1) A=' 0 initially, and (2) a very small learning rate (e.g.. 5e-7) is usually adopted during alignment training.
Now, we proceed with the main proof:
Lemma B.5. Under Condition 1, the difference in hidden representations A = h.L4 -rf can be approximated by:
AhLL
hhg hich
(17)
J.L. Lemma B.6. The difference in logits A related the difference in sparse activations Act by:
where Act. = c0
Az' K&ct where K =W WT
(18)
Lemma B.7, For small A the KL divergence can be bounded by:
DKL (prer (1p)) - ||
(19)
Proof. Using second-order Taylor expansion and noting that the maximum eigenvalue of the Hessian of KL divergence concerning logits is Amax H)) 1:
DRL (RI. 2
(A=')" H (=fet)
(20)
Amax (H)||
(21)
5 | 11
(22)
Combining these lemmas:
DKI(Pherllo)) S
||
(23)
5 1K 12
(24)
M²
- 2
|
11
(25)
The right-hand side proportional the MSE of the sparse activations:
|
J.L. J.L. Coo ref,i)
t.b. Ab MSE
ref
(26)
Let Jm be the set of indices corresponding to the top m activations. Then:
2
(1pb) S
Σ
(chlt - Gof ,22,)2
(27)
2
NE
2
MSE
(28)
Therefore, minimizing the MSE of sparse activation leads to minimizing an upper bound on DKL (p Pp)).
13

---

Constrain Alignment with Sparse Autocncoders
C. Concrete Examples of Feature-Level Representations vs. Token-Level Embeddings
This section provides concrete examples and visualizations to highlight the differences between feature-level representations and token-level embeddings in our framework.
C.1. Definitions and Intuitions
Token-level embeddings correspond directly to the token output probabilities (logits) generated by model. These embeddings are nigh-dimensional vectors representing each token the model's vocabulary. For sequence z [11, ........ .T). the token-level embeddings at position are computed as:
Token-Level Embeddings:
h4 fuken (x+) E R
where the vocabulary size, and finken is the output projection from the model's hidden state.
Feature-Level Representations: Feature-level representations, the other hand, are high-level abstractions derived from the model's intermediate layers. These representations capture patterns and salient features across sequences. Using Sparse Autoencoder (SAE), the hidden state I at layer l can be transformed into sparse activations ef, defined as:
BenUt + b),
where Wene Rmxdd b R. and m V. This sparse activation ensures only subset of features is active, making the representation interpretable and efficient.
C.2. Concrete Example: A Mathematical Query Consider the input query:
"What is the derivative of x.2 + 3x + 5?"
Token-Level Embedding: The token-level output probabilities for each token the response sequence, such "The derivative is 2x + 3.", involve logits for every token:
logits
[log P("The", log P('derivative"), P('is"),...)
Feature-Level Representation: Using SAE on the 25th layer, the sparse feature representation for the same sequence might activate specific features corresponding mathematical operations or semantic groupings:
=
|activation, (Polynomial). activation (AritioontAiiittettii......
D. Experiments on Additional Baselines and Ablation Studies
In response to reviewer feedback, we conducted additional experiments to address their concerns and validate our methodol- ogy. These include comparisons with the SimPO+KL baseline and ablations on multi-layer sparse autoencoders (SAEs).
D.1. Comparison with SimPO+KL
This subsection provides a direct comparison of our method against SimPO+KL. We implemented SimPO+KL. following the same experimental settings in Section 4. Specifically, tested on the Gemma-2-2B model using the AlpacaEval-22 dataset, evaluating both winning rate (WR) and Length-controlles winning rate (WR-L). Results summarized in Table 6.
The results show that FPO achieves comparable or better performance than SimPO+KL in both WR and WR-L metrics. This highlights the effectiveness of feature-level constraints in maintaining both alignment quality and diversity, with a competitive computational cost.
Discussion:
D.2. Ablation Study on Multi-Layer SAEs
To find out the effect extending SAEs across multiple layers, we conducted experiments adding SAEs at different layer combinations. Table presents the performance metrics when SAEs were applied to various combinations of shallow, middle, and deep layers.
14

---

Constrain Alignment with Sparse Autoencoders
Table 6: Comparison of FPO with SimPO+KL on the AlpacaEval-2 dataset. (Entropy), WR (%), and WR-L (%).
Metrics include Accuracy (%), Diversity
Method
Accuracy (%) f Diversity (Entropy)
WR (%) f
WR-L (%) f
FPO (Ours)
64.1
1.68
51.8
50.2
SimPO+KI.
63.6
1.66
50.8
SimPO
63.4
1.64
50.2
49.8
TDPO-2
64.2
1.68
50.0
50.0
Table Ablation study on using SAEs multiple layers in FPO, Metrics include Accuracy (%), Diversity (Entropy), WR (%), and WR-L (%)
SAE Layers
Accuracy (%) 1 Diversity (Entropy)
WR (%) f
WR-I (%) f
Single Layer (Layer 25)
64.1
1.68
51.8
50.2
Layers 0, 25
62.1
1.70
47.2
48.8
Layers 12, 25
61.9
1.64
48.4
49.5
Layers 24, 25
1.66
48.6
Layers 0, 12, 25
51.4
1.66
47.8
49.8
Results indicate that adding multiple SAE layers does not consistently improve performance and may even degrade alignment metrics (e.g.. accuracy and WR). The best results were achieved with single SAE layer (Layer 25), confirming that simplicity in feature extraction leads to more stable alignment.
Discussion:
15


Reviews:
[{"review_id": "Y0Vu846CZQ", "reviewer": "Reviewer_f5gf", "summary": "This work proposes a feature-level direct preference optimization algorithm, FPO, for LLM preference learning. Specifically, it revised the token-level KL reguralization of TDPO into the feature level KL regularization where features are obtained from a pre-trained SAE.\n\nThe reference features and tokens are both precomputed and stored offline to reduce the memory footprint.\n\nThe results demonstrate that FPO can achieve comparable performance with TDPO2 while with fewer memory usage.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: Why the title is Accurately? The key advantages of this work are its efficiency and reduced memory usage.**\n\nWe appreciate the reviewer’s question regarding the title’s use of “Accurately” and their observation that the work’s key advantages seem to lie in efficiency and reduced memory usage. While efficiency and memory benefits are indeed significant outcomes of FPO, we argue that “Accurately” reflects the precision in controlling model capabilities. Our work leverages SAEs to decompose model representations into features, offering a level of granularity that token-level methods lack. By applying MSE constraints to specific features e.g., features that controlling the json format output, FPO enables precise manipulation of model alignment. To substantiate this, we present experiments to accurately regulate specific capabilities during alignment, leaving others unaffected while achieving strong overall performance. We assessed four critical domains: Instruction Following with IFEval, Multilingual Capability with MultiAlpaca and WildChat, Safety with Jailbreak Bench and AdvBench, and Sentiment with Twitter Financial News Sentiment. Detailed experimental settings and results are provided: https://github.com/FPO-code/FPO_code/blob/main/accurate_control.md\n\n**Q2: Why there is no related work?**\n\nWe are sorry for missing the related work section due to the page limit. Considering our work lie on the intersection of alignment and mechanistic interpretability,  we’d like to give a series of important related works of these two fields. We will add the related work section to the paper in the following versions.\n\n**Q3: The performance ... versus TDPO includes: (1) length normalization, and (2) computing compressed feature-level KL (at the last layer) instead of token-level KL ... contribute to the performance?**\n\nWe analyze these two factors with a new ablation study in https://github.com/FPO-code/FPO_code/blob/main/tdpo_vs_fpo.md. Results show that (1) FPO performs better at length norm and (2)  computing compressed feature-level KL shows no obvious effect on the downstream performance (but more efficient).\n\n**Q4: The FPO results in Table 2 are not presented. I cannot see the effectiveness of the proposed method. But from the descriptions, the FPO seems cannot scale to larger models.**\n\nΔScores in Table 2 as the margin between FPO and other methods (as described in Line 296, Evaluation Benchmarks). As for the scaling results, we provides FPO in different architectures (LLaMA, Gemma) and parameters from 2B - 9B: https://github.com/FPO-code/FPO_code/blob/main/LLaMA_Experiments.md\n\n**Q5: Figure 1 presents ambiguous visual information ... refine this figure and amplify the caption.**\n\nWe thanks the reviewer's valuable feedback, we will refine this figure in the following version.\n\n**Q6: Lacks some implementational details. Is the pre-trained SAE is fixed? ... cause higher reconstruction error along the training?**\n\nSAE are commonly robust under alignment, considering the model architectures are fixed and alignment process's data amount is much smaller comparing to pretraining. We test the reconstruction error to verify this: https://github.com/FPO-code/FPO_code/blob/main/reconstruction_error.md.\n\n**Q7: It may seem weird that the author claims the alignment stability ... How stable are TDPO, SimPO, and FPO which tradeoffs?**\n\nWe thanks the reviewer's suggestions about the alignment stability. Here we perform comparison of token log-probability difference: https://github.com/FPO-code/FPO_code/tree/main/FPO-code-main. We will add this into the appendix.\n\n**Q8: What motivates you to change the base model in SimPO evaluation setting ... on one arbitrary base model in that setting?**\n\nWe changes the SimPO evaluation settings because we found the current settings are better than the SimPO paper proposed. Here is the comparison: https://github.com/FPO-code/FPO_code/blob/main/simpo.md.\n\n**Q9: In Figure 2 left, is this an example instance or statistics? What is the analysis setting?**\n\nThis is the globally activated feature statistics on 1024 examples from Ultrafeedback dataset for illustrating the sparsity of SAE features. We use the Gemma2-2b SAE with width 16k. \n\n**Q10: In Fig. 3, how KL of different methods are computed? Since FPO’s regularization is feature-level KL, is it compared at the feature or token level KL?**\n\nFeature-level KL is a more accurate and efficient way comparing to token-level KL. Here we propose a detailed comparison: https://github.com/FPO-code/FPO_code/blob/main/tdpo_vs_fpo.md.\n\n**Q11: What does uniqueness mean in line 84?**\n\nThis means that FPO supposed a specific control of the alignment process (refer to Q1).\n\n**Q12: What is the wall-clock computational time of FPO vs TDPO?**\n\nData loading does need more times. However, the efficiency in training saves more times (and memory). Details can be found at: https://github.com/FPO-code/FPO_code/blob/main/time.md.", "claims_and_evidence": "Why the title is *Accurately*? The key advantages of this work are its efficiency and reduced memory usage.", "methods_and_evaluation_criteria": "Please see Other Strengths And Weaknesses.", "theoretical_claims": "The FPO feature distance bounds the feature level KLD.", "experimental_designs_or_analyses": "- Need to present ablation on the revised loss components.\n- Others please see **Other Strengths And Weaknesses**.", "relation_to_broader_scientific_literature": "This proposed approach can help reduce the memory footprint while maintaining equivalent alignment performance.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "strength:\n\n- The writing is clear and easy to follow.\n- The proposed method is novel and neat. FPO can significantly reduce computational overhead and memory footprint.\n\nweaknesses:\n\n- Why there is no related work?\n- The performance (alignment acc and diversity) of the proposed method almost has no difference compared with TDPO2. Technically, the change versus TDPO includes: (1) length normalization, and (2) computing compressed feature-level KL (at the last layer) instead of token-level KL. However, in ablation experiments, not both factors are decomposed for analysis. How does each factor contribute to the performance?\n- The FPO results in Table 2 are not presented. I cannot see the effectiveness of the proposed method. But from the descriptions, the FPO seems cannot scale to larger models.\n- Figure 1 presents ambiguous visual information. It is difficult for readers to get the key idea of this paper by viewing this figure with its caption. I would suggest the author refine this figure and amplify the caption.\n- Lacks some implementational details. Is the pre-trained SAE is fixed? Considering the pre-trained SAE is fixed, will the distribution shift of one layer’s output features (input of SAE) cause higher reconstruction error along the training?\n- It may seem weird that the author claims the alignment stability everywhere while not including any quantitative comparisons. Better to include (in the appendix) statistics on token log-probability difference as [1] does. This could help readers to realize the significance of this problem. How stable are TDPO, SimPO, and FPO which tradeoffs?\n\n[1] Understanding Reference Policies in Direct Preference Optimization\n\nquestions:\n\n- What motivates you to change the base model in SimPO evaluation setting, considering you inherit most of it? If no special reason, could you please provide results on one arbitrary base model in that setting? It helps readers to get a sense of the advantage of the FPO method with comprehensive comparisons and its generalizability to various architectures.\n- In Figure 2 left, is this an example instance or statistics? What is the analysis setting?\n- In Fig. 3, how KL of different methods are computed? Since FPO’s regularization is feature-level KL, is it compared at the feature or token level KL?\n- What does *uniqueness* mean in line 84?\n- What is the wall-clock computational time of FPO vs TDPO? I think the data loading process will also reduce the advantage of FPO.", "other_comments_or_suggestions": "None.", "questions_for_authors": "Please see Other Strengths And Weaknesses for questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4nQnI8Lx3K", "reviewer": "Reviewer_axev", "summary": "The paper proposes Feature-level constrained Preference Optimization (FPO), a novel method designed to improve the alignment of LLMs. FPO utilizes sparse features activated in a trained sparse autoencoder and employs feature-level offline reference to maintain the quality of sequential KL divergence. Experiment results demonstrate that FPO achieves a 5% absolute improvement in win rate compared to baselines while requiring lower computational resources.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: The paper doesn’t have a related work section and doesn’t offer a sufficiently comprehensive overview of key prior work in LLM alignments**\n\nWe are sorry for missing the related work section due to the page limit. Considering our work lie on the intersection of alignment and mechanistic interpretability, we’d like to give a series of important related works of these two fields. We will add the related work section to the paper in the following versions.\n\n**Q2: According to Table 1 left, the performance gains of FTO appear marginal. For example, on the Gemma-2-2B model evaluated with AlpacaEval-2 (805 questions), the method surpasses DPO by only 13 or 14 questions and surpasses simPO by 8 or 9 questions. On the Gemma-2-9B model, the performance gap even becomes smaller, improving only 3 or 4 over SimPO. In Table 1 right, TDPO-2 achieves better performance compared with FTO.**\n\nWe thank the reviewer for their observation regarding the seemingly marginal performance gains of FPO in Table 1 and the comparison with TDPO-2. We acknowledge that the initial results on Gemma showed modest improvements. To address this, we extended our experiments to include **LLaMA** models and tested an enhanced approach using **Multi-Layer Sparse Autoencoders (MLSAE)** to extract cross-layer global features, which yielded superior performance. MLSAE extracts features from residual streams in all layers and thus make the constraint deeper. We also update the efficiency tests for this new method and show that MLSAE does not increase the wall-time computational cost. Results can be found at: https://github.com/FPO-code/FPO_code/blob/main/LLaMA_Experiments.md\n\n\nAlso, we propose experiments regarding the **accuracy** of FPO in controlling model capabilities. Our work leverages SAEs to decompose model representations into features, offering a level of granularity that token-level methods lack. By applying constraints to specific features e.g., features that controlling the json format output, FPO enables precise manipulation of model alignment, which is beyond merely improved performance. The experimental details and results can be found at: https://github.com/FPO-code/FPO_code/blob/main/accurate_control.md\n\n**Q3: Although FTO can reduce GPU memory consumption, there is no theoretical guarantee that the proposed MSE constraints is better than the KL divergence constraints**\n\nWe provide the bounding KL Divergence with MSE of sparse activation in Appendix B. And here we also give additional broad experimental verification of the comparison between MSE constraints and the original KL divergence constraints: https://github.com/FPO-code/FPO_code/blob/main/mse_vs_kl.md.\n\n**Q4: The hyperparameters in FTO include the SAE layer, hyperparameters  and stop-gradient operator. The hyperparameter searching process is time-consuming and challenging.**\n\nRegarding the SAE layer search, we have updated our approach with a MLSAE which achieves superior performance while eliminating the need to identify a specific layer for alignment. This enhancement streamlines the process and boosts effectiveness. As for the stop-gradient operator and hyperparameter settings, we conducted a preliminary search and confirmed that the stop-gradient operator remains essential across different models, consistent with TDPO’s findings. Consequently, we believe the cost of hyperparameter tuning is less burdensome than the reviewer might assume, as our initial exploration efficiently validates these key components.\n\n**Q4: The hyperparameters in FTO include the SAE layer, hyperparameters  and stop-gradient operator. The hyperparameter searching process is time-consuming and challenging.**\n\nRegarding the SAE layer search, we have updated our approach with a MLSAE which achieves superior performance while eliminating the need to identify a specific layer for alignment (refer to Q2). This enhancement streamlines the process and boosts effectiveness. As for the stop-gradient operator and hyperparameter settings, we conducted a preliminary search and confirmed that the stop-gradient operator remains essential across different models, consistent with TDPO’s findings. Consequently, we believe the cost of hyperparameter tuning is less burdensome than the reviewer might assume, as our initial exploration efficiently validates these key components.\n\n**Q5: The example code is not provided, which limits the reviewers' ability to verify reproducibility.**\n\nWe thanks the reviewer's feedback regarding the transparency and reproducibility. Here we provide the code and additional experimental settings and results: https://github.com/FPO-code/FPO_code.", "claims_and_evidence": "The claims are generally fine.", "methods_and_evaluation_criteria": "The logical flow of the method looks good to me. For the concerns, please refer to Strengths and Weaknesses.", "theoretical_claims": "I checked the proof in Appendix B.", "experimental_designs_or_analyses": "For the experimental design, I have the following concerns:\n\nThe experiments are only conducted in Gemma model. What about other open-source models like Llama or Mistral?", "relation_to_broader_scientific_literature": "FPO combines existing alignment methods including DPO SimPO, and TDPO. Instead of using the KL divergence constraints, FPO uses SAEs to project LLM’s internal representations onto a sparse feature space to get feature vectors. Then FPO uses the MSE constraints of the feature vectors to replace the KL divergence constraints in TDPO.", "essential_references_not_discussed": "The paper doesn’t have a related work section and doesn’t offer a sufficiently comprehensive overview of key prior work in LLM alignments. For example, [1], [2], [3], [4], [5] and [6].\n\n[1] Azar, Mohammad Gheshlaghi, et al. \"A general theoretical paradigm to understand learning from human preferences.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2024.\n\n[2] Rafailov, Rafael, et al. \"From $ r $ to $ q^* $: Your language model is secretly a q-function.\" arXiv preprint arXiv:2404.12358 (2024).\n\n[3] Kong, Lingkai, et al. \"Aligning large language models with representation editing: A control perspective.\" Advances in Neural Information Processing Systems 37 (2024): 37356-37384.\n\n[4] Ji, Xiang, et al. \"Self-play with adversarial critic: Provable and scalable offline alignment for language models.\" arXiv preprint arXiv:2406.04274 (2024).\n\n[5] Cen, Shicong, et al. \"Value-incentivized preference optimization: A unified approach to online and offline rlhf.\" arXiv preprint arXiv:2405.19320 (2024).\n\n[6] Richemond, Pierre Harvey, et al. \"Offline regularised reinforcement learning for large language models alignment.\" arXiv preprint arXiv:2405.19107 (2024).", "other_strengths_and_weaknesses": "Strengths:\n1. The paper is well-written and easy to follow.\n2. FTO can reduce GPU memory consumption compared to baseline alignment algorithms.\n3. By using SAE, FTO provides a potentially more interpretable and computationally efficient alternative to token-level alignment methods.\n\nWeaknesses:\n1. According to Table 1 left, the performance gains of FTO appear marginal. For example, on the Gemma-2-2B model evaluated with AlpacaEval-2 (805 questions), the method surpasses DPO by only 13 or 14 questions and surpasses simPO by 8 or 9 questions. On the Gemma-2-9B model, the performance gap even becomes smaller, improving only 3 or 4 over SimPO. In Table 1 right, TDPO-2 achieves better performance compared with FTO.\n2. Although FTO can reduce GPU memory consumption, there is no theoretical guarantee that the proposed MSE constraints is better than the KL divergence constraints\n3. The hyperparameters in FTO include the SAE layer, hyperparameters $\\alpha$ and stop-gradient operator. The hyperparameter searching process is time-consuming and challenging.\n4. The example code is not provided, which limits the reviewers' ability to verify reproducibility.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "Please refer to the weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LCKlRgAKsz", "reviewer": "Reviewer_jmiE", "summary": "This paper proposed a DPO variant that is both good performing and computationally efficient. They replace expensive per token KL divergence regularization w/ an SAE-based regularization -- hinging the efficiency on the SAE sparsity.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: Only 2 models are used of the same family.**\nWe apologize for the lack of evaluation across other model families. During our experiments with FPO, the Gemma model series was the only one with a relatively complete set of SAE models. Consequently, we limited our training and testing to the Gemma model. Following the release of LLaMA-scope w/ LLaMA-3 8B, we conducted subsequent experiments. The following details outline our experimental setup, including baseline selection, benchmark selection, experimental parameters, and results: https://github.com/FPO-code/FPO_code/blob/main/LLaMA_Experiments.md.\n\n**Q2: Without including SAE training in the computational efficiency evaluation, it is still questionable if this method is really more compationally efficient than the baselines. Without the computational efficiency edge, the claim is questionable.**\n\nWe maintain that our methods offer significant efficiency advantages. We present two primary lines of evidence: \n\n(1) SAEs are pre-trained on extensive datasets and can be reused across multiple downstream tasks (and can be generalized to base and chat model). Therefore, their training cost can be considered a one-time investment, rather than a recurring expense factored into each task's efficiency evaluation. This approach aligns with common practices in machine learning, such as leveraging pre-trained word embeddings without recalculating their training cost for every application. In such scenarios, efficiency assessments typically focus on the downstream task (e.g., alignment), as pre-training is treated as a standardized preliminary step. We provide a detailed analysis of the wall-clock time and computational costs associated with the combined SAE training and alignment process.\n\n(2) To further enhance efficiency, we have introduced a Multi-Layer Sparse Autoencoder (MLSAE) approach, which not only outperforms the original single-layer SAE method but also reduces training costs. By capturing cross-layer global features, MLSAE eliminates the need for extensive layer-specific searches, streamlining the process while improving alignment accuracy. Our updated efficiency analysis, now incorporating MLSAE, demonstrates lower computational overhead compared to baselines like DPO and TDPO-2, even when accounting for pre-training. These results reinforce our claim of computational efficiency, alongside the method’s superior performance, as evidenced in our revised experiments on Gemma and LLaMA models. Experimental details can be found at https://github.com/FPO-code/FPO_code/blob/main/LLaMA_Experiments.md.\n\n**Q3: I would be interested to see if the use of SAE can enable even a more precise control over the LLMs, like increasing only certain traits?**\n\nOur work leverages SAEs to decompose model representations into features, offering a level of granularity that token-level methods lack. By applying MSE constraints to specific features e.g., features that controlling the json format output, FPO enables precise manipulation of model alignment. To substantiate this, we present experiments to accurately regulate specific capabilities during alignment, leaving others unaffected while achieving strong overall performance. We assessed four critical domains: Instruction Following with IFEval, Multilingual Capability with MultiAlpaca and WildChat, Safety with Jailbreak Bench and AdvBench, and Sentiment with Twitter Financial News Sentiment. Detailed experimental settings and results are provided: https://github.com/FPO-code/FPO_code/blob/main/accurate_control.md", "claims_and_evidence": "Somewhat -- the theory part connecting FPO (theirs) loss w/ the previous DPO variants loss is convincing (how their loss makes it more stable and efficient). However, the empirical result is a bit lacking for the following reasons:\n- Only 2 models are used of the same family (Gemma-2)\n- The efficiency claim is somewhat questionable -- SAE training is not considered in the Efficiency experiment (5.1)", "methods_and_evaluation_criteria": "yes, for the most part. Again, I am not convinced of the efficiency evaluation -- the SAE training is not included.", "theoretical_claims": "Yes, the relationship between FPO loss w/ DPO's and other variants. Looks correct to me.", "experimental_designs_or_analyses": "Yes, everything. As previously mentioned, i have 2 big issues w/ the experimental design (see above)", "relation_to_broader_scientific_literature": "They propose a DPO loss variant that is stable in training, computationally efficient, and high performing. Previous work in the literature only satisfy 2 out of the three.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "My biggest worry is on the 2 issues I have stated before:\n1. Without including SAE training in the computational efficiency evaluation, it is still questionable if this method is really more compationally efficient than the baselines. Without the computational efficiency edge, the claim is questionable.\n2. Using more model family is needed, to show that this method works well for other models as well", "other_comments_or_suggestions": "I would be interested to see if the use of SAE can enable even a more precise control over the LLMs, like increasing only certain traits?", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IgkTDItq5x", "reviewer": "Reviewer_7zgf", "summary": "The author proposes Feature-level Constrained Preference Optimization (FPO). FPO uses a SimPO objective plus a regularizer that compares features in a lower-dimensional space, rather than token probabilities in the high-dimensional vocabulary space. These features are obtained by a sparse autoencoder. In this way, FPO can still regularize the training relative to a reference model, thus improving stability relative to SimPO, while computing the regularizing term in a lower-dimensional space, thus improving efficiency relative to TDPO.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer's suggestions. We note that the reviewer has raised three valuable concerns, and we would like to address them.\n\n**Q1: To me, a shortcoming of the paper .... doesn't seem to be relevant to this regularization.** \n\nWe acknowledge this gap in our original submission and address it here by providing both motivation and experimental evidence. SAEs decompose activations into sparse, interpretable features. Unlike traditional token-level constraints, which apply broad regularization, SAEs allow us to pinpoint and manipulate individual abilities, by selecting target features and constraining them. In our FPO method, we leverage this by the weighting parameter, beta, to these features. A higher beta strengthens the constraint, reducing the model’s ability in that domain, while a lower beta preserves or enhances it. \n\nWe conducted large-scale experiments using the Gemma2-2B model with four tasks: (1) Instruction Following on IFEval with tasks in JSON Format, Capitalize, Highlight Text, Lowercase, Bullet List, and Quotation. (2) Multilingual Capability. 1,000 entries from MultiAlpaca and WildChat, limited to French, Canadian French, German, and Italian, with English questions from MKQA. (3) Safety. 300 question-answer pairs from Jailbreak Bench and AdvBench, completed and verified with DeepseekV3. (4) Sentiment. 1,000 samples from Twitter Financial News Sentiment, split evenly (500 positive, 500 negative). For each domain, we identified globally activated features by averaging activations across all tokens (sparsity details in Q2). Result tables can be found here: \nhttps://github.com/FPO-code/FPO-code\n\nSAEs presents **strong and accurate** control in specific domains e.g., enhancing model's ability to generate in json format while avoiding it to use bullet list ( please refer to our experiment tables ). We believe that FPO is the first method that expanding the alignment process to feature-level, which is a relatively important contribution to the community.\n\n**Q2: Furthermore, the authors use pooled sparse activations in Equation (10), which are no longer sparse.** \n\nWhile the initial step of average pooling across tokens in the provided text results in pooled activations that are no longer strictly sparse, the subsequent application of the top-k function restores a form of sparsity. The top-k operation is both meaningful and critical to our approach. This operation aligns with common practices in mechanistic interpretability, such as difference-in-means analysis, where the focus is on the most significant features. This is particularly effective because our alignment datasets exhibit global features—consistent patterns that persist across tokens, such as structured response formatting (e.g., JSON formatting) or safety-related behaviors (e.g., avoiding unsafe responses). \n\nThe top-k operation ensures that our regularization focuses on these dominant features, allowing us to control alignment constraints at a high level while maintaining computational efficiency. Without this step, the MSE would be computed over all components of $\\bar{c}^\\ell$, including noise from less significant activations, diluting the regularization’s effectiveness.\n\nHere we provide some examples of extracted features in Q1 selected with our top-k algorithm (refer to 1. Can SAEs Identify Relevant Features from These Datasets?):\nhttps://github.com/FPO-code/FPO_code/blob/main/accurate_control.md.\n\n**Q3: It is therefore possible that the SAE's role might simply be to generate embeddings for regularization, ... simply minimize the weight difference between the unaligned and aligned models.**\n\nWe appreciate these insightful suggestions and agree that exploring alternative regularization methods is valuable. Indeed, approaches like PReLU-based autoencoders, pooled activations, or weight difference minimization could provide some regularization benefits. However, our choice of SAEs is not arbitrary—it is motivated by their unique ability to decompose model representations into monosemantic, interpretable features, enabling precise and targeted control over specific capabilities during alignment. As detailed in Q1, this granularity surpasses what token-level constraints or broader weight-based methods can achieve. \n\nTo test the reviewer’s hypothesis, we compared our SAE-based FPO method with an alternative: TDPO with activation-level alignment, where regularization is applied directly to pooled activations (akin to the second suggestion). We used the same experimental setup as in Q1 and evaluated performance across the four domains: Instruction Following, Multilingual Capability, Safety, and Sentiment. For TDPO, we pooled activations from the model’s last layer and applied a constraint to minimize their divergence during alignment, simulating an embedding-like regularization. Result tables are here:\nhttps://github.com/FPO-code/FPO_code/blob/main/accurate_control.md", "claims_and_evidence": "The claims and evidences are in general clear.", "methods_and_evaluation_criteria": "Based on the evaluation, FPO's performance is better or comparable to TDPO's performance, but with less cost. Essentially, FPO has 2 changes relative to TDPO:\n\n(i) length normalization in the log probability difference and in the margin\n(ii) SAE feature regularization\n\nIdeally, it would be great to understand the roles these choices play separately. For example, if we use only (i) but still use token-level regularization, can we increase the performance relative to TDPO? Or is (i) only useful with (ii)? Or maybe it suffices to use only (ii) without (i)?", "theoretical_claims": "The paper has no proofs for theoretical claims.", "experimental_designs_or_analyses": "The experimental designs are in general sound. See my comment in the \"Methods And Evaluation Criteria\" for suggestions on ablating the design choice relative to TDPO.", "relation_to_broader_scientific_literature": "Based on the evaluation, FPO's performance is better or comparable to TDPO's performance, but with less cost. Essentially, FPO has 2 changes relative to TDPO:\n\n(i) length normalization in the log probability difference and in the margin\n(ii) SAE feature regularization\n\nThe approach is motivated by the merits and shortcomings of SimPO and TDPO:\n\n- SimPO is computationally efficient as it does not use a reference model. However, it is sometimes unstable.\n- TDPO is numerically more stable than SimPO thanks to the token-level KL regularizer. However, it is computationally expensive, especially when the models have a large vocabulary size.\n\nThe proposed FPO can be seen as an extension of TDPO: it regularizes not in the high-dimensional token probability space but in the low-dimensional feature space.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "The main strength of the paper, in my opinion, is that it demonstrates it's unnecessary to regularize an aligned model in token space. Regularization in feature space—specifically, using features from an SAE as the authors do—is more cost-effective and yields a similar effect.\n\nTo me, a shortcoming of the paper is that it doesn't adequately motivate why an SAE is a natural choice for producing the features to be regularized. SAEs are primarily used in interpretability studies, but their interpretability doesn't seem to be relevant to this regularization. Furthermore, the authors use pooled sparse activations in Equation (10), which are no longer sparse. It is therefore possible that the SAE's role might simply be to generate embeddings for regularization, and the specific model (SAE or others) may not be crucial. A simple test is to, replace the ReLU in the SAE with a PReLU; this results in a standard autoencoder for embedding token features, and may work just as well. Another option would be to use the pooled activations directly---without using the SAE---for regularization. Alternatively, one could simply minimize the weight difference between the unaligned and aligned models. I'm not suggesting the authors explore all these variations, as the possibilities are endless. I'm raising these examples because the SAE doesn't appear to be the most obvious choice, and it would be beneficial if the authors could clarify this and provide evidence supporting the advantage of using an SAE.", "other_comments_or_suggestions": "NA", "questions_for_authors": "See \"Methods And Evaluation Criteria\" and \"Other Strengths And Weaknesses\".", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes Feature-level constrained Preference Optimization (FPO), a method for aligning large language models using sparse autoencoders (SAEs) to enforce alignment constraints in a lower-dimensional feature space. FPO aims to improve computational efficiency and stability by replacing token-level KL divergence with feature-level mean-squared error constraints derived from pretrained SAEs. The authors demonstrate that FPO achieves competitive alignment performance with lower memory overhead and add empirical results showing the method’s ability to selectively control model capabilities. 

This paper presents a novel and promising direction for more interpretable and efficient LLM alignment through feature-level regularization. While not without weaknesses, the core idea is sound, the empirical results are encouraging, and the rebuttal demonstrates a strong commitment to addressing reviewer concerns. The ability to constrain specific capabilities via sparse features is especially compelling and opens up new avenues in fine-grained alignment. Given the positive trajectory of the discussion and the potential impact of the approach, the paper is recommended for acceptance once the authors incorporate all the reviewer's feedback in the final version.
Decision: Accept (poster)