SparsePO: Controlling Preference Alignment of LLMs via Sparse Token Masks

TL;DR

SparsePO introduces token-level weighting for Direct Preference Optimization by learning sparse masks that regulate per-token rewards and KL penalties. It offers two mask-generation schemes—activation-based and learnable sparse masks—and supports common or independent masking for rewards and KL terms, enabling flexible control over token influence. Across sentiment, summarization, and dialogue tasks, SparsePO yields favorable reward-KL frontiers and improved win rates relative to strong PO baselines, with mask sparsity adapting to the target preference and beta settings. The approach demonstrates scalable improvements across model sizes and domains, though performance is task-dependent, and limitations exist for highly structured outputs like code.

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 harmlessness, and summary quality. Our method obtains +10% and +3% win rate points in summarization and dialogue scenarios, respectively, without compromising model reasoning or the relevancy and faithfulness of the summary response.

Figures (22)

• 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.
• Figure 2: Pareto frontier of expected reward and response-level KL divergence w.r.t. the reference model, for a sentiment control scenario over the IMDB dataset. Solid lines estimate the frontier for each system, and points represent hyper-parameter variations.
• Figure 3: Sparsity levels in the reward mask ($m_{u}$, left) and the token-level KL divergence mask ($m_d$, middle), as well as token-level KL divergence of chosen responses during training (over IMDB), for increasing values of $\beta$.
• Figure 4: Token-level heatmaps for chosen responses for TDPO-v2 SparsePO. Darker color indicates higher values. All scores are scaled in $[0,1]$ for comparison.
• Figure 5: Win rates of system responses against chosen responses in Anthropic HH single-turn dialogue.