Implicit biases in multitask and continual learning from a backward error analysis perspective

Overparameterized neural networks trained to interpolate are able to generalize surprisingly well in spite of the high complexity of their hypothesis space [1]. One key concept to understand this phenomenon is that of implicit regularization or implicit training biases, which are quantities that are not explicitly regularized in the loss during training but by other mechanisms, guiding the network toward simpler solutions [2, 3]. Several groups [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] have recently used a technique called Backward Error Analysis (BEA) to compute implicit biases as a measure of the discrepancy between an optimizer iterates and the solutions of Gradient Flow (GF), which are the unique continuous paths of steepest descent. Because of its flexibility BEA has been used to compute optimizer implicit biases in many settings: Gradient Descent (GD) [4], Stochastic Gradient Descent (SGD) [5, 10], Momentum [16], Adam and RMSProp [14], GAN’s [7, 12, 11], and diffusion processes [9], among others [6, 8, 13, 15].

Our contribution: We add to this body of work by computing implicit biases in multitask learning [17] and continual learning [18, 19] settings optimized with SGD. In both cases, the output of BEA is a modified loss implicitly minimized by the optimizer. It consists of the original loss plus additional terms, which can be split in two parts: 1) a beneficial implicit flatness regularizer proportional to the learning rate and already observed in single-task learning (in [4, 5] using BEA as well as with other approaches in [20, 21, 22]), and 2) a conflict term, due to the presence of several tasks, and which can be detrimental to both convergence and implicit flatness regularization. In multitask, the conflict term is the inner product between the task gradients, creating an implicit propensity in the learning dynamics to seek misaligned task gradients, which is known to be detrimental and needs to be mitigated [23, 24, 25]. In continual learning, the conflict term is the Lie bracket [26] between the task gradients whose non-vanishing may possibly be related to catastrophic forgetting [18, 19] where the performance of previous tasks degrades as new ones are learned. We hope to foster interest on Lie brackets in optimization, which is one of the basic tools in differential geometry [26].

To illustrate BEA, we now derive an implicit bias of SGD after a single mini-batch update by adapting the derivation for full-batch GD from [4]. Consider the loss $L_{B}(\theta)$ computed on a batch of data $B$ from a dataset $D$. At a given step, the SGD iterate is $\theta^{\prime}=\theta-h\nabla L_{B}(\theta)$, while the solution of GF (which exactly minimizes the batch loss) is the curve $\theta(t)$ solving the differential equation $\dot{\theta}(t)=-\nabla L_{B}(\theta(t))$ with $\theta(0)=\theta$. The discretization drift is the difference between the two, i.e., $\|\theta(h)-\theta^{\prime}\|$, and it is of order $\mathcal{O}(h^{2})$ for GD (see [27] for details). BEA proposes to compute a modified equation, in the form of GF plus corrections in terms of powers of the learning rate

so that the solution $\tilde{\theta}(t)$ of the modified equation exactly coincides with the GD iterate: $\theta^{\prime}=\tilde{\theta}(h)$. The idea of BEA is to use the continuous modified equation to analyze the discrete optimizer. Note that if we truncate the modified equation at order $n$ (i.e. removing the terms of order $h^{n}$ and higher), the discretization drift becomes of order only $\mathcal{O}(h^{n+1})$. Let us compute $f_{1}$ following [4]: First, we expand the solution of the modified equation in a Taylor series:

For $\tilde{\theta}(h)$ to coincide with $\theta^{\prime}=\theta-h\nabla L_{B}(\theta)$, we need all the terms in power of $h^{2}$ or higher to vanish. This gives us the first correction (and recursively the higher ones too as needed; see [27, 4, 6, 13]): $f_{1}(\theta)=-\nabla\left(\frac{1}{4}\|\nabla L_{B}(\theta)\|^{2}\right)$. Therefore the gradient update $\theta^{\prime}$ follows a GF with drift only of order $\mathcal{O}(h^{3})$ but for a modified loss $\widetilde{L}_{B}$, since the modified equation is of the form

The second term in this modified loss is a flatness bias, called Implicit Gradient Regularization (IGR) in [4], which prefers optimization paths with shallower slopes (i.e. lower gradients) guiding the trajectory toward flatter regions, very much in line with other flatness biases in SGD found by different means [21, 22, 28, 29]. We now turn to applying BEA to multitask and continual settings.

Multitask learning trains a neural network jointly on several tasks hoping that knowledge gained from each task will transfer to the other tasks, helping generalization, and useful in case of data scarcity [30, 31]. However, it has been observed at times that learning multiple tasks at once can be detrimental, a circumstance attributed to the loss gradients for each task being misaligned [23, 24, 25, 32]. BEA shed some theoretical light on this since the implicit multitask dynamics of SGD given by its modified equation has a term, the conflict term, with propensity to guide the training in regions with misaligned gradients. In terms of losses, the simplest multitask setting corresponds to having two losses $L_{1}(\phi_{1},\theta)$ for the first task and $L_{2}(\phi_{2},\theta)$ for the second task. The parameters $\theta$ correspond to the part of the network that is shared between the two tasks, while $\phi_{1}$ and $\phi_{2}$ are the parameters corresponding to the two task heads. The training setup is to devise a global loss

consisting on a weighted average of the two losses and update the network parameters with

Note that the update above can be considered either as a full-batch GD update as in [4], or a single-step batch update in SGD within an epoch as in Section 2.

Continual learning is concerned with learning from a data distribution that is changing over time with tasks corresponding to locally stationary phases of the evolution [18, 19]. One of its major issues is catastrophic forgetting, when the updates from latter tasks degrade the performance on earlier tasks. While catastrophic forgetting is an issue for all modern approaches [18, 19, 33], its causes remain unclear. As we will see, the BEA modified equation for continual learning may help shed some new light on this issue. Namely, consider the continual learning setting where we perform two successive SGD updates $\theta_{1}=\theta_{0}-h\nabla L_{1}(\theta_{0})$ and $\theta_{2}=\theta_{1}-h\nabla L_{2}(\theta_{1})$, with the two losses computed on two successive batches from a changing data distribution. Using BEA, we want to first compute a modified loss whose continuous minimization approximates well the two successive updates. Then we want to identify possibly detrimental terms in the modified equation that may be responsible for a decreased performance on the first batch by the second update. It turns out that such a detrimental term pops up, controlled by the Lie bracket of the two batch gradients:

Lie brackets are fundamental tools in differential geometry [26]. They help quantify how flows intertwine. For instance, if the Lie bracket between loss gradients for different tasks vanish, i.e., $[\nabla L_{1},\nabla L_{2}]=0$, this implies that their gradient flows commute: Following the gradient flow of first $L_{1}$ and then $L_{2}$ yields the same result as the reverse [26], with their flows somehow spanning "non-interacting" subspaces. The next theorem states that when this happens two consecutive SGD updates as above can be approximated by GF for a modified loss of the form:

where the IGR terms encourage the learning trajectory toward flatter regions for each task. Note that flatness preservation between tasks seems helpful to combat catastrophic forgetting [34, 35]. However, when $[\nabla L_{1},\nabla L_{2}]\neq 0$, a term of order $h$ in the modified equation (theorem below) and proportional to the Lie bracket can potentially disrupt that implicit flatness regularization induced by the modified loss above. Since it is the only term of order $h$ that can do so, we conjecture that the non-vanishing of the Lie bracket between loss gradients pertaining to different tasks may be linked to catastrophic forgetting in continual learning. The following theorem gives an exact description of how this Lie bracket affects the implicit gradient regularization dynamics:

We computed implicit biases in multitask and continual learning optimized with SGD using backward error analysis. These biases are local, measuring the discrepancy between one step of SGD and gradient flow on the batch loss. In both cases we found a beneficial flatness bias proportional to the learning rate and preferring smaller slopes on the loss surface for each task along the learning trajectories similar to single-task supervised learning [4, 5]. We also found a detrimental implicit bias in both cases (due to the presence of several tasks and which we called conflict term) that has the potential to steer the learning dynamics away from the flatter regions with higher generalization power. For multitask learning, the detrimental implicit bias is controlled by the inner product of the task loss-gradients $\langle\nabla L_{1},\nabla L_{2}\rangle$, which is a known key quantity in multitask learning already (e.g., [23, 24, 25, 32]). For continual learning the detrimental bias is a new quantity, the Lie bracket $[\nabla L_{1},\nabla L_{2}]$ between the task loss-gradients measuring how much their respective gradient flows span independent regions of the parameter space. Despite their wide use in many areas of mathematics, Lie brackets are new to deep learning optimization to the best of our knowledge. We hope this work will help foster the use of backward error analysis in deep learning, and serve as a theoretical motivation to devise methods relying on Lie brackets in continual learning.

We would like to thank Mihaela Rosca, Maxim Neumann, Michael Munn, Javier Gonzalvo, David Barrett, and Hossein Mobahi for helpful discussions and feedback. We would also like to thank Patrick Cole for his support.