Lagrangian Duality and Constrained Optimisation
How a constraint becomes a penalty with a price attached, why the KL-regularised objective at the heart of RLHF has a closed-form solution, and where duality quietly fails.
Every alignment method built on reinforcement learning solves the same shaped problem: improve a reward without straying too far from a reference model. Written as a constraint it is awkward, and written as a penalty it is trivial, and the machinery that moves between the two forms is Lagrangian duality. Knowing that machinery is what turns the DPO derivation from magic into arithmetic.
The price of a constraint
Given a problem \(\min_x f(x)\) subject to \(g(x) \le 0\), form the Lagrangian
The multiplier \(\lambda\) is a price: it converts one unit of constraint violation into units of objective. The dual function \(d(\lambda) = \min_x \mathcal{L}(x, \lambda)\) is a lower bound on the constrained optimum for every \(\lambda \ge 0\), which is weak duality and holds unconditionally. When the problem is convex and mildly regular, the best lower bound equals the true optimum, strong duality, and the KKT conditions characterise the solution: stationarity, feasibility, dual feasibility, and complementary slackness \(\lambda\, g(x) = 0\). Complementary slackness has an interpretation worth internalising: a constraint that is not tight has price zero, and a constraint with a positive price is exactly binding.
The closed form that alignment runs on
Consider the objective every RLHF pipeline optimises, maximise expected reward under a KL penalty toward a reference policy:
This is a constrained problem in disguise: \(\pi\) must be a probability distribution, so \(\sum_y \pi(y \mid x) = 1\). Attach a multiplier to that constraint, differentiate the Lagrangian with respect to each \(\pi(y \mid x)\), set to zero, and the solution falls out:
with \(Z(x)\) the normaliser enforcing the constraint. The optimal policy is the reference reweighted exponentially by reward, temperature \(\beta\). Rearranging for \(r\) expresses the reward in terms of the optimal policy and the reference, and that inversion is the whole of DPO: a reward model was never needed because the policy already is one (Rafailov et al., NeurIPS 2023, arXiv:2305.18290). \(\beta\) is not a tuning knob invented for convenience; it is the dual price of the trust region.
The same structure appears in trust-region policy optimisation, where a KL constraint on the policy update is handled through its Lagrangian and solved with a natural-gradient step, and PPO's clipped objective is a cheap heuristic surrogate for that constrained problem rather than a different idea.
When duality misleads
Non-convexity opens a gap. Neural network objectives are non-convex, so the best dual bound can sit strictly below the primal optimum. Reasoning about a deep model's training as if strong duality held is a common and quiet error; what actually justifies the RLHF derivation above is that the optimisation is over the distribution \(\pi\), which is a convex problem in \(\pi\) even though the parameterisation is not convex in \(\theta\).
Fixed penalties are not fixed constraints. Choosing \(\beta\) once and leaving it is not the same as constraining KL to a budget. As training moves, the KL corresponding to a fixed \(\beta\) drifts, sometimes by an order of magnitude, which is why adaptive-\(\beta\) controllers exist in RLHF implementations and why reported KL values matter more than reported \(\beta\) values.
Constraint qualification is not automatic. Strong duality needs a regularity condition such as Slater's, a strictly feasible point. Problems with equality constraints on the boundary, or with constraints that are only satisfiable at a single point, can have a duality gap even when convex.
The multiplier is only interpretable at the optimum. Reading \(\lambda\) as a sensitivity, how much the optimum improves per unit of constraint relaxation, is valid at a solution. Reading mid-optimisation multipliers as prices is how people talk themselves into believing a partially trained model is trading off correctly.
12 flashcards for this concept
Click a card to reveal the answer.