TAB-PO: Preference Optimization with a Token-Level Adaptive Barrier for Token-Critical Structured Generation

TL;DR

Token-Adaptive Barrier Preference Optimization (TAB-PO), which augments DPO with token-weighted, reference-adjusted advantages that prioritize high-value semantic tokens, and a conditional token-level barrier that regularizes under-confident tokens balancing SFT-anchored likelihood and preference-driven separation in low-separation, importance-skewed regimes is introduced.

Abstract

Direct Preference Optimization is an offline post-SFT method for aligning language models from preference pairs, with strong results in instruction following and summarization. However, DPO's sequence-level implicit reward can be brittle for token-critical structured prediction settings such as medical annotation, which often exhibit (i) low-separation preference pairs, where chosen and rejected completions differ by minimal edit distance (often 1-3 tokens), and (ii) token-importance skew, where sparse semantic tokens (hierarchical labels and evidence Spans) carry disproportionate task importance relative to high-frequency structural tokens (JSON scaffolding). In this regime, standard DPO suffers from margin collapse (insufficient log-probability separation between near-identical preferences), likelihood squeezing (the margin objective shifts the absolute likelihoods of both completions together), and gradient dilution, where uniform sequence-level weighting diffuses learning signal across shared scaffolding while rare, confusable label tokens receive weak, noisy updates. We introduce Token-Adaptive Barrier Preference Optimization (TAB-PO), which augments DPO with token-weighted, reference-adjusted advantages that prioritize high-value semantic tokens, and a conditional token-level barrier that regularizes under-confident tokens balancing SFT-anchored likelihood and preference-driven separation in low-separation, importance-skewed regimes. We evaluate TAB-PO on medical communication annotation, a task requiring joint prediction of hierarchical labels and evidence Spans from patient-provider messages. TAB-PO achieves a ~ 4% relative improvement in micro-F1 over SFT and consistently outperforms recent preference-optimization baselines.

Figures (10)

• Figure 1: Overview of the TAB-PO training pipeline: Codebook Creation$\rightarrow$ Manual Annotation $\rightarrow$ Prompt Engineering $\rightarrow$ Supervised Fine-tuning $\rightarrow$ Preference Optimization with TAB-PO.
• Figure 2: Performance comparison of preference optimization techniques across Code, Sub-code, and Span F1 scores.
• Figure 3: Confusion matrices comparing top 5 misclassifications between SFT and TAB-PO models.
• Figure 4: TAB-PO ablation study on Qwen2.5-1.5B-Instruct. Results are reported in F1 and averaged over five runs with different random seeds. LN is length normalization , CB is class balancing, and TW is token weighting.
• Figure 5: Comparison of F1 scores between baseline and prompt-engineered instructions across two language models (Llama-3.1-8B and Qwen2.5-14B). Prompt engineering significantly improves performance.