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 degree of freedom: a unit budget $\v 1^\top\v w=1$ that decides which coefficients receive the strongest update. This part derives the optimal allocation and shows why it improves the MM step.

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

A convergence race

the optimal allocation reaches the solution in fewer iterations than the baseline allocations.

Consider an ordinary lasso with $p=12$ correlated, standardized predictors and $n=300$ observations. Three coefficients generate the response; at $\lambda=0.45\lambda_{\max}$, two remain active in the representative solution. The MM update is identical in every run; only the weight policy changes.

Optimality gap $g(\v\gamma^{(t)})-g^*$ on one representative instance. Across 20 replications, median iterations to a $10^{-6}$ gap were: optimal weights 7.5, round-robin 26.5, uniform 30, and random 31.5.

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$ leaves the objective and current iterate fixed, but redistributes the curvature of the upper bound and therefore changes the proposed next step. All the surrogates touch the orange objective at the current iterate. The blue curve uses the optimal weights: its minimum is the lowest in the family, so it gives the strongest guaranteed one-step improvement.

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. The star marks the optimal weight and the open circle the uniform allocation.

Part 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 into the surrogate and minimize 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} $$

The cost contributed by coordinate $j$ is 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, so $c_j$ is its deactivation point. Beyond that point, additional weight no longer changes its active status and has diminishing marginal value, $(\widetilde\gamma_j\tau_j)^2/w_j^2$.

A direct KKT interpretation. At zero allocated weight, $-\phi_j'(0)=\xi_j^2$, so an eligible coordinate's squared KKT discrepancy is exactly its initial marginal contribution to improving the surrogate bound. The optimal-weight problem therefore prioritizes the coordinates furthest from KKT balance. In particular, for a zero coefficient that violates $\tau_j\le\alpha$, this priority is $(\tau_j-\alpha)^2$: the larger the violation, the greater the initial value of assigning weight to that coordinate.

The allocation problem in one picture

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

This is a general separable resource-allocation problem. Each curve $\phi_j(w_j)$ tells us the surrogate cost remaining after coordinate $j$ receives weight $w_j$. Every curve would prefer more weight, but the coordinates share one budget:

$$ \underset{\v w\ge\v0}{\operatorname{minimize}} \quad \sum_{j=1}^p\phi_j(w_j) \qquad\text{subject to}\qquad \sum_{j=1}^p w_j=1. $$
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 actual costs with common $\alpha=5$. 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 unit budget on coordinate $j$ and lowers the objective by approximately $-\phi_j'(w_j)\,dw_j$. The vertical cost values themselves do not equal water heights. The allocation geometry lives in marginal-slope space: in the vessel picture, $w_j$ is the fill height and $-\phi_j'(w_j)$ is the vessel width.

The connection to the lasso KKT gap is at the left edge of the plot. For every eligible coordinate, before any budget is allocated,

$$ -\phi_j'(0)=\xi_j^2=(\alpha-\tau_j)^2. $$

Thus the squared auxiliary KKT gap is the coordinate's initial marginal value: it is the amount by which an infinitesimal unit of weight lowers the surrogate bound. As weight is added, that marginal value stays flat until $c_j$ and then decreases. The optimal-weight procedure repeatedly gives the next increment to the currently largest $-\phi_j'(w_j)$. It stops when $\sum_jw_j=1$; at that point all funded coordinates have the same marginal value $\nu^*$, while every unfunded coordinate starts at or below it.

The common marginal value

The KKT conditions for the simplex problem introduce one multiplier $\nu^*$. Every funded coordinate has the same final marginal value, while an eligible but unfunded coordinate starts 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. $$

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 is assigned first where it gives the greatest marginal increase in transmission rate. As power is added, these marginal returns decrease until every funded channel has the same marginal return. A channel whose initial marginal return is already below that common level receives no power.

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

Because each $-\phi_j'(w_j)$ starts at height $\xi_j^2$ and only decreases as $w_j$ grows, the KKT conditions above lead to a classic water-filling procedure. Picture one vessel per coordinate whose width at fill height $w_j$ is $-\phi_j'(w_j)$ — base width $\xi_j^2$, narrowing past the knot $c_j$. Pour the unit budget into whichever vessel is currently widest: funded surfaces then always share one common width — the water level $\nu$ — which only descends. When the budget is exhausted the level rests at $\nu^*$: vessels whose base $\xi_j^2$ never rose above the final level stay dry. Termination lands in one of two cases — either $\nu^*$ equals some coordinate's base width $\xi_{(k)}^2$ (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)$ — negligible next to the linear solve in each MM step.

Explore the optimal-weight procedure

Each eligible coordinate is now a vessel. Its base width is the squared auxiliary KKT gap $\xi_j^2$, its fill height is the allocated weight $w_j$, and its width at height $w_j$ is the current marginal value $-\phi_j'(w_j)$. The budget always enters the vessel with the widest available surface.

The colored fill rises as the unit budget is allocated. A horizontal white bar marks $c_j$, the point where coefficient $j$ becomes inactive; above it the vessel narrows because additional weight has diminishing marginal value. An empty vessel is not excluded from the model: its base is simply narrower than the current surface width $\nu$, so it has not started receiving budget. Widths are drawn on a square-root scale so that all five vessels remain visible; the labels give the exact squared KKT gaps $\xi_j^2$.

Follow the KKT allocation

The vessel picture has already supplied the intuition. The animation below keeps only the two views needed to follow the algorithm: the marginal curves show which coordinate should receive the next increment of budget, while the cost curves show the resulting decrease in the separable objective. At each step the algorithm increases a coordinate with the largest available value of $-\phi_j'(w_j)$. It stops when $\sum_j w_j=1$ and the KKT conditions hold. The controls compare the two ways that this stopping point can occur.

How the scalar-lasso MM algorithm distributes one unit of update budget across coefficients. The animation assigns the next increment to the eligible coordinate with the greatest marginal improvement to the surrogate bound.
Show the full allocation derivation

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

Predictor \(\mathbf x_j\) is centred 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 minimiser \(\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\). For an active coefficient, \(\xi_j^2\) is its initial marginal value. A zero coefficient with \(\xi_j<0\) violates its KKT condition and remains eligible. A zero coefficient with \(\xi_j\ge0\) is already correct and is removed before allocation. Hence, among eligible coordinates, a larger \(\xi_j^2\) means a larger KKT discrepancy and a greater initial benefit from assigning weight there. For a violating zero coefficient, \(\xi_j^2=(\tau_j-\alpha)^2\).

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.3\)
\(w_j\)
\(\phi_j(w_j)\)
An example \(\phi_j\) with \(\tilde\gamma_j=0.4,\ \tau_j=3,\ \alpha=7\): raw gap \(\xi_j=4\), knot at \(c_j = \tilde\gamma_j\tau_j/\xi_j = 0.3\). 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)
\(w_j\)
\(\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)\)
\(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 lasso KKT conditions identify which coefficients are active and which zero coefficients already satisfy optimality. The auxiliary scales turn those conditions into a convex problem over $\bb{R}_+^p$. The matrix bound supplies a concurrent closed-form MM update, and the optimal-weight procedure selects the member of that surrogate family with the smallest one-step upper bound in $\c O(p\log p)$ time.

With exact linear solves and a unique auxiliary minimizer, the theory gives global convergence to the lasso KKT point. The full theory, implementation and experiments are in the paper.

← Part 4: A Separable Bound Back to blog →