Optimization

Online Scaling for Operator Splitting without a Smooth Envelope

Posted by Wanyu Zhang on September 10, 2026

OSGM uses online learning to choose preconditioners for gradient methods. Our recent work, Operator Splitting Methods with Online Scaling (OSOP), extends this idea to operator splitting methods. This is joint work with Wenzhi Gao and Madeleine Udell.

The starting point is the envelope: an auxiliary objective that turns a splitting step into a gradient step in a suitable metric. Writing the envelope as $E$, the relation is $T(z)=z-B(z)\nabla E(z)$ under suitable assumptions, with $B(z)$ positive definite. Minimizing the envelope recovers a solution of the original problem. See my earlier post for related background.

The idea of OSOP is simple: apply OSGM with some adaptations to these envelopes. It tries $z-P(z-T(z))$ and learns $P$ from a designed feedback function. When the envelope is smooth and convex, we naturally have convergence guarantees. We also establish more general sufficient conditions to ensure smooth and convex envelopes.

However, envelope can be nonsmooth

For PDHG on LPs, the envelope can have kinks. Consider

\[\min_{x\geq0}\;c^\top x \qquad\text{subject to}\qquad Ax=b,\]

with Lagrangian $L(x,y)=c^\top x+y^\top(b-Ax)$. For fixed $\tau,\sigma>0$ with $\tau\sigma|A|_2^2<1$, PDHG takes the step

\[x^+=[x-\tau(c-A^\top y)]_+,\qquad y^+=y+\sigma\bigl(b-A(2x^+-x)\bigr).\]

Here $[\cdot]_+$ replaces negative coordinates by zero. Using the DRS reformulation of PDHG, the associated envelope simplifies to

\[E_{\tau,\sigma}(x,y) ={}c^\top x^+ +y^\top(b-Ax^+) -\tfrac{1}{2\tau}\|x^+-x\|^2 +\tfrac{\sigma}{2}\|b-A(2x^+-x)\|^2.\]

Unfortunately, this envelope is nonsmooth in general, so the smooth-envelope analysis no longer applies. Then a natural question is:

Can we accelerate operator splitting methods such as PDHG when the envelope is nonsmooth, or even a general fixed-point method?

Residual feedback (contractive $T$)

For a fixed-point method $z^+=T(z)$, define

\[r(z):=z-T(z),\qquad \ell_k(P):=\tfrac{\|r(z^k-Pr(z^k))\|}{\|r(z^k)\|}.\]

The residual vanishes exactly at a fixed point; the feedback measures its relative reduction. If $T$ is a contraction with factor $\rho<1$, the identity scaling guarantees progress:

\[\ell_k(I) =\tfrac{\|T(z^k)-T^2(z^k)\|}{\|r(z^k)\|} \leq\rho.\]

For PDHG, we obtain a guarantee from a block of ordinary steps.

Block residual feedback (general $T$ with sharpness)

Let $Z^\star$ be the nonempty fixed-point set and $d(z):=\operatorname{dist}(z,Z^\star)$ in a fixed inner-product norm. Assume:

  • Firm nonexpansiveness:

    \[\|T(z)-T(w)\|^2+\|r(z)-r(w)\|^2\leq\|z-w\|^2.\]

  • Sharpness: for a uniform $0<\alpha<1$ along the accepted iterates and their ordinary inner steps,

    \[\alpha d(z)\leq\|r(z)\|.\]

Block contraction. With $q:=\sqrt{1-\alpha^2}$, these assumptions give

\[\begin{aligned} d(T(z))^2&\leq d(z)^2-\|r(z)\|^2\leq q^2d(z)^2,\\ \|r(T^m(z))\|&\leq d(T^m(z)) \leq\tfrac{q^m}{\alpha}\|r(z)\|. \end{aligned}\]

Here we also used $|r(z)|\leq d(z)$, another consequence of firm nonexpansiveness. Choose $m$ so that $\rho_m:=q^m/\alpha<1$.

Block update. Choose a compact convex scaling set $\mathcal P$ containing $I$, and initialize $P_0=I$. At each outer iteration:

\[\begin{aligned} b_k&:=T^m(z^k),\qquad s_k:=z^k-b_k,\\ u_k(P)&:=z^k-Ps_k,\\ \ell_k(P)&:=\tfrac{\|r(u_k(P))\|}{\|r(z^k)\|},\\ z^{k+1}&\in\operatorname*{argmin}_{v\in\{u_k(P_k),\,b_k\}}\|r(v)\|,\\ P_{k+1}&:=\Pi_{\mathcal P}(P_k-\eta_kG_k). \end{aligned}\]

Stop when the residual is zero. Take $G_k$ from the trial loss, using a branch derivative at a projection kink, and update the learner even if its trial is rejected. Hold $s_k$ fixed when differentiating.

Linear convergence. Regardless of the learner’s choices,

\[\boxed{ \begin{aligned} \|r(z^K)\|&\leq\rho_m^K\|r(z^0)\|,\\ d(z^K)&\leq\tfrac{\rho_m^K}{\alpha}\|r(z^0)\|. \end{aligned} }\]

Proof. The block comparison gives

\[\|r(z^{k+1})\|\leq\|r(b_k)\|\leq\rho_m\|r(z^k)\|.\]

Multiply over iterations and apply sharpness. Moreover, nonincreasing ordinary residuals give

\[\|s_k\|\leq m\|r(z^k)\|, \qquad \|z^{k+1}-z^k\|\leq Bm\|r(z^k)\|,\]

where $B:=\max_{P\in\mathcal P}|P|$. Summing the geometric tail proves linear convergence to a fixed point.

Choosing $m$. Without knowing $\alpha$, run the ordinary block until $|r(b_k)|\leq\theta|r(z^k)|$ for a fixed $0<\theta<1$. Uniform sharpness bounds the block length; acceptance gives contraction by $\theta$ per block.

Application on PDHG for LP

For the LP above, let $z=(x,y)$ and $T(z)=(x^+,y^+)$. Use the metric

\[M:=\begin{pmatrix}\tau^{-1}I&A^\top\\ A&\sigma^{-1}I\end{pmatrix}, \qquad \|v\|_M:=\sqrt{v^\top Mv}.\]

PDHG is firmly nonexpansive in this metric (Lu and Yang). For an LP with solutions, residual sharpness holds on bounded regions (Proposition 2). Assume the trajectory is bounded, and choose a uniform $\alpha$ valid on a bounded neighborhood containing it. Keep $m$ fixed for the two-stage bound below.

A practical choice is diagonal scaling. For $b_k=(\bar x,\bar y)$,

\[u_k(p)= \begin{pmatrix} x^k-p_x\odot(x^k-\bar x)\\ y^k-p_y\odot(y^k-\bar y) \end{pmatrix}, \qquad p_{k+1}=\operatorname{clip}_{[0,p_{\max}]} \bigl(p_k-\eta_k g_k\bigr),\]

with $p_0=\mathbf1$ and $p_{\max}>1$. The feedback and acceptance rule use the same $M$-norm.

Stage I: active set identification

The block bound gives

\[\|r(z^k)\|_M\leq\rho_m^k\|r(z^0)\|_M, \qquad z^k\longrightarrow z^\star.\]

Assume the limit is strictly complementary:

\[x_i^\star>0 \quad\text{or}\quad c_i-A_i^\top y^\star>0 \qquad\text{for every }i,\]

where $A_i$ is column $i$ of $A$. Every projection input then has a nonzero limit. Since the scalings are bounded and $s_k\to0$, there is a finite $K_0$ after which the iterates, ordinary inner steps, and all allowed trials share one projection pattern. On this neighborhood,

\[T(u)=Hu+h,\qquad r(u)=(I-H)(u-z^\star).\]

Stage II: a faster local rate

A better comparator exists after active-set identification. Let $\mathcal R:=\operatorname{range}(I-H)$ be the residual subspace. It is the $M$-orthogonal complement of the fixed directions. Firm nonexpansiveness and local sharpness imply

\[\|Hv\|_M^2 \leq\|v\|_M^2-\|(I-H)v\|_M^2 \leq q^2\|v\|_M^2, \qquad v\in\mathcal R.\]

Since $r(T^m(z))=H^m r(z)$ in this neighborhood,

\[\ell_k(I)\leq q^m=\alpha\rho_m<\rho_m, \qquad k\geq K_0.\]

Thus the best fixed scaling along the local trajectory satisfies

\[\rho_\star:= \min_{P\in\mathcal P}\sup_{k\geq K_0}\ell_k(P) \leq q^m<\rho_m.\]

After identification,

\[\ell_k(P)= \tfrac{\|(I-H)(z^k-Ps_k)-h\|_M}{\|r(z^k)\|_M}\]

is convex in $P$. Its subgradients are bounded because $|s_k|_M/|r(z^k)|_M\leq m$. Projected online subgradient descent with $\eta_k\propto(k+1)^{-1/2}$ therefore gives

\[\sum_{k=K_0}^{K_0+N-1}\ell_k(P_k) \leq N\rho_\star+R_N, \qquad R_N=O(\sqrt N).\]

The acceptance rule and the arithmetic-geometric mean inequality now give

\[\tfrac{\|r(z^{K_0+N})\|_M}{\|r(z^{K_0})\|_M}\leq\prod_{k=K_0}^{K_0+N-1}\ell_k(P_k)\leq\left(\rho_\star+\tfrac{R_N}{N}\right)^N.\]

Combining the two stages,

\[\boxed{ \begin{aligned} \|r(z^{K_0+N})\|_M &\leq\|r(z^0)\|_M\, \underbrace{\rho_m^{K_0}}_{\text{Stage I}}\\ &\quad{}\times \underbrace{\left(\rho_\star+O(N^{-1/2})\right)^N}_{\text{Stage II}}, \qquad \rho_\star<\rho_m. \end{aligned} }\]

The strict improvement here is over the Stage I bound. Relative to local PDHG, the learner competes with the best allowed fixed scaling; a diagonal family need not give a strict gain.

This gives a theoretical counterpart to my earlier empirical work with Haihao Lu, Enhanced PDHG for Linear Programming with Online Preconditioning. That method learns internal preconditioners using Lagrangian feedback; the result here concerns residual feedback with a block comparison. Trial evaluations cost extra work, so a better iteration rate still needs a runtime evaluation.

References