Constrain Alignment with Sparse Autoencoders

We greatly appreciate the reviewer's suggestions. We note that the reviewer has raised three valuable concerns, and we would like to address them.

Q1: To me, a shortcoming of the paper .... doesn't seem to be relevant to this regularization.

We 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.

We 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: https://github.com/FPO-code/FPO-code

SAEs 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.

Q2: Furthermore, the authors use pooled sparse activations in Equation (10), which are no longer sparse.

While 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).

The 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.

Here 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?): https://github.com/FPO-code/FPO_code/blob/main/accurate_control.md.

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.

We 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.

To 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: https://github.com/FPO-code/FPO_code/blob/main/accurate_control.md

Q1: Only 2 models are used of the same family. We 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.

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.

We maintain that our methods offer significant efficiency advantages. We present two primary lines of evidence:

(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.

(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.

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?

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

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

We 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.

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.

We 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

Also, 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

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

We 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.

Q4: The hyperparameters in FTO include the SAE layer, hyperparameters and stop-gradient operator. The hyperparameter searching process is time-consuming and challenging.

Regarding 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.

Q4: The hyperparameters in FTO include the SAE layer, hyperparameters and stop-gradient operator. The hyperparameter searching process is time-consuming and challenging.

Regarding 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.

Q5: The example code is not provided, which limits the reviewers' ability to verify reproducibility.

We 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.

Q1: Why the title is Accurately? The key advantages of this work are its efficiency and reduced memory usage.

We 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

Q2: Why there is no related work?

We 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.

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?

We 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).

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.

Δ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

Q5: Figure 1 presents ambiguous visual information ... refine this figure and amplify the caption.

We thanks the reviewer's valuable feedback, we will refine this figure in the following version.

Q6: Lacks some implementational details. Is the pre-trained SAE is fixed? ... cause higher reconstruction error along the training?

SAE 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.

Q7: It may seem weird that the author claims the alignment stability ... How stable are TDPO, SimPO, and FPO which tradeoffs?

We 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.

Q8: What motivates you to change the base model in SimPO evaluation setting ... on one arbitrary base model in that setting?

We 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.

Q9: In Figure 2 left, is this an example instance or statistics? What is the analysis setting?

This 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.

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?

Feature-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.

Q11: What does uniqueness mean in line 84?

This means that FPO supposed a specific control of the alignment process (refer to Q1).

Q12: What is the wall-clock computational time of FPO vs TDPO?

Data 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.