Source-linked AI summary

Automatic differentiation of nonsmooth iterative algorithms

Jérôme Bolte, Edouard Pauwels, Samuel Vaiter

arXiv:2206.00457v1math.OCcs.LG

TL;DR

The paper studies the asymptotics of piggyback automatic differentiation for nonsmooth iterative algorithms and asks whether its limit has derivative meaning. It uses conservative derivatives to characterize a set-valued limiting recursion, establishing almost-everywhere convergence of classical derivatives while identifying a divergence case for momentum methods.

  • Problem

    For nonsmooth recursions, the asymptotics of piggyback derivatives and their interpretation as derivatives of the fixed point are open questions addressed by the paper.

    Source: §1 · p. 2
  • Method

    Under suitable nonexpansivity conditions, the paper uses conservative derivatives to define a set-valued nonsmooth piggyback recursion and characterize its limiting fixed point.

    Source: §1 · p. 2
  • Results

    The limiting set is conservative for the fixed-point map, and for almost all parameters classical piggyback Jacobians converge to the classical Jacobian of the fixed point.

    Source: §1 · p. 2
  • Takeaways & Limitations

    The framework applies to proximal splitting methods in nonsmooth convex optimization, while nonsmooth piggyback derivatives for momentum methods may diverge despite linearly converging iterates.

    Source: §1 · pp. 2–3

    The convergence results rely on a contraction condition, which may hold locally when the iterates converge to a unique fixed point; Heavy-Ball piggyback derivatives may nevertheless diverge.

    Source: Assumption 1 (The conservative… · p. 4, §5.2 · p. 10

Abstract

from arXiv · show

Differentiation along algorithms, i.e., piggyback propagation of derivatives, is now routinely used to differentiate iterative solvers in differentiable programming. Asymptotics is well understood for many smooth problems but the nondifferentiable case is hardly considered. Is there a limiting object for nonsmooth piggyback automatic differentiation (AD)? Does it have any variational meaning and can it be used effectively in machine learning? Is there a connection with classical derivative? All these questions are addressed under appropriate nonexpansivity conditions in the framework of conservative derivatives which has proved useful in understanding nonsmooth AD. For nonsmooth piggyback iterations, we characterize the attractor set of nonsmooth piggyback iterations as a set-valued fixed point which remains in the conservative framework. This has various consequences and in particular almost everywhere convergence of classical derivatives. Our results are illustrated on parametric convex optimization problems with forward-backward, Douglas-Rachford and Alternating Direction of Multiplier algorithms as well as the Heavy-Ball method.

1 Introduction

The paper asks how nonsmooth piggyback derivatives behave asymptotically and whether their limit can serve as a derivative of an algorithm’s fixed point. Under suitable nonexpansivity assumptions, it develops a conservative-derivative framework for these questions and reports both convergence results and a divergence case for momentum methods.

1 Introduction

  • Piggyback differentiation propagates iterate Jacobians by the chain rule; the paper considers full Jacobians, although AD frameworks often compute vector-Jacobian or Jacobian-vector products.The results therefore apply to both reverse and forward modes.Source: §1 · p. 2
  • The paper studies the chain rule’s asymptotics for nonsmooth recursions and how to interpret the limit as a derivative of the fixed point.These questions arise for iterative algorithms used in differentiable programming.Source: §1 · p. 2
  • Set-valued nonsmooth piggyback recursions have a well-defined limit that is conservative for the fixed-point map, under suitable nonexpansivity assumptions.The paper presents this as a nonsmooth infinite chain rule for automatic differentiation.Source: §1 · p. 2
  • For almost all parameters, the classical-Jacobian piggyback recursion is well defined and converges to the classical Jacobian of the fixed point.The result has implications for both forward- and reverse-mode automatic differentiation.Source: §1 · p. 2
  • Nonsmooth piggyback derivatives for momentum methods such as Heavy-Ball may diverge even when the iterates converge linearly.This contrasts with the reported convergence results for proximal splitting methods.Source: §1 · p. 3, Figure 2 · p. 3

2 Nonsmooth piggyback differentiation

The paper extends piggyback differentiation to nonsmooth recursions using conservative Jacobians and a set-valued recursion. Under a contraction condition, the iterates converge linearly to a unique fixed point, and the framework also permits a local version of the condition.

2 Nonsmooth piggyback differentiation

  • Conservative Jacobians provide a nonsmooth extension of piggyback iterations compatible with automatic differentiation.The framework uses set-valued derivatives for path-differentiable functions and their calculus properties.Source: §2 · pp. 3–4

Assumption 1 (The conservative Jacobian of the iteration mapping is a contraction)

  • The standing condition requires a locally Lipschitz, path-differentiable iteration map whose conservative Jacobian’s state block has operator norm at most ρ<1.This makes each parameter-fixed iteration map a strict contraction.Source: Assumption 1 (The conservative… · p. 4
  • Under this condition, the iterates converge linearly to the unique fixed point of the iteration mapping.The paper also allows the operator-norm condition to hold only at the fixed point when the iterates converge to it.Source: Assumption 1 (The conservative… · p. 4
  • The set-valued piggyback recursion forms each next Jacobian as AJ+B using [A, B] from the conservative Jacobian and a current Jacobian J.A recursive selection from this set-valued update matches operations implemented in nonsmooth AD frameworks.Source: Assumption 1 (The conservative… · p. 4

3 Asymptotics of nonsmooth piggyback differentiation

Under contraction conditions, set-valued piggyback iterations converge to a conservative Jacobian for the recursion’s fixed-point map. The results also establish almost-everywhere convergence of classical derivatives, convergence guarantees for JVPs and VJPs, and a linear rate under an additional selection-function condition.

3.1 Fixed point of affine iterations

  • A unique nonempty compact fixed-point set exists for the set-valued affine iteration, which iterates compact matrix sets using contractions in the matrix blocks.The iteration is defined through the Hausdorff-distance framework.Source: §3.1 · pp. 4–5

3.2 An infinite chain rule and its consequences

  • The limiting set-valued piggyback mapping is a conservative Jacobian for the fixed-point map, and the set-valued piggyback iterations converge to it.For almost all parameters, the classical Jacobians of the iterates are well defined and converge to the classical Jacobian of the fixed point; when F is nondifferentiable, the implicit conservative-Jacobian inclusion may be strict.Source: §3.2 · pp. 5–6

3.3 Consequence for algorithmic differentiation

  • For almost all parameters, forward-mode JVPs converge to the fixed-point derivative applied to the input direction; reverse-mode VJPs converge when the terminal adjoints converge.A C1 loss provides one example of convergent terminal adjoints: the gradient evaluated at the iterates.Source: §3.3 · p. 6

3.4 Linear convergence rate for semialgebraic piecewise smooth selection function

  • With a Lipschitz gradient selection structure, piggyback derivatives converge linearly; semialgebraicity is sufficient, but not necessary if the required properties hold by other means.The paper notes that a broader class of definable functions could be considered when those properties are guaranteed.Source: §3.4 · p. 7

4 Application to proximal splitting methods in convex optimization

The section applies the paper’s convergence framework to proximal splitting methods for parametric convex optimization. Under stated regularity and contraction conditions, it establishes when forward–backward, Douglas–Rachford, and ADMM iterations can be differentiated, and illustrates the results on several optimization problems.

4 Application to proximal splitting methods in convex optimization

  • Under Assumption 2, semialgebraic convex objectives and locally Lipschitz gradient and proximal operations provide sufficient conditions for the paper’s convergence results to apply.The section states that these conditions cover a broad range of practical gradient and proximal operations.Source: §4 · p. 7

4.1 Splitting algorithms

  • Forward–backward iterations are strict contractions when either the smooth term or the nonsmooth term is strongly convex, under the stated step-size and regularity conditions.The proposition requires 0 < α < 2 L and Assumption 2.Source: §4.1 · p. 8
  • Douglas–Rachford iterations are strict contractions when f is strongly convex; with a minimizer, the iterates converge to a fixed point yielding an optimization solution.The contraction result is stated under Assumption 2 and α > 0.Source: §4.1 · p. 8
  • ADMM is treated as Douglas–Rachford splitting on the Fenchel dual, and its derivatives can be combined when φθ is strongly convex with Lipschitz gradient and Aθ is injective.Under those conditions, ADMM converges linearly.Source: §4.1 · p. 8

4.2 Numerical illustration.

  • The examples span Ridge, Lasso, sparse inverse covariance selection, and trend filtering, with iterate and derivative linear rates reported as medians and deciles over 100 runs.For Lasso, a generalized-support qualification condition yields the contraction properties needed for the qualitative convergence results.Source: §4.2 · p. 9
  • The numerical illustrations apply the methods to problems including Ridge, Lasso, sparse inverse covariance selection, and trend filtering.For sparse inverse covariance selection, the section notes that f is locally strongly convex, so the convergence results apply qualitatively.Source: §4.2 · p. 9

5 Failure of automatic differentiation for inertial methods

The section contrasts Heavy-Ball’s smooth-case derivative convergence with a nonsmooth counterexample. In that example, Heavy-Ball converges, yet piggyback differentiation can diverge even though the true derivative is zero.

5 Failure of automatic differentiation for inertial methods

  • For a C2 objective, Heavy-Ball derivative propagation converges to the solution derivative, but a C1,1 strongly convex parametric objective can produce diverging derivative vectors.The counterexample’s gradient is path differentiable; implicit differentiation or differentiating gradient descent avoids the stated divergence.Source: §5 · pp. 9–10

5.1 Heavy-ball algorithm and global convergence

  • For strongly convex objectives with Lipschitz gradient, an appropriate step size and momentum bound give Heavy-Ball global linear convergence; with a C2 objective, derivative propagation converges to the solution derivative.The cited convergence conditions include α = 1/L and the stated bound on β.Source: §5.1 · p. 10

5.2 A diverging Jacobian accumulation

  • With α = 1 and β = 3/4, the example’s Heavy-Ball iterates converge globally to 0, while a four-step product of Jacobian matrices has eigenvalue −9/8.The objective is piecewise quadratic, and its gradient is path differentiable.Source: §5.2 · p. 10
  • At θ = 0, alternating two steps on equivalent implementations of the objective makes Heavy-Ball’s differentiated iterates diverge, although the true derivative is 0.The underlying iterate sequence is constant in this example.Source: §5.2 · p. 10

6 Conclusion

The conclusion presents a framework for interpreting piggyback differentiation of nonsmooth recursions and applies it to widely used convex-optimization solvers. It identifies broader recursion classes as a future direction, particularly beyond strict contraction.

6 Conclusion

  • The framework gives a rigorous meaning to differentiating nonsmooth solvers and applies to forward–backward, Douglas–Rachford, and ADMM methods.It is also presented as relevant to abstract recursions for variational inequalities and saddle-point problems.Source: §6 · pp. 10–11
  • Figure 3 contrasts stable derivative propagation for fixed-step gradient descent with unstable propagation for Heavy-Ball, with both methods initialized at the optimum.Source: Figure 3 · p. 10
  • Relaxing Assumption 1 to cover a wider class of methods, including recursions where F is not a strict contraction, is left for future work.Source: §6 · p. 11

A.1 Banach–Picard theorem: Proof of Theorem 1

The appendix proves convergence of the set-valued iteration through contraction in Hausdorff distance and develops stability results for its fixed set. It also establishes regularity of the piggyback fixed-point mapping under the paper’s contraction assumptions.

A.1 Banach–Picard theorem: Proof of Theorem 1

  • When the matrix action has contraction factor ρ < 1, the set-valued iteration is a contraction in Hausdorff distance, and Banach–Picard gives a unique fixed set.The proof also bounds iterates to keep the recursion in a compact set.Source: §A.1 · pp. 15–16

A.2 Properties of the fixed-set mapping

  • The fixed-set mapping is locally Lipschitz continuous on Cρ, and the appendix also states a monotonicity property for fixed sets.The proof bounds the distance from fix(J0) to J(fix(J0)) in terms of the distance between J0 and J.Source: §A.2 · pp. 16–17

A.3 Perturbed iterations

  • Perturbed iterations remain close to the unperturbed fixed set under the stated contraction and perturbation bounds; vanishing perturbations yield convergence to that set.The appendix gives an explicit distance bound and a corollary for perturbations whose magnitudes decay geometrically.Source: §A.3 · pp. 17–18

B.1 Regularity of J pb

  • Under Assumption 1, the piggyback fixed-point mapping is nonempty-valued, locally bounded, and has a closed graph.The assumption bounds the operator norm of the state-matrix component of each conservative Jacobian by ρ < 1.Source: §B.1 · p. 19

B.2 Proof of Theorem 2

The proofs establish conservativity of the limiting set-valued piggyback derivative and connect it to classical derivatives almost everywhere. They also relate the iteration to implicit differentiation and verify conditions for several convergence results.

B.2 Proof of Theorem 2

  • Measurable selections and dominated convergence establish that the limiting set-valued piggyback derivative is conservative for the fixed-point map.The approximating conservative Jacobians converge to the limit in Hausdorff distance.Source: §B.2 · p. 20

B.3 Proof of Corollary 1

  • For almost all parameters, the conservative Jacobian limit is a singleton equal to the classical Jacobian.The proof uses the fact that conservative Jacobians equal classical Jacobians almost everywhere.Source: §B.3 · p. 20

B.4 Proof of Corollary 3

  • A full-measure set makes the approximating Jacobians singletons, so their convergence to the piggyback limit is classical rather than set-valued.The convergence to the limit is in Hausdorff distance.Source: §B.4 · pp. 20–21

B.5 Proof of Proposition 1

  • Forward- and reverse-mode propagation implement the same Jacobian recursion, whose iterates converge to the classical Jacobian of the fixed-point map.The forward-mode recursion follows by setting the propagated tangent to J_i times the parameter direction.Source: §B.5 · pp. 21–22

C Connection with implicit differentiation

  • With classical Jacobians, the piggyback iteration converges to a classical implicit derivative when the iteration map is continuously differentiable.The stated inclusion between the piggyback limit and the implicit-derivative set can be strict.Source: C Connection with implicit… · p. 22

We set

D.1 Lipschitz property of conservative Jacobians of selections

  • For continuous semialgebraic functions with a Lipschitz gradient selection, conservative Jacobians of selections satisfy a local Lipschitz-like bound.The bound uses the Lipschitz constant associated with the selection structure.Source: §D.1 · pp. 22–23

D.2 Proof of Corollary 3

The proof derives the result by combining linear convergence with earlier bounds.

D.2 Proof of Corollary 3

  • The proof combines linear convergence of the iteration with Lemma 3 and Corollary 4 to establish the result.The recursion considered is x_{k+1} = F(x_k, θ).Source: §D.2 · p. 23

E.1 Proof of Proposition 2

The appendix verifies contraction conditions for forward–backward and Douglas–Rachford maps, establishes the ADMM/dual Douglas–Rachford equivalence, and examines Heavy-Ball differentiation. Its example shows that forward derivative propagation can diverge despite stable primal convergence.

E.1 Proof of Proposition 2

  • For forward–backward iterations, strong convexity of either function makes the corresponding map a strict contraction under the stated assumptions.The proof establishes this through the Jacobian factors when f or g is strongly convex.Source: §E.1 · pp. 23–24

E.2 Proof of Proposition 3

  • For Douglas–Rachford, strong convexity of f makes the reflected proximal map a strict contraction, yielding the required contraction condition for the iteration.The proof uses the eigenvalues of the Jacobian factors for the proximal and reflected maps.Source: §E.2 · p. 24

E.3 Equivalence between ADMM and dual Douglas–Rachford

  • The appendix derives the equivalence between ADMM and dual Douglas–Rachford, using a lemma that preserves convexity, lower semicontinuity, and closedness.The text notes that the equivalence is known and includes the derivation for completeness.Source: §E.3 · pp. 24–25

F Inertial methods

  • With a strongly convex objective and Lipschitz gradient, Heavy-Ball converges globally at a linear rate to the unique solution.The iteration map includes an inertial parameter β.Source: F Inertial methods · p. 26

F.1 Failure of Forward differentiation for C1,1 objectives

  • For a C1,1 objective, forward derivative propagation can diverge geometrically even when Heavy-Ball converges linearly and its fixed point is differentiable.The example's fixed point is constant in θ, so implicit differentiation gives derivative 0; the divergence does not occur for C2 objectives.Source: §F.1 · pp. 26–27

G Experiments details

The appendix specifies the experimental protocol and the problem settings, and lists the software assets used.

G Experiments details

  • Experiments were repeated 100 times; plots show the median and first and last deciles, with 2000 iterations run and the first 1000 plotted.All methods were differentiated in forward mode using JAX's jacfwd.Source: G Experiments details · p. 27
  • The experiments cover ridge regression and lasso with forward–backward, sparse inverse covariance selection with Douglas–Rachford, and trend filtering with ADMM.The stated dimensions are n=500, p=300 for ridge; n=50, p=500 for lasso; n=50 for covariance selection; and n=p=75 for trend filtering.Source: G Experiments details · p. 27

H Assets used

  • The listed assets are NumPy [29], Matplotlib, and JAX, each with its stated software license.The licenses listed are BSD-3, PSF, and Apache-2.0, respectively.Source: H Assets used · p. 28
Loading 2206.00457v1…