Stochastic Modified Equations for Stochastic Gradient Descent in Infinite-Dimensional Hilbert Spaces

Abstract: Inverse problems in scientific computing often require optimization over infinite-dimensional Hilbert spaces. A commonly used solver in such settings is stochastic gradient descent (SGD), where gradients are approximated using randomly sampled sub-objective functions. In this work we study the continuous-time limit of SGD in the small step-size regime. We show that the discrete dynamics can be approximated by a stochastic differential equation (SDE) driven by cylindrical Brownian motion. The analysis extends diffusion-approximation results previously established in Euclidean spaces to the infinite-dimensional setting. Two analytical difficulties arise in this extension. First, the cylindrical nature of the noise requires establishing well-posedness of the resulting stochastic evolution equation through appropriate structural conditions on the covariance operator. Second, since the randomness in SGD originates from discrete sampling while the limiting equation is driven by Gaussian noise, the comparison between the two dynamics must be carried out in a weak sense. We therefore introduce a suitable class of smooth functionals on the Hilbert space and prove that the discrepancy between SGD and the limiting SDE, when evaluated through these functionals, is of second order in the step size. Numerical experiments confirm the predicted convergence behavior.

• The paper introduces a stochastic modified equation that approximates SGD in infinite-dimensional Hilbert spaces with a notable second-order weak error bound.
• It rigorously derives the SPDE limit under cylindrical noise conditions and validates the approach through numerical experiments on inverse problems.
• The work establishes well-posedness of the SPDE model and elucidates the practical connection between discrete SGD and continuous-time stochastic dynamics.

Stochastic Modified Equations for Stochastic Gradient Descent in Infinite-Dimensional Hilbert Spaces

Introduction and Motivation

This work addresses a central challenge in the analysis of optimization algorithms for inverse problems and scientific computing, where parameters and unknowns are naturally modeled as elements of infinite-dimensional Hilbert spaces. It provides a rigorous treatment of stochastic gradient descent (SGD) in such infinite-dimensional settings, advancing the theoretical connection between discrete-time stochastic optimization algorithms and their continuous-time stochastic partial differential equation (SPDE) analogs—specifically, the so-called Stochastic Modified Equations (SMEs).

The analysis extends prior work on diffusion approximations of SGD in finite-dimensional Euclidean spaces to the regime of separable, infinite-dimensional Hilbert spaces, which is technically non-trivial due to the subtlety of "cylindrical" noise and the structural properties required for well-posedness of infinite-dimensional SDEs.

Mathematical Framework

The paper considers stochastic optimization problems of the form: $$\min_{\phi \in H}\ \mathcal{F}(\phi) := \mathbb{E}_\gamma\,\mathcal{F}_\gamma(\phi)$$ where $H$ is a real, separable Hilbert space, and $\mathcal{F}_\gamma$ is a family of sub-objective functionals. SGD iterates are therefore: $$\phi_{n+1} = \phi_n - \eta \nabla \mathcal{F}_{\gamma_n}(\phi_n)$$ where the stochasticity originates from i.i.d. sampling of $\gamma_n$.

A key conceptual advance is a weak error analysis for the discrete-time SGD dynamics versus the continuous-time SME, in the small $\eta$ (step-size) regime.

Cylindrical Noise and Covariance Operators

An essential part of the analysis is the infinite-dimensional nature of noise. The SGD noise term, when mapped to the limit, manifests as cylindrical Brownian motion, requiring the specification of a covariance operator with careful control to ensure the corresponding SPDE is well-posed. The noise is only defined as a cylindrical process in $H$ and may not correspond to a Hilbert-Schmidt operator unless trace-class assumptions are imposed, a technical hurdle absent in finite dimensions.

The covariance operator $Q(\phi)$ associated to the stochastic gradient at $\phi$, given by: $$Q(\phi) = \mathbb{E}_\gamma \left[ \left(\nabla \mathcal{F}_\gamma(\phi) - \nabla \mathcal{F}(\phi)\right) \otimes \left(\nabla \mathcal{F}_\gamma(\phi) - \nabla \mathcal{F}(\phi)\right) \right]$$ is central both in the SDE modeling and the validation of the weak error expansion.

Derivation of the SPDE Limit

Under the classical regime of vanishing step size, the discrete SGD dynamics are approximated, in a weak sense, by the solution of the following SPDE: $H$0 where $H$1 is cylindrical Brownian motion and $H$2 is constructed so $H$3. The SME thereby captures the asymptotic "diffusive" effect of stochasticity in SGD.

A technical novelty lies in the explicit construction of the diffusion $H$4 in terms of a mapping from $H$5 (the space of squares integrable functions over the randomness) into $H$6 and the demonstration that it is Hilbert-Schmidt under suitable moment assumptions.

Well-Posedness Results

The authors rigorously establish the existence, uniqueness, and $H$7-boundedness of solutions for the resulting SME using standard results for infinite-dimensional SDEs. The argument requires local Lipschitz continuity and polynomial growth conditions for both the drift and the diffusion, all of which are related back to regularity of the sub-objective gradients and their covariance structure.

Second-Order Weak Convergence Analysis

The core result is a second-order weak error expansion between the law of the SGD iterates and the discrete SME with the same step size $H$8: $H$9 uniformly for $\mathcal{F}_\gamma$0 and any sufficiently smooth test functional $\mathcal{F}_\gamma$1.

This is a noteworthy claim: Classical numerical analysis for the Euler-Maruyama scheme of SDEs in both finite and infinite dimensions indicates only $\mathcal{F}_\gamma$2 weak accuracy for general functionals. However, the SGD and the discretized SME align to higher order, $\mathcal{F}_\gamma$3, in the weak sense—without requiring high-order numerical schemes. The main technical tools are the careful perturbative expansion, the structure of the SPDE, and uniform control on multi-linear transformations (Fréchet derivatives) of functionals on $\mathcal{F}_\gamma$4.

The proof proceeds via a telescoping expansion and discrete-time induction, separating one-step error estimates from the test functional regularity propagation along the stochastic flow. The class of test functionals $\mathcal{F}_\gamma$5 considered is composed of functionals with uniformly bounded derivatives up to third order in the operator topology.

Numerical Validation

Comprehensive numerical experiments are presented in two classes of problems:

• Quadratic (homogeneous) noise models: Error scaling between (i) SGD, (ii) discretized SME, and (iii) exact SME expectations, as predicted, matches $\mathcal{F}_\gamma$6 for weak errors when the test function is a polynomial in the $\mathcal{F}_\gamma$7-norm, confirming the high-fidelity matching of discrete-time and SPDE models.
• State-dependent noise (inverse problems): The analysis is extended to infinite-dimensional inverse problems arising in imaging (e.g., optical tomography), with the SGD noise explicitly dependent on the state. The numerical results consistently show $\mathcal{F}_\gamma$8 convergence, robust to truncation in the number of basis functions, illustrating the dimension-independent nature of the analysis.

Broader Implications and Future Directions

The results have significant theoretical and practical ramifications for the mathematical foundation of SGD and its stochastic approximations in infinite dimensions. Practically, this justifies the use of SME-based continuous-time models, even in high- or infinite-dimensional settings, for understanding the fluctuation and statistical properties of SGD. The higher-order weak accuracy in infinite dimensions indicates that “vanilla” SGD is particularly well-suited for sampling or approximate Bayesian inference over function spaces.

The approach opens the way to more sophisticated analytical tools for infinite-dimensional stochastic optimization methods. It suggests new avenues, such as the extension to momentum and adaptive variants of SGD, ergodic/long-time asymptotics for the limiting SPDE, and the interplay with varying continuous discretization and projection strategies that naturally arise in inverse problems. Relaxation of smoothness and trace-class noise conditions, or the study of irregular functionals, poses further mathematical challenges.

Conclusion

This paper establishes a rigorous stochastic diffusion analysis for SGD in infinite-dimensional Hilbert spaces, demonstrating well-posedness of the SME and a second-order weak error bound between the true SGD iterates and the discretized SME. The methodology and results substantially extend the repertoire of stochastic analysis for optimization in Hilbert and Banach spaces and have direct implications for large-scale inverse problems and function-space learning frameworks.

Reference:

"Stochastic Modified Equations for Stochastic Gradient Descent in Infinite-Dimensional Hilbert Spaces" (2604.10860)