Every Sample Counts: Supervised Fine-Tuning of Language Models with Pointwise Constraints
TL;DR
This work proposes a novel alignment framework that addresses the gap between per-sample constraints and average loss by enforcing per-sample constraints while still minimizing an average loss, and develops an augmented Lagrangian approach tailored to this formulation.
Abstract
Fine-tuning language models often requires enforcing constraints on individual inputs without compromising downstream performance. Existing constrained alignment methods impose constraints on average, which can induce undesirable disparities across inputs or users. We propose a novel alignment framework that addresses this gap by enforcing per-sample constraints while still minimizing an average loss. To mitigate the impact of overly restrictive constraints and outliers, we introduce a learned, sample-dependent relaxation that minimally adjusts the constraints, trading off a user-defined relaxation cost with the training objective. To address practical duality and optimization challenges, we develop an augmented Lagrangian approach tailored to this formulation. We demonstrate the flexibility of the framework by instantiating it under distinct small language-model fine-tuning tasks and constraints: safety in instruction following, preferences in function calling and length in re-ranking. Across these settings, our approach reduces tail constraint violations while largely preserving the model's performance.