PaperlayerView exact revision on arXiv

Source-linked AI summary

Preference Data Selection for Mitigating the Alignment Tax in Large Language Models

Minsu Kim, Jianxun Lian, Xing Xie, Steven Euijong Whang

arXiv:2608.24192v1cs.AIcs.CL

TL;DR

LLM preference alignment can incur an alignment tax by degrading pretrained general capabilities, while the data characteristics driving this degradation are underexplored. BALIGN analyzes preference-optimization gradients and selects data using a composite risk score over three features. Across the reported evaluations, it preserves general capabilities while maintaining alignment performance and reaches the stated optimal balance.

  • Problem

    Preference alignment can degrade pretrained general capabilities, but the preference-data characteristics associated with this alignment tax remain underexplored.

  • Method

    BALIGN selects preference samples using a composite risk score combining reference margin, token length difference, and TF–IDF similarity.

  • Results

    BALIGN consistently achieves the optimal balance between general capability and alignment, with the best general capability and best or second-best alignment performance.

  • Takeaways & Limitations

    The results support BALIGN as a data-centric strategy for minimizing catastrophic forgetting while preserving alignment signals.

Abstract

from arXiv · show

Aligning large language models to human preferences is crucial for real-world deployment but frequently incurs an alignment tax, leading to the catastrophic forgetting of pre-trained general capabilities. While previous works primarily frame this problem as an optimization or architectural challenge, the inherent characteristics of preference data that drive this degradation remain largely underexplored. In this paper, we propose BALIGN, a balanced data selection strategy that explicitly mitigates catastrophic forgetting while optimizing alignment efficacy. Through theoretical and empirical analyses of the preference optimization gradient, we identify three key data-centric features that dictate parameter drift: the reference model's log-probability margin, the token length difference between chosen and rejected responses, and the TF-IDF similarity to general capability corpora. By aggregating these orthogonal features into a unified composite risk score, BALIGN systematically filters out high-risk preference samples that disrupt intrinsic model parameters or provide minimal alignment utility. Extensive experiments on standard human preference datasets demonstrate that BALIGN strongly preserves foundational capabilities without compromising alignment gains, consistently achieving the optimal Pareto frontier with minimal computational overhead.

Introduction

LLM alignment can cause an alignment tax that degrades general capabilities, motivating a data-centric approach that balances stability and plasticity. BALIGN selects preference data using three risk-related features and reports preserved capabilities alongside alignment performance.

  • Alignment can cause catastrophic forgetting of general knowledge, instruction following, reasoning, code, and math capabilities.
  • Preserving general capabilities while integrating human preferences requires balancing stability and plasticity.
  • Prior work largely treats catastrophic forgetting as an optimization or architectural challenge, leaving preference-data characteristics underexplored.
  • BALIGN identifies reference margin, token length difference, and TF–IDF similarity as three data-centric features for selection.
  • Selecting minimal margins, small length differences, and preference samples distinct from general SFT data mitigates forgetting while improving alignment.
  • BALIGN assigns feature-specific and composite risk scores to filter samples that degrade stability or plasticity.

Related Work

Related work frames alignment, catastrophic forgetting, and mitigation methods through established optimization and continual-learning approaches. However, existing methods have mainly targeted supervised fine-tuning, leaving the joint dynamics of forgetting and preference alignment underexplored.

  • LLM alignment adjusts pretrained models so outputs conform to human preferences concerning helpfulness, harmlessness, honesty, and safety.
  • RLHF combines supervised fine-tuning, a learned reward model, and PPO, while GRPO reduces critic-model memory overhead but remains engineering-heavy and expensive.
  • Alignment-related catastrophic forgetting is the degradation of previously acquired knowledge or pretrained capabilities while optimizing for human preferences.
  • Mitigation research includes experience replay, regularization, self-synthesized rehearsal, self-distillation, and model merging.
  • These approaches primarily focus on supervised fine-tuning for task-specific performance, leaving the joint dynamics of catastrophic forgetting and alignment underexplored.

Problem Definition

The paper formulates preference alignment from a pretrained reference model using preference pairs and DPO, while evaluating general capability retention and preference accuracy. Data selection is posed as a bi-level optimization balancing these objectives.

  • Notation and Preliminaries: The reference model is a pretrained language model, and each preference pair contains a chosen response preferred over a rejected response for a prompt.
  • Alignment Objective: DPO optimizes the policy to satisfy observed preferences without explicitly learning a reward model.
  • Evaluation: General capability measures retention of previously acquired knowledge across benchmarks disjoint from the alignment distribution.
  • Evaluation: Catastrophic forgetting is defined as F(πθ) = C(πref) − C(πθ), with positive values indicating degraded general capability relative to the reference.
  • Evaluation: Alignment quality is measured by preference accuracy, the proportion of held-out samples where the chosen response receives higher log-probability than the rejected response.
  • Data Selection: Preference-data selection is framed as a bi-level optimization problem that chooses a subset of a specified size to balance alignment and catastrophic forgetting.

Method

BALIGN selects preference data using three features that characterize stability and plasticity during alignment: reference-model margin, response-length difference, and TF–IDF similarity. It converts their normalized risks into a composite score for filtering high-risk samples.

  • BALIGN derives three data-centric features—reference margin, token length difference, and TF–IDF similarity—to characterize stability-plasticity trade-offs during alignment.Reference margin and length difference indicate stability, while TF–IDF similarity proxies plasticity.
  • Reference Margin: The reference margin compares average per-token log-probabilities of chosen and rejected responses under the reference model, with large absolute margins treated as risky.A large positive margin can indicate that the reference model already prefers the chosen response, so further updating may add little alignment signal while increasing drift.
  • Token Length Difference: The token length difference measures chosen-minus-rejected response length, because large asymmetry can make the margin signal length-driven and inflate alignment updates.The corresponding risk score accounts for asymmetric risk through a hyperparameter γ.
  • TF–IDF Similarity: TF–IDF similarity averages cosine similarity between each preference sample and a general-capability corpus, with higher similarity indicating greater gradient-interference risk.Preference samples are represented using their prompt and chosen response, and the vectorizer is fit on preference and general data together.
  • Empirical Analysis: Bucket experiments show complementary roles: reference-margin and length risks primarily track catastrophic forgetting, whereas TF–IDF risk captures preference-learning failure.Experiments use HH-RLHF, Llama-3.1-8B-Instruct as the reference model, five equal-sized risk buckets, and evaluations including MMLU, IFEval, ARC-Challenge, HumanEval, and GSM8K.
  • BALIGN Data Selection Strategy: BALIGN aggregates min–max-normalized risks into a composite score and selects samples under a budget to reduce expected stability or plasticity risks.The score combines reference-margin, length-difference, and TF–IDF risks; λ controls the balance among them.

Experiments

Experiments evaluate BALIGN’s general-capability retention, alignment performance, controllability, ablation behavior, and selection efficiency against multiple baselines on HH-RLHF.

  • Evaluation setup: Experiments measure both absolute general capability and alignment, plus catastrophic forgetting and alignment gain relative to the base reference model.These metrics represent stability and plasticity, respectively.
  • Evaluation setup: The evaluation uses HH-RLHF for helpfulness and harmlessness alignment and five general-capability datasets spanning knowledge, instruction following, reasoning, coding, and mathematics.The five datasets are MMLU, IFEval, ARC-Challenge, HumanEval, and GSM8K.
  • General capability and alignment results: Figure 2 reports preserved diverse general capabilities alongside alignment performance comparable to training on the full preference dataset.The comparison concerns the base model’s capabilities and full-dataset training.
  • General capability and alignment results: BALIGN achieves the best general capability and the best or second-best alignment performance, consistently balancing capability retention with alignment.Its alignment performance is only marginally below the top performer when it is not best.
  • Controllability of stability and plasticity: Varying DPO β controls the stability-plasticity trade-off: lower β generally favors plasticity, higher β favors stability, and β=0.1 provides a strong balance.Across β values from 0.01 to 0.5, BALIGN traces the Pareto frontier and occupies the top-right region.
  • Ablation study: Removing any one of BALIGN’s three risk scores harms equilibrium: reference-margin and length risks worsen general capability, while TF-IDF risk reduces alignment gain.The ablation supports complementary and synergistic contributions from all three scores.
  • Computation time analysis: BALIGN requires lower selection overhead because its three risk scores need at most one forward pass, unlike baselines using per-sample gradients or auxiliary training.Figure 4 compares total computation time during data selection.

Conclusion

BALIGN addresses the stability-plasticity dilemma in LLM alignment through efficient, data-centric preference selection. Its composite risk score removes high-risk samples, improving the capability-alignment trade-off with lower computational overhead than existing baselines.

  • Conclusion: BALIGN jointly optimizes model stability and plasticity through a data-centric preference-selection framework.The framework targets catastrophic forgetting and preference alignment together.
  • Conclusion: Extreme reference margins and length asymmetries exacerbate catastrophic forgetting, while lexical overlap with general-capability domains governs alignment gains.These findings motivate the three-feature risk design.
  • Conclusion: BALIGN uses a unified composite risk score to remove high-risk preference samples before optimization.The selection strategy integrates the identified data-centric signals.
  • Conclusion: Evaluations show superior capability retention and preference alignment with significantly lower computational overhead than existing baselines.The conclusion characterizes BALIGN as a scalable preprocessing foundation for aligned LLMs.

More Details on Empirical Analysis

The empirical analysis shows that BALIGN’s three risk scores capture complementary forms of model degradation. Two identify samples associated with catastrophic forgetting, while the third isolates weak preference learning.

  • Empirical analysis: As risk rises from B1 to B5, catastrophic forgetting increases for s∆pref and s∆ℓ, reaching 3.48% and 2.19%, respectively.These two scores act as primary indicators of forgetting.
  • Empirical analysis: The highest-risk sτ bucket substantially reduces alignment gain while keeping forgetting relatively stable.This score primarily captures failure of preference learning.
  • Empirical analysis: The distinct behaviors of the three scores support integrating them for a robust alignment framework.The analysis describes their effects as complementary.

More Details on Experimental Settings

The experiments use fixed computational and training settings to evaluate BALIGN on preference alignment and general capabilities. They compare against adapted data-selection baselines using full-parameter DPO training.

  • Experimental settings: Training and evaluation run on Intel Xeon Silver 4210R CPUs and NVIDIA RTX A6000 GPUs.These are the reported hardware settings.
  • Experimental settings: Evaluation maximizes alignment gain while minimizing catastrophic forgetting, using deterministic greedy decoding and log-likelihood scoring.Greedy decoding is used for generative benchmarks and log-likelihood for multiple-choice tasks.
  • Experimental settings: HH-RLHF provides separate training and test splits for helpfulness and harmlessness alignment.The helpfulness and harmlessness training subsets contain 43,835 and 42,537 samples, respectively.
  • Experimental settings: Experiments use Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct as base reference models, restricted to open instruction-tuned models in the 7B–8B range.The selection reflects a balance between base capabilities and GPU tractability; closed-weight models cannot be used.
  • Baselines: BALIGN is compared with naive, SFT-oriented, stability- and plasticity-aware SFT, and plasticity-aware alignment baselines.SFT-specific methods are adapted by treating each chosen response as the ground-truth target.
  • Training configuration: DPO training uses full-parameter fine-tuning with batch size 128, learning rate 5 × 10^-6, cosine scheduling, β=0.1, three epochs, and bfloat16 precision.Training runs under DeepSpeed ZeRO-3.

More Details on General Capability and Alignment Results

Tables 4–7 provide detailed general-capability and alignment results across helpfulness and harmlessness on HH-RLHF, using Llama-3.1-8B-Instruct or Qwen2.5-7B-Instruct as base reference models.

  • The paper directs readers to Tables 4–7 for more detailed general-capability and alignment results.
  • Table 4 reports helpfulness results on HH-RLHF with Llama-3.1-8B-Instruct as the base reference model.
  • Table 5 reports helpfulness results on HH-RLHF with Qwen2.5-7B-Instruct as the base reference model.
  • Table 6 reports harmlessness results on HH-RLHF with Llama-3.1-8B-Instruct as the base reference model.
  • Table 7 reports harmlessness results on HH-RLHF with Qwen2.5-7B-Instruct as the base reference model.
Loading 2608.24192v1…