S parse PO : Controlling Preference Alignment of LLM s via Sparse Token Masks (2025)

Abstract: Direct alignment algorithms have proven an effective step for aligning language models to human-desired behaviors. Current variants of the Direct Preference Optimization objective have focused on a strict setting where all tokens are contributing signals of KL divergence and rewards to the loss function.However, human preference is not affected equally by each word in a sequence but is often dependent on specific words or phrases, e.g. existence of toxic terms leads to non-preferred responses. Based on this observation, we argue that not all tokens should be weighted equally during PO and propose a flexible objective termed SparsePO, that aims to automatically learn to weight the KL divergence and reward corresponding to each token during PO training. We propose two different variants of weight-masks that can either be derived from the reference model itself or learned on the fly. Notably, our method induces sparsity in the learned masks, allowing the model to learn how to best balance reward and KL divergence contributions at the token level, learning an optimal level of mask sparsity.Extensive experiments illustrate the effectiveness of our approach at aligning to preference proxies, including sentiment control, helpfulness and harmless, and summary quality.Our method obtains +10\% and +3\% win-rate points in summarization and dialogue scenarios, respectively,without compromising the reasoning capabilities of the model, or the relevancy and faithfulness of the summary response.

## introduction the rise of employing large language models (llms) as conversational agents has increased the importance of aligning them with human preferences. preference optimization (po), i.e. the training paradigm that aims to steer models to a desired behavior (typically related to human perception), * equal contribution is considered one of the most important step in the pipeline of llm training for producing accurate, harmless and controllable models. reinforcement learning from human feedback (rlhf; @xcite ) was the primary method for obtaining such a behavior. however, due to it's inherent complexity it has been overpowered by direct preference optimization (dpo; @xcite ), a simpler, offline approach that produces a policy model that fits the preference data without the need for reinforcement learning. dpo performs at the sequence level, optimizing rewards and measuring kl divergence for complete responses. however, various studies have shown that signals from specific tokens are primarily responsible for learning desired behaviors, both during pre-training @xcite and preference optimization @xcite . in particular, in domains where the preference is determined by a specific aspect (e.g. sentiment, toxicity) or when the decision relies on certain subsequences @xcite , it is necessary to consider more finegrained updates. to further illustrate this point, figure 1 shows that dpo is already learning implicitly to assign different token-level rewards, with higher values on a few tokens with positive/negative polarity (e.g. pretty, weak). however, noting the various lone tokens with high rewards, dpo's reward distribution seems inconsistent, and we posit that it would benefit from a more explicit signal. aligned with prior work, we argue that not all tokens are important in preference optimization. we further propose that in order to have more diverse responses, and flexible optimization, we should allow only certain tokens to be close to the reference model so that the rest are able to grow beyond it-dismissing the need for measuring kl divergence on all tokens. as such, in this work we propose sparse token-level preference optimization (sparsepo), a method that learns automatically during training inherently sparse masks figure 1 : token-level rewards for chosen (top) and rejected (bottom) responses given an input prompt. after a gpt2-large model is trained with dpo on the imdb dataset to generate positive movies reviews, these rewards are calculated as the log ratio of token probabilities between policy (dpo) and reference model (original gpt2-large). denser values indicate higher probability score assigned to a token by the policy than the reference, implying importance towards that preference. our contributions include (1) a flexible framework, termed sparsepo, for weighting token-level reward and kl contributions tailored to the offline preference optimization objective, (2) analyses over the induced masks' sparsity and reward frontier and how they correlate with controlled kl divergence, (3) quantitative and qualitative gains when employing our proposed approach to different domains with explicit or implicit preference indicators. extensive experiments demonstrate the effectiveness of sparsepo, obtaining +10% and +3% win rate points in summarization and dialogue tasks, respectively, without compromising the reasoning capabilities of the model, or the relevancy, lexical diversity, and faithfulness of the summary response. ## experiments in this section, the effectiveness of sparsepo is investigated in both proxy-preference and humanpreference setups. proxy-preference setups are analyzed through sentiment control, and summarization, whereas human-preference setup is analyzed through single-turn dialogue tasks. see appendix b and c for further details on experimental setup and complementary results, respectively. ## discussion based on the controlled experiments we conducted in the previous section, here we briefly discuss our overall findings. firstly, based on the sentiment control analysis, sparsepo allows larger kl divergence at little to no cost in expected ground-truth reward. the β value is able to control sparsity in both masks, across domains, with values between 0.6 to 4 leading to mid-range sparsity levels, depending on the domain and target preference proxy. we found that higher sparsity was present in sentiment control, highlighting a certain triviality of the task as the sft model seems able to already identify words that are important for the target preference. on the other end, for summarization, lower sparsity between 0.2 and 0.4 seemed best in terms of alignment accuracy as summary correctness is a less well-defined preference proxy. for helpfulness control, optimal sparsity was found instead between 0.6 and 0.8, possibly as existence of toxic terms renders response dispreferred. we argue that the mask works in tandem with beta and we observed that the range of betas that are effective with sparsepo is generally higher than dpo (with best values between 0.4-1). * furthermore, sparsepo avoids reward hacking by maintaining a meaningful distribution of values in masks for reward and token-level kl terms, as revealed by our analysis in § 3.2 and § c.5. from our analysis over dpo, tdpo and their variants, it is important to note that, although restricting divergence at the response or tokenlevel proves effective at maintaining the model indomain, this does not guarantee better ground-truth rewards or better downstream task performance. for cases in which the preference proxy is complex, such as 'helpfulness', 'summary quality', this plain control can even hinder performance. in contrast, we devise a training procedure in which a model can learn to enhance or suppress the reward and kl divergence for each token independently. our qualitative analysis shows that indeed for trivial tasks, tokens that are considered important for the preference proxy get high rewards and low kl divergence, meaning they need to be close to the reference predictions to maintain preference. finally, regarding the scalability of our method, we provided a principled comparison across a wide range of mode sizes (albeit considered relatively small for current industrial standards) going from 500m (gpt2) to 1.4b (pythia) and 6b (gptj), as well as dataset sizes going from 10k (imdb) to 150k (hh). the results showed that sparsepo can scale effectively to model size and training data size across a variety of domains. ## conclusion we introduced sparse token-level preference optimization (sparsepo), a novel lm alignment strategy that learns to weight the reward and kl divergence for each particular token in a response during po training. we proposed two masking strategies, obtaining model activation-based masks from the reference model and learning mask representations either commonly for both reward and divergence terms or independently. by allowing masks to be learned along with preference, we observed that they converged to a non-trivial level of sparsity which can be controlled with well-studied hyperparameters in preference optimization, while being dependent on target preference proxy. extensive experiments across several tasks and domains, reveal that our method consistently outperforms strong baselines that model preference at the response and token-level, while assigning higher rewards and lower kl values to tokens that are important for inferring target preference.