Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

On-Average Stability of Multipass Preconditioned SGD and Effective Dimension

Simon Vary, Tyler Farghly, Ilja Kuzborskij, Patrick Rebeschini

Abstract

We study trade-offs between the population risk curvature, geometry of the noise, and preconditioning on the generalisation ability of the multipass Preconditioned Stochastic Gradient Descent (PSGD). Many practical optimisation heuristics implicitly navigate this trade-off in different ways -- for instance, some aim to whiten gradient noise, while others aim to align updates with expected loss curvature. When the geometry of the population risk curvature and the geometry of the gradient noise do not match, an aggressive choice that improves one aspect can amplify instability along the other, leading to suboptimal statistical behavior. In this paper we employ on-average algorithmic stability to connect generalisation of PSGD to the effective dimension that depends on these sources of curvature. While existing techniques for on-average stability of SGD are limited to a single pass, as first contribution we develop a new on-average stability analysis for multipass SGD that handles the correlations induced by data reuse. This allows us to derive excess risk bounds that depend on the effective dimension. In particular, we show that an improperly chosen preconditioner can yield suboptimal effective dimension dependence in both optimisation and generalisation. Finally, we complement our upper bounds with matching, instance-dependent lower bounds.

On-Average Stability of Multipass Preconditioned SGD and Effective Dimension

Abstract

We study trade-offs between the population risk curvature, geometry of the noise, and preconditioning on the generalisation ability of the multipass Preconditioned Stochastic Gradient Descent (PSGD). Many practical optimisation heuristics implicitly navigate this trade-off in different ways -- for instance, some aim to whiten gradient noise, while others aim to align updates with expected loss curvature. When the geometry of the population risk curvature and the geometry of the gradient noise do not match, an aggressive choice that improves one aspect can amplify instability along the other, leading to suboptimal statistical behavior. In this paper we employ on-average algorithmic stability to connect generalisation of PSGD to the effective dimension that depends on these sources of curvature. While existing techniques for on-average stability of SGD are limited to a single pass, as first contribution we develop a new on-average stability analysis for multipass SGD that handles the correlations induced by data reuse. This allows us to derive excess risk bounds that depend on the effective dimension. In particular, we show that an improperly chosen preconditioner can yield suboptimal effective dimension dependence in both optimisation and generalisation. Finally, we complement our upper bounds with matching, instance-dependent lower bounds.
Paper Structure (42 sections, 20 theorems, 154 equations, 1 figure)

This paper contains 42 sections, 20 theorems, 154 equations, 1 figure.

Table of Contents

  1. INTRODUCTION
  2. Our contributions
  3. Smooth strongly convex losses.
  4. On-average stabiliy for smooth PL losses.
  5. Lower bounds.
  6. Notation and terminology.
  7. Proof Sketch and Technical Challenges
  8. Generalisation Geometry via On-Average Multipass Stability with Correlated Iterates
  9. Spectral Alignment under Geometric Mismatch
  10. Preliminaries
  11. Relative smoothness & strong convexity.
  12. Generalised co-coercivity.
  13. Excess risk bounds of PSGD via on-average stability
  14. On-average stability and risk bounds for strongly convex smooth losses
  15. Risk bounds for non-convex losses under PL-property
  16. ...and 27 more sections

Key Result

lemma 1

[lemma]lemma:cocoercivity Let $f$ be $\alpha$-strongly convex and $\beta$-smooth w.r.t. $\| \cdot \|_H$ and $P$ is $C_{\ell, P}$-spectrally aligned with $\ell(\cdot,z)$, i.e., $\kappa(PH) < \rho_\ell^2$ in def:relcond_bound. Then for all $x, y \in \mathbb{R}^d$:

Figures (1)

  • Figure 1: Illustration of model misspecification. The geometry of the expected loss curvature $\nabla^2 f$ differs from the geometry of the gradient noise ($\Sigma$). While setting $P \approx \Sigma^{-1}$ whitens the noise, it may result in unstable updates along high-curvature directions.

Theorems & Definitions (41)

  • definition 1: Smoothness w.r.t. $\|\cdot\|_H$
  • definition 2: Strong convexity w.r.t. $\|\cdot\|_H$
  • definition 3: Spectrally aligned preconditioner
  • lemma 1: Co-coercivity of spectrally aligned PSGD updates
  • lemma 2
  • lemma 3: On-average parameter stability of PSGD
  • proposition 1: Risk bounds in geometry defined by $P^{-1}$)
  • remark 1: Approximate NGD under misspecification
  • proposition 2: Risk bounds in geometry defined by $H$
  • proposition 3: Excess risk bounds for PL-losses
  • ...and 31 more