Conceptio › Archive › arXiv CS

arXiv CSopen access

Decoupled Descent: Exact Test Error Tracking Via Approximate Message Passing

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

neural-networks

machine learning, deep learning, neural networks

DECOUPLED DESCENT: EXACT TEST ERROR TRACKING VIA APPROXIMATE MESSAGE PASSING

MAX LOVIG

arXiv:2604.27883v1 [math.ST] 30 Apr 2026

Statistics and Data Science, Yale University; [email protected] Abstract. In modern parametric model training, full-batch gradient descent (and its variants) suffers due to progressively stronger biasing towards the exact realization of training data; this drives the systematic “generalization gap”, where the train error becomes an unreliable proxy for test error. Existing approaches either argue this gap is benign through complex analysis or sacrifice data to a validation set. In contrast, we introduce decoupled descent (DD), a novel theory-based training algorithm that satisfies a train-test identity—enforcing the train error to asymptotically track the test error for stylized Gaussian mixture models. Within this specific regime, leveraging approximate message passing theory, DD iteratively cancels the biases due to data reuse, rigorously demonstrating the feasibility of zero-cost validation and 100% data utilization. Moreover, DD is governed by a low-dimensional state evolution recursion, rendering the dynamics of the algorithm transparent and tractable. We validate DD on XOR classification, yielding superior performance compared to GD; additionally, we implement noisy MNIST and non-linear probing of CIFAR-10, demonstrating that even when our stylized assumptions are relaxed, DD narrows the generalization gap compared to GD.

CONTENTS A.1. Notation

1. Introduction

2. The Train-Test Identity And Its Algorithmic Consequences

A.2. A Method Of Descent Under The Train-Test Identity 15

A.3. Deferred Examples

2.1. Correcting Full-batch GD

A.4. Main Assumption

2.2. The Train-Test Identity

A.5. The Asymptotic Test Error

2.3. Natural Algorithmic Consequences

A.6. Full Definition Of State Evolution 17

A.7. Approximate Message Passing And Proving Lemma 3.6

A.8. The Proof Of Theorem 3.8

3. Decoupled Descent 3.1. Main Result: The State Evolution Of Decoupled Descent 4. Algorithm Design 4.1. Pure Decoupled Descent 5. Applications 5.1. Improved Training On The XOR Model

8 8 8 9

A.9. The State Evolution Of Pure DD 30 A.10. Taylor Expansion Of (Pure) DD 31 A.11. The Benefits of Damping For Test Error Critical Points Appendix B. Deferred Applications

36 38

B.1. Faithful Test Error Tracking In Signal-less Learning

B.2. The XOR Training Iterations

5.3. Training A MLP Head On CIFAR10 ResNet Embeddings For Classifying Cats And Dogs 10

B.3. A Run-time Analysis For The MNIST Problem

B.4. Implementation of MNIST Problem With Nine Hidden Layers

5.2. Varying Layer Widths For Classifying Zeros And Eights For MNIST Data

6. Discussion And Conclusion

Acknowledgements

References

B.5. Whitening Technique For CIFAR10 Embedding Vectors 42

Appendix A. Deferred Theory

B.6. Data Collection and Attribution

1. Introduction Consider training data (X, y) = (xi , yi )i∈[n] ∈ Rn×d where xi ∈ Rd and yi ∈ R are i.i.d. (xi , yi ) ∼ Px,y . We model this data using a parametric function Mβ : Rd → R with parameters β ∈ Rp under the loss function L(Mβ (x), y) : R × R → R. We learn β by descending the objective n

1X L(β) = L(Mβ (xi ), yi ), n

(1.1)

i=1

using some variant of gradient descent (GD). The desired goal, however, is not to directly minimize (1.1) but to minimize the test error, i.e. with (x̌, y̌) ∼ Px,y , minimize Ľ(β) = Ex̌,y̌ [L(Mβ (x̌), y̌)] .

(1.2)

In the classical regime of n → ∞ with d, p fixed, L(β) is consistent for Ľ(β). Unfortunately, as seen below, in the modern regime of n, d → ∞ with n ≍ d ≍ p, the convergence L(β) → Ľ(β) fails. Example 1.1. Consider high-dimensional linear regression, for n i.i.d. samples xi ∼ N (0, Idd /d) with labels yi = 0. We fit Mβ (x) = x⊤ β for β ∈ Rd under MSE loss L(ŷ, y) = (ŷ − y)2 = (x⊤ β)2 . Initializing β1 independent of (X, y) and setting M = X ⊤ X, the GD update β2 = (Id − ηX ⊤ X)β1 has train (1.1) and test (1.2) errors, L(β2 ) = β1⊤ (Id − ηM )M (Id − ηM )β1 /n,

Ľ(β2 ) = β1⊤ (Id − ηM )(Id − ηM )β1 /d.

If n → ∞ with d bounded then M/n → Idd /d almost surely so L(β2 ) → Ľ(β2 ); however, when n, d → ∞ with n ≍ d, M no longer concentrates to the identity and thus L(β2 ) − Ľ(β2 ) ̸→ 0. Consequently, in modern scaling regimes, a decrease in the training objective for GD may not represent a decrease in test error. Previous work on GD must either justify that this mismatch is benign or sacrifice a portion of the training data to validate test error performance. In this work, we construct a class of algorithms termed decoupled descent (DD). For a set of stylized high-dimensional problems DD satisfies a train-test identity—any mismatch between (1.1) and (1.2) is forbidden. Consequently, this method allows for validation without sacrificing any training data. Theorem 1.2 (Informal version of Theorem 3.8). For n i.i.d. draws (xi )i∈[n] from a d-dimensional Gaussian mixture with finitely many modes and responses (yi )i∈[n] , any DD iterate βt satisfies Training Error(βt ) − Test Error(βt ) → 0,

almost surely as n, d → ∞ with n/d → α.

As a benefit of Theorem 1.2, DD enables zero-cost validation, meaning the training iterate with the lowest train error also has the lowest test error asymptotically. By eliminating the need for the validation set, users achieve 100% data utilization. Moreover, techniques such as hyperparameter tuning and early stopping on the train error immediately translate performance improvements to the population-level test error. DD originates from a carefully designed approximate message passing (AMP) algorithm [7, 10, 13]. Thus, DD admits a tractable low-dimensional law for the train and test error. This permits a simple physical interpretation for DD’s dynamics and gives rigorous asymptotic guarantees on the behavior of the test error, enabling a principled way to implement descent algorithms. While modern practice uses stochastic gradients, eliminating the full-batch generalization gap is a fundamental first step. To conclude this section, we outline our contributions below. – We introduce a family of training algorithms, termed decoupled descent, for a set of stylized learning problems with Gaussian mixture data, parametric models and general loss functions, covering a wide range of supervised learning tasks.

– In contrast to GD, where the training error is a biased proxy that worsens over time, we show that parameters trained by any DD algorithm enforce the train error to asymptotically equal the test error at each iterate. – We derive a sequence of low-dimensional recursions that track the algorithm’s dynamical trajectory. This provides insight into how DD algorithms explicitly control the test error. – We illustrate the effectiveness of DD using various applications: – XOR Classification: Validates DD under stylized model assumptions and demonstrates better test error performance compared to GD. – Noisy MNIST (0 vs. 8): Demonstrates that DD is robust to different noise distributions and that the train-test identity is consistent across varying ratios of parameter count to data dimension. – CIFAR-10 Probing: Applies DD to (possibly whitened) ResNet-18 embeddings. Although the train-test identity is not exact, DD narrows the generalization gap compared to GD despite the lack of Gaussian structure. We provide an overview of the notation for this paper in Section A.1. 2. The Train-Test Identity And Its Algorithmic Consequences 2.1. Correcting Full-batch GD. Let the i-th row (or element) of X ∈ Rn×d and y ∈ Rn have (xi , yi ) ∼ Px,y where Px is a Gaussian mixture with modes µ1 , . . . , µJ ; conditioned on mode j, we have yi ∼ Pj independent of the Gaussian realization (see (3.1) for more details). Consider parameter ′ β = (θ, a) where θ ∈ Rd×L and a ∈ RL (with L, L′ bounded) alongside a parametric function Mθ,a (x) = Ma (x⊤ θ) with Ma : RL → R. We use the following shorthand. Definition 2.1. Given a parametric model Mθ,a (x) = Ma (x⊤ θ) = Ma (h) and loss function L, let ′ Ψ(h, y, a) = L(Ma (h), y, a) and further define its derivatives, ∇h Ψ(h, y, a) : RL × R × RL → RL , ′ ′ ′ ∇a Ψ(h, y, a) : RL × R × RL → RL , and Hessian ∇2h Ψ(h, y, a) : RL × R × RL → RL×L . Moreover, when h ∈ Rn×L and y ∈ Rn , each of the above functions is its corresponding row-wise application, for example Ψ(h, y, a)i = Ψ(hi , yi , a). We then run full-batch GD, depending on learning rate parameters η, γ > 0, by the iteration ht = Xθt ,

ĥt = ∇h Ψ ht , y, at ,

θt+1 = θt − ηX ⊤ ĥt ,

at+1 = at −

γX ∇a Ψ(ht , yi , at ). n i=1

As explained in the introduction, this algorithm produces iterates θt , at where the training error does not track the test error. We correct this behavior using the following pure decoupled descent iteration, ht = Xθt +η

t−1 X

ĥs ,

ĥt = ∇h Ψ ht , y, at ,

s=1 n 1X θt+1 = θt −η X ⊤ ĥt −α ∇2h Ψ ht,i , yi , at θt , n i=1

n γX at+1 = at − ∇aΨ ht,i , yi , at . n

(2.1)

i=1

P P n 1 2Ψ h , y , a DD adds the correction terms η t−1 ĥ , ηα ∇ applied to the pre-activation t,i i t s=1 s i=1 h n Xθt and gradient of θt , respectively. These terms cancel correlations due to reuse of data matrix X, avoiding the train-test disconnect that GD suffered from in the introduction.

2.2. The Train-Test Identity. Next, we codify the notion of train and test error equality from Theorem 1.2. Consider the algorithm A(X, y) 7→ (h1 , . . . , hT , θ1 , . . . , θT , a1 , . . . , aT ) where, for each ′ t ∈ [T ], (ht , θt , at ) ∈ Rn×L × Rd×L × RL . When A represents training for a parametric model, ht represents the pre-activations of our n samples input into Mat , depending on the parameters at . Then, θt represents the parameter we apply to new test examples at time t of training. The algorithm A satisfies the train-test identity if the following holds: Let (X, y) have i.i.d. rows (xi , yi ) ∼ Px,y specified in Section 2.1, (x̌, y̌) ∼ Px,y and L be a loss function. For all t ∈ [T ], the following limit holds almost surely with respect to the data (X, y), n

1X L(Mat (ht,i ), yi ) − Ex̌,y̌ [L(Mat (x̌⊤ θt ), y̌)] = 0. n,d→∞ n lim

i=1

Notably, the train-test identity does not require the loss to decrease. Assuming the train-test identity holds, Appendix A.2 provides a method to monotonically decrease the test error. 2.3. Natural Algorithmic Consequences. Algorithms satisfying the train-test identity (for example pure DD) can implement a “zero-cost validation” phenomenon. Standard practice saves an ϵ proportion of the data to estimate the test error with the average of the trained model’s loss on the unseen validation set. This introduces a frustrating tradeoff where one wants ϵ to be small to maximize the data involved in training but not too small where the estimator for the test error becomes unreliable. Thus, a zero-cost validation method, i.e. when all data points are used directly in training, is desirable. This suggests the train-test identity as an algorithmic principle for training parameters; some benefits of this principle are given below. P – Early Stopping (Online): Denoting Et = n1 ni=1 L(Mat (ht,i ), yi ), let Dt = Et+1 − ET and D̃t = log(Et+1 ) − log(Et ). Stopping when Dt ≥ ϵ or D̃t ≤ log(1 − ϵ) provides a certificate that the subsequent update will not improve the test error by the specified threshold. In particular, ϵ = 0 guarantees the test error is non-increasing during training. – Early Stopping (Offline): Save a subsequence of iterates T ⊆ [T ] and select the specific parameters (θt∗ , at∗ ) that minimize the train loss Et across the saved time steps. – Initializations/Hyperparameter Tuning: Conduct many parallel runs with varying initializations or hyperparameters (e.g. η, γ) and select the configuration that minimizes the train error, and by the train-test identity, minimizes the test error. – Architecture Search: Optimize model complexity (dimension L or architecture Ma ) by selecting the run with the lowest train error; this equivalently identifies the optimal test error architecture. We expect an online method for architecture selection is also possible. 3. Decoupled Descent Data Model and Parametric Function. For bounded J ∈ N, signal vectors µ1 , . . . , µJ , response laws P1 , . . . , PJ , and class probabilities p1 , . . . , pJ , consider the Gaussian mixture classification model, µj Idd , ⊗ Pj with probability pj . (3.1) (xi , yi ) ∼ N d d We then define the matrix-vector pair X ∈ Rn×d , y ∈ Rn where the i-th row of X and y are i.i.d. draws xi and yi . We consider the limit n, d → ∞ with n/d → α. The data matrix X is commonly seen in mean field analysis of Gaussian mixture models [20]. Specific statistical problems using data (3.1) are provided in Example A.1. Remark 3.1. This work requires that J is bounded, we expect this can be relaxed to a J → ∞ limit after the limit n, d → ∞. This naturally follows by designing finer and finer discrete distribution approximations to a low-dimensional data-generating process. Such a low-dimensional process is

seen in Example A.1 (3) where we could instead consider signal vectors µ ∼ cv with v ∈ Rd and c is endowed some prior distribution supported on a compact interval. Although the finite mixture assumption on the signal is stylized, it establishes a rigorous framework for isolating and correcting data-reuse bias in high-dimensional dynamics. ′

For parameters θ ∈ Rd×L and a ∈ RL , with L, L′ bounded, define Mθ,a (x) : Rd → R with Mθ,a (x) = Ma (x⊤ θ), where Ma : RL → R uses parameter a. To evaluate the performance of this model, we consider a loss function L(ŷ, y) : R × R → R. Specific parametric model and loss pairs following this format are given in Example A.2. Training Protocol. Now, we introduce the general family of DD algorithms. ′

′

Definition 3.2. Let functions g : RL × R × RL → RL , f : RL × R × RL → RL and hyperparameters η0 , η1 , γ0 , γ1 parameterize decoupled descent. When h ∈ Rn×L and y ∈ Rn , we denote g(h, y, a)i = g(hi , yi , a) for each i ∈ [n] as the row wise application of g; consider the iteration, t−1 X (t−1)−s ht = Xθt + η1 η0 ĥs ,

ĥt = g(ht , y, at )

s=1 ⊤

θ̃t = X ĥt − α

! n 1X ∇h g(ht,i , yi , at ) θt n

(3.2)

i=1

θt+1 = η0 θt − η1 θ̃t ,

at+1 = γ0 at − γ1

1X f (ht,i , yi , at ). n i=1

To provide a streamlined analysis of DD algorithms, we make strong assumptions on the above setting. We describe them informally below, see Appendix A.4 for the formal statement. Assumption 3.3 (Informal Version Of Assumption A.3). Assume that: (1) The rows of (X, y) are i.i.d. from distribution (3.1) where n, d → ∞ with n/d → α ∈ (0, ∞). Moreover, we assume that d−1 µ⊤ j µk → χj,k almost surely as n, d → ∞. (2) Initializations θ1, a1 are data-independent with bounded limiting norms and signal alignment. (3) Uniformly over y, the functions f, g (and ∇h f, ∇h g) are Lipschitz and bounded in h and a. (4) Ψ is (C 4 , C 2 ) ∩ Lip in (h, a) with bounded expected derivatives under distribution (3.1). Remark 3.4 (Universality and more general activation functions). As DD is a designed AMP algorithm, we conjecture our results are robust to relaxations of Assumption 3.3. Indeed, AMP is known to exhibit universality and we empirically demonstrate this robustness in Section 5, for example: (1) Previous works allow one to replace the Gaussian noise in (3.1) with mean zero, variance 1/d, independent (but not identical) sub-Gaussian noise [26, 5, 3]. We expect our results to hold under this change in noise. (2) Previous works relax the bounded and Lipschitz requirement for the activation/loss functions f, g and Ψ [19, 23, 6], although truncating a desired function and considering a bounded Lipschitz extension can give a suitable analysis. 3.1. Main Result: The State Evolution Of Decoupled Descent. We begin by providing an asymptotically representation of the test error.

7 ′

Definition 3.5. Given mj,θ ∈ RL for j ∈ [J], Ωθ ∈ RL×L and ā ∈ RL , define test(m1,θ , . . . , mJ,θ , Ωθ , ā) =

J X j=1

pj EŽθ ∼N (0,Ωθ ) [L(Mā (mj,θ + Žθ ), Y̌j )]. Y̌j ∼Pj

−1 ⊤ It is easy to see that for trained parameters θ, a with almost sure limits d−1 µ⊤ j θ → mj,θ , d θ θ → Ωθ , a → ā, we have that limd→∞ Ex̌,y̌ [L(Mθ,a (x̌), y̌)] = test(m1,θ , . . . , mJ,θ , Ωθ , ā) almost surely. This is proven in Appendix A.5.

We now present the main technical lemma of this work, the state evolution of decoupled descent. For simplicity, we defer the full low-dimensional system of recursive equations describing the state evolution parameters to Appendix A.6. ′

Lemma 3.6. If Assumption 3.3 holds and ϕ : RL × R × RL → R, f, g are suitably regular (see Assumption A.7), then almost surely for (ht , at , θt )t∈[T ] from (3.2), n

X 1X ϕ(ht,i , yi , at ) = pj E[ϕ(mj,t + Gt , Yj , āt )], n→∞ n lim

i=1

Gt ∼ N (0, Ωt [t, t]),

Yj ∼ Pj ,

j=1

where Ωt , Σt , Ξt , mj,t , āt are given by the following almost sure limits as n, d → ∞ with n/d → α, 1 ⊤ 1 ⊤ 1 ⊤ 1 ⊤ µj θt → mj,t , θt θt → Ωt [t, t], θt θ̃t → Ξt [t, t], θ̃ θ̃t → αΣt [t, t], at → āt . d d d d t Remark 3.7 (Finite Sample Guarantees). Again, basing our descent algorithm in approximate message passing has the benefits of appealing to prior AMP literature, specifically one can endow the above lemma with finite sample guarantees using works [24, 17, 23, 2]. The proof of Lemma 3.6 follows by a change of variables to traditional AMP algorithms and is deferred to Appendix A.7. As a consequence, we have the following theorem that confirms DD satisfies the train-test identity. Theorem 3.8. Consider the state evolution variables in Definition 3.6 and test from Definition 3.5. If Assumption 3.3 holds and g, f are suitably regular (Assumption A.7), the following holds almost surely for the iterates (ht , at , θt )t∈[T ] from (3.2) and (x̌, y̌) are drawn from distribution (3.1), n

1X lim L(Mat (ht,i ), yi ) = test(m1,t , . . . , mJ,t , Σt [t, t], āt ), n→∞ n i=1

as a consequence, the train-test identity holds for DD, i.e. almost surely, n 1X L(Mat (ht,i ), yi ) − lim Ex̌,y̌ [L(Mat (x̌⊤ θ), y̌)] = 0. lim d→∞ n,d→∞ n i=1

The proof of this statement is deferred to Appendix A.8. We sketch the proof of Theorem 3.8 below. Proof Sketch: We rewrite DD, via a change of variables, as a non-separable matrix-valued AMP algorithm [4, 19]. This establishes that the coordinate-wise average of ϕ applied to ht is governed by a low-dimensional set of recursive equations (Appendix A.6). Choosing ϕ = Ψ, this average −1 ⊤ (with respect to iterate θt ) only depends on the almost sure limits of d−1 µ⊤ j θt and d θt θt . These moments evolve identically as if we trained on the data X (t) = S + Z (t) at time t where Z (t) is a fresh Gaussian matrix and S has rows independently assigned µj /d with probability pj . Due to this behavior, the training error must be the test error since DD essentially evaluates on a fresh independent data set at each update.

Remark 3.9. A dynamical mean field theory (DMFT) analysis of overfitting in full-batch GD is given in [21, Equations (C.7)-(C.10)]. This description contains complex “response function” terms represented by the integration of a two-time correlation matrix which accounts for the reuse of the data matrix. By implementing DD with AMP, however, our algorithm is self correcting in the sense that the response function is zero, admitting a simpler analysis than DMFT methods. This simplicity also allows for better physical interpretation and insight into algorithmic design. 4. Algorithm Design The choice of g, f in algorithm (3.2) controls the state evolution variables, and in turn, the asymptotic test error testt = test(m1,t , . . . , mJ,t , ΩT [t, t], āt ) by Theorem 3.8. This means that given g, f we can calculate testt exactly by simulations. Thus, given candidate function classes G and F, with a budgeted run-time T , the optimal g ∗ , f ∗ is found by running |F||G| total simulations and selecting the run that minimized mint∈[T ] testt . Unfortunately, to run these simulations, we require exact knowledge of the data generating processes (3.1). Therefore, it is desirable to select g, f which produce low test error independent of the data generating process. 4.1. Pure Decoupled Descent. Using (3.2), we recover pure DD from Section 2 by selecting g = ∇h Ψ, f = ∇a Ψ, η0 = γ0 = 1 and η1 = γ1 = η. See (2.1) for the exact iteration and Appendix A.9 for the corresponding state evolution. Traditionally in the analysis of GD, we consider a Taylor expansion on the train error when the learning rate η → 0. Because pure DD satisfies the train-test identity, in contrast to GD, we can directly expand the asymptotic test error. With Gt defined by the state evolution in Appendix A.9 and Yj ∼ Pj from (3.1), let Gt = E[∇a Ψ(mj,t + Gt , Yj , āt )] , Ut = E[∇h Ψ(mj,t + Gt , Yj , āt )] , j∈[J]

j∈[J]

and denote χ = (χj,k )j,k∈[J] and p = (pj )j∈[J] such that (diag(p)χdiag(p))j,k = pj χj,k pk . Theorem 4.1. Let pure DD satisfy the conditions of Assumption 3.3 and Assumption A.7 for sufficiently small η up to bounded time T ∈ N. As η → 0, testt+1 satisfies the following Taylor expansion with state evolution parameters corresponding to the state evolution recursion (A.17), testt+1 = testt − η∥Gt p∥22 − ηα∥Ut (diag(p)χdiag(p))1/2 ∥2F

(4.1)

−

E ηX D pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], Ξt [t, t] + Ξt [t, t]⊤ + ϵt , 2

(4.2)

j=1

where supt∈[T ] |ϵt | ≤ Cη 2 as η → 0 (where C depends on constants in Assumption A.7). Notice the terms in (4.1) are non-positive, meaning that they directly lead to a decrease in the test error. The remaining term in (4.2), however, may be positive. We describe methods to bound the impact of this term in Remark A.18. The proof of Theorem 4.1 is deferred to Appendix A.10 alongside a general Taylor expansion for generic choices of g, f and learning rate parameters. To give insight on the global convergence properties of DD, we describe fixed points of DD algorithms and compare them to critical points of the test error in Appendix A.11. 5. Applications Below we present three applications of decoupled descent, a fourth application on the signal-less regression problem from the introduction is deferred to Appendix B.1. Experiments were conducted on a M4 Mac mini (2024, 10-core CPU, 10-core GPU) with 16GB of unified memory.

Figure 1. Summary statistics for 20 XOR runs (n = d = 1000): GD (left) vs. DD (right) with η = 0.05 and SNR λ = 1 (upper left), λ = 4 (upper right) and λ = 8 (bottom). Blue/red denote train/test error; solid lines are medians, shaded areas are interquartile ranges, and dotted lines show min/max. Low SNR (λ = 1): GD overfits (low train/high test error); DD stabilizes both near log(2), reflecting the non-informative regime. Medium SNR (λ = 4): DD maintains train-test parity, outperforming GD’s overfit solution. High SNR (λ = 8): Signal dominates noise; both algorithms achieve similar performance. 5.1. Improved Training On The XOR Model. We consider a high-dimensional variant of the XOR problem (J = 4, pj = 1/4) with n i.i.d. data points with the following signal vectors dependent on fixed v ∈ Rd/2 , µ1 = [v, v], µ2 = [−v, −v], µ3 = [−v, v], µ4 = [v, −v], where yj = 0 for j ∈ {1, 2} and yj = 1 for j ∈ {3, 4}. We fit model Ma,θ1 ,θ2 (x) = σ(a(x⊤ θ1 )(x⊤ θ2 )), where θ1 ∈ Rd , θ2 ∈ Rd , a ∈ R and σ(x) = 1+e1−x , under cross entropy loss (i.e. L(ŷ, y) = −y log(ŷ) − (1 − y) log(1 − ŷ)). These parameters are initialized at θ11 ∼ N (0, Idd ), θ12 ∼ N (0, Idd ) and a1 ∼ N (0, 1) independently. Both GD and DD iterations for this problem are given in Appendix B.2. Figure 1 plots the train and test errors for GD and DD when each vi = λ with λ ∈ {1, 4, 8} as a signal to noise ratio (SNR) parameter. In all cases, DD maintains the train-test identity. In low SNR regimes (λ ∈ {1, 4}), DD outperforms GD and maintains the train-test identity while GD rapidly overfits; for high SNR regimes, DD matches GD’s performance because, under the λ → ∞ limit, the signal dominates the noise and the effect of memorization is negligible. We also provide an example of zero-cost validation hyperparameter tuning from Section 2. We fix the value of at = 1 in Ma and train θ1 and θ2 with a damped variant of pure DD given in Appendix B.2; this training model incorporates varying η0 and represents weight regularization for the parameters. In Figure 2, we run 20 shared replications of the XOR data and plot the train-test error for each choice of η0 ∈ {1, 0.9, 0.8} (recalling η0 = 1 is pure DD). We can see that our zero-cost validation method is successful, each run demonstrates the train-test identity, and we observe that some weight regularization is beneficial in this case. 5.2. Varying Layer Widths For Classifying Zeros And Eights For MNIST Data. We apply DD to a simple MNIST binary classification test (i.e. zeros vs. eights). For the train data, we sample n = 800 digits (grayscale images of dimension d = 784) from all zeros and eights in the MNIST data

η0 = 1.0

η0 = 0.9

η0 = 0.8

Figure 2. Summary statistics for 50 XOR runs (n = d = 1000) of damped DD (defined in Appendix B.2): η = 0.05, fixed at = 1, λ = 4 and η0 = 1 (left), η0 = 0.9 (middle) and η0 = 0.8 (right). Blue/red colors and line type are equivalent to Figure 1. No Regularization: η0 = 1 (pure DD) gives an initial decrease in the loss but suffers from inflated values of Ωt [t, t] for later iterations (note, the early stopping technique of Section 2 would be helpful in this case). Mild Regularization: η0 = 0.9 finds a good balance in weight regularization which controls variance inflation while achieving good test error and training stability. High Regularization: η0 = 0.8 regularizes too strongly and leads to an estimator that is too conservative, leading to a larger train-test error. set, rescale by SNR λ/d, and inject mean zero, variance 1/d, discrete noise, δ−√2/√d wp 1/4,

δ√2/√d wp 1/4,

δ0 wp 1/2,

(5.1)

(instead of the standard Gaussian noise from distribution (3.1)). The test set is all zero and eight digits in MNIST with identical processing; this experiment verifies if the train-test identity of decoupled descent demonstrates universality phenomena similar to AMP algorithms and is robust when the data is not immediately represented by J mixture modes. Additionally, to see if the train-test identity for decoupled descent degrades as L increases in size, we train a two-layer network (hidden widths L ∈ {3, 9, 27}, tanh activation, sigmoid output) using cross entropy loss and learning rate η = 1. We directly implemented this training procedure in Pytorch using a special MLP class (see Appendix B.4 for details) and thus utilized the default PyTorch initializations for our parameters. Figure 3 plots the train-test error for these three models, we see the train-test identity still holds regardless of the size of the hidden layer, the discrete (instead of Gaussian) noise and a possibly more complex signal distribution. Remark 5.1 (On computational complexity). Recall implementing pure DD required two correction P terms. Term η t−1 but no computational overhead, trivial compared to the s=1 ĥs adds O(n) storageP size of the data. The Hessian term ηα n1 ni=1 ∇2h Ψ(ht,i , yi , at ) requires a per-iteration Jacobian. This Hessian overhead is a negligible, d independent constant—the necessary mechanism for DD to maintain the train-test identity. It is also worth noting that using standard auto-differentiation allows the process of deriving the Hessian automatically, avoiding the process of deriving the correction term on a case by case basis. For completeness, we provide some comparative run-times between GD and pure DD for the L = 9 case in Appendix B.3. 5.3. Training A MLP Head On CIFAR-10 ResNet Embeddings For Classifying Cats And Dogs. We finally consider a modern CIFAR-10 classification task. The goal of this task is to classify images as either pictures of cats (CIFAR-10 class 3) or dogs (CIFAR-10 class 5). We note that the DD algorithms in this section were also natively implemented using Pytorch’s torch.autograd function (See Appendix B.4 for the MNIST example, the CIFAR-10 example is similar).

Figure 3. MNIST zeros vs. eights train/test errors (d = 784, n = 800, λ = 30) over 20 replications: GD (left) vs. DD (right) with discrete noise from distribution (5.1). Blue/red denote train/test error; solid lines are medians, shaded areas are IQRs, dotted lines are min/max. We run a two layer network with hidden layers (L ∈ {3, 9, 27}). The train-test error identity continues to hold as L grows and when the Gaussian noise is replaced with this discrete counterpart. The data matrix for this application is wildly different from Assumption A.3. Instead of noisy images, we use pre-trained ResNet-18 embeddings of dimension d = 512. We sample n = 800 random embedded vectors from the full pool of embedded cat and dog vectors from CIFAR-10 for both the training and testing set independently. Then, we train a two layer MLP classification head with hidden layer width five and tanh activation using both GD and pure DD. The motivation for using DD on these vectors is from Gaussian equivalence theory [12] which postulates one can analyze the intermediate layers of a trained network with Gaussian surrogates. We consider increasingly stronger “whitening” over our training and testing data. In all cases, we standardize the embedding vectors by subtracting their coordinate-wise mean and dividing by the coordinate-wise standard deviation (with respect to only the training distribution). A detailed explanation of the whitening √ procedures are deferred to Appendix B.5, informally we have (1) Vanilla, rescaling features by 1/ d; (2) ZCA (Train), applying whitening computed solely on training data [14]; and (3) Joint ZCA, whitening across the combined train-test pool to prevent covariance drift. These procedures should be interpreted as increasingly stronger processing steps to make the embedded vectors of the train and test set seem more Gaussian. We can see the training and test curves for both GD and pure DD in Figure 4. Although DD has reduced the generalization gap, it is larger than previous examples as we have not injected random noise in the data set and therefore are unlikely to satisfy Assumption 3.3 completely. This suggests a more general method to apply DD to models with non-independent noise structures. 6. Discussion And Conclusion To conclude, we presented a novel machine learning training algorithm and demonstrated it solves the train-test error disconnect present in full-batch GD. We additionally provided some initial guidance on how to design such algorithms and gave empirical validation of their success.

Figure 4. Train/test errors for CIFAR-10 (cats vs. dogs). GD (left) vs. DD (right). Blue/red denote train/test error; solid lines are medians, shaded areas cover the inter-quartile range, dotted lines show min/max, and title is the method of whitening (see Section 5.3). DD reduces overfitting effects compared to GD for training a classification head to ResNet-18 embeddings.

Limitations While DD guarantees an exact asymptotic train-test identity, the practical deployment faces several constraints. The proofs assume Gaussianity on the data X, although some relaxation is expected (Remark 3.4 (1)). Real-world data, however, contains structured correlations that do not satisfy our assumptions. Our results are asymptotic and would benefit from specified finite sample rates (Appendix 3.7). We require full-batch training which rules out stochastic methods like SGD; this additionally introduces DD memory and computational overhead due to the correction terms which can result in ≈ L× longer step times (Appendix B.3). We have only considered models with a wide first layer and finite subsequent layers; over-parameterization of later layers may break the train-test identity, although we do not see this in simulations (Section 5.2). Finally, we do not implement a mechanism in DD to represent signal drift, a common second source of overfitting. Broader Impacts When Assumption 3.3 approximately holds, DD offers three benefits to practitioners: (1) Our zero-validation principle maximizes data utilization, vital in domains where data is scarce or expensive to collect. (2) The number of training runs can be limited by embedding hyperparameter tuning into training dynamics. (3) Exact test error tracking gives a honest assessment of generalization, preventing the deployment of overfitted models. Future Directions We present three future directions. Orthogonally Invariant Noise: Can DD extend to data with dependent elements [29, 8, 18] with symmetry properties? It remains open whether the standard Onsager term’s partial success on CIFAR-10 is a coincidence or a general property of embedding vectors. Stochastic DD: To reduce the full-batch computational overhead, can recent DMFT derivations for SGD [9, 22] inform a mini-batch DD iteration, greatly expanding its practical applicability. Large-Width Networks: DD currently requires finite-width subsequent layers, adapting deep AMP/DMFT frameworks [28, 11, 27] to DD with multiple d-scaling layers would increase the applicability of these results.

Acknowledgements Thanks to Zhou Fan for the many helpful discussions and support on this project. References [1] Z. D. Bai and Y. Q. Yin. Limit of the smallest eigenvalue of a large dimensional sample covariance matrix. The Annals of Probability, 21(3), July 1993. ISSN 0091-1798. DOI: https://doi.org/10.1214/aop/1176989118. URL http://dx.doi.org/10.1214/aop/1176989118. [2] Zhigang Bao, Qiyang Han, and Xiaocong Xu. A leave-one-out approach to approximate message passing. The Annals of Applied Probability, 35(4), August 2025. ISSN 1050-5164. DOI: https://doi.org/10.1214/25-aap2186. URL http://dx.doi.org/10.1214/25-aap2186. [3] Mohsen Bayati, Marc Lelarge, and Andrea Montanari. Universality in polytope phase transitions and message passing algorithms. The Annals of Applied Probability, 25(2), apr 2015. DOI: https://doi.org/10.1214/14-aap1010. [4] Raphaël Berthier, Andrea Montanari, and Phan-Minh Nguyen. State evolution for approximate message passing with non-separable functions. Information and Inference: A Journal of the IMA, 9(1):33–79, January 2019. ISSN 2049-8772. DOI: https://doi.org/10.1093/imaiai/iay021. URL http://dx.doi.org/10.1093/imaiai/iay021. [5] Wei-Kuo Chen and Wai-Kit Lam. Universality of approximate message passing algorithms. Electronic Journal of Probability, 26(none), January 2021. ISSN 1083-6489. DOI: https://doi.org/10.1214/21-ejp604. URL http://dx.doi.org/10.1214/21-ejp604. [6] Yatin Dandi, David Gamarnik, Francisco Pernice, and Lenka Zdeborová. Sequential dynamics in ising spin glasses, 2025. [7] David L. Donoho, Arian Maleki, and Andrea Montanari. Message passing algorithms for compressed sensing: I. motivation and construction. In IEEE Information Theory Workshop 2010 (ITW 2010), page 1–5. IEEE, January 2010. DOI: https://doi.org/10.1109/itwksps.2010.5503193. URL http://dx.doi.org/10.1109/itwksps.2010.5503193. [8] Zhou Fan. Approximate message passing algorithms for rotationally invariant matrices. The Annals of Statistics, 50(1), February 2022. ISSN 0090-5364. DOI: https://doi.org/10.1214/21aos2101. URL http://dx.doi.org/10.1214/21-aos2101. [9] Zhou Fan and Leda Wang. High-dimensional learning dynamics of multi-pass stochastic gradient descent in multi-index models, 2026. [10] Oliver Y. Feng, Ramji Venkataramanan, Cynthia Rush, and Richard J. Samworth. A Unifying Tutorial on Approximate Message Passing. Now Publishers, 2022. ISBN 9781638280057. DOI: https://doi.org/10.1561/9781638280057. URL http://dx.doi.org/10.1561/9781638280057. [11] Cédric Gerbelot and Raphaël Berthier. Graph-based approximate message passing iterations. Information and Inference: A Journal of the IMA, 12(4):2562–2628, September 2023. ISSN 2049-8772. DOI: https://doi.org/10.1093/imaiai/iaad020. URL http://dx.doi.org/10.1093/imaiai/iaad020. [12] Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Marc M’ezard, and Lenka Zdeborov’a. The gaussian equivalence of generative models for learning with shallow neural networks. In Mathematical and Scientific Machine Learning, 2020. URL https://api.semanticscholar.org/CorpusID:235165686.

[13] A. Javanmard and A. Montanari. State evolution for general approximate message passing algorithms, with applications to spatial coupling. Information and Inference, 2(2): 115–144, October 2013. ISSN 2049-8772. DOI: https://doi.org/10.1093/imaiai/iat004. URL http://dx.doi.org/10.1093/imaiai/iat004. [14] Agnan Kessy, Alex Lewin, and Korbinian Strimmer. Optimal whitening and decorrelation. The American Statistician, 72(4):309–314, January 2018. ISSN 1537-2731. DOI: https://doi.org/10.1080/00031305.2016.1277159. URL http://dx.doi.org/10.1080/00031305.2016.1277159. [15] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009. [16] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998. [17] Gen Li and Yuting Wei. A non-asymptotic framework for approximate message passing in spiked models, 2023. [18] Songbin Liu and Junjie Ma. Unifying amp algorithms for rotationally-invariant models, 2024. [19] Max Lovig, Tianhao Wang, and Zhou Fan. On universality of non-separable approximate message passing algorithms, 2025. [20] Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification*. Journal of Statistical Mechanics: Theory and Experiment, 2021(12):124008, December 2021. ISSN 17425468. DOI: https://doi.org/10.1088/1742-5468/ac3a80. URL http://dx.doi.org/10.1088/17425468/ac3a80. [21] Andrea Montanari and Pierfrancesco Urbani. Dynamical decoupling of generalization and overfitting in large two-layer networks, 2025. [22] Sota Nishiyama and Masaaki Imaizumi. High-dimensional limit of stochastic gradient flow via dynamical mean-field theory, 2026. [23] Galen Reeves. Dimension-free bounds for generalized first-order methods via gaussian coupling, 2025. [24] Cynthia Rush and Ramji Venkataramanan. Finite sample analysis of approximate message passing algorithms. IEEE Transactions on Information Theory, 64(11):7264–7286, nov 2018. DOI: https://doi.org/10.1109/tit.2018.2816681. [25] Sophia Sklaviadis, Thomas Moellenhoff, Andre F. T. Martins, Mario A. T. Figueiredo, and Mohammad Emtiyaz Khan. A stein identity for q-gaussians with bounded support, 2026. [26] Tianhao Wang, Xinyi Zhong, and Zhou Fan. Universality of approximate message passing algorithms and tensor networks. The Annals of Applied Probability, 34(4), August 2024. ISSN 1050-5164. DOI: https://doi.org/10.1214/24-aap2056. URL http://dx.doi.org/10.1214/24aap2056. [27] Yizhou Xu, TianQi Hou, ShanSuo Liang, and Marco Mondelli. Approximate message passing for multi-layer estimation in rotationally invariant models. In 2023 IEEE Information Theory Workshop (ITW), page 294–298. IEEE, April 2023. DOI: https://doi.org/10.1109/itw55543.2023.10160238. URL http://dx.doi.org/10.1109/itw55543.2023.10160238.

[28] Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning (ICML), pages 11727–11737. PMLR, 2021. [29] Xinyi Zhong, Tianhao Wang, and Zhou Fan. Approximate message passing for orthogonally invariant ensembles: multivariate non-linearities and spectral initialization. Information and Inference: A Journal of the IMA, 13(3), July 2024. ISSN 2049-8772. DOI: https://doi.org/10.1093/imaiai/iaae024. URL http://dx.doi.org/10.1093/imaiai/iaae024. Appendix A. Deferred Theory A.1. Notation. Let [n] = {1, . . . , n}; R and N denote the real and natural numbers, with (·)k as the k-fold Cartesian product. For M ∈ Rn×L , Mi ∈ RL denotes the i-th row. For a block matrix M , M [r, s] is its (r, s)-th block. vec(·) and mat(·) represent standard vectorization and matricization. ∇k(·) is the k-fold partial derivative. Norms ∥ · ∥2 , ∥ · ∥∞ , ∥ · ∥F , ∥ · ∥op are the Euclidean, infinity, Frobenius, and operator norms. 11{·} is the indicator function and δx is the Dirac distribution at x ∈ R. ⊗, ⊙, ⟨·, ·⟩ denote Kronecker, Hadamard, and inner products (where ⊗ additionally means the independent product of two probability measures). eB ∈ Rn is (eB )i = 11{i ∈ B}. O, o, Θ represent standard asymptotic notation with n-dependent growth. A.2. A Method Of Descent Under The Train-Test Identity. Suppose an algorithm A satisfies the train-test identity and generates the iterate (θt , ht , at ). In addition, there exists a number of (k) (k) (k) updates (θt+1 , ht+1 , at+1 )k∈[K] generated by K different algorithms (A(k) )k∈[K] , each maintaining the train-test identity. Each algorithm A(k) produces an update with training error n 1X (k) L(Ma(k) (ht+1,i ), yi ). (A.1) n t+1 i=1

And, by the train-test identity, (A.1) is an asymptotically consistent estimator for the test error of (k) (k) (k) (k) update (θt+1 , ht+1 , at+1 ), i.e. Ex̌,y̌ [L(Ma(k) (x̌⊤ θt+1 ), y̌)]. Therefore, reducing the test error at step t+1

t + 1 follows by choosing the update that minimizes (A.1), so long as one exists. Other algorithms which do not satisfy the train-test identity can fundamentally never offer such a guarantee due to the train-test disconnect from the introduction. A.3. Deferred Examples. Example A.1. Many common learning problems can be realized by the above data model. (1) The signal-less regression model from the introduction is given with J = 1, µj = 0 and Pj = δ0 (i.e. a Dirac measure at zero). (2) A simple classification problem with flipped class noise has J = 2, µ1 = v and µ2 = −v, and for some ϵ ∈ (0, 1/2)) we set P1 = (1 − ϵ)δ−1 + ϵδ1 and P2 = (1 − ϵ)δ1 + ϵδ−1 with class probabilities p1 = p2 = 1/2. (3) A discrete-valued regression problem has J being some large constant, say 100. We let µj = jv with v ∈ Rd . Then let Pj be the convolution of δc1 j+c2 and Pnoise where c1 , c2 ∈ R and Pnoise is the response’s noise distribution, finally we can let p1 , . . . , pJ be some arbitrary prior over data examples. (4) An XOR classification problem has J = 4, µ1 = [v, v], µ2 = −[v, v], µ3 = [−v, v], µ4 = [v, −v] with v ∈ Rd/2 , Pj = δ0 for j ≤ 2 and Pj = δ1 otherwise. Let each pj = 1/4. Example A.2. Many architectures utilized in practice can be used under our parametric models. Each example below is trained using either MSE loss 21 (ŷ − y)2 or MAD loss |ŷ − y| for simplicity. (1) A simple linear regression model is given by Mθ,a (x) = x⊤ θ.

(2) A generalized linear model is given by Mθ,a (x) = σ(x⊤ θ) for a link function σ : R → R. (3) Let L = L′ , and let σ, ϕ : R → R be applied element-wise, a two-layer network is given by Mθ,a (x) = ϕ(σ(x⊤ θ)a). Deeper networks of finite width are given by setting a = (vec(W1 ), . . . , vec(Wk−1 ), wk ) where Wi ∈ RL×L for i ∈ [k − 1], wk ∈ RL and selecting Mθ,a (x) = ϕ(σ(· · · σ(σ(x⊤ θ)W1 ) · · · Wk−1 )wk ). (4) Corresponding to Example A.1 (4), we could consider a model for the XOR problem with L = 2, θ = (θ1 , θ2 ), L′ = 1 and selecting Mθ,a (x) = a(x⊤ θ1 )(x⊤ θ2 ). A.4. Main Assumption. Assumption A.3. The following hold when both n, d → ∞, with n/d → α ∈ (0, ∞) fixed: (1) [Data Composition] In distribution (3.1), for each j, k ∈ [J], Pj has bounded moments of all n×d × Rn orders, pj ∈ [0, 1] and limd→∞ d−1 µ⊤ j µk exists. Further, each row of data (X, y) ∈ R is drawn i.i.d. from distribution (3.1). (2) [Initialization Limits] Initialization θ1 is independent of the data (X, y) and limd→∞ d−1 θ1⊤ θ1 = θ̄2 almost surely. For each j ∈ [J], limd→∞ d−1 θ1⊤ µj = mj,1 almost surely. The norm of the initialization ∥a1 ∥2 is uniformly bounded. (3) [DD Lipschitzness] Uniformly over y ∈ R, the functions f (h, y, a), g(h, y, a) and their derivatives ∇h f (h, y, a), ∇h g(h, y, a) are Lipschitz and bounded in h and a. (4) [Model And Loss Smoothness] Let Ψ(h, y, a) from Definition 2.1, uniformly over y ∈ R: (a) Ψ is Lipschitz, four times continuously differentiable with respect to h and twice continuously differentiable with respect to a. (b) Let ω ∈ RL×L be a covariance matrix with bounded operator norm, vectors m ∈ RL ′ and a ∈ RL with ∥m∥2 , ∥a∥2 bounded, and let h = m + ω 1/2 G where G ∼ N (0, IdL ). For each class j ∈ [J] with Yj ∼ Pj , there exists a bounded constant C > 0 where, max ∥E[∇2a Ψ(h, Yj , a)]∥op , ∥E[∇h ∇a Ψ(h, Yj , a)]∥op , ∥E[∇2h Ψ(h, Yj , a)]∥op ≤ C max ∥E[∇2h ∇a Ψ(h, Yj , a)]∥F , ∥E[∇3h Ψ(h, Yj , a)]∥F , ∥E[∇4h Ψ(h, Yj , a)]∥F ≤ C A.5. The Asymptotic Test Error. We prove the following proposition. ⊤ Proposition A.4. If Assumption A.3 holds and limd→∞ µ⊤ j θ/d = mj,θ , limd→∞ θ θ/d = Ωθ , limn,d→∞ a = ā almost surely for trained parameters (θ, a), then there exists a deterministic function test depending only on the aforementioned limits such that,

lim Ex̌,y̌ [L(Mθ,a (x̌), y̌)] = test(m1,θ , . . . , mJ,θ , Ωθ , ā).

d→∞

Proof. Consider the test error from (1.2) with β = (θ, a) denoting Ľ(θ, a) = Ex̌,y̌ [L(Mθ,a (x̌), y̌)] where (x̌, y̌) are drawn from distribution (3.1), we write lim Ľ(θ, a) = lim

d→∞

J X

pj E[L(Ma (µ⊤ j θ/d + Žθ ), Y̌j )],

j=1

⊤ where Žθ ∼ N (0, θ⊤ θ/d) and Y̌j ∼ Pj . Under the assumption that µ⊤ j θ/d, θ θ/d and a have almost sure limits mj,θ (for each j ∈ [J]), Ωθ and ā respectively, we have that

lim Ľ(θ, a) =

d→∞

J X j=1

pj EŽθ ∼N (0,Ωθ ) [L(Mā (mj,θ + Žθ ), Y̌j )],

(A.2)

Y̌j ∼Pj

by dominated convergence as Assumption 3.3 assumes Ψ(h, y, a) = L(Ma (h), yj ) is Lipschitz. Noting the right hand side of (A.2) is from Definition 3.5 gives the proof. ■

A.6. Full Definition Of State Evolution. As we see momentarily, DD can be equivalently written as an AMP algorithm, thus off-the-shelf AMP results immediately provide a distributional character′ ization of the pre-activations h1 , . . . , ht ∈ Rn×L and parameters θ1 , . . . , θt ∈ Rd×L , a1 , . . . , at ∈ RL . This description is given by a set of low-dimensional recursive equations termed state evolution. Informally, we derive a recursion of variables Ωt , Σt , Ξt , m1,t , . . . , mJ,t , āt which represent the almost sure limits of (self-)overlaps between θt , µj and θ̃t alongside the limiting value of at . Definition A.5. Let (Pj )j∈[J] , (χj,k )j,k∈[J] , (pj )j∈[J] be from Assumption A.3 (1), (mj,1 )j∈[J] , θ̄2 be the limits from Assumption A.3 (2) and let g, f be the functions from Assumption A.3 (3). Define the following state evolution parameters mj,t ∈ RL for j ∈ [J], t ∈ [T ], block matrices ′ Σt , Ωt ∈ (RL×L )(t+1)×(t+1) , Ξt ∈ (RL×L )t×t and āt ∈ RL for t ∈ [T ] recursively as follows. Recall the indexing [r, s] denotes the (r, s)-th L × L block of a given matrix, we have Σt [r, s] =

J X

pj E[g(Gr + mj,r , Yj , ār )g(Gs + mj,s , Yj , ās )⊤ ]

j=1

ℓj,t = pj E[g(Gt + mj,t , Yj , āt )]

Ξt [r + 1, s] = η0 Ξt−1 [r, s] − η1 αΣt [r, s] + α

J X

χj,k ℓj,r ℓ⊤ k,s

j,k=1

Ωt+1 [r + 1, s + 1] = η02 Ωt [r, s] − η0 η1

Ξt [r, s] + Ξt [s, r]

⊤

+ η12

αΣt [r, s] + α

J X

χj,k ℓj,r ℓ⊤ k,s

j,k=1

mj,t+1 = η0 mj,t − η1 α

J X

χj,k ℓk,t

k=1

āt+1 = γ0 āt − γ1

J X

pj E[f (Gt + mj,t , Yj , āt )]

j=1

with indices r, s ∈ [t] and j ∈ [J]. The expectations h are taken over i Yj ∼ Pj , the random vectors Ω [r,r] Ω [r,s]

Gt ∼ N (0, Ωt [t, t]), and the pairs (Gr , Gs ) ∼ N 0, Ωtt [s,r] Ωtt [s,s] . The system is initialized with P ā1 = a1 , Ωt [1, 1] = Ω1 [1, 1] = θ̄2 , Ξt [1, s] = α Jj=1 mj,1 ℓ⊤ j,s , and Ωt+1 [1, s + 1] = η0 Ωt [1, s] − ⊤ η1 Ξt [1, s] = Ωt+1 [s + 1, 1] .

Remark A.6. As a special case of Definition A.5, we have the recursion,  Ωt+1 [t + 1, t + 1] = η02 Ωt [t, t] − η1 η0 (Ξt [t, t] + Ξt [t, t]⊤ ) + η12 αΣt [t, t] + α2

J X

  χj,k ℓj,t ℓ⊤ k,t .

j,k=1

A.7. Approximate Message Passing And Proving Lemma 3.6. For Lemma 3.6 and Theorem 3.8, we require the following assumption. ′

Assumption A.7. Uniformly over y ∈ R, function ϕ(h, y, a) : RL × R × RL → R is Lipschitz in h and a. Moreover, there exists a bounded constant C > 0, independent of t, where the matrices from −1 Definition A.5 satisfy maxt∈[T ] max(∥Σ−1 t ∥op , ∥Ωt ∥op , ∥ΣT [t, t]∥op , ∥ΩT [t, t]∥op ) ≤ C. A.7.1. Relating Back To AMP. Consider the original DD algorithm dependent on the activation ′ ′ functions g : RL × R × RL → RL and f : RL × R × RL → R and hyperparameters η0 , η1 , γ0 , γ1 from

Equation (3.2), we repeat the algorithm below for convenience below, t−1 X (t−1)−s η0 ĥs ht = Xθt + η1 s=1

ĥt = g(ht , y, at ) θ̃t = X ⊤ ĥt − α

! n 1X ∇h g(ht,i , yi , at ) θt n

(A.3)

i=1

θt+1 = η0 θt − η1 θ̃t n

at+1 = γ0 at − γ1

1X f (ht,i , yi , at ). n i=1

Note, the above algorithm P is invariant to a permutation of the rows of X, thus we fix a representation with X = d1 S + Z = d1 Jj=1 eBj µ⊤ j + Z where B1 , . . . , BJ partition [n] and Z is an element-wise independent Gaussian matrix with Zij ∼ N (0, 1/d). Moreover, by Assumption A.3 (1), we have the almost sure limits limn→∞ |Bj |/n = pj . Therefore, we can equivalently write (A.3) as, ht =

J X j=1

eBj

t−1 X µ⊤ j θt (t−1)−s + Zθt + η1 η0 ĥs d s=1

ĥt = g(ht , y, at ) θ̃t =

J X j=1

⊤ n eBj ĥt + Z ⊤ ĥt − α µj d n

! n 1X ∇h g(ht,i , yi , at ) θt n i=1

θt+1 = η0 θt − η1 θ̃t n

at+1 = γ0 at − γ1

1X f (ht,i , yi , at ). n i=1

We prove Lemma 3.6 in three steps. µ⊤ θt

e⊤ B ĥt

′

⊤ 1×L and ā ∈ RL respectively. (1) We replace jd , dj and at by pre-specified vectors m⊤ t j,t , ℓj,t ∈ R Moreover, with Z̃ ∈ Rn×d , we couple Z = α1/2 Z̃ and replace Z by α1/2 Z̃ noting that Z̃ has independent N (0, 1/n) entries.

We then analyze the following frozen algorithm, recalling that α = n/d, hfrozen = t

J X

1/2 eBj m⊤ Z̃θtfrozen + η1 j,t + α

t−1 X (t−1)−s frozen η0 ĥs s=1

j=1

ĥfrozen = g(hfrozen , y, āt ) t t θ̃tfrozen = α

J X

1/2 ⊤ frozen µj ℓ⊤ Z̃ ĥt −α j,t + α

j=1

! n 1X ∇h g(hfrozen , yi , āt ) θtfrozen t,i n i=1

frozen θt+1 = η0 θtfrozen − η1 θ̃tfrozen .

The state evolution for this algorithm is given in Appendix A.7.2.

(A.4)

19 e⊤ B ĥt

µ⊤ θt

j j ⊤ (2) We prove that if m⊤ and at respectively, j,t , ℓj,t and āt are the almost sure limits of d , d ′ n×(L·T ) n then for a suitable class of test functions, say represented by ϕ : R × R × RL ·T → Rn and ′ φ : Rd×(L·T ) × RL ·T → Rn , we have that both,

i=1

1X 1X ϕ(h1 , . . . , hT , y, a1 , . . . , aT )i − ϕ(hfrozen , . . . , hfrozen , y, ā1 , . . . , āT )i = 0 1 T n→∞ n n lim

and d

i=1

1X 1X lim φ(θ̃1 , . . . , θ̃T , a1 , . . . , aT )i − φ(θ̃1frozen , . . . , θ̃Tfrozen , ā1 , . . . , āT )i = 0, n→∞ d d almost surely. This is the content of Appendix A.7.3. (3) We derive the desired almost sure limits mj,t , ℓj,t and āt in Appendix A.7.4. Combining these three steps proves Lemma 3.6. A.7.2. A Re-parameterized State Evolution. For this subsection, we drop the superscript frozen for notational convenience. We continue with our analysis of (A.4), consider the change of variables   J J X X −1/2  ⊤ 1/2 zt = α ht − eBj mj,t , ht = α zt + eBj m⊤ j,t j=1

 wt = α−1/2 θ̃t − α

j=1

J X

 1/2  µj ℓ⊤ wt + α j,t , θ̃t = α

(A.5)

J X

µ j ℓ⊤ j,t .

j=1

Then, to analyze algorithm (A.4), it suffices to analyze the following algorithm and undo the above change of variables, zt = Z̃θt + α

−1/2

t−1 X (t−1)−s η1 η0 ĥs s=1

 ĥt = g α1/2 zt +

J X

  eBj m⊤ j,t , y, āt

j=1

 wt = Z̃ ⊤ ĥt − α1/2 

θt+1 = η0t θ1 − η1

t X

n 1X

 ∇h g α1/2 wt + α

J X

i=1

j=1





η0t−s α1/2 ws + α

s=1

J X





(A.6)

 , yi , āt  θt µ j ℓ⊤ j,t i

 µ j ℓ⊤ j,s ,

j=1

where the final equality comes from unrolling the recursion of θt+1 = η0 θt − η1 θ̃t and changing variables. Then, we denote ! J X 1/2 ⊤ G(zt ) = g α zt + eBj mj,t , y, āt j=1

Ft (w1 , . . . , wt ) = η0t θ1 − η1

t X s=1

η0t−s

1/2

ws + α

J X j=1

! µj ℓ⊤ j,s

(A.7)

Now, treating y, µ1 , . . . , µJ , eB1 , . . . , eBj , θ1 as fixed vectors, we can equivalently write (A.6) as, zt = Z̃θt −

t X

ĥs B⊤ t,s

s=1

(A.8)

ĥt = G(zt ) wt = Z̃ ⊤ ĥt − θt C⊤ t θt+1 = Ft (wt ), where, denoting the derivative with respect to a specified variable (·) as ∂(·) , d

Bt,s =

X 1X 1 ∂ws,i Ft−1 (w1,i , . . . , wt−1,i ) = α−1 ∂ws,i Ft−1 (w1,i , . . . , wt−1,i ) n d

i=1 (t−1)−s −1/2 = −η1 η0 α Id n X

Ct =

1 n

i=1

∂zt,i G(zt,i ) = α1/2 n−1

i=1

n X

⊤ ∇h g(α1/2 zt,i + m⊤ j(i),t , yi , āt ) ,

i=1

with j(i) mapping a coordinate i ∈ [n] to its corresponding block index of B1 , . . . , BJ . Note, we can immediately check that G is Lipschitz as g is Lipschitz in zt (see Assumption A.3 (3)) where,     J J X X  − g α1/2 z̃t +  ∥G(zt ) − G(z̃t )∥F = g α1/2 zt + eBj m⊤ eBj m⊤ j,t , y, āt j,t , y, āt j=1

≤ L∥α

1/2

j=1

(zt − z̃t )∥F

≤ (Lα1/2 )∥zt − z̃t ∥F , and Ft is clearly Lipschitz in w1 , . . . , wt as it is a linear function of w1 , . . . , wt with bounded coefficients. We can then see that algorithm (A.8) is a full-history matrix-valued extension of the non-separable AMP theory from [4] (or see [19] for a full-history vector-valued AMP algorithm). Using the standard construction of reformatting the vector amp iterates (temporarily overloading the definition of the iterates zt , wt ) (z1 , w1 ), . . . , (zt·L , wt·L ) ∈ R2×n into matrix amp iterates ztmat and wtmat with ztmat = [z(t−1)+1 , . . . , zt·L−1 ],

wtmat = [w(t−1)+1 , . . . , wt·L−1 ],

we can conclude that (A.8) uses the correct Onsager correction terms and is a valid rectangular AMP algorithm. We leave further details of matrix valued AMP iterates to the survey [10, Section 6.7] and [13]. We then have the following low dimensional prescription of the iterates (zt , wt )t∈[T ] from (A.8). Definition A.8. Given the functions G, Ft from (A.7) and deterministic vector θ1 , with initialization α−1 ⊤ z w w z z d θ1 θ1 = Ω1 , we recursively define the sequence of matrices Σ1 , . . . , ΣT , Ω2 , . . . , ΩT +1 with 1 E[G(Zr )⊤ G(Zs )], r, s ∈ [t] n z Ωt [r, r] Ωzt [r, s] (Zr , Zs ) ∼ N 0, z ⊗ Idn Ωt [s, r] Ωzt [s, s] 1 Ωzt+1 [r + 1, s + 1] = E[Fr (W1 , . . . , Wr )⊤ Fs (W1 , . . . , Ws )], r, s ∈ [t] n α−1 = E[Fr (W1 , . . . , Wr )⊤ Fs (W1 , . . . , Ws )] d Σw t [r, s] =

α−1 ⊤ θ E[Fs (W1 , . . . , Ws )] = Ωzt+1 [s + 1, 1]⊤ , d 1 Ωzt+1 [1, 1] = Ωz1

Ωzt+1 [1, s + 1] =

s ∈ [t]

(W1 , . . . , Wt ) ∼ N (0, Σw t ⊗ Idd ) . Lemma A.9 ([19], Theorem 3.3 (extended to matrix-valued AMP algorithms)). Let ϕ1 , ϕ2 : Rn×(L·T ) → Rn , φ1 , φ2 : Rd×(L·T ) → Rd each be uniformly Lipschitz function in n, d. Let (zt , wt ) be the iteration (A.8), and denote (Z1 , . . . , ZT ) ∼ N (0, ΩzT ⊗ Idn ), (W1 , . . . , WT ) ∼ N (0, Σw T ⊗ Idd ) w )−1 ∥ , ∥Ωz ∥ , ∥(Ωz )−1 ∥ ) is bounded almost ∥ , ∥(Σ from Definition A.8. If maxt∈[T ] max(∥Σw op op t op t t op t surely, then the following limits hold almost surely, n

1X 1X ϕ(z1 , . . . , zT )i − lim E[ϕ(Z1 , . . . , ZT )i ] = 0 n→∞ n n→∞ n lim

i=1

lim

d 1X

d→∞ d

i=1

φ(w1 , . . . , wT )i − lim

d 1X

d→∞ d

i=1

E[φ(W1 , . . . , WT )i ] = 0,

i=1

where ϕ(·)i = ϕ1 (·)i ϕ2 (·)i and φ(·)i = φ1 (·)i φ2 (·)i . Definition A.10. Let Assumption A.3 hold with (Pj )j∈[J] from Assumption A.3 (1), θ̄2 from Assumption A.3 (2), and g be from Assumption A.3 (3). Given initialization Ω̌1 = θ̄2 , sequences (mj,t )j∈[J],t∈[T ] , (ℓj,t )j∈[J],t∈[T ] and (āt )t∈[T ] , we recursively define the sequence of matrices Σ̌1 , . . . , Σ̌T , Ω̌2 , . . . , Ω̌T +1 with Σ̌t [r, s] =

J X

pj E[g(Žr + mj,r , Yj , ār )g(Žs + mj,s , Yj , ās )⊤ ],

r, s ∈ [t]

j=1

Ω̌t [r, r] Ω̌t [r, s] (Žr , Žs ) ∼ N 0, Ω̌t [s, r] Ω̌t [s, s] J X 2 ⊤ χj,k ℓj,r ℓk,s , Ξ̌t [r + 1, s] = η0 Ξ̌t−1 [r, s] − η1 αΣ̌t [r, s] + α

r ∈ [t − 1], s ∈ [t]

j,k=1

Ξ̌t [1, s] = α

J X

mj,1 ℓ⊤ j,s ,

s ∈ [t]

j=1

Ω̌t+1 [r + 1, s + 1] = η02 Ω̌t [r, s]−η0 η1

Ξ̌t [r, s]+Ξ̌t [s, r]

⊤

+η12

αΣ̌t [r, s]+α

J X

χj,k ℓj,r ℓ⊤ k,s

j,k=1

r, s ∈ [t] Ω̌t+1 [1, s + 1] = η0 Ω̌t [1, s] − η1 Ξ̌t [1, s] = Ω̌t+1 [s + 1, 1]⊤ ,

s ∈ [t]

Ω̌t+1 [1, 1] = θ̄2 , Corollary A.11. Let (ht , ĥt , θ̃t , θt ) be the iteration from (A.4) and let Assumption A.3 hold. The following limits hold almost surely, for j ∈ [J], r, s, t ∈ [T ], J

X 1 ⊤ 1 1/2 lim µj θt − µ ⊤ W̌t + α µ j ℓ⊤ j E[(α j,t )] = 0 d→∞ d d j=1

lim

1 ⊤ θr θs = Ω̌T [r, s]

d→∞ d

1 ⊤ θr θ̃s = Ξ̌[r, s] d→∞ d 1 lim ĥ⊤ ĥs = Σ̌T [r, s], n→∞ n r lim

where (W̌1 , . . . , W̌T ) ∼ N (0, Σ̌T ⊗ Idd ) from Definition A.10 and Yj ∼ Pj . ′

Moreover, if ϕ̄(h, y, a, j) : RL × R × RL × [J] → R is a Lipschitz function in h uniformly over y, a, j, then the following limit also holds almost surely, n

X 1X lim ϕ̄(ht,i , yi , āt , j(i)) = pj E[ϕ̄(mj,t + Žt , Yj , āt , j)], n,d→∞ n i=1

j∈[J]

where j(i) maps coordinate i to its block Bj with i ∈ Bj and (Ž1 , . . . , ŽT ) ∼ N (0, Ω̌T ) from Definition A.10. Proof. First, we prove the following inductive claim, for all t ∈ [T ], there exists almost sure limits z w Ω̃zt , Σ̃w t of Ωt , Σt as d → ∞ which are recursively defined by the following system initialized at z −1 2 Ω̃1 = α θ̄ (with θ̄2 from Assumption A.3) where Yj ∼ Pj , Σ̃w t [r, s] =

J X

pj E[g(α1/2 Z̃r + mj,r , Yj , ār )g(α1/2 Z̃s + mj,s , Yj , ās )⊤ ],

r, s ∈ [t]

j=1

z Ω̃t [r, r] Ω̃zt [r, s] (Z̃r , Z̃s ) ∼ N 0, z Ω̃t [s, r] Ω̃zt [s, s] J X 2 ⊤ Ξ̃t [r + 1, s] = η0 Ξ̃t−1 [r, s] − η1 αΣ̃w [r, s] + α χ ℓ ℓ j,k j,r k,s , t

r, s ∈ [t]

j,k=1

Ξ̃t [1, s] = α

J X

mj,1 ℓ⊤ j,s ,

(A.9)

s ∈ [t]

j=1 J X 2 ⊤ Ω̃zt+1 [r + 1, s + 1] = Ω̃zt [r, s]−α−1 Ξ̃t [r, s]−α−1 Ξ̃t [s, r]⊤ +α−1 αΣ̃w [r, s]+α χ ℓ ℓ j,k j,r k,s , t j,k=1

r, s ∈ [t] Ω̃zt+1 [1, s + 1] = η0 Ω̃zt [1, s] − α−1 η1 Ξ̃t [1, s] = Ω̃t+1 [s + 1, 1]⊤ ,

s ∈ [t]

Ω̃zt+1 [1, 1] = α−1 θ̄2 . We proceed by induction. For each t ∈ [T ], we first prove the almost sure limit of Ωzt and then the almost sure limit of Σw t , deriving the limits of intermediate variables in Definition A.8 between these θ⊤ θ

steps. The base case follows immediately as limd→∞ Ωz1 = limd→∞ α−1 1d 1 = α−1 θ̄2 by Assumption A.3. Then, with Z1 ∼ N (0, Ωz1 ⊗ Idn ), we write that   ⊤   J J X X 1   1/2   g α1/2 Z1 + E g α Z1 + eBj m⊤ eBj m⊤ Σw 1 = j,1 , y, ā1 j,1 , y, ā1  n j=1

j=1

J X |Bj | 1 X ⊤ ⊤ E[g(α1/2 Z1,i + mj(i),1 , yi , ā1 )g(α1/2 Z1,i + mj(i),1 , y, ā1 )⊤ ]. n |Bj | j=1

i∈Bj

Then, applying the strong law of large numbers on the set of coordinates i ∈ Bj for the random variables yi ∼ Pj , the continuous mapping theorem for the almost sure limit limn→∞ |Bj |/n = pj , and the continuous mapping theorem applied to the variance Z1,i (i.e. Ωz1 with almost sure limit Ω̃z1 ) using that g is a continuous function, gives that the almost sure limit of Σw 1 is Σ̃w 1 =

J X

pj E[g(α1/2 Z̃1 + mj,1 , Yj , ā1 )g(α1/2 Z̃1 + mj,1 , Yj , ā1 )⊤ ],

j=1

where Z̃1 ∼ N (0, Ω̃z1 ). Now, we assume the inductive claim holds up to time t − 1, i.e. the matrices Ω̃zt−1 and Σ̃w t−1 are the almost sure limits of the matrices Ωzt−1 and Σw respectively. t−1 Recall the choice of functions G, Ft from (A.7), we can then rewrite Ωzt , with r, s ∈ [t − 1], as α Ωzt [r + 1, s + 1] =

−1

η0r θ1 −η1

!!⊤ !!# r J s J X X X X ⊤ s ⊤ r−p 1/2 s−q 1/2 η0 θ1 −η1 η0 α Wp +α µj ℓj,p η0 α Wq +α µk ℓk,q p=1

α−1 = E d

η0r−1 θ1 −η1

η0

r−1 X

(r−1)−p η0

Wp +α

µj ℓ⊤ j,p

−η1 α

1/2

Wr +α

j=1

J X

!)⊤

µj ℓ⊤ j,r

j=1

η0s−1 θ1 −η1

k=1

J X

}

Ar−1

(

}

!)# !! s−1 J J X X X (s−1)−q ⊤ 1/2 1/2 ⊤ −η1 α Ws +α µk ℓk,s η0 α Wq +α µk ℓk,q q=1

k=1

| = α−1

1/2

p=1

| × η0

q=1

j=1

}

As−1

}

1 ⊤ ⊤ 2 ⊤ E[η02 A⊤ r−1 As−1 − η0 η1 (Ar−1 Bs + Br As−1 ) + η1 Br Bs ] d

where (W1 , . . . , Wt−1 ) ∼ N (0, Σw t−1 ⊗ Idd ). Moreover, we have the following simplifications, α−1 z E[A⊤ (As Ar−1 is Fr−1 (W1 , . . . , Wr ) and similarly for As−1 ) r−1 As−1 ] = Ωt−1 [r, s], d  ⊤  r−1 X (r−1)−p 1 1  r−1  E[A⊤ η0 Bp  Bs  , r−1 Bs ] = E  η0 θ1 − η1 d d p=1





 s−1 X 1 1 (s−1)−q E[Br⊤ As−1 ] = E Br⊤ η0s−1 θ1 − η1 η0 Bq  , d d q=1   J J J X X X 1 1 3/2 2 ⊤ E[Br⊤ Bs ] = αE[Wr⊤ Ws ]+α3/2 E[Wr ]⊤ µk ℓ⊤ ℓj,r µ⊤ ℓj,r µ⊤ j E[Ws ]+α j µk ℓk,s k,r +α d d j=1

k=1

=α

X µ⊤ E[Wr⊤ Ws ] j µk ⊤ + α2 ℓj,r ℓ d d k,s j,k=1

2 = αΣw t−1 [r, s] + α

J X j,k=1

ℓj,r

µ⊤ j µk ⊤ ℓ d k,s

j,k=1

as E[Wr ] = 0. This in turn implies that, Ωzt [r + 1, s + 1] = η02 Ωzt−1 [r, s] − α−1 η0 η1 (Ξzt−1 [r, s] + Ξzt−1 [s, r]⊤ )   J ⊤ X µj µk ⊤ 2 + α−1 η12 αΣw ℓj,r ℓ , t−1 [r, s] + α d k,s

(A.10)

j,k=1

where, with initialization Ξzt [1, s] = d1 θ1⊤ E[Bs ] and r ∈ [t − 2], s ∈ [t − 1], we recursively define,  ⊤  r X r−p 1   Ξzt−1 [r + 1, s] = E η0 θ1 − η1 ηθ,0 Bp  Bs  d p=1

   ⊤  r−1 X (r−1)−p 1   = E η0 η0 θ1 − η1 ηθ,0 Bp  − η1 Br  Bs  d p=1

i 1 h = η0 Ξzt−2 [r, s] − η1 E Br⊤ Bs d

J X

2 = η0 Ξzt−2 [r, s] − η1 αΣw t−1 [r, s] + α

j,k=1

 µ⊤ µ j k ⊤  ℓ . ℓj,r d k,s

w By the inductive hypothesis, we have the covariance matrix Σw t−1 has the almost sure limit Σ̃t−1 and the covariance matrix Ωzt−1 has the almost sure limit Ω̃zt−1 . Thus, under equation (A.10), we apply the strong law of large numbers on the coordinates i ∈ Bj for the independent random variables yi ∼ Pj , the continuous mapping theorem for the almost sure limit limn→∞ |Bj |/n = pj , limit µ⊤ µk

θ⊤ µ

limd→∞ jd = χj,k from Assumption A.3 (1) and limit limd→∞ 1d j = mj,1 from Assumption A.3 (2) to give that Ωzt has the following almost sure limit, with r, s ∈ [t − 1], Ω̃zt [r + 1, s + 1] = η02 Ω̃zt−1 [r, s] − α−1 η0 η1 (Ξ̃t−1 [r, s] + Ξ̃t−1 [s, r]⊤ )   J X 2  χj,k ℓj,r ℓ⊤ + α−1 η12 αΣ̃w t−1 [r, s] + α k,s , j,k=1

where we recursively define, with r ∈ [t − 2] and s ∈ [t − 1],  2 Ξ̃t−1 [r + 1, s] = η0 Ξ̃t−2 [r, s] − η1 αΣ̃w t−1 [r, s] + α

J X

  χj,k ℓj,r ℓ⊤ k,s ,

j,k=1

with Ξ̃t−1 [1, s] = α

⊤ j=1 mj,1 ℓj,s .

Moreover, we can write "  s−1 −1 X α (s−1)−q Ωzt [1, s + 1] = θ1⊤ E η0 η0s−1 θ1 − η1 η0 d q=1 !# J X 1/2 ⊤ − η 1 α Ws + α µk ℓk,s k=1

= η0 Ωzt−1 [1, s] − α−1 η1 Ξzt−1 [1, s]

α1/2 Wq + α

J X k=1

!  µk ℓ⊤ k,q

as E[Wq ] = 0, which can be argued identically as above to have the almost sure limit, Ω̃zt [1, s + 1] = η0 Ω̃zt−1 [1, s] − α−1 η1 Ξ̃t−1 [1, s]. As a final step, we can use the aforementioned argument on the convergence of Ωz1 to give that Ωzt [1, 1] has the almost sure limit Ω̃zt [1, 1] = Ω̃z1 [1, 1]. Then, we can write Σw t as,      J J X X 1  g α1/2 Zs +  Σw E g α1/2 Zr + eBj m⊤ eBj m⊤ t [r, s] = j,r , y, ār j,s , y, ās n j=1

J X j=1

j=1

|Bj | 1 X E[g(α1/2 Zr,i + mj(i),r , yi , ār )⊤ g(α1/2 Zs,i + mj(i),s , yi , ās )⊤ ] n |Bj | i∈Bj

where (Z1 , . . . , Zt ) ∼ N (0, Ωzt ⊗ Idn ). An application of the strong law of large numbers over the block Bj where each yi ∼ Pj independently and the continuous mapping theorem on the covariance of Zr,i and Zs,i , using that g is a continuous function, then gives that Σw t has the almost sure limit, Σ̃w t [r, s] =

J X

pj E[g(α1/2 Z̃r + mj,r , Yj , ār )g(α1/2 Z̃s + mj,s , Yj , ās )⊤ ],

j=1

z Ω̃ [r, r] Ω̃zt [r, s] where (Z̃r , Z̃s ) ∼ N 0, zt . This proves the inductive claim. Ω̃t [s, r] Ω̃zt [s, s] Recalling system (A.9), we can relate this system back to Definition A.10 by considering the change of variables Žt = α1/2 Z̃t , Ω̌t+1 = αΩ̃zt+1 , Ξ̌t = Ξ̃t and Σ̌t = Σ̃w t . Recall the reparameterization (A.5), where, for any two functions ϕ, φ given in Lemma A.9, we can write   n n J J X X X X 1 1 1/2  , ϕ(h1 , . . . , hT )i = ϕ α1/2 z1 + eBj m⊤ zT + eBj m⊤ j,1 , . . . , α j,T n n i=1

i=1

or d

i=1

j=1



j=1



X X 1 X  1/2 1X 1/2  . φ(θ̃1 , . . . , θ̃T ) = φ α w1 + α µ j ℓ⊤ wT + α µ j ℓ⊤ j,1 , . . . , α j,T d d i

Notice, if both ϕ1 , ϕ2 and φ1 , φ2 are L-Lipschitz in their first argument then the following functions are Lα1/2 Lipschitz with respect to either z1 , . . . , zT or w1 , . . . , wT , for both i ∈ {1, 2},   J J X X 1/2  (z1 , . . . , zT ) 7→ ϕi α1/2 z1 + eBj m⊤ zT + eBj m⊤ j,1 , . . . , α j,T j=1

j=1

J X

 (w1 , . . . , wT ) 7→ φi α1/2 w1 + α

1/2 µ j ℓ⊤ wT + α j,1 , . . . , α

j=1

 . µ j ℓ⊤ j,T

j=1

Therefore, by Lemma A.9 and the aforementioned change of variables between system (A.9) and Definition A.10, we have almost surely that   n n J J X X 1X 1X   = 0, (A.11) lim ϕ(h1 , . . . , hT )i − lim ϕ Ž1 + eBj m⊤ eBj m⊤ j,1 , . . . , ŽT + j,T n→∞ n n→∞ n i=1

i=1

j=1

or   d d J J X X 1X 1 X  1/2 1/2  = 0, lim φ(θ̃1 , . . . , θ̃T )i − lim φ α W̌1 +α µj ℓ⊤ WT +α µj ℓ⊤ j,1 , . . . , α j,T d→∞ d d→∞ d i=1

i=1

j=1

(A.12) where (Ž1 , . . . , ŽT ) ∼ N (0, Ω̌T ⊗ Idn ) and (W̌1 , . . . , W̌T ) ∼ N (0, Σ̌T ⊗ Idd ). Each of the limits in the statement of the lemma follow immediately from the following choice of ϕ and φ, using equations (A.11) or (A.12) and using the identicality of the row distribution of Ž1 , . . . , ŽT and W̌1 , . . . , W̌T (recalling that φ(·)i = φ1 (·)i φ2 (·)i and similarly for ϕ): 1 ⊤ µ θ̃t , φ(θ̃1 , . . . , θ̃T )i = µj,i θ̃t,i d j 

  s r X X 1 ⊤ η0s−q θ̃q,i  η0r−p θ̃p,i  η0s θ1 − η1 θ θs , φ(θ̃1 , . . . , θ̃T )i = η0r θ1 − η1 d r q=1 p=1   r X 1 ⊤ θr θ̃s , φ(θ̃1 , . . . , θ̃T )i = η0r θ1 − η1 η0r−p θ̃p,i  θ̃s d p=1

1 ⊤ ĥ ĥs , ϕ(h1 , . . . , hT )i = g(hr,i , yi , ār )g(hs,i , yi , ās ) (A.13) n r and, for the first three limits, invoking the law of large numbers on row-wise independent laws of W̌1 , . . . , W̌t , Ž1 , . . . , Žt to relate back to Definition A.10. Note, (A.13), and more generally the P P |B | P statement on n1 ni=1 ϕ̄(ht,i , yi , āt , j) = Jj=1 nj i∈Bj ϕ̄(ht , y, āt , j)i (where ϕ̄ is applied row-wise) in the statement of this corollary follows by selecting (ϕ1 )i = ϕ̄(·, yi , āt , j) (or g(·, yi , āt )) for each i ∈ [n] and ϕ2 being the all ones vector (or g(·, yi , āt )) and using equation (A.11). Then an application of the continuous mapping theorem for the almost sure limit limn→∞ |Bj |/n, the identicality of the row distribution of Ž1 , . . . , ŽT , alongside the strong law of large numbers applied to the block Bj where each yi ∼ Pj independently proves the claim. ■ We conclude this section by observing that Definition A.10, with the exception of the updates on the frozen state evolution parameters, is equivalent Definition A.5. ⊤

A.7.3. Convergence of Lower-order Recursions Of mj,t , ℓj,t And āt . Let m̃⊤ j,t = consider the following algorithm equivalent to (A.3), ht =

J X

eBj m̃⊤ j,t + Zθt + η1

t−1 X (t−1)−s η0 ĥs s=1

j=1

ĥt = g(ht , y, at ) θ̃t =

J X

⊤ αµj ℓ̃⊤ j,t + Z ĥt − α

j=1

! n 1X ∇h g(ht,i , yi , at ) θt n i=1

θt+1 = η0 θt − η1 θ̃t n

1X at+1 = γ0 at − γ1 f (ht,i , yi , at ). n i=1

We prove the following.

e ĥt µ⊤ j θt ⊤ = Bj , ℓ̃ and j,t d d

Lemma A.12. Assume that almost surely the following limits exist, limd→∞ m̃j,t = mj,t ∈ RL , ′ limd→∞ ℓ̃j,t = ℓj,t ∈ RL and limd→∞ at = āt . If ϕ1 , ϕ2 : Rn×(L·T ) × Rn × RL ·T → Rn , φ1 , φ2 : ′ Rd×(L·T ) × RL ·T → Rn are Lipschitz in their first argument (and third argument for ϕ1 , ϕ2 , both uniformly over the second argument y ∈ Rn ), then we have that almost surely, n

i=1

1X 1X ϕ(h1 , . . . , hT , y, a1 , . . . , aT )i − ϕ(hfrozen , . . . , hfrozen , y, ā1 , . . . , āT )i = 0 1 T n→∞ n n lim

and d

i=1

1X 1X lim φ(θ̃1 , . . . , θ̃T , a1 , . . . , aT )i − ϕ(θ̃1frozen , . . . , θ̃Tfrozen , ā1 , . . . , āT )i = 0, n→∞ d n where ϕ(·)i = ϕ1 (·)i ϕ2 (·)i and φ(·)i = φ1 (·)i φ2 (·)i . Proof. We proceed by induction on t ∈ [T ] to prove the iterates of the original and frozen algorithm have a vanishing difference in two-norm almost surely as n, d → ∞ with n/d → α. Specifically, we prove the following limits hold almost surely, 1 1 lim ∥ht − hfrozen ∥22 = 0, lim ∥ĥt − ĥfrozen ∥22 = 0, (A.14) t t n→∞ n n→∞ n 1 1 frozen 2 ∥2 = 0. (A.15) lim ∥θ̃t − θ̃tfrozen ∥22 = 0, lim ∥θt+1 − θt+1 d→∞ d d→∞ d Base Case (t = 1): By initialization, θ1 = θ1frozen . Evaluating the pre-activations gives,     J J X X frozen  − h1 − hfrozen = eBj m̃⊤ eBj m⊤ 1 j,1 + Zθ1 j,1 + Zθ1 j=1

J X

j=1

eBj (m̃j,1 − mj,1 )⊤ .

j=1

Since limd→∞ m̃j,1 = mj,1 almost surely, and that limn→∞ ∥eBj ∥22 /n = pj ∈ [0, 1] from Assumption A.3 (1) (with J being n-independent), we have limn→∞ n1 ∥h1 − hfrozen ∥22 = 0 almost surely. Thus, as 1 the function g is Lipschitz with respect to input h1 , a1 (for some bounded constant C > 0), we have that 1 1 ∥ĥ1 − ĥfrozen ∥22 = ∥g(h1 , y, a1 ) − g(hfrozen , y, ā1 )∥22 1 1 n n 1 frozen 2 2 ∥h1 − h1 ∥2 + ∥a1 − ā1 ∥2 . ≤C n Therefore, as an immediate consequence of limn→∞ n1 ∥h1 − hfrozen ∥22 = 0 and limd→∞ a1 = ā1 , we 1 1 frozen 2 have limn→∞ n ∥ĥ1 − ĥ1 ∥2 = 0 almost surely. Inductive Step: Assume equations (A.14) and (A.15) hold for all s ≤ t. We first consider the update for θ̃t : θ̃t − θ̃tfrozen = α

J X

µj (ℓ̃j,t − ℓj,t )⊤ + Z ⊤ (ĥt − ĥfrozen ) t

j=1

D E frozen − α ⟨∇h g(ht , y, at )⟩ θt − ∇h g(hfrozen , y, ā ) θ , t t t where ⟨·⟩ is a shorthand for the empirical average over n coordinates. We bound the two-norm of each of these three terms separately:

(1) Signal Term: Since limd→∞ ℓ̃j,t = ℓj,t and d1 ∥µj ∥22 = O(1) by Assumption A.3 (1), we P immediately have limd→∞ d1 ∥α Jj=1 µj (ℓ̃j,t − ℓj,t )⊤ ∥22 = 0 almost surely. (2) Noise Term: Since Z ∈ Rn×d is a matrix with independent N (0, 1/d) entries, the operator norm ∥Z∥op is almost surely bounded [1] as n, d → ∞. Thus, almost surely, ∥Z∥2op 1 ⊤ 1 frozen 2 lim ∥Z (ĥt − ĥt )∥2 ≤ lim ∥ĥt − ĥfrozen ∥22 = lim α∥Z∥2op ∥ĥt − ĥfrozen ∥22 = 0, t t d→∞ d d→∞ d→∞ d n where the final equality applied the inductive hypothesis on ĥt . , y, āt )⟩θt , we have by the triangle (3) Onsager Term: Adding and subtracting ⟨∇h g(hfrozen t inequality and Cauchy-Schwarz, 2 α2 ⟨∇h g(ht , y, at )⟩θt − ⟨∇h g(hfrozen , y, āt )⟩θtfrozen t d 2 2 ∥θt ∥2 2 ≤ 2α2 ⟨∇h g(ht , y, at )⟩ − ⟨∇h g(hfrozen , y, āt )⟩ t d op 2 ∥θt − θ frozen ∥2 t 2 , y, ā )⟩ . + 2α2 ⟨∇h g(hfrozen t t d op By Assumption A.3 (3), ∇h g is Lipschitz in both h and a. Thus, the operator norm difference 2 ⟨∇h g(ht , y, at )⟩ − ⟨∇h g(hfrozen , y, āt )⟩ op vanishes almost surely by the inductive hypothesis t on ht and the almost sure convergence of limn→∞ at = āt . Since d1 ∥θt ∥22 = O(1) almost surely (by the inductive hypothesis and Corollary A.11), the first term vanishes. The second term vanishes because d1 ∥θt − θtfrozen ∥22 → 0 by the inductive hypothesis, and the empirical average of ∇h g is bounded.

Therefore, d1 ∥θ̃t − θ̃tfrozen ∥22 → 0. frozen ∥2 → 0. By the linear update θt+1 = η0 θt − η1 θ̃t , we trivially obtain d1 ∥θt+1 − θt+1 2

Similarly, evaluating the update for ht+1 , ht+1 − hfrozen = t+1

J X

frozen ) eBj (m̃j,t+1 − mj,t+1 )⊤ + Z(θt+1 − θt+1

j=1

+ η1

t X

). η0t−s (ĥs − ĥfrozen s

s=1

Applying the assumed limits for m̃j,t+1 , the bounded operator norm of Z, and the inductive hypothesis 2 on the past ĥs iterates, we find limn→∞ n1 ∥ht+1 − hfrozen t+1 ∥2 = 0 almost surely. Convergence of Test Functions: We now prove that almost sure convergence of the iterates implies the almost sure convergence of the test functions. Consider ϕ(·)i = ϕ1 (·)i ϕ2 (·)i ′ where ϕ1 , ϕ2 : Rn×(L·T ) × Rn × RL ·T → Rn are Lipschitz functions satisfying the statement of the Lemma. For notational simplicity, we abbreviate ϕ(h1 , . . . , hT , y, a1 , . . . , aT ) as ϕ(h) and . ϕ(hfrozen , . . . , hfrozen , y, ā1 , . . . , āT ) as ϕ(hfrozen ). 1 T We then have that, n

1 X 1 (ϕ(h)i − ϕ(hfrozen )i ) = ϕ1 (h)⊤ ϕ2 (h) − ϕ1 (hfrozen )⊤ ϕ2 (hfrozen ) n n i=1 1 ≤ ϕ1 (h)⊤ ϕ2 (h) − ϕ1 (h)⊤ ϕ2 (hfrozen ) + ϕ1 (h)⊤ ϕ2 (hfrozen ) − ϕ1 (hfrozen )⊤ ϕ2 (hfrozen ) n (A.16)

By the Cauchy-Schwarz inequality, the first term is bounded by, 1/2 1/2 1 1 2 frozen 2 ∥ϕ1 (h)∥2 ∥ϕ2 (h) − ϕ2 (h )∥2 . n n Because ϕ1 is Lipschitz (for some constant C > 0) in h, a we have that 1 1 ∥ϕ1 (h)∥22 = ∥ϕ1 (h) − ϕ1 (hfrozen ) + ϕ1 (hfrozen )∥22 n n 2 2 ≤ ∥ϕ1 (h) − ϕ1 (hfrozen )∥22 + ∥ϕ1 (hfrozen )∥22 n n ! T X 1 2 ≤ 2C ∥h − hfrozen ∥2F + ∥at − āt ∥22 + ∥ϕ1 (hfrozen )∥22 n n t=1

and using the inductive hypothesis, the assumed almost sure convergence of at and Corollary A.11 on n2 ∥ϕ1 (hfrozen )∥22 implies that it is almost surely bounded by a constant. Moreover, using that ϕ2 is Lipschitz (for some bounded Lipschitz constant C > 0) in both argument h1 , . . . , hT and a1 , . . . , aT , we have that 1 1 ∥ϕ2 (h) − ϕ2 (hfrozen )∥22 = ∥ϕ2 (h, y, a1 , . . . , aT ) − ϕ2 (hfrozen , y, ā1 , . . . , āT )∥22 n n ! T X 1 ∥h − hfrozen ∥2F + ∥at − āt ∥22 . ≤C n t=1

As we have previously shown that ( n1 ∥h − hfrozen ∥2F ) converges to 0 almost surely, then the first term of (A.16) converges to zero almost surely. Applying the symmetric argument to the second term in (A.16) gives the desired almost sure limit for ϕ. An identical argument gives the exact same result for φ(θ̃1 , . . . , θ̃T , a1 , . . . , aT ). This completes the proof. ■ A.7.4. The Almost Sure Limits mj,t , ℓj,t And āt . To avoid confusion with the function ϕ used in the appendix, we denote the test function in Lemma 3.6 as ϕ̄. Combining Corollary A.11 and Lemma A.12, the proof of Lemma 3.6 follows after observing that Assumption A.7 can allow ϕ̄ to be utilized in both Corollary A.11 and Lemma A.12 with the identification of ϕ1 = ϕ̄ (applied row-wise) and ϕ2 being the all ones vector. As Definition A.8 and Definition A.5 are identical up to the frozen state evolution parameters, and have proven we can replace m̃j,t , ℓ̃j,t and at by the frozen sequence of their respective almost sure limits, we are done once we identify the recursions for the almost sure limits of m̃j,t , ℓ̃j,t and at . For the case of m̃j,t , we have that, m̃j,t+1 =

⊤ µ θt+1 θ̃⊤ µj j = η0 m̃j,t − η1 t , d d

recursively plugging in the almost sure limit of d1 µ⊤ j θ̃t from Corollary A.11 and that E[W̌t ] = 0, we have that for each s ∈ [T ], J

k=1

X µ⊤ µj X θ̃⊤ µj k lim s = lim α ℓk,s = α χj,k ℓk,s , d→∞ d d→∞ d

30 µ⊤ µk

µ⊤ µ

by the assumed limit of jd = kd j in Assumption A.3 (1). Thus, we have that almost surely that limd→∞ m̃j,t = mj,t defined by the recursion mj,t+1 = η0 mj,t − η1

J X

αχj,k ℓk,t .

k=1

Similarly, using that ℓ̃j,t = almost sure limit,

e⊤ B ĥt j

|Bj | 1 P |Bj | 1 P i∈Bj ĥt,i = n |Bj | i∈Bj g(ht,i , yi , at ), we have the n |Bj |

ℓj,t = pj E[g(Zt + mj,t , Yj , āt )], where Zt ∼ N (0, Ω̌t [t, t]). P P |B | As a final note, using that at+1 = γ0 at − γ1 Jj=1 nj |B1j | i∈Bj f (ht,i , yi , at ), and thus by Corollary A.11 with ϕ1 = f and ϕ2 being the all ones vector (recalling from Assumption A.3 (3) that f is Lipschitz in its first argument), we have the following almost sure limit, āt+1 = γ0 āt − γ1

J X

pj E[f (Gt + mj,t , Yj , āt )],

j=1

noting that ā1 = a1 and we assumed in Assumption A.3 (2) that ∥a1 ∥2 is bounded and that f is bounded from Assumption A.3 (3), meaning that each limit has ∥āt ∥2 bounded for t ∈ [T ]. Combining these almost sure limits with the state evolution recursion from Definition A.10, we obtain the state evolution in Definition A.5 and thus conclude Lemma 3.6. A.8. The Proof Of Theorem 3.8. Proof Of Theorem 3.8. Under Assumption A.3, we have that Assumption A.7 holds for the choice of ϕ(h, y, a) = L(Ma (h), y). Invoking Lemma 3.6 gives the following almost surely, n

X 1X L(Mat (ht,i ), yi ) = pj E[L(Māt (mj,t + Gt ), Yj )], Gt ∼ N (0, Ωt [t, t]), Yj ∼ Pj , n→∞ n i=1 j=1 P it is then immediate that limn→∞ n1 ni=1 L(Mat (ht,i ), yi ) = test(m1,t , . . . , mJ,t , Ωt [t, t], āt ) from Definition 3.5, invoking Proposition A.4 concludes the proof. ■ lim

A.9. The State Evolution Of Pure DD. Using Definition A.5, we can specify the state evolution for this algorithm. We abbreviate ∇h Ψr,j = ∇h Ψ(mj,r + Gr , Yj , ār ) where G1 , . . . , Gt ∼ N (0, Ωt ), P Yj ∼ Pj and define Vr,s = αΣt [r, s] + α2 Jj,k=1 χj,k pj pk E[∇h Ψr,j ]E[∇h Ψs,k ]⊤ . With identical initializations from Definition A.5 and indices r, s ∈ [t], we recursively define, Σt [r, s] =

J X

pj E[∇h Ψr,j ∇h Ψ⊤ s,j ]

j=1

Ξt [r + 1, s] = Ξt−1 [r, s] − ηVr,s Ωt+1 [r + 1, s + 1] = Ωt [r, s] − η(Ξt [r, s] + Ξt [s, r]⊤ ) + η 2 Vr,s mj,t+1 = mj,t − ηα

J X

(A.17) χj,k pk E[∇h Ψt,k ]

k=1

āt+1 = āt − η

J X j=1

pj E[∇a Ψ(mj,t + Gt , Yj , āt )].

A.10. Taylor Expansion Of (Pure) DD. Lemma A.13. Let µ ∈ RL , Σ ∈ RL×L and c ∈ R. Then for X ∼ N (cµ, Σ), the following identities hold, ∇µ E[f (X)] = cE[∇x f (x)|x=X ] 1 ∇Σ E[f (X)] = E[∇2x f (x)|x=X ] 2 ∇2µ E[f (X)] = c2 E[∇2x f (x)|x=X ] 1 ∇2Σ E[f (X)] = E[∇4x f (x)|x=X ] 4 c ∇µ ∇Σ E[f (X)] = E[∇3x f (x)|x=X ]. 2 ■

Proof. See [25, Equation (2)].

Lemma A.14. Suppose Assumption A.3 holds. Let T be the space of all θ = (ā, m1 , . . . , mJ , ω) where ∥ā∥2 ≤ C, ∥mj ∥2 ≤ C, and ∥ω −1 ∥op ≤ C, ∥ω∥op ≤ C for some constant C > 0. If G ∼ N (0, IdL ) and Yj ∼ Pj , then the expectations of the derivatives of Ψ(mj + ω 1/2 G, Yj , ā) with respect to ā, mj , and ω exist, and the expectation operator can be interchanged with the derivative operator. Proof. To interchange the expectation and derivative, it suffices to show that the partial derivatives of Ψ(mj + ω 1/2 G, Yj , ā) with respect to ā, mj , and ω are uniformly bounded by an integrable function for all θ ∈ T . Let x = mj + ω 1/2 G. By Assumption A.3 (4) (a), Ψ is Lipschitz with respect to h and a. Therefore, the gradients ∇h Ψ and ∇a Ψ exist almost everywhere and are uniformly bounded by constant L > 0. Thus, each of the following derivatives satisfy ∥∇ā Ψ∥2 = ∥∇a Ψ∥2 ≤ L ∥∇mj Ψ∥2 = ∥∇h Ψ∥2 ≤ L ∥∇ω Ψ∥F ≤ ∥∇h Ψ∥2 ∥∇ω (ω 1/2 G)∥op ≤ L · C ′ ∥G∥2 , where C ′ > 0 is a constant bounding the derivative of the matrix square root uniformly over the domain T ; such a constant exists by the assumption that max(∥ω −1 ∥op , ∥ω∥op ) ≤ C. We construct the dominating function H(G) = max(L, L · C ′ ∥G∥2 ). Because G ∼ N (0, IdL ), the second moment ∥G∥2 is bounded. Thus, H(G) is integrable and bounds the parameter derivatives independently of θ ∈ T , allowing the interchange of expectation and differentiation. ■ Lemma A.15. Let Assumption A.3 (4) hold for a bounded constant C > 0, the operator norm of the block Hessian matrix ∇2a,m,vec(ω) E[Ψ(m + ω 1/2 G, Yj , a)] is bounded by a constant C ′ > 0 depending only on C. Proof. We can expand the matrix ∇2a,m,vec(ω) E[Ψ] into the following 3 × 3 block matrix, ∇2a E[Ψ(m+ω 1/2 G, y, a)]  ∇m ∇a E[Ψ(m+ω 1/2 G, y, a)] ∇vec(ω) ∇a E[Ψ(m+ω 1/2 G, y, a)] 

∇a ∇m E[Ψ(m+ω 1/2 G, y, a)] ∇2m E[Ψ(m+ω 1/2 G, y, a)] ∇vec(ω) ∇m E[Ψ(m+ω 1/2 G, y, a)]

 ∇a ∇vec(ω) E[Ψ(m+ω 1/2 G, y, a)] ∇m ∇vec(ω) E[Ψ(m+ω 1/2 G, y, a)] . ∇2vec(ω) E[Ψ(m+ω 1/2 G, y, a)]

Applying the derivative identities from Lemma A.13, each block in the above matrix becomes one of E[∇2a Ψ], E[∇h ∇a Ψ], E[∇2h Ψ], 12 E[∇2h ∇a Ψ], 12 E[∇3h Ψ], and 14 E[∇4h Ψ], where we have suppressed the input for notational simplicity.

Assumption A.3 (4) guarantees the operator or Frobenius norm of each of these derivatives is bounded by C. Because the operator norm of any block matrix is bounded by the summed operator norms of its block matrices and the Frobenius norm strictly upper-bounds the operator norm, the operator norm of the entire block matrix is bounded by a finite linear combination of constant C, concluding the proof. ■ Theorem A.16. Let T ∈ N finite. Consider algorithm (3.2); recall the state evolution variables (specifically Ωt , mj,t and āt ) from Definition A.5 and testt = test(m1,t , . . . , mJ,t , Ωt [t, t], āt ) from Definition 3.5. Let ϵ = max(|1 − η0 |, |1 − γ0 |, η1 , γ1 ). If Assumption A.3 and Assumption A.7 (with respect to g, f ) hold for sufficiently small ϵ > 0, then in the limit ϵ → 0, testt satisfies the following expansion with Gt ∼ N (0, Ωt [t, t]), testt+1 = testt − (1 − γ0 )

J X

D E pj E[∇a Ψ(mj,t + Gt , Yj , āt )], āt

(A.18)

j=1 J X

− γ1

D E pj pj ′ E[∇a Ψ(mj,t + Gt , Yj , āt )], E[f (mj ′ ,t + Gt , Yj ′ , āt )]

(A.19)

j,j ′ =1 J X

D E pj E[∇h Ψ(mj,t + Gt , Yj , āt )], mj,t

(A.20)

J E D X −η1 α pj χj,j ′ pj ′ E[∇hΨ(mj,t +Gt , Yj , āt )], E[g(mj ′ ,t +Gt , Yj ′ , āt )]

(A.21)

− (1 − η0 )

j=1

j,j ′ =1

− η0 (1 − η0 )

J X

D E pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], Ωt [t, t]

(A.22)

j=1 J E D X 1 − η0 η1 pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], Ξt [t, t] + Ξt [t, t]⊤ 2

(A.23)

j=1

+ δt , where supt∈[T ] |δt | ≤ Cϵ2 as ϵ → 0. Remark A.17 (Physical Interpretation Of The First Order Expansion). The generalization error expansion from Theorem A.16 decomposes the local behavior of DD optimization into distinct physical terms. Connecting to works in spin glasses, one can interpret the order parameters mj,t and Ωt [t, t] as the magnetization (with respect to external field vectors µ1 , . . . , µJ ) and self-overlap of the parameter θt respectively. In addition, each of terms (A.18)-(A.23) has the following physical interpretation: – Equations (A.18) and (A.20) correspond to weight decay due to the introduction of a non-zero damping factor of 1 − γ0 and 1 − η0 respectively. – Equation (A.19) enforces the representation of the model Ma to align the low-dimensional signal a with the loss landscape through the use of function f . – Equation (A.21) drives the signal acquisition of each of the vectors µ1 , . . . , µJ by θ using the function g. Notice, the presence of χj,j ′ and α incorporates the overlap structure and data aspect ratio natively into the path of the optimization trajectory.

– Equations (A.22) and (A.23) govern the variance inflation of the relative noise in the directions perpendicular to µ1 , . . . , µJ . The first term represents the effect of damping. The second term, through unrolling the definition of Ξt [t, t], can be seen as cumulative correlation of the current gradient direction with the direction of past gradient evaluations. Each of these effects is compared to the relative scale of the expected Hessian of the post-activations corresponding to the iterate θt , eliciting the intuition that flat regions of the loss landscape suffer from little inflation while steep regions are subject to strong inflation. Proof Of Theorem A.16. By Definition 3.5, the value of testt+1 can be written as, testt+1 =

J X

pj E[L(Māt+1 (mj,t+1 + (Ωt+1 [t+1, t+1])1/2 G), Yj )],

j=1

where G ∼ N (0, IdL ) and Yj ∼ Pj . For convenience, we abbreviate Ωs [s, s] = ωs for all s ∈ [t + 1]. Using the recursive definitions of mj,t+1 , āt+1 from Definition A.5 and ωt+1 from Remark A.6, we have that testt+1 =

J X

pj E[L(Māt +∆ā ((mj,t + ∆mj ) + (ωt + ∆ω)1/2 G), Yj )]

j=1

J X

pj E[Ψ((mj,t + ∆mj ) + (ωt + ∆ω)1/2 G, Yj , āt + ∆ā)],

j=1

where we define (recalling ℓk,t from Definition A.5) ∆ā = −(1 − γ0 )āt − γ1 A, ∆mj = −(1 − η0 )mj,t − η1 Bj ,  ∆ω = −(1 − η02 )ωt − η0 η1 (Ξt [t, t] + Ξt [t, t]⊤ ) + η12 αΣt [t, t] + α2

J X

  χk,k′ ℓk,t ℓ⊤ k′ ,t

k,k′ =1

J X

1/2

pj E[f (mj,t + ωt G, Yj , āt )],

j=1

Bj = α

J X

1/2

χj,k pk E[g(mk,t + ωt G, Yk , āt )]

k=1

ℓk,t = pj E[g(mk,t + mk,t , Yk , āt )], with (G1 , . . . , Gt ) ∼ N (0, Ωt ) implicit in the definition of Σt and Ξt . By Assumption A.3 (3), the functions f, g are bounded, therefore each of A, (Bj )j∈[J] , (ℓj,t )j∈[J] are element-wise bounded which further implies that ∥A∥2 , ∥Bj ∥2 , ∥ℓj,t ∥2 are bounded (dependent on dimension L) for all j ∈ [J]. Moreover, by Assumption A.7 the operator norm of the matrices ωt = Ωt [t, t] = ΩT [t, t] and Σt [t, t] = ΣT [t, t] have bounded operator norm independent of t, and therefore the values of ∥ωt ∥F , ∥Σt [t, t]∥F θ⊤ θ̃

are bounded independent of t ∈ [T ]. Invoking Lemma 3.6, we have that limd→∞ td t = Ξ[t, t], then using that v s ! u ⊤ ⊤ ⊤ θ̃ θt θt u θt θ̃t θt θ̃t θ̃ t t tTr √ = Tr ≤ √ d d d d d F F F

34 θ⊤ θ

θ̃⊤ θ̃

and the limiting values of limd→∞ td t = ωt and limd→∞ td t = αΣt [t, t], we have that ∥Ξt [t, t]∥F is also bounded independent of t. Therefore, each term of ∆ā, ∆mj (for each j ∈ [J]) and ∆ω is bounded by Cϵ (for some constant C > 0 when ϵ is sufficiently small) uniformly over t ∈ [T ]. Now, define G(ā, m1 , . . . , mJ , ω) =

J X

pj E[Ψ(mj + ω 1/2 G, Yj , ā)].

j=1

By Taylor’s theorem, there exists a point ξ = (ξā , ξ(mj )j∈[J] , ξω ) on the line connecting (āt , (mj,t )j∈[J] , ωt ) to (āt+1 , (mj,t+1 )j∈[J] , ωt+1 ) such that, G(āt+1 , m1,t+1 , . . . , mJ,t+1 , ωt+1 ) = G(āt , m1,t , . . . , mJ,t , ωt ) + ∇G(āt , m1,t , . . . , mJ,t , vec(ωt ))⊤ ∆ā, ∆m1 , . . . , ∆mJ , vec(∆ω) ⊤ 1 + ∆ā, ∆m1 , . . . , ∆mJ , vec(∆ω) ∇2 G(ξ) 2 × ∆ā, ∆m1 , . . . , ∆mJ , vec(∆ω) , (A.24) where both the gradient and Hessian are with respect to (ā, m1 , . . . , mJ , vec(ω)). Notice the matrix ∇2 G(ξ) is a finite linear combination of the Hessian matrices from Lemma A.15. Recall from Definition A.5 that each state evolution parameter is a derived by a recursion of expectations over the functions g, f . Using the boundedness of these functions and the boundedness of the initialization ∥a1 ∥2 = ∥ā1 ∥2 , we can therefore conclude that each of ∥āt ∥2 , ∥āt+1 ∥2 , maxj∈[J] ∥mj,t ∥2 , maxj∈[J] ∥mj,t+1 ∥2 , ∥ωt ∥op , ∥ωt+1 ∥op are bounded. Thus, the corresponding values for ξ are also bounded due to being a linear interpolation of the above iterates. Then, applying the triangle equality on the operator norm of ∇2 G(ξ) to each j-term and using Assumption A.3 (4) (b) to invoke Lemma A.15, we have that ∥∇2 G(ξ)∥op ≤ C for some constant C > 0. Moreover, as we previously concluded that (∆ā, ∆m1 , . . . , ∆mJ , ∆ω) is bounded by Cϵ for sufficiently small ϵ uniformly over t, each second order term in (A.24) is of order ϵ2 uniformly over t ∈ [T ]. It remains to compute the first order Taylor expansion. Using Lemma A.14 to interchange expectation and derivative, we can calculate, ∇ā G(ā, m1 , . . . , mJ , ω) =

J X

h i pj E ∇ā Ψ(mj + ω 1/2 G, Yj , ā) ,

(A.25)

j=1

h i ∇mj G(ā, m1 , . . . , mJ , ω) = pj E ∇h Ψ(mj + ω 1/2 G, Yj , ā) , J h i 1X ∇ω G(ā, m1 , . . . , mJ , ω) = pj E ∇2h Ψ(mj + ω 1/2 G, Yj , ā) , 2

(A.26) (A.27)

j=1

where the latter two derivatives follow from Lemma A.13. Using (A.25), the first order contribution from ∆ā gives, ⟨∇ā G, ∆ā⟩ = −(1 − γ0 )

J X j=1

D E pj E[∇a Ψ(mj,t + Gt , Yj , āt )], āt

− γ1

J X

D E pj pk E[∇a Ψ(mj,t + Gt , Yj , āt )], E[f (mk,t + Gt , Yk , āt )] .

j,k=1

This produces terms (A.18) and (A.19). Using (A.26), the first order contribution from ∆mj gives, J X

⟨∇mj G, ∆mj ⟩ = −(1 − η0 )

j=1

J X

D E pj E[∇h Ψ(mj,t + Gt , Yj , āt )], mj,t

j=1

− η1 α

J X

D E pj χj,k pk E[∇h Ψ(mj,t + Gt , Yj , āt )], E[g(mk,t + Gt , Yk , āt )] .

j,k=1

This produces terms (A.20) and (A.21). Using (A.27), The first order contribution from ∆ω gives, J D E X 1 2 ⟨∇ω G, ∆ω⟩ = − (1 − η0 ) pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], ωt 2 j=1

1 − η0 η1 2 1 + η12 2

J X

E D pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], Ξt [t, t] + Ξt [t, t]⊤

j=1

J X

pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], αΣt [t, t] + α2

χk,k′ ℓk,t ℓ⊤ k′ ,t

k,k′ =1

j=1

= −η0 (1 − η0 )

J X

E D pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], ωt

j=1 J D E X 1 − η0 η1 pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], Ξt [t, t] + Ξt [t, t]⊤ 2 j=1

J J E X 1 X D χk,k′ ℓk,t ℓ⊤ pj E[∇2h Ψ(mj,t +Gt , Yj , āt )], αΣt [t, t]+α2 + η12 ′ k ,t 2 ′ j=1

(A.28)

k,k =1

E X D 1 − (1 − η0 )2 pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], ωt 2

(A.29)

j=1

= −η0 (1 − η0 )

J X

D E pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], ωt

j=1 J

E X D 1 − η0 η1 pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], Ξt [t, t] + Ξt [t, t]⊤ + δt , 2 j=1

where supt∈[T ] |δt | ≤ Cϵ2 for some constant C > 0 for sufficiently small ϵ. The final inequality in the above display follows by rewriting (A.29) using (1 − η02 ) = (1 − η0 )2 + 2η0 (1 − η0 ), and using the inequality ⟨M1 , M2 ⟩ ≤ ∥M1 ∥F ∥M2 ∥F to absorb term (A.28) using the η12 prefactor. This second argument is allowed by the assumed boundedness of ∥E[∇2h Ψ(mj,t + Gt , Yj , āt )]∥F (Assumption A.3 (4) (b)), the boundedness of ∥ωt ∥F , ∥Σt [t, t]∥F independent of time t ∈ [T ] and recognizing that ∥ℓk,t ℓ⊤ k′ ,t ∥F ≤ ∥ℓk,t ∥F ∥ℓk′ ,t ∥F is bounded as ℓk,t is an expectation over g, a bounded function. This produces terms (A.22) and (A.23). Combining all first order contributions and bounding the second order remainder by the ϵ2 order remainder δt (say by enlarging its bound by 2Cϵ2 ) completes the proof. ■

Proof of Theorem 4.1. First, we invoke Theorem A.16 with the choice of η0 , γ0 = 1, η1 = γ1 = η, g = ∇h Ψ, f = ∇a Ψ. This gives the Taylor expansion testt+1 = testt −η

J X

D E pj pj ′ E[∇a Ψ(mj,t + Gt , Yj , āt )], E[∇a Ψ(mj ′ ,t + Gt , Yj ′ , āt )]

j,j ′ =1

− ηα

J X

D E pj χj,j ′ pj ′ E[∇h Ψ(mj,t + Gt , Yj , āt )], E[∇h Ψ(mj ′ ,t + Gt , Yj ′ , āt )]

j,j ′ =1 J

−

E ηX D pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], Ξt [t, t] + Ξt [t, t]⊤ + ϵt , 2 j=1

where supt∈[T ] |ϵt | ≤ Cη 2 for constant C > 0 for sufficiently small η > 0. Using the notation given in the statement of Theorem 4.1 and simplifying the above equation gives the proof. ■ Used as an important step in the proof of Theorem A.16, we present a high level observation that allows one control the norm of the matrix Ξ[t, t] in the Taylor expansion of Theorem 4.1. Remark A.18. Observe, the first term on line (4.2) need not be negative and may increase the test error. Moreover, unrolling Ξt [t, t] from Definition A.5, this term is equal to, η

J X

J t−1 J D E X X X 2 ⊤ pj E[∇2h Ψ(mj,t + Gt , Yj , āt )], α mk,1 ℓ⊤ − η αΣ [s, t] + α χ ℓ ℓ . t a,s a,b k,t b,t

j=1

k=1

s=1

a,b=1

It is not clear the above summation is bounded independent of time t. Thankfully, the assumption that Σt [t, t] and Ωt [t, t] have bounded operator norms (and thus bounded Frobenius norms) from Assumption A.7 allows for such control. Indeed, using the limits from Lemma 3.6 and θt⊤ θ̃t d

θt θ̃ √t ≤ √ , d d F F F

we have that ∥Ξt [t, t]∥F , and thus the term in (4.2) by Assumption A.3 (4), is bounded in the n, d limit. A.11. The Benefits of Damping For Test Error Critical Points. Further tuning hyperparameters in (3.2) can give favorable global behavior as well. Theorem A.19. Let g = ∇h Ψ and f = ∇a Ψ in Algorithm (3.2), with η, γ ∈ (0, 1), consider the one dimensional hyperparameter subspace over c ∈ R+ , η0 = 1 − η,

η1 = cη,

γ0 = 1,

γ1 = γ, .

Any fixed point (m∗1 , . . . , m∗J , Ω∗ , ā∗ ) of Definition A.5, must satisfy the equations ! J X ∗ ∗ ∗ ∗ mj = −cα χj,k pk E[g(mk + G , Yk , ā )]

(A.30)

k=1

0=−

J X

pj E[f (m∗j + G∗ , Yj , ā∗ )],

(A.31)

j=1

where Yj ∼ Pj and G∗ ∼ N (0, Ω∗ ). Moreover, for any deterministic vector v ∈ Rd perpendicular to (µj )j∈[J] , any iterate θ produced by the above fixed point must satisfy limd→∞ d−1 v ⊤ θ = 0.

Proof. Assume the fixed point given in the statement of the theorem. Using the equations of mj,t+1 and āt+1 for the given fixed point, we have that (m∗j )j∈[J] , Ω∗ , ā∗ satisfy the following system, m∗j = η0 m∗j − η1 α ā∗ = γ0 ā∗ − γ1

J X

χj,k pk E[g(m∗k + G∗ , Yk , ā∗ )]

k=1 J X

pj E[f (m∗j + G∗ , Yj , ā∗ )],

j=1

where Yj ∼ Pj and G∗ ∼ N (0, Ω∗ ). Rearranging these equations and simplifying gives the fixed point condition in the theorem. Note, by adding a dummy signal vector µJ+1 = v, pJ+1 = 0 where v ⊥ µj for j ∈ [J] into distribution (3.1), we can also conclude the second statement in the theorem, i.e. that m∗J+1 = 0, as χJ+1,k = 0 for all k ∈ [J] and η0 ∈ (0, 1). ■ PJ 2 ⊤ ∗ Let (θ∗ , a∗ ) be a critical point for the test error, assume that H ∗ = j=1 pj E[∇h Ψ(µj θ /d + Z ⊤ θ∗ , Yj , a)] is invertible and isotropic when Z ∼ N (0, 1/d). Then, there exists a c ∈ R+ where (A.30) and (A.31) are equivalent to projected critical point conditions on (θ∗ , a∗ ). A test error critical point satisfies ∇a Ex̌,y̌ L(Ma∗ (x̌⊤ θ∗ ), y̌) = 0 and ∇θ Ex̌,y̌ L(Ma∗ (x̌⊤ θ), y̌) = 0 with (x̌, y̌) drawn from distribution (3.1). Expanding (3.1) in terms of the mixture on µ1 , . . . , µJ , interchanging expectation and derivative by Lemma A.14, we equivalently write this system as, J X

∗ ⊤ ∗ ∗ pj E[∇a Ψ(µ⊤ j θ /d + Z θ , Yj , a )] = 0,

(A.32)

∗ ⊤ ∗ ∗ ∗ pj µj E[∇h Ψ(µ⊤ j θ /d + Z θ , Yj , a )] = θH ,

(A.33)

j=1

−

J X j=1

where Z ∼ N (0, Idd /d). As H ∗ is invertible, for any deterministic vector v ∈ Rd , (A.33) implies that, J X v ⊤ µj v⊤ θ∗ ∗ ⊤ ∗ ∗ −1 =− E[∇h Ψ(µ⊤ (A.34) pj j θ /d + Z θ , Yj , a))] (H ) . d d j=1

∗ There are two cases of interest: (1) If v = µj for some j ∈ [J], assuming that limd→∞ µ⊤ j θ /d = mj,θ∗ , limd→∞ (θ∗ )⊤ θ∗ /d = Ωθ∗ , limn,d→∞ a∗ = āθ∗ , and replacing (H ∗ )−1 = (c∗ )−1 Id for some constant c∗ > 0 by the isotropic assumption, asymptotically solving (A.34) is equivalent to,

mj,θ = −c∗

J X

χj,k pk E[∇h Ψ(mk,θ∗ + G, Yj , āθ∗ ))],

G ∼ N (0, Ωθ∗ );

(A.35)

k=1 ⊤

(2) If v is perpendicular to the span of (µj )j∈[J] , then assuming that limd→∞ v d θ = m⊥ almost surely, we have the asymptotic equation m⊥ = 0.

(A.36)

⊤ Next, under the same limits a, µ⊤ j θ/d and θ θ/d, (A.32) is asymptotically equivalent to, J X j=1

pj E[∇ā Ψ(mj,θ∗ + G, Yj , āθ∗ )] = 0,

G ∼ N (0, Ωθ∗ ) .

(A.37)

Comparing (A.35), (A.36) and (A.37) with mj,θ∗ = m∗j for each j ∈ [J], Ωθ∗ = Ω∗ and āθ∗ = ā∗ , and selecting c = (c∗ α)−1 , the statement of Theorem A.19 is identical to the set of equations (A.32) and (A.35). Remark A.20 (Adaptive Hessian estimation). Note, in the case where H ∗ is not isotropic, one naturally would incorporate the Hessian Ht into a DD algorithm, the data-based estimate of n

Ĥt =

1X 2 ∇h Ψ(ht,i , yi , at ), n i=1

is already calculated for the correction term in both pure DD and is expected to be an asymptotically exact estimate of Ht . But, this requires a time-inhomogeneous choice of g = gt ; such an extension to our results is possible but is left to future work. Appendix B. Deferred Applications B.1. Faithful Test Error Tracking In Signal-less Learning. Recall this problem from the introduction, consider n i.i.d. data points xi ∼ N (0, Idd /d) and yi = 0, generating the data (X, y). Let Mθ (x) = x⊤ θ with θ ∈ Rd initialized at θ1 ∼ N (0, Idd ) and trained under loss L(ŷ, y) = 12 (ŷ −y)2 (rescaled for convenience), realized by distribution (3.1) with J = 1 and µ1 being the zero vector. We run the following algorithms: GD

(Damped) Pure DD

ht = Xθt

t−1 X (t−1)−s ht = Xθt + η1 η0 ĥs

θ̃t = X ⊤ ht θt+1 = θt − η θ̃t

s=1 ⊤

θ̃t = X ht − αθt θt+1 = η0 θt − η1 θ̃t

Figure 5 plots the train and test errors for GD and pure DD (η0 = 1, η1 = η). We immediately see the train-test disconnect in GD is corrected in both DD algorithms. Note, for pure DD, the train-test error explodes. This is predicted from the variance inflation terms in Ωt from Definition A.5 and occurs because this problem has no signal. Specifically, we can use Remark A.6 to conclude the variance recursion (and thus the test error recursion) takes form, Ωt+1 [t + 1, t + 1] = η02 Ωt [t, t] − η1 η0 (Ξt [t, t] + Ξt [t, t]⊤ ) + η12 αΩt [t, t], which may or may not be a contraction depending on the choices of η0 , η1 . This naturally leads into relying on the zero-cost validation feature of DD. We consider the following one dimensional subspace dependent on η, η0 = 1 − η,

η1 = cη.

(B.1)

We then sweep the value of c over some specified range for a fixed data realization. The results of many individual replications of this experiment are given in Figure 6. Again, for each simulation the train-test identity holds perfectly as expected and the trajectory of the test error is tracked by the state evolution prediction. Moreover, since each of these replications used the same dataset, we can optimize over the hyperparameter c by selecting whichever algorithm gets the optimal training (and thus test) error. B.2. The XOR Training Iterations.

Figure 5. Summary stats for 100 signal-less regression runs (n = 200, d = 800): GD (top) vs. DD (bottom) with η = 0.05. Blue/red denote train/test error; solid lines are medians, shaded areas are interquartile ranges, and dotted lines show min/max. Notice, by design, the trajectories of the train and test error are identical for pure DD while they significantly diverge for GD.

Pure DD (h1t , h2t ) = X(θt1 , θt2 ) + η

(h1t , h2t ) = X(θt1 , θt2 ) ŷt = σ(at (h1t ⊙ h2t )) ĥ1t = (ŷt − y) ⊙ (at h2t ) ĥ2t = (ŷt − y) ⊙ (at h1t ) 1 θt+1 = θt1 − ηX ⊤ ĥ1t 2 θt+1 = θt2 − ηX ⊤ ĥ2t

at+1 = at − η(ŷt − y)⊤ (h1t ⊙ h2t )/n

t−1 X

s=1 1 2 ŷt = σ(at (ht ⊙ ht )), ĥ1t = (ŷt − y) ⊙ (at h2t ) ĥ2t = (ŷ − y) ⊙ (at h1t ), n X

1 Ht = n

(ĥ1s , ĥ2s ),

ŷt,i (1− ŷt,i )a2t

(h2t,i )2 h1t,i h2t,i h1t,i h2t,i (h1t,i )2

i=1

! 0 at +(ŷt,i −yi ) at 0 θ̃t = X ⊤ (ĥ1t , ĥ2t ) − α(θt1 , θt2 )Ht 1 2 (θt+1 , θt+1 ) = (θt1 , θt2 ) − η θ̃t

at+1 = at − η(ŷt − y)⊤ (h1t ⊙ h2t )/n

Figure 6. An example of hyperparameter tuning with DD over 100 individual replications, in each figure is the train and test error for damped pure DD algorithm in the signal-less regression problem with n = 800, d = 800, η = 0.05 and we sweep c over the sub-space given in (B.1). Blue lines refer to train error and red lines refer to test error; the solid line is the median error, the shaded region is the range from the 25-th to 75-th quartiles and the dotted lines are the minimum and maximum error runs. This represents an application where multiple instances of DD are run on the same data set and, because the train-test identity, we get a zero-cost validation method to selecting good hyperparameters that aid in training. Damped Variant Of Pure DD (Fixed at = 1) t−1 X (t−1)−s 1 2 (h1t , h2t ) = X(θt1 , θt2 ) + η1 η0 (ĥs , ĥs ), s=1

ŷt = σ(h1t ⊙ h2t ), ĥ1t = (ŷt − y) ⊙ (h2t ) ĥ2t = (ŷt − y) ⊙ (h1t ), ! 2 2 n 1X (ht,i ) h1t,i h2t,i 0 1 ŷt,i (1− ŷt,i ) 1 2 +(ŷt,i −yi ) Ht = 1 0 ht,i ht,i (h1t,i )2 n

i=1 θ̃t = X ⊤ (ĥ1t , ĥ2t ) − α(θt1 , θt2 )Ht 1 2 (θt+1 , θt+1 ) = η0 (θt1 , θt2 ) − η1 θ̃t

Figure 7. Empirical density estimates for runtime on a width-9 two-layer neural network (25 replications, 500 epochs). Left: Distribution of per-epoch clock time (ms) for GD vs. DD. Right: Total training time (s) per replication for the MNIST zeros/eights problem. B.3. A Run-time Analysis For The MNIST Problem. Figure 7 tracks the “wall clock” time per epoch (i.e. a single full batch) for our width-nine hidden layer model at both the per epoch level and for total training time. Although our implementation of DD suffers from a ≈ L× constant factor overhead, we believe the practical benefits of DD (including use of the full training set) may outweigh this computational penalty. A variant of DD mimicking SGD may further reduce this overhead (Section 6; future directions). B.4. Implementation of MNIST Problem With Nine Hidden Layers. Below is our implementation of the MNIST zeros/eights problem for a width-nine hidden layer model. 1 2 3

import time, math, torch, numpy as np import torch.nn as nn, torch.nn.functional as F from torchvision import datasets, transforms

4 5 6 7 8

def get_data(train): d = datasets.MNIST("./data", train=train, download=True, transform=transforms.ToTensor()) m = (d.targets == 0) | (d.targets == 8) return d.data[m].float() / 255.0, (d.targets[m] == 8).float()

9 10 11

x_tr, y_tr = get_data(True) x_test, y_test = get_data(False)

12 13 14 15 16

N, T, H, lr, lam, reps = 800, 500, 9, 1, 30, 20 device = "mps" if torch.backends.mps.is_available() else "cpu" x_test, y_test = x_test.unsqueeze(1).to(device), y_test.to(device) alpha = N / 784

17 18 19 20 21 22

def inject_noise(x, lam): D = x.shape[-1] * x.shape[-2] p = torch.rand_like(x) z = torch.where(p < 0.25, -math.sqrt(2/D), torch.where(p < 0.5, math.sqrt(2/D), 0.0)) return (lam / D) * x + z.to(x.device)

23 24 25 26 27 28 29 30

class DecoupledMLP(nn.Module): def __init__(self): super().__init__() self.fc1, self.fc2 = nn.Linear(784, H, bias=False), nn.Linear(H, 1, bias=False) def forward(self, x, mem=0): z = self.fc1(x.view(x.size(0), -1)) + mem return self.fc2(torch.tanh(z)).squeeze(1), z

42 31 32 33

dd_err = np.zeros((2, reps, T)) dd_time = []

34 35 36 37 38

for r in range(reps): idx = torch.randperm(len(x_tr))[:N] x_n, y_n = inject_noise(x_tr[idx].unsqueeze(1).to(device), lam), y_tr[idx].to(device) x_te_n = inject_noise(x_test, lam)

39 40 41 42 43 44 45

m_dd, mem = DecoupledMLP().to(device), torch.zeros(N, H, device=device) t0 = time.time() for t in range(T): logits, z = m_dd(x_n, mem) loss = F.binary_cross_entropy_with_logits(logits, y_n) h_hat = torch.autograd.grad(loss, z, create_graph=True)[0]

46 47

J_bar = torch.stack([torch.autograd.grad(h_hat[:, i].sum(), z, ,→ retain_graph=True)[0].mean(dim=0) for i in range(H)])

48 49 50 51 52 53

with torch.no_grad(): m_dd.fc2.weight -= lr * torch.autograd.grad(loss, m_dd.fc2.weight, ,→ retain_graph=True)[0] theta = m_dd.fc1.weight m_dd.fc1.weight.copy_(theta - lr * (h_hat.t() @ x_n.view(N, -1) - alpha * (J_bar @ ,→ theta))) mem += lr * h_hat.detach()

54 55 56 57

dd_err[0, r, t] = loss.item() with torch.no_grad(): dd_err[1, r, t] = F.binary_cross_entropy_with_logits(m_dd(x_te_n)[0], ,→ y_test).item() dd_time.append(time.time() - t0)

B.5. Whitening Technique For CIFAR-10 Embedding Vectors. We considered the following pre-processing steps, which approach the conditions of Assumption A.3 to varying degrees: √ – Vanilla (no whitening): Features are rescaled by a 1/ d factor (assuming the embedding vectors have Gaussian entries). This leaves correlations between coordinates of the embedding vectors which our current DD algorithm does not account for. – ZCA (whitening) on the train data: We compute the empirical covariance matrix Σtrain = 1 T N −1 Xtrain Xtrain exclusively on the training set and compute the eigen-decomposition Σtrain = T V ΛV . We construct the Zero-phase Component Analysis (ZCA) [14] matrix Wtrain = ′ V (Λ)−1/2 V T , which is then applied to both √ the train and test data using X = XWtrain and ′ X̌ = X̌Wtrain . Finally, we apply the 1/ d normalization from the Vanilla bullet. – Joint ZCA (whitening) on train and test data: To prevent drift in the covariance between the train and test sets, we concatenate both data sets to form a joint covariance matrix Σjoint . The resulting Wjoint is applied to all data with, ′ X XWjoint = . X̌ ′ X̌Wjoint √ Finally, we apply the 1/ d normalization from the Vanilla bullet. B.6. Data Collection and Attribution. This work utilized both the MNIST and CIFAR-10 datasets, standard benchmarks in the high-dimensional statistics and machine learning literature. – MNIST: A collection of 70, 000 handwritten digit images [16]. The dataset is available under the Creative Commons Attribution-Share Alike 3.0 license.

– CIFAR-10: Consists of 60, 000 32 × 32 color images in 10 classes [15]. The dataset is used in accordance with the researchers’ terms at the University of Toronto.

Record · ID 146951 · SHA-256 88ff050a0a11421c
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.