Implicit differentiation

We derive the linear system \(M_{AA}\,d\beta^\star_A/d\alpha = -r\) that sparho.implicit_forward() solves, starting from either of two equivalent fixed points: KKT stationarity of the inner problem on the active set, or the proximal-gradient fixed point. The two viewpoints give the same system; we present the KKT view first (cleaner, no prox-Jacobian) and then sketch the prox view for context.

Setup

The inner problem is

\[ \beta^\star(\alpha) \;=\; \arg\min_{\beta \in \mathbb{R}^p} \; F(\beta;\alpha) \;:=\; L(X\beta, y) \;+\; R(\beta; \alpha), \]

with \(L\) smooth and convex, \(R\) convex but generally non-smooth. We assume:

  1. Inner uniqueness. \(F(\cdot;\alpha)\) has a unique minimizer \(\beta^\star(\alpha)\). This holds for SquaredLoss + L1 with full column rank on the active columns; for LogisticLoss \(L\) is strictly convex.

  2. Strict activity on \(A\). At \(\beta^\star\), the subgradient of \(R\) at every \(j \in A\) is single-valued — i.e. no coordinate sits exactly at a kink (\(\beta^\star_j = 0\) for separable L1-type, or \(\|\beta^\star_{G_k}\| = 0\) for GroupL1). This is the measure-zero genericity assumption discussed in Active sets and why we restrict.

  3. Local \(C^1\) regularity of \(\alpha \mapsto \beta^\star(\alpha)\) on \(A\). Under (1)+(2) the inner problem on \(A\) reduces to a smooth strongly-convex problem and the classical implicit function theorem [Krantz and Parks, 2013] applies. See [Bertrand et al., 2022] for a careful treatment in this exact setting, and [Bolte et al., 2021] for a generalized nonsmooth IFT covering the boundary cases.

Under (1)–(3), \(\beta^\star_I \equiv 0\) on the inactive set \(I = \{0,\dots,p-1\} \setminus A\) in a neighborhood of \(\alpha\) (see Active sets and why we restrict), so \(d\beta^\star_I/d\alpha = 0\) identically and the implicit-diff machinery only needs to track \(A\).

KKT view

On \(A\), the subdifferential \(\partial R\) is single-valued; the inner KKT stationarity condition is

\[ \nabla_{\!A}\, L(X\beta^\star, y) \;+\; s_A(\beta^\star_A; \alpha) \;=\; 0, \qquad s_A(\beta_A;\alpha) := \partial_\beta R(\beta;\alpha)\big|_A. \]

For L1, \(s_A = \alpha \operatorname{sign}(\beta_A)\), locally constant on \(A\) (its derivative in \(\beta_A\) is zero); for WeightedL1 it is \(\alpha_A \odot \operatorname{sign}(\beta_A)\); for ElasticNet \(s_A = \alpha (\rho \operatorname{sign}(\beta_A) + (1-\rho)\beta_A)\); for GroupL1 \(s_{G_k} = \alpha w_k \, \beta_{G_k}/\|\beta_{G_k}\|\) — a block-structured map enumerated explicitly in Penalties: prox, Jacobian, \partial_\beta s, \partial_\alpha s.

Both sides are \(C^1\) in \((\alpha, \beta_A)\) on a neighborhood of the solution (by strict activity), so we differentiate in \(\alpha\):

\[ \underbrace{\Big(\nabla_{\!\beta_A \beta_A}^{\!2}\,L \;+\; \partial_\beta s_A(\beta_A;\alpha)\Big)}_{=:\,M_{AA}} \;\frac{d\beta^\star_A}{d\alpha} \;+\; \underbrace{\partial_\alpha s_A(\beta_A;\alpha)}_{=:\,r} \;=\; 0. \]

Rearranged:

\[ \boxed{\;M_{AA} \, \frac{d\beta^\star_A}{d\alpha} \;=\; -\,r.\;} \]

This is the system sparho.implicit_forward() builds and inverts. The two operators on its left:

  • \(H_{L,AA} = \nabla^2_{\beta_A \beta_A} L(X\beta^\star, y)\) — data-side Hessian. For SquaredLoss, \(H_{L,AA} = \tfrac{1}{n} X_A^\top X_A\); for LogisticLoss, \(H_{L,AA} = X_A^\top \operatorname{diag}(w)\, X_A\) with \(w_i = \sigma(z_i)(1-\sigma(z_i))\), \(z = X\beta^\star\).

  • \(\partial_\beta s_A\) — penalty curvature on \(A\). Zero for L1 / WeightedL1; uniform diagonal \(\alpha(1-\rho) I\) for ElasticNet; block-diagonal \((\alpha w_k / \|\beta_{G_k}\|)\,(I - u_k u_k^\top)\) per active group for GroupL1, with \(u_k = \beta_{G_k} / \|\beta_{G_k}\|\).

The right-hand side \(r = \partial_\alpha s_A\) is what sparho.implicit_forward() calls the “penalty α-Jacobian on the active set” — sign(β_A) for L1, ρ sign(β_A) + (1-ρ)β_A for ElasticNet, etc. See Penalties: prox, Jacobian, \partial_\beta s, \partial_\alpha s for the per-variant list.

Chain rule for \(dC/d\alpha\)

The outer criterion \(C\) enters only through its gradient \(\partial C/\partial \beta\) at \(\beta^\star\). By chain rule (with \(\beta^\star_I \equiv 0\) near \(\alpha\)):

\[ \frac{dC}{d\alpha} \;=\; \Big(\frac{\partial C}{\partial \beta}\Big)^{\!\top} \frac{d\beta^\star}{d\alpha} \;=\; \Big(\frac{\partial C}{\partial \beta_A}\Big)^{\!\top} \frac{d\beta^\star_A}{d\alpha} \;=\; -\,\Big(\frac{\partial C}{\partial \beta_A}\Big)^{\!\top} M_{AA}^{-1}\, r. \]

The standard adjoint trick avoids materializing \(M_{AA}^{-1}\): solve once for \(v = M_{AA}^{-1}\,\partial C/\partial \beta_A\) (symmetric system, \(M_{AA} = M_{AA}^\top\)), then return \(-v^\top r\). This is exactly the structure of sparho.implicit_forward():

v, info = cg(op, grad_C_A, ...)         # M_AA · v = ∂C/∂β_A
return -np.dot(jac_alpha, v)             # = -rᵀ v

where op is the matrix-free operator wrapping \(H_{L,AA}+\partial_\beta s_A\) and jac_alpha is \(r\). For WeightedL1 \(r\) is a vector lifted back to \(\mathbb{R}^p\) rather than scalar-summed.

Proximal-gradient fixed-point view

The same linear system can be derived from the proximal-gradient fixed point

\[ \beta^\star \;=\; \operatorname{prox}_{\gamma R(\cdot;\alpha)} \!\Big(\beta^\star - \gamma\,\nabla L(X\beta^\star, y)\Big), \]

valid for any step \(\gamma \in (0, 2/\!\operatorname{Lip}(\nabla L))\). Differentiating in \(\alpha\) and using

\[ \frac{d\beta^\star}{d\alpha} \;=\; J_z\!\Big(\beta^\star - \gamma\nabla L; \alpha\Big) \Big(\frac{d\beta^\star}{d\alpha} - \gamma H_L \frac{d\beta^\star}{d\alpha}\Big) \;+\; J_\alpha\!\Big(\beta^\star - \gamma\nabla L; \alpha\Big), \]

where \(J_z\) and \(J_\alpha\) are the prox Jacobians w.r.t. its input and \(\alpha\) respectively (see Penalties: prox, Jacobian, \partial_\beta s, \partial_\alpha s and crates/sparho-core/src/prox.rs), one arrives at the same restricted system after using the fact that \(J_z\) is the orthogonal projector onto \(A\) on the active set (and zero on the inactive set, hence killing inactive rows automatically). sparho takes the KKT route because it avoids materializing the prox Jacobian inside hypergrad.py — the Rust kernels in crates/sparho-core/src/prox.rs expose prox_jacobian_* for checking the math in unit tests and for any future hypergradient mode that wants them, but sparho.implicit_forward() itself does not call them.

The two derivations are equivalent and both appear in the sparse-ho line of work; the KKT-restricted version is the practical formulation introduced in [Bertrand et al., 2020] and developed in [Bertrand et al., 2022].

Why \(M_{AA}\) is SPD (generically)

Inside CG we need \(M_{AA}\) to be symmetric positive definite (SPD) so that conjugate gradients converges and the system has a unique solution.

  • Symmetry. \(H_{L,AA}\) is symmetric since \(L\) is twice-differentiable (and is itself a Gram matrix for SquaredLoss / a weighted Gram for LogisticLoss). \(\partial_\beta s_A\) is symmetric for each penalty variant: zero / scalar diagonal for L1 / ElasticNet, and \((I - u_k u_k^\top)\) is symmetric for GroupL1.

  • Positive semidefiniteness. \(H_{L,AA}\) is PSD by convexity of \(L\); \(\partial_\beta s_A\) is PSD because \(R\) is convex and we evaluated on the smooth branch (for L1 / WL1 it is identically zero; for ElasticNet it is a non-negative scalar diagonal; for GroupL1 \(I - u_k u_k^\top\) is the orthogonal projector onto \(u_k^\perp\), PSD).

  • Strict definiteness. Generic active sets give \(X_A\) full column rank (\(|A| \leq n\)), so \(H_{L,AA} \succ 0\) for SquaredLoss; the weighted Gram \(X_A^\top \operatorname{diag}(w) X_A\) is positive definite for Logistic when \(w_i > 0\) (always true) and the same rank condition holds. The penalty curvature is PSD on top — ElasticNet adds a positive scalar shift; GroupL1 adds a PSD block — so \(M_{AA}\) inherits strict positive definiteness from \(H_{L,AA}\).

Ridge stabilization

When \(X_A\) is near-rank-deficient (highly collinear features, common on dense designs with \(|A|\) approaching \(n\)), \(H_{L,AA}\) becomes ill-conditioned and CG either stalls or returns a non-finite iterate. sparho stabilizes by replacing \(M_{AA}\) with \(M_{AA} + \varepsilon I\):

\[ (M_{AA} + \varepsilon I)\,v_\varepsilon \;=\; \partial C/\partial \beta_A. \]

The induced bias in \(v\), and hence in \(dC/d\alpha\), is bounded as follows. Let \(M_{AA}\) have eigendecomposition \(M_{AA} = U \Lambda U^\top\) with \(\Lambda = \operatorname{diag}(\lambda_i)\), \(\lambda_i > 0\). Then

\[ \|v - v_\varepsilon\| \;\leq\; \max_i \frac{\varepsilon}{\lambda_i (\lambda_i + \varepsilon)} \;\|\partial C/\partial \beta_A\|. \]

For directions whose eigenvalue \(\lambda_i \gg \varepsilon\) the bias is \(O(\varepsilon/\lambda_i^2) \cdot \lambda_i = O(\varepsilon/\lambda_i)\) — negligible. For directions where \(\lambda_i \approx \varepsilon\) the ridge effectively replaces the answer with a soft-pseudoinverse, which is the desired behavior: the original problem was undefined there.

sparho auto-scales \(\varepsilon\) to the operator’s natural diagonal magnitude:

\[ \varepsilon \;=\; 10^{-10} \cdot \frac{\operatorname{tr}(M_{AA})}{|A|}, \]

so on well-conditioned problems CG returns results bit-identical to \(\varepsilon = 0\) across eight orders of magnitude (verified at v0.1.0 release on the libsvm Lasso benchmarks). Pass ridge=0.0 to sparho.implicit_forward() to disable; pass an explicit value to override the auto-scaling.

What happens when CG fails

hypergrad.implicit_forward treats scipy.sparse.linalg.cg failure (non-zero info, or non-finite output) as a hard miss: it emits a RuntimeWarning and returns a zero hypergradient for that outer iteration. The outer loop then takes a zero step, the next iteration retries with potentially-different inner state (warm-start drift, HOAG’s tolerance schedule reducing inner tolerance, etc.). This is safer than propagating a NaN — see also the NaN-guards in sparho.grad_search() and sparho.hoag_search.

What this enables

A single CG solve of an \(|A| \times |A|\) SPD system per outer iteration replaces the \(O(p)\)-dimensional fixed-point iterations that a naïve unrolled hypergradient would need. For sparse problems where \(|A| \ll p\) — the regime sparse-ho was designed for — this turns hyperparameter tuning into a tractable bilevel problem, which is the core contribution of the implicit-differentiation line [Bertrand et al., 2020, Bertrand et al., 2022, Pedregosa, 2016].

See Active sets and why we restrict for the assumptions under which the active set is locally constant, Penalties: prox, Jacobian, \partial_\beta s, \partial_\alpha s for the per-variant breakdown of \(\partial_\beta s\) and \(\partial_\alpha s\), and Convergence: HOAG outer loop for the HOAG outer-loop analysis.