The Lasso and MM Algorithms - 5 of 5

Part 5: Optimal Weights

Anant Mathur · July 2026

1 · Lasso + CD 2 · Auxiliary KKT 3 · MM Algorithms 4 · A Separable Bound 5 · Optimal Weights

The bound in Part 4 leaves one design choice: how to split a unit budget $\v 1^\top\v w=1$, which decides which coefficients receive the strongest update. This final part works out the best possible allocation and shows what it buys.

Section A
Why the allocation matters
Start with the computational consequence, then inspect the geometry behind it.

A convergence race

the optimal weights reach the solution in a fraction of the iterations the baselines need.

The figure compares how quickly several weight policies reduce the optimality gap on the same lasso problem. Each run uses the same MM update. The only thing that changes is how the weight budget is allocated.

Optimality gap $g(\v\gamma^{(t)})-g^*$ on one representative instance. The dots on the dashed line mark the iteration at which each policy first reaches the $10^{-6}$ target. The inset bars grow with the race and stop at each policy's median iteration count to that target over 20 replications: optimal weights 7.5, round-robin 26.5, uniform 30, random 31.5. On this particular instance uniform and random do not reach the target within 60 iterations.

Uniform and random allocations spend most of the budget on coordinates that do not need it. Round-robin eventually visits the right coordinates, but must wait for its turn. The optimal-weight rule instead removes zero coefficients that already satisfy $\tau_j\le\alpha$, then equalizes the marginal value among the eligible coordinates it funds.

Reproducible experiment. The published Python script generates the chart data and replication summary.

Explore the surrogate family

drag the white point on the simplex to choose $\v w$ and compare its surrogate with the optimal choice in blue.

Every point on the simplex defines a valid MM surrogate for the same lasso objective. Changing $\v w$ does not touch the objective or the current iterate; it only moves curvature around in the upper bound. That is enough to change the proposed next step. All the surrogates touch the orange objective at the current iterate, but the blue curve, the one built from the optimal weights, has the lowest minimum in the family and hence the strongest guaranteed one-step improvement.

jump to

Peach curves sample the surrogate family, blue is the optimal choice, and the dashed curve follows the white simplex point. Vertical guides mark the current iterate, the optimal MM step and the selected surrogate's minimum on this slice. On the right, the heatmap and bars show each allocation's full guaranteed decrease; the star marks $\widehat{\v w}$ and the open circle the uniform allocation.

Section B
Computing the optimal weights
Derive the separable weight problem, then compute its minimizer with a water-filling procedure.

From the surrogate to a budget problem

choose the allocation that gives the smallest guaranteed upper bound after the next MM step.

Substitute the closed-form update of Part 4 back into the surrogate. The bound then reduces to a sum of coordinate costs $\phi_j(w_j)$, one per coefficient, and the best allocation minimizes that sum over the unit simplex:

$$ \widehat{\v w} =\argmin_{\v w\ge\v0,\;\v1^\top\v w=1} \sum_{j=1}^p\phi_j(w_j). $$

At the current iterate write $\tau_j:=|\v x_j^\top \m G(\widetilde{\v\gamma})\v y|$ and $\xi_j:=\alpha-\tau_j$. A zero coordinate with $\xi_j\ge0$ already satisfies its lasso KKT condition, so it is removed before allocation. For every remaining coordinate, define the deactivation threshold

$$ c_j:= \begin{cases} \widetilde\gamma_j\tau_j/\xi_j, & \xi_j>0,\\ +\infty, & \xi_j\le0. \end{cases} $$

In this notation, the cost for coordinate $j$ works out to the continuously differentiable convex function

$$ \phi_j(w_j)= \begin{cases} \dfrac{(\widetilde\gamma_j\tau_j)^2}{w_j} +2\widetilde\gamma_j\tau_j^2, & w_j\ge c_j\text{ and }c_j\in(0,1],\\[0.8em] 2\widetilde\gamma_j\alpha\tau_j-w_j\xi_j^2, & \text{otherwise.} \end{cases} $$

Before the knot, one unit of weight improves the bound at the constant marginal rate $-\phi_j'(w_j)=\xi_j^2$. At $w_j=c_j$, the next update sets coefficient $j$ to zero; this is the deactivation point. Past it, extra weight cannot change the coefficient's active status, and its marginal value decays as $(\widetilde\gamma_j\tau_j)^2/w_j^2$.

A direct KKT interpretation. At zero weight, $-\phi_j'(0)=\xi_j^2=(\alpha-\tau_j)^2$: an eligible coordinate's squared KKT gap is exactly its initial marginal value. The weight problem therefore funds first the coordinates furthest from satisfying their KKT conditions.

The allocation problem in one picture

minimize a sum of convex coordinate costs while distributing one unit of weight among them.

This is a classic separable resource-allocation problem. Each curve $\phi_j(w_j)$ is the surrogate cost left after coordinate $j$ receives weight $w_j$. Every curve would prefer more weight, but the coordinates share one unit of budget.

The classical picture. Boyd and Vandenberghe, Example 5.2 and Figure 5.7, use the same simplex-constrained allocation problem with $\phi_j(w_j)=-\log(w_j+a_j)$. The figure below replaces those logarithms by the actual MM coordinate costs from one lasso iteration.

Four of the five coordinates from the running example ($\alpha=5$), the same instance animated in the vessel and KKT-allocation views below. Filled circles mark the funded optima $\widehat w_j$; the open circle is an unfunded coordinate. The short dashed tangents at the funded points are parallel because their slopes all equal $-\nu^*$.

How to read the plot. Moving right by $dw_j$ spends that much of the budget on coordinate $j$ and lowers the objective by about $-\phi_j'(w_j)\,dw_j$. What matters is slope, not height: the allocation is decided by the marginal values $-\phi_j'(w_j)$, which the vessel picture below makes visible.

The common marginal value

The KKT conditions for the simplex problem introduce one multiplier $\nu^*$. Every funded coordinate ends with marginal value exactly $\nu^*$, and every eligible but unfunded coordinate starts at or below it:

$$ -\phi_j'(\widehat w_j)=\nu^*\quad\text{if }\widehat w_j>0, \qquad \xi_j^2\le\nu^*\quad\text{if }\widehat w_j=0. $$

Each marginal value is flat at $\xi_j^2$ until the knot $c_j$ and only decreases after it, so these conditions can be met greedily: hand each small increment of budget to whichever coordinate currently has the largest $-\phi_j'(w_j)$, and stop when $\sum_jw_j=1$.

Water-filling computes the optimal weights

Background. The water-filling algorithm is best known from digital communications, where it allocates a fixed power budget across parallel channels in multicarrier systems. Power goes first to the channel where it buys the largest gain in transmission rate. Those returns shrink as power accumulates, until every funded channel offers the same marginal gain; a channel that starts below that common level gets nothing.

Here, channels become eligible coordinates, power becomes the unit budget of MM weights, and transmission-rate gain becomes reduction of the surrogate bound.

That greedy rule is water-filling. Picture one vessel per coordinate whose width at fill height $w_j$ is $-\phi_j'(w_j)$: the base has width $\xi_j^2$, and the vessel narrows past the knot $c_j$. Pour the budget into whichever vessel is currently widest. Funded surfaces always share one common width, the water level $\nu$, which only descends; when the budget runs out it rests at $\nu^*$, and vessels whose base never exceeded that level stay dry. The procedure stops in one of two ways. Either $\nu^*$ equals some coordinate's base width $\xi_{(k)}^2$ and that pivot coordinate takes the leftover budget, or $\nu^* = (\sum_{j \le k} \widetilde\gamma_j \tau_j)^2$ sits strictly between two base widths and the funded weights split proportionally to $\widetilde\gamma_j \tau_j$. Sorting the $\xi_j^2$ once gives the exact solution in $\c O(p\log p)$ time, a negligible cost next to the linear solve in each MM step.

Explore the optimal-weight procedure

Here is that procedure on the running example, one vessel per eligible coordinate.

The colored fill rises as the budget is poured. A white bar marks the knot $c_j$, where coefficient $j$ becomes inactive and the vessel starts to narrow. An empty vessel is still eligible: its base is narrower than the current level $\nu$, so it has not yet received budget. Widths use a square-root scale so all five vessels stay visible; the labels give the exact $\xi_j^2$.

Follow the KKT allocation

The last animation follows the same allocation in two views: panel (a) shows the marginal values that decide where the next increment goes, and panel (b) shows the costs $\phi_j$ falling as weight arrives. The case buttons compare the two ways the procedure can stop. The collapsed derivation collects every formula from this section in one place.

Show the full allocation derivation

Setup & notation: \(j=1,\dots,p\) indexes coefficients

Predictor \(\mathbf x_j\) is centered and scaled to \(\|\mathbf x_j\|_2=1\), and \(\tilde\gamma_j\ge0\) is its scale at the current MM iterate. Two quantities determine coordinate \(j\)'s update: \[ \alpha:=n\lambda \;\;\text{(common lasso threshold)}, \qquad\quad \tau_j:=\bigl|\mathbf x_j^\top \mathbf G(\tilde{\boldsymbol\gamma})\mathbf y\bigr| \;\;\text{(current residual signal)}. \]

KKT conditions of the auxiliary problem

At the global minimizer \(\widehat{\boldsymbol\gamma}\), \[ \tau_j = \alpha \;\;\text{ for every active coefficient } (\widehat\gamma_j > 0), \qquad\quad \tau_j \le \alpha \;\;\text{ for every zero coefficient } (\widehat\gamma_j = 0). \] Write \(\xi_j:=\alpha-\tau_j\). A zero coefficient with \(\xi_j\ge0\) already satisfies its condition and is removed before allocation. Every other coordinate is eligible, and its initial marginal value \(\xi_j^2\) measures how far it is from its KKT condition.

The weight problem and the deactivation threshold \(c_j\)

\[ \widehat{\boldsymbol w} \;=\; \mathop{\mathrm{argmin}}\limits_{\;\boldsymbol w \,\ge\, 0,\; w_1 + \cdots + w_p \,=\, 1}\;\; \sum_{j=1}^{p} \phi_j(w_j), \qquad\quad c_j := \frac{\tilde\gamma_j\,\tau_j}{\xi_j}\quad\text{when }\xi_j>0 \;\;\text{(smallest weight that switches coefficient }j\text{ off)}. \] When \(\xi_j>0\), giving coordinate \(j\) weight \(w_j\ge c_j\) makes the next update set \(\gamma_j\to0\); smaller weights leave it active. Thus \(c_j\) is the point where coefficient \(j\) becomes inactive. Past it, further weight has diminishing marginal value because \(-\phi_j'(w_j)\) decreases as \(w_j^{-2}\). When \(\xi_j\le0\), no finite deactivation threshold exists and no weight can switch the coefficient off; \(\phi_j\) stays linear over the whole budget.

Coordinate \(j\)'s contribution to the bound and its marginal value

linear part - coefficient j stays active
for \(w_j < c_j\) (and whenever \(c_j \notin (0,1]\)):
\[ \phi_j(w_j) = 2\,\tilde\gamma_j\alpha\tau_j - w_j\,\xi_j^{2}, \qquad -\,\phi_j'(w_j) = \xi_j^{2} \;\;\text{(constant slope)} \]
curved part - coefficient j is deactivated
for \( w_j \ge c_j \) when \( c_j \in (0,1] \):
\[ \phi_j(w_j) = \frac{(\tilde\gamma_j\tau_j)^2}{w_j} + 2\,\tilde\gamma_j\tau_j^2, \qquad -\,\phi_j'(w_j) = \frac{(\tilde\gamma_j\tau_j)^2}{w_j^{2}} \;\;\text{(decays as } w_j \text{ grows)} \]
\(\text{linear part: slope } -\xi_j^2\)
\(\text{nonlinear part: } \dfrac{(\tilde\gamma_j\tau_j)^2}{w_j} + 2\tilde\gamma_j\tau_j^2\)
\(\text{knot } c_j = 0.375\)
\(w_j\)
\(\phi_j(w_j)\)
Coordinate 3 of the running example: \(\tilde\gamma_j=0.25,\ \tau_j=3,\ \alpha=5\), so \(\xi_j=2\) and the knot is at \(c_j = \tilde\gamma_j\tau_j/\xi_j = 0.375\). The linear part (constant slope \(-\xi_j^2\)) meets the nonlinear part smoothly at the knot; \(\phi_j\) is continuously differentiable. At \(c_j\), coefficient \(j\) becomes inactive; beyond it, further weight has diminishing marginal value.

\(\nu^*\): the optimal KKT multiplier

The multiplier \(\nu:=\max_j-\phi_j'(w_j)\) is the greatest marginal value still available; funded coordinates sit exactly on it. In the vessel picture this is the water level. When the budget runs out, \(-\phi_j'(\widehat w_j)=\nu^*\) for funded coordinates, while eligible unfunded ones satisfy \(\xi_j^2\le\nu^*\). Where \( \nu^* \) lands decides the KKT regime:

The two termination regimes  (eligible coordinates sorted by \(\xi_{(1)}^2\ge\xi_{(2)}^2\ge\cdots\))

Boundary threshold: ν* matches one squared KKT gap \[ \nu^* = \xi_{(k)}^2 \quad\text{for some } k \] \[ \widehat w_{(j)} = \begin{cases} \tilde\gamma_{(j)}\tau_{(j)} / \sqrt{\nu^*}, & j < k \;\;\text{(nonlinear part)}\\[0.6ex] 1 - \sum_{i<k} \widehat w_{(i)}, & j = k \;\;\text{(boundary coordinate: residual on the flat branch)}\\[0.6ex] 0, & j > k \end{cases} \]
Coordinates before \(k\) lie on their curved branches. Coordinate \(k\) remains on its flat marginal segment and receives exactly the residual needed to make the weights sum to one. Winner-takes-all is the endpoint \(k=1\): the preceding sum is empty, so \(\widehat w_{(1)}=1\) and every other weight is zero.
Interior common slope: ν* lies between squared KKT gaps \[ \xi_{(k)}^2 \;>\; \nu^* = \Bigl(\sum_{j \le k} \tilde\gamma_{(j)}\tau_{(j)}\Bigr)^{2} \;>\; \xi_{(k+1)}^2 \] \[ \widehat w_{(j)} = \begin{cases} \tilde\gamma_{(j)}\tau_{(j)} / \sqrt{\nu^*}, & j \le k \;\;\text{(all funded, nonlinear part)}\\[0.6ex] 0, & j > k \end{cases} \]
All \(k\) funded coordinates lie beyond their knots. Equal marginal values and the unit-budget constraint give a split proportional to \(\tilde\gamma_j\tau_j\).

KKT allocation across coordinates

\(-\phi_j'(w_j)\) (√ scale)
\(\nu\) = —
(a) Marginal values determine the next allocation. Coordinate \(j\)'s marginal value \(-\phi_j'(w_j)\) is flat at \(\xi_j^2\), then decays after the knot. The next increment goes to a dot at the greatest available height. The dashed line is the common KKT multiplier \(\nu\): funded coordinates lie on it, while unfunded coordinates wait at \((0,\xi_j^2)\) below it.
\(\phi_j(w_j)\)
(b) The separable objective decreases. Coordinate \(j\)'s contribution \(\phi_j(w_j)\) to the bound. All funded coordinates descend with the same slope \(-\nu\) at the solution (parallel white tangent strokes); ○ marks the knot \(c_j\) where the linear part hands over to the nonlinear part.
\(\nu\) = —
KKT multiplier (marginal value of budget)
\(\textstyle\sum_j w_j\) · budget used

Closing the loop

The pieces of this series now fit together. The lasso KKT conditions sort coefficients into active ones and zeros that already satisfy optimality; the auxiliary scales turn those conditions into a convex problem over $\bb{R}_+^p$. The matrix bound of Part 4 then gives a closed-form MM update for every coordinate at once, and the weight problem solved here picks out the best member of that surrogate family for $\c O(p\log p)$ extra work per iteration.

With exact linear solves and a unique auxiliary minimizer, the iterates converge globally to the lasso KKT point. The full statements and proofs, along with the implementation and experiments, are in the paper.

← Part 4: A Separable Bound Back to blog →