Stacking as Accelerated Gradient Descent

We consider a fairly general supervised learning setting. Denote the input space by $\mathcal{X}$ and the output space by $\mathcal{Y}$. Examples $(x,y)\in\mathcal{X}\times\mathcal{Y}$ are drawn from a distribution $D$ (which may simply be the empirical data distribution in the case of empirical risk minimization). We aim to model the input-output relationship via predictions in $\mathbb{R}^{d}$, for some dimension parameter $d$. Given an example $(x,y)\in\mathcal{X}\times\mathcal{Y}$ the quality of a prediction $\widehat{y}\in\mathbb{R}^{d}$ is measured via a loss function $\ell:\mathbb{R}^{d}\times\mathcal{Y}\to\mathbb{R}$. Predictions are computed using functions $f:\mathcal{X}\to\mathbb{R}^{d}$. We will assume that the predictor functions $f$ are square integrable with respect to $D$, i.e. $\operatorname*{\mathbb{E}}_{(x,y)\sim D}[\|f(x)\|^{2}]<\infty$. The space of such functions forms a Hilbert space, denoted $\mathcal{L}_{2}$, with the inner product defined as $\langle f,g\rangle=\operatorname*{\mathbb{E}}_{(x,y)\sim D}[\langle f(x),g(x)\rangle]$. Unless specified otherwise, all functions in the subsequent discussion will be assumed to be in $\mathcal{L}_{2}$. The loss function can then naturally be extended to predictor functions $f\in\mathcal{L}_{2}$ by defining, with some abuse of notation, $\ell(f):=\operatorname*{\mathbb{E}}_{(x,y)\sim D}[\ell(f(x),y)]$. The goal of training is to obtain a function $f\in\mathcal{L}_{2}$ that minimizes $\ell(f)$.

In the rest of this section, we perform the analysis in a purely functional setting, which affords a convenient analysis. However, we note that, in practice, functions are parameterized (say by neural networks) and hence update rules for functions may not always be realizable via the specific parameterization used. The functional setting allows us to sidestep realizability issues and to focus on the conceptual message that stacking initialization enables accelerated updates.

We now define a general ensemble learning setup within the above setting. In this setup, we aim to approximate the minimizer of $\ell$ on $\mathcal{L}_{2}$ via an ensemble, which is a sequence of functions $(f_{1},f_{2},\ldots,f_{T})$, where $T>0$ is a given parameter defining the size of the ensemble. The functions in the ensemble are typically “simple” in the sense that they are chosen from a class of functions that is easy to optimize over. A predictor function can be obtained from an ensemble $(f_{1},f_{2},\dots,f_{T})$ by aggregating its constituent functions into a single function $F_{T}:\mathcal{X}\to\mathbb{R}^{d}$. The loss of an ensemble can then be defined (again with some abuse of notation) in terms of its aggregation as $\ell((f_{1},f_{2},\ldots,f_{T})):=\ell(F_{T})$. Two specific aggregation operators we consider are the following:

1. 1.

   Addition: (E.g. boosting.) This is a summation over ensemble outputs: $F_{T}=f_{1}+f_{2}+\cdots f_{T}$.

2. 2.

   Residual composition: (E.g. deep residual neural networks.) This is a composed function $F_{T}=(I+f_{T})\circ(I+f_{T-1})\circ\cdots\circ(I+f_{1})$, where the domain is $\mathcal{X}=\mathbb{R}^{d}$ and $I:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is the identity mapping.

Stagewise training is a simple greedy procedure to train ensembles in a progressive manner. Suppose we have already obtained a (partial) ensemble $(f_{1},f_{2},\ldots,f_{t})$. Then, the next function in the ensemble, $f_{t+1}$, is ideally obtained by minimizing the loss of the new ensemble, i.e. $f_{t+1}=\arg\min_{f}\ell((f_{1},f_{2},\ldots,f_{t},f))$.

However, in practice, this ideal is hard to implement, and instead two heuristics are commonly used: (a) the new function to be trained is initialized in some carefully chosen manner, and (b) the optimization above is done using early stopping, i.e. a few steps of gradient descent, which ensures that the new function stays close to initialization. We analyze these heuristics in a functional optimization setting as follows.

First, we assume that the function $f_{t+1}$ to be trained is initialized at some carefully chosen value $f_{t+1}^{0}$. For notational convenience, we denote the aggregation of the ensemble $(f_{1},f_{2},\ldots,f_{t},f_{t+1}^{0})$ by $F_{t+1}^{0}$ and that of the generic ensemble $(f_{1},f_{2},\ldots,f_{t},f)$ by $F$.

Next, we note that an exact analysis for early stopping quickly becomes technically intractable. Instead, for a theoretical analysis, we model the heuristic of early stopping by using $\ell_{2}$ regularization around the initialization and linearizing the loss near the initialization, as follows. It is known (see, e.g. [13, Section 7.8]) that early stopping acts as a form of $\ell_{2}$ regularization which ensures that the trained function $f_{t+1}$ remains close to its initialization, $f_{t+1}^{0}$, which implies that $F_{t+1}$ remains close to $F_{t+1}^{0}$. Thus, early stopping can be modeled as minimizing $\ell(F)+\frac{\lambda}{2}\|F-F_{t+1}^{0}\|^{2}$, for some regularization parameter $\lambda$. Further, since the trained function remains close to the initialization, we also approximate $\ell(F)$ by its linearization around the initialization: $\ell(F)\approx\ell(F_{t+1}^{0})+\langle\nabla\ell(F_{t+1}^{0}),F-F_{t+1}^{0}\rangle$. Here, $\nabla\ell(\cdot)$ is the Fréchet derivative, and $\langle\cdot,\cdot\rangle$ denotes the inner product in $\mathcal{L}_{2}$. These considerations lead to the following key modeling assumption.

We can now consider specific initialization strategies (i.e. zero initialization, random initialization, and stacking initialization) in the context of additive and residual compositional models and see how these initializations lead to various forms of functional gradient descent.

First, consider stagewise training where functions are initialized to be zero functions, i.e. $f_{t+1}^{0}=0$. It is easy to see that with this initialization, for both additive and residual compositional models, we have $F_{t+1}^{0}=F_{t}$. Thus, from (1), we have that the updated ensemble’s predictor can be written as

This exactly describes functional gradient descent with step size $\frac{1}{\lambda}$. In the additive setting this is well-known: indeed, boosting can be seen as a functional gradient descent [24]. The result for the residual compositional setting appears to be new.

We now consider stagewise training where functions are initialized randomly, i.e. $f_{t+1}^{0}$ is a randomly drawn function, independent of all randomness up to stage $t$. In the following, we will assume that $\operatorname*{\mathbb{E}}[f_{t+1}^{0}]=0$, where the $0$ on the RHS denotes the zero function. With this initialization, for both additive and residual compositional models, we have $F_{t+1}^{0}=F_{t}+g_{t}$, where $g_{t}=f_{t+1}^{0}$ for additive models, and $g_{t}=f_{t+1}^{0}\circ F_{t}$ for residual compositional models. In either case, note that $\operatorname*{\mathbb{E}}[g_{t}]=0$. Now define the loss functional $\ell_{t}(F):=\operatorname*{\mathbb{E}}[\ell(F+g_{t})]$. Since $\operatorname*{\mathbb{E}}[g_{t}]=0$, we can interpret $\ell_{t}$ as a randomized smoothing of $\ell$, similar to convolving with a Gaussian. Then, from (1), we have that the updated ensemble’s predictor can be written as

Now, note that

Or in other words, the above update can be seen as a stochastic functional gradient descent step on the smoothed loss function $\ell_{t}$.

We now consider stagewise training where functions are initialized in a stacking-like fashion with $f_{t+1}^{0}=f_{t}$, which we will refer to as the stacking initialization. When the ensemble aggregation operator is addition, we have $F_{t+1}^{0}=f_{t}+F_{t}=F_{t}+(F_{t}-F_{t-1})$ and hence (1) implies that the updated ensemble’s predictor is

The above formula essentially describes Nesterov’s accelerated gradient descent, which has the following update rule:

Here, $\beta\in[0,1)$ is a constant that can depend on $t$. In fact, we can exactly recover Nesterov’s accelerated gradient descent if we modify the stacking initialization to $f_{t+1}^{0}=\beta f_{t}$. Thus, stacking enables accelerated descent for training additive models.

When the ensemble aggregation operator is residual composition, stagewise training with the stacking initialization $f_{t+1}^{0}=f_{t}$ results in

Equation (1) therefore implies the updated ensemble’s predictor is

In contrast, Nesterov’s update rule (2), and the fact that for residual compositional models,

yields the following equation for $F_{t+1}$:

Comparing (3) and (4), barring the minor difference in $\beta$ parameters, which can be easily rectified as in the case of the additive models by setting $f_{t+1}^{0}=\beta f_{t}$, the major difference is that $f_{t}\circ F_{t}$ replaces $f_{t}\circ F_{t-1}$. Although possibly intractable to prove formally, we believe that the updates in (3) also provide an accelerated convergence rate, since we expect $F_{t-1}$ to be close to $F_{t}$ as iterates converge to the optimal function.

In the following section, we show that in certain deep linear networks, the above intuition is indeed correct and provide a rigorous proof that stacking provides an accelerated convergence rate.

Consider again the general supervised learning setting from Section 2 and suppose, as is often the case in modern neural networks, that examples consist of inputs $x\in\mathbb{R}^{d}$ and outputs $y\in\mathcal{Y}$. The loss function $\ell:\mathbb{R}^{d}\times\mathcal{Y}\to\mathbb{R}$ is assumed to be convex in the first argument. Let the samples be drawn from a distribution $D$ over $\mathbb{R}^{d}\times\mathcal{Y}$. Then the expected loss of the linear predictor $x\mapsto Wx$ for a matrix $W\in\mathbb{R}^{d\times d}$ is (with some abuse of notation) $\ell(W):=\operatorname*{\mathbb{E}}_{(x,y)\sim D}[\ell(Wx,y)]$. In the following, we will assume the expected loss $\ell(W)$ is $L$-smooth and $\mu$-strongly convex in $W$, by which we mean that the following inequalities hold for any $W,V\in\mathbb{R}^{d\times d}$:

Here, for matrices $W,V\in\mathbb{R}^{d}$, $\langle W,V\rangle=\text{Tr}(W^{\top}V)$, and $\left\lVert W\right\rVert$ is the Frobenius norm of $W$. The condition number $\kappa$ of the loss is defined as $\kappa:=\frac{L}{\mu}$.

The deep residual neural networks we consider have $t$ layers with weight matrices $w_{1},w_{2},\ldots,w_{t}$, and the function they compute is $x\mapsto W_{t}x$, where

Here, $I\in\mathbb{R}^{d\times d}$ is the identity matrix providing the residual connection. The expected loss of the neural network described above on the data is $\ell(W_{t})$.

Suppose we train the deep residual linear network described above using stacking initialization, but incorporating $\beta$-scaling: i.e., to train the $(t+1)$-th layer, its weight matrix is initialized to $w_{t+1}^{0}=\beta w_{t}$, for some constant $\beta\in[0,1]$, and then trained. Following the exact same steps as in the derivation of stacking updates in the functional setting of Section 2, we end up with the following formula for $W_{t+1}$:

When $W_{t-1}$ is non-singular, we have $w_{t}=W_{t}W_{t-1}^{-1}-I=(W_{t}-W_{t-1})W_{t-1}^{-1}$, so the above equation can be rewritten as

As previously noted, (5) differs from Nesterov’s AGD method in form: Nesterov’s AGD updates would be

Despite the differences with Nesterov’s method, we can show that stacking notably still yields a provably accelerated convergence rate. Let $W^{*}:=\arg\min_{W}\ell(W)$. Suppose that $W^{*}$ is non-singular, i.e. its smallest singular value $\sigma_{\min}(W^{*})>0$. Theorem 3.1 shows that as long as the first two layers are initialized so that $W_{1}$ and $W_{2}$ are close to optimal, stacking results in a suboptimality gap of $\exp(-\widetilde{\Omega}(T/\sqrt{\kappa}))$ after $t$ stages of stacking. This is of the same order as the rate obtained by Nesterov’s acceleration; note that, in comparison, stagewise training with zero initialization results would result in a suboptimality gap of $\exp(-\widetilde{\Omega}(T/\kappa))$.

The $\widetilde{\Omega}(\cdot)$ notation above hides polylogarithmic dependence on the problem parameters for clarity of presentation. Precise expressions can be found in the proof. The primary insight behind the result of Theorem 3.1 is that Nesterov’s accelerated gradient method is relatively robust to perturbations in its update rules. This robustness is formalized below in Lemma 3.2. The lemma is described in a fairly general, standalone setting since it may be of independent interest.

We will apply Lemma 3.2 using the correspondence $x_{0}=y_{0}=W_{0}=V_{0}$, and for all $t\geq 1$, $y_{t}=W_{t}$ and $x_{t}=W_{t}+\beta(W_{t}-W_{t-1})W_{t-1}^{-1}W_{t}$. Lemma 3.2 implies that even though stagewise training with stacking differs from Nesterov’s method, their similar form allows us to express the former as a perturbation of the latter. In particular, if we write our stacking update (5) for deep residual linear networks as a perturbation of Nesterov’s method, as in (7), the perturbation term is exactly

for all $t\in[2,T-1]$. That is, if we examine the sequence of iterates $W_{1},\dots,W_{T}$ produced by stagewise training with stacking initialization, the term $\Delta_{t}$ is a measure of the disagreement between the realized iterate $W_{t+1}$ and what Nesterov’s method says the iterate $W_{t+1}$ should be conditioned on the iterates $W_{1},\dots,W_{t}$ from previous timesteps. We note that, even when these perturbation terms $\Delta_{1},\dots,\Delta_{T-1}$ are small in norm, it is possible for Nesterov’s method to describe an iterate sequence that diverges significantly in norm from the iterates realized by stagewise training.

Rewriting (8) as $\Delta_{t}=\beta(W_{t}-W_{t-1})W_{t-1}^{-1}(W_{t}-W_{t-1})$, we immediately see that in order to satisfy the requirement of Lemma 3.2 that $\left\lVert\Delta_{t}\right\rVert\in O(\kappa^{-2}\left\lVert W_{t}-W_{t-1}\right\rVert)$, we simply need $\left\lVert\beta W_{t-1}^{-1}(W_{t}-W_{t-1})\right\rVert\in O(\kappa^{-2})$. This is satisfied when $W_{t}\approx W_{t-1}$ and $W_{t-1}$ is reasonably non-singular. A sufficient condition for this is that the iterates are sufficiently close to the ground-truth solution $W^{*}$ and $W^{*}$ is non-singular, which explains the conditions of Theorem 3.1.

We now prove Theorem 3.1 formally.

To prove the main result, it suffices to show that $\left\lVert\Delta_{t}\right\rVert\leq\alpha\left\lVert W_{t}-W_{t-1}\right\rVert$ for all $t\in[T-1]$. With this claim Lemma 3.2 immediately gives a convergence rate of

This gives the desired result $\ell(W_{T})-\ell(W^{*})\in\exp(-\Omega(\tfrac{T}{\sqrt{\kappa}}+\log(\tfrac{T}{\sqrt{\kappa}})))$ as

where $C=\left(\tfrac{L{}}{2}+\tfrac{8\sqrt{\kappa}\mu}{4(4\sqrt{\kappa}-3)}\right)\left\lVert W_{1}-W^{*}\right\rVert^{2}\in\exp(-\Omega(\log(\tfrac{T}{\sqrt{\kappa}})))$.

Thus, we need to show that $\left\lVert\Delta_{t}\right\rVert\leq\alpha\left\lVert W_{t}-W_{t-1}\right\rVert$ for all $t\in[T-1]$. For this, we use the following claim, where $\eta\coloneqq\sqrt{\frac{16\delta}{\mu}}$:

We can now complete the proof by showing the following claim via induction on $t$:

Note that $W_{0}=V_{0}$. We have

which implies that $\left\lVert W_{0}-W^{*}\right\rVert\leq\sqrt{\frac{2\delta}{\mu}}\leq\eta$. Furthermore, since $W_{1}=W_{0}-\frac{1}{L}\nabla\ell(W_{0})$ and $\ell$ is $L$-smooth, we have $\ell(W_{1})\leq\ell(W_{0})$, which implies, exactly as above, that $\left\lVert W_{1}-W^{*}\right\rVert\leq\eta$. Thus, by Claim 3.3, we conclude that $\left\lVert\Delta_{1}\right\rVert\leq\alpha\left\lVert W_{1}-W_{0}\right\rVert$.

Fixing $t\in[2,T-1]$, and assume, as our inductive hypothesis, that $\left\lVert\Delta_{\tau}\right\rVert\leq\alpha\left\lVert W_{\tau}-W_{\tau-1}\right\rVert$ and $\left\lVert W_{\tau}-W^{*}\right\rVert\leq\eta$ for all $\tau\in[t-1]$. Due to this inductive hypothesis, we can invoke Lemma 3.2 to observe that for all $\tau\in[T-1]$:

Thus, we have $\|W_{t}-W^{*}\|\leq\sqrt{\frac{16\delta}{\mu}}=\eta$. Since $\left\lVert W_{t-1}-W^{*}\right\rVert\leq\eta$ due to our induction hypothesis, we can now use Claim 3.3 to conclude that $\left\lVert\Delta_{t}\right\rVert\leq\alpha\left\lVert W_{t}-W_{t-1}\right\rVert$. This completes the inductive proof.

∎

∎

Theorem 3.1 presents a local convergence result: it assumes that the initial two layers put the network in the vicinity of the optimal solution $W^{*}$. We can also provide a global convergence result by using a small warmup period phase to stacking—that is, by training the first few (only a constant number) stages of the deep linear network without stacking.

We now turn to proving Lemma 3.2.