Truncated Back-propagation for Bilevel Optimization

INTRODUCTION

Bilevel optimization has been recently revisited as a theoretical framework for designing and analyzing algorithms for hyperparameter optimization and meta learning . Mathematically, these problems can be formulated as a stochastic optimization problem with an equality constraint (see Section 1.1):

where $w$ and $\lambda$ are the parameter and the hyperparameter, $F$ and $f_{S}$ are the expected and the sampled upper-level objective, $g_{S}$ is the sampled lower-level objective, and $S$ is a random variable called the context. The notation $\approx_{\lambda}$ means that $\hat{w}_{S}^{*}(\lambda)$ equals the unique return value of a prespecified iterative algorithm (e.g. gradient descent) that approximately finds a local minimum of $g_{S}$. This algorithm is part of the problem definition and can also be parametrized by $\lambda$ (e.g. step size). The motivation to explicitly consider the approximate solution $\hat{w}_{S}^{*}(\lambda)$ rather than an exact minimizer $w_{S}^{*}$ of $g_{S}$ is that $w_{S}^{*}$ is usually not available in closed form. This setup enables $\lambda$ to account for the imperfections of the lower-level optimization algorithm.

Solving the bilevel optimization problem in (1) is challenging due to the complicated dependency of the upper-level problem on $\lambda$ induced by $\hat{w}_{S}^{*}(\lambda)$. This difficulty is further aggravated when $\lambda$ and $w$ are high-dimensional, precluding the use of black-box optimization techniques such as grid/random search and Bayesian optimization .

Recently, first-order bilevel optimization techniques have been revisited to solve these problems. These methods rely on an estimate of the Jacobian $\nabla_{\lambda}\hat{w}_{S}^{*}(\lambda)$ to optimize $\lambda$. Pedregosa and Gould et al. assume that $\hat{w}_{S}^{*}(\lambda)=w_{S}^{*}$ and compute $\nabla_{\lambda}\hat{w}_{S}^{*}(\lambda)$ by implicit differentiation. By contrast, Maclaurin et al. and Franceschi et al. treat the iterative optimization algorithm in the lower-level problem as a dynamical system, and compute $\nabla_{\lambda}\hat{w}_{S}^{*}(\lambda)$ by automatic differentiation through the dynamical system. In comparison, the latter approach is less sensitive to the optimality of $\hat{w}_{S}^{*}(\lambda)$ and can also learn hyperparameters that control the lower-level optimization process (e.g. step size). However, due to superlinear time or space complexity (see Section 2.2), neither of these methods is applicable when both $\lambda$ and $w$ are high-dimensional .

Few-step reverse-mode automatic differentiation and few-step forward-mode automatic differentiation have recently been proposed as heuristics to address this issue. By ignoring long-term dependencies, the time and space complexities to compute approximate gradients can be greatly reduced. While exciting empirical results have been reported, the theoretical properties of these methods remain unclear.

In this paper, we study the theoretical properties of these truncated back-propagation approaches. We show that, when the lower-level problem is locally strongly convex around $\hat{w}_{S}^{*}(\lambda)$, on-average convergence to an $\epsilon$-approximate stationary point is guaranteed by $O(\log 1/\epsilon)$-step truncated back-propagation. We also identify additional problem structures for which asymptotic convergence to an exact stationary point is guaranteed. Empirically, we verify the utility of this strategy for hyperparameter optimization and meta learning tasks. We find that, compared to optimization with full back-propagation, optimization with truncated back-propagation usually shows competitive performance while requiring half as much computation time and significantly less memory.

The goal of hyperparameter optimization is to find hyperparameters $\lambda$ for an optimization problem $P$ such that the approximate solution $\hat{w}^{*}(\lambda)$ of $P$ has low cost $c(\hat{w}^{*}(\lambda))$ for some cost function $c$. In general, $\lambda$ can parametrize both the objective of $P$ and the algorithm used to solve $P$. This setup is a special case of the bilevel optimization problem (1) where the upper-level objective $c$ does not depend directly on $\lambda$. In contrast to meta learning (discussed below), $c$ can be deterministic . See Section 4.2 for examples.

Many low-dimensional problems, such as choosing the learning rate and regularization constant for training neural networks, can be effectively solved with grid search. However, problems with thousands of hyperparameters are increasingly common, for which gradient-based methods are more appropriate .

Another important application of bilevel optimization, meta learning (or learning-to-learn) uses statistical learning to optimize an algorithm $\mathcal{A}_{\lambda}$ over a distribution of tasks $\mathcal{T}$ and contexts $S$:

It treats $\mathcal{A}_{\lambda}$ as a parametric function, with hyperparameter $\lambda$, that takes task-specific context information $S$ as input and outputs a decision $\mathcal{A}_{\lambda}(S)$. The goal of meta learning is to optimize the algorithm’s performance $c_{\mathcal{T}}$ (e.g. the generalization error) across tasks $\mathcal{T}$ through empirical observations. This general setup subsumes multiple problems commonly encountered in the machine learning literature, such as multi-task learning and few-shot learning .

Bilevel optimization emerges from meta learning when the algorithm computes $\mathcal{A}_{\lambda}(S)$ by internally solving a lower-level minimization problem with variable $w$. The motivation to use this class of algorithms is that the lower-level problem can be designed so that, even for tasks $\mathcal{T}$ distant from the training set, $\mathcal{A}_{\lambda}$ falls back upon a sensible optimization-based approach . By contrast, treating $\mathcal{A}_{\lambda}$ as a general function approximator relies on the availability of a large amount of meta training data .

In other words, the decision is $\mathcal{A}_{\lambda}(S)=(\hat{w}_{S}^{*}(\lambda),\lambda)$ where $\hat{w}_{S}^{*}(\lambda)$ is an approximate minimizer of some function $g_{S}(w,\lambda)$. Therefore, we can identify

BILEVEL OPTIMIZATION

2 Computing the hypergradient

Like , we treat the iterative optimization algorithm that solves the lower-level problem as a dynamical system. Given an initial condition $w_{0}=\Xi_{0}(\lambda)$ at $t=0$, the update rule can be written asFor notational simplicity, we consider the case where $w_{t}$ is the state of (4); our derivation can be easily generalized to include other internal states, e.g. momentum.

in which $\Xi_{t}$ defines the transition and and $T$ is the number iterations performed. For example, in gradient descent, $\Xi_{t+1}(w_{t},\lambda)=w_{t}-\gamma_{t}(\lambda)\nabla_{w}g(w_{t},\lambda)$, where $\gamma_{t}(\lambda)$ is the step size.

By unrolling the iterative update scheme (4) as a computational graph, we can view $\hat{w}^{*}$ as a function of $\lambda$ and compute the required derivative $\text{d}_{\lambda}f$ . Specifically, it can be shown by the chain ruleNote that this assumes $g$ is twice differentiable.

where $A_{t+1}=\nabla_{w_{t}}\Xi_{t+1}(w_{t},\lambda)$, $B_{t+1}=\nabla_{\lambda}\Xi_{t+1}(w_{t},\lambda)$ for $t\geq 0$, and $B_{0}=\text{d}_{\lambda}\Xi_{0}(\lambda)$.

The computation of (5) can be implemented either in reverse mode or forward mode . Reverse-mode differentiation (RMD) computes (5) by back-propagation:

and finally $\text{d}_{\lambda}f=h_{-1}$. Forward-mode differentiation (FMD) computes (5) by forward propagation:

TRUNCATED BACK-PROPAGATION

In this paper, we investigate approximating (5) with partial sums, which was previously proposed as a heuristic for bilevel optimization ( Eq. 3, Eq. 2). Formally, we perform $K$-step truncated back-propagation ($K$-RMD) and use the intermediate variable $h_{T-K}$ to construct an approximate gradient:

This approach requires storing only the last $K$ iterates $w_{t}$, and it also saves computation time. Note that $K$-RMD can be combined with checkpointing for further savings, although we do not investigate this.

We first establish some intuitions about why using $K$-RMD to optimize $\lambda$ is reasonable. While building up an approximate gradient by truncating back-propagation in general optimization problems can lead to large bias, the bilevel optimization problem in (1) has some nice structure. Here we show that if the lower-level objective $g$ is locally strongly convex around $\hat{w}^{*}$, then the bias of $h_{T-K}$ can be exponentially small in $K$. That is, choosing a small $K$ would suffice to give a good gradient approximation in finite precision. The proof is given in Appendix A.

Assume $g$ is $\beta$-smooth, twice differentiable, and locally $\alpha$-strongly convex in $w$ around $\{w_{T-K-1},\dots,w_{T}\}$. Let $\Xi_{t+1}(w_{t},\lambda)=w_{t}-\gamma\nabla_{w}g(w_{t},\lambda)$. For $\gamma\leq\frac{1}{\beta}$, it holds

where $M_{B}=\max_{t\in\{0,\dots,T-K\}}\|B_{t}\|$. In particular, if $g$ is globally $\alpha$-strongly convex, then

Note $0\leq(1-\gamma\alpha)<1$ since $\gamma\leq\frac{1}{\beta}\leq\frac{1}{\alpha}$. Therefore, Proposition 3.1 says that if $\hat{w}^{*}$ converges to the neighborhood of a strict local minimum of the lower-level optimization, then the bias of using the approximate gradient of $K$-RMD decays exponentially in $K$. This exponentially decaying property is the main reason why using $h_{T-K}$ to update the hyperparameter $\lambda$ works.

Next we show that, when the lower-level problem $g$ is second-order continuously differentiable, $-h_{T-K}$ actually is a sufficient descent direction. This is a much stronger property than the small bias shown in Proposition 3.1, and it is critical in order to prove convergence to exact stationary points (cf. Theorem 3.4). To build intuition, here we consider a simpler problem where $g$ is globally strongly convex and $\nabla_{\lambda}f=0$. These assumptions will be relaxed in the next subsection.

Let $g$ be globally strongly convex and $\nabla_{\lambda}f=0$. Assume $g$ is second-order continuously differentiable and $B_{t}$ has full column rank for all $t$. Let $\Xi_{t+1}(w_{t},\lambda)=w_{t}-\gamma\nabla_{w}g(w_{t},\lambda)$. For all $K\geq 1$, with $T$ large enough and $\gamma$ small enough, there exists $c>0$, s.t. $h_{T-K}^{\top}\text{d}_{\lambda}f\geq c\|\nabla_{\hat{w}^{*}}f\|^{2}$. This implies $h_{T-K}$ is a sufficient descent direction, i.e. $h_{T-K}^{\top}\text{d}_{\lambda}f\geq\Omega(\|\text{d}_{\lambda}f\|^{2})$.

The full proof of this non-trivial result is given in Appendix B. Here we provide some ideas about why it is true. First, by Proposition 3.1, we know the bias decays exponentially. However, this alone is not sufficient to show that $-h_{T-K}$ is a sufficient descent direction. To show the desired result, Lemma 3.2 relies on the assumption that $g$ is second-order continuously differentiable and the fact that using gradient descent to optimize a well-conditioned function has linear convergence . These two new structural properties further reduce the bias in Proposition 3.1 and lead to Lemma 3.2. Here the full rank assumption for $B_{t}$ is made to simplify the proof. We conjecture that this condition can be relaxed when $K>1$. We leave this to future work.

2 Convergence

With these insights, we analyze the convergence of bilevel optimization with truncated back-propagation. Using Proposition 3.1, we can immediately deduce that optimizing $\lambda$ with $h_{T-K}$ converges on-average to an $\epsilon$-approximate stationary point. Let $\nabla F(\lambda_{\tau})$ denote the hypergradient in the $\tau$th iteration.

Suppose $F$ is smooth and bounded below, and suppose there is $\epsilon<\infty$ such that $\|h_{T-K}-\text{d}_{\lambda}f\|\leq\epsilon$. Using $h_{T-K}$ as a stochastic first-order oracle with a decaying step size $\eta_{\tau}=O(1/\sqrt{\tau})$ to update $\lambda$ with gradient descent, it follows after $R$ iterations,

That is, under the assumptions in Proposition 3.1, learning with $h_{T-K}$ converges to an $\epsilon$-approximate stationary point, where $\epsilon=O((1-\gamma\alpha)^{-K})$.

We see that the bias becomes small as $K$ increases. As a result, it is sufficient to perform $K$-step truncated back-propagation with $K=O(\log 1/\epsilon)$ to update $\lambda$.

Next, using Lemma 3.2, we show that the bias term in Theorem 3.3 can be removed if the problem is more structured. As promised, we relax the simplifications made in Lemma 3.2 into assumptions 2 and 3 below and only assume $g$ is locally strongly convex.

Under the assumptions in Proposition 3.1 and Theorem 3.3, if in addition

$g$ is second-order continuously differentiable

$B_{t}$ has full column rank around $w_{T}$

$\nabla_{\lambda}f^{\top}(\text{d}_{\lambda}f+h_{T-K}-\nabla_{\lambda}f)\geq\Omega(\|\nabla_{\lambda}f\|^{2})$

the problem is deterministic (i.e. $F=f$)

then for all $K\geq 1$, with $T$ large enough and $\gamma$ small enough, the limit point is an exact stationary point, i.e. $\lim_{\tau\to\infty}\|\nabla F(\lambda_{\tau})\|=0$.

Theorem 3.4 shows that if the partial derivative $\nabla_{\lambda}f$ does not interfere strongly with the partial derivative computed through back-propagating the lower-level optimization procedure (assumption 3), then optimizing $\lambda$ with $h_{T-K}$ converges to an exact stationary point. This is a very strong result for an interesting special case. It shows that even with one-step back-propagation $h_{T-1}$, updating $\lambda$ can converge to a stationary point.

This non-interference assumption unfortunately is necessary; otherwise, truncating the full RMD leads to constant bias, as we show below (proved in Appendix E).

There is a problem, satisfying all but assumption 3 in Theorem 3.4, such that optimizing $\lambda$ with $h_{T-K}$ does not converge to a stationary point.

Note however that the non-interference assumption is satisfied when $\nabla_{\lambda}f=0$, i.e. when the upper-level problem does not directly depend on the hyperparameter. This is the case for many practical applications: e.g. hyperparameter optimization, meta-learning regularization models, image desnosing , data hyper-cleaning , and task interaction .

3 Relationship with implicit differentiation

The gradient estimate $h_{T-K}$ is related to implicit differentiation, which is a classical first-order approach to solving bilevel optimization problems . Assume $g$ is second-order continuously differentiable and that its optimal solution uniquely exists such that $w^{*}=w^{*}(\lambda)$. By the implicit function theorem , the total derivative of $f$ with respect to $\lambda$ can be written as

Here we show that, in the limit where $\hat{w}^{*}$ converges to $w^{*}$, $h_{T-K}$ can be viewed as approximating the matrix inverse in (11) with an order-$K$ Taylor series. This can be seen from the next proposition.

Under the assumptions in Proposition 3.1, suppose $w_{t}$ converges to a stationary point $w^{*}$. Let $A_{\infty}=\lim_{t\to\infty}A_{t}$ and $B_{\infty}=\lim_{t\to\infty}B_{t}$. For $\gamma<\frac{1}{\beta}$, it satisfies that

By Proposition 3.6, we can write $\text{d}_{\lambda}f$ in (11) as

That is, $h_{T-K}$ captures the first $K$ terms in the Taylor series, and the residue term has an upper bound as in Proposition 3.1.

Given this connection, we can compare the use of $h_{T-K}$ and approximating (11) using $K$ steps of conjugate gradient descent for high-dimensional problems . First, both approaches require local strong-convexity to ensure a good approximation. Specifically, let $\kappa=\frac{\beta}{\alpha}>0$ locally around the limit. Using $h_{T-K}$ has a bias in $O((1-\frac{1}{\kappa})^{K})$, whereas using (11) and inverting the matrix with $K$ iterations of conjugate gradient has a bias in $O((1-\frac{1}{\sqrt{\kappa}})^{K})$ . Therefore, when $w^{*}$ is available, solving (11) with conjugate gradient descent is preferable. However, in practice, this is hardly true. When an approximate solution $\hat{w}^{*}$ to the lower-level problem is used, adopting (11) has no control on the approximate error, nor does it necessarily yield a descent direction. On the contrary, $h_{T-K}$ is based on Proposition 3.1, which uses a weaker assumption and does not require the convergence of $w_{t}$ to a stationary point. Truncated back-propagation can also optimize the hyperparameters that control the lower-level optimization process, which the implicit differentiation approach cannot do.

EXPERIMENTS

This deterministic problem satisfies all of the assumptions in the previous section, particularly those of Theorem 3.4: $g$ is $1$-smooth and $\frac{1}{2}$-strongly convex, with

and $B_{0}=0$. Although $f$ is somewhat complicated, with many saddle points, it satisfies the non-interference assumption because $\nabla_{\lambda}f=0$.

Figure 1 visualizes Proposition 3.1 by plotting the approximation error $\|h_{T-K}-\text{d}_{\lambda}f\|$ and the theoretical bound $\frac{(1-\gamma\alpha)^{K}}{\gamma\alpha}\|\nabla_{\hat{w}^{*}}f\|M_{B}$ at $\lambda=(1,1)$. For this problem, $\alpha=\frac{1}{2}$, $M_{B}=\|\gamma G\|=\gamma$, and $\nabla_{\hat{w}^{*}}f$ can be found analytically from $\hat{w}^{*}=Cw_{0}+(I-C)\lambda$, where $C=(I-\gamma G)^{T}$. Figure 4 (left) plots the iterates $\lambda_{\tau}$ when optimizing $f$ using $1$-RMD and a decaying meta-learning rate $\eta_{\tau}=\frac{\eta_{0}}{\sqrt{\tau}}$. Because $\|h_{T-K}\|$ varies widely with $K$, we tune $\eta_{0}$ to ensure that the first update $\eta_{1}h_{T-K}(\lambda_{1})$ has norm $0.6$. In comparison with the true gradient $\text{d}_{\lambda}f$ at these points, we see that $h_{T-1}$ is indeed a descent direction. Figure 2 (left) visualizes this in a different way, by plotting $h_{T-K}^{\top}\text{d}_{\lambda}f/\|\text{d}_{\lambda}f\|^{2}$ for various $K$ at each point $\lambda_{\tau}$ along the $K=1$ trajectory. By Lemma 3.2, this ratio stays well away from zero.

To demonstrate the biased convergence of Theorem 3.3, we break assumption 3 of Theorem 3.4 by changing the upper objective to $\widetilde{f}(\hat{w}^{*},\lambda):=f(\hat{w}^{*},\lambda)+5\|\lambda-(1,0)\|^{2}$ so that $\nabla_{\lambda}\widetilde{f}\neq 0$. The guarantee of Lemma 3.2 no longer applies, and we see in Figure 2 (right) that $h_{T-K}^{\top}\text{d}_{\lambda}f/\|\text{d}_{\lambda}f\|^{2}$ can become negative. Indeed, Figure 3 shows that optimizing $\widetilde{f}$ with $h_{T-1}$ converges to a suboptimal point. However, it also shows that using larger $K$ rapidly decreases the bias.

For the original objective $f$, Theorem 3.4 guarantees exact convergence. Figure 4 shows optimization trajectories for various $K$, and a log-scale plot of their convergence rates. Note that, because the lower-level problem cannot be perfectly solved within $T$ steps, the optimal $\lambda$ is offset from the origin. Truncated back-propagation can handle this, but it breaks the assumptions required by the implicit differentiation approach to bilevel optimization.

2 Hyperparameter optimization problems

We optimize the lower-level problem $g$ through $T=100$ steps of gradient descent with $\gamma=1$ and consider how adjusting $K$ changes the performance of $K$-RMD. See Appendix G.1 for more experimental setup. Our hypothesis is that $K$-RMD for small $K$ works almost as well as full RMD in terms of validation and test accuracy, while requiring less time and far less memory. We also hypothesize that $K$-RMD does almost as well as full RMD in identifying which samples were corrupted . Because our formulation of the problem is unconstrained, the weights $\sigma(\lambda_{i})$ are never exactly zero. However, we can calculate an F1 score by setting a threshold on $\lambda$: if $\sigma(\lambda_{i})<\sigma(-3)\approx 0.047$, then the hyper-cleaner has marked example $i$ as corrupted.F1 scores for other choices of the threshold were very similar. See Appendix G.1 for details.

Table 2 reports these metrics for various $K$. We see that $1$-RMD is somewhat worse than the others, and that validation loss (the outer objective $f$) decreases with $K$ more quickly than generalization error. The F1 score is already maximized at $K=5$. These preliminary results indicate that in situations with limited memory, $K$-RMD for small $K$ (e.g. $K=5$) may be a reasonable fallback: it achieves results close to full backprop, and it runs about twice as fast.

From a theoretical optimization perspective, we wonder whether $K$-RMD converges to a stationary point of $f$. Data hypercleaning satisfies all of the assumptions of Theorem 3.4 except that $B_{t}$ is not full column rank (since $M<N$). In particular, the validation loss $f$ is deterministic and satisfies $\nabla_{\lambda}f=0$. Figure 5 plots the norm of the true gradient $\text{d}_{\lambda}f$ on a log scale at the $K$-RMD iterates for various $K$. We see that, despite satisfying almost all assumptions, this problem exhibits biased convergence. The limit of $\|\text{d}_{\lambda}f\|$ decreases slowly with $K$, but recall from Table 2 that practical metrics improve more quickly.

2.2 Task interaction

We next consider the problem of multitask learning . Similar to , we formulate this as a hyperparameter optimization problem as follows. The lower-level objective $g(w,\{C,\rho\})$ learns $V$ different linear models with parameter set $w=\{w_{v}\}_{v=1}^{V}$:

where $l(w)$ is the training loss of the multi-class linear logistic regression model, $\rho$ is a regularization constant, and $C$ is a nonnegative, symmetric hyperparameter matrix that encodes the similarity between each pair of tasks. After $100$ iterations of gradient descent with learning rate $0.1$, this yields $\hat{w}^{*}$. The upper-level objective $c(\hat{w}^{*})$ estimates the linear regression loss of the learned model $\hat{w}^{*}$ on a validation set. Presumably, this will be improved by tuning $C$ to reflect the true similarities between the tasks. The tasks that we consider are image recognition trained on very small subsets of the datasets CIFAR-$10$ and CIFAR-$100$. See Appendix G.2 for more details.

From an optimization standpoint, we are most interested in the upper-level loss on the validation set, since that is what is directly optimized, and its value is a good indication of the performance of the inexact gradient. Figure 6 plots this learning curve along with two other metrics of theoretical interest: norm of the true gradient, and cosine similarity between the true and approximate gradients. In CIFAR100, the validation error and gradient norm plots show that $K$-RMD converges to an approximate stationary point with a bias that rapidly decreases as $K$ increases, agreeing with Proposition 3.1. Also, we find that negative values exist in the cosine similarity of $1$-RMD, which implies that not all the assumptions in Theorem 3.4 hold for this problem (e.g. $B_{t}$ might not be full rank, or the the inner problem might not be locally strong convex around $\hat{w}^{*}$.) In CIFAR10, some unusal behavior happens. For $K>1$, the truncated gradient and the full gradient directions eventually become almost the same. We believe this is a very interesting observation but beyond the scope of the paper to explain.

In Table 3, we report the testing accuracy over 10 trials. While in general increasing the number of back-propagation steps improves accuracy, the gaps are small. A thorough investigation of the relationship between convergence and generalization is an interesting open question of both theoretical and practical importance.

3 Meta-learning: One-shot classification

The aim of this experiment is to evaluate the performance of truncated back-propagation in multi-task, stochastic optimization problems. We consider in particular the one-shot classification problem , where each task $\mathcal{T}$ is a $k$-way classification problem and the goal is learn a hyperparameter $\lambda$ such that each task can be solved with few training samples.

In each hyper-iteration, we sample a task, a training set, and a validation set as follows: First, $k$ classes are randomly chosen from a pool of classes to define the sampled task $\mathcal{T}$. Then the training set $S=\{(x_{i},y_{i})\}_{i=1}^{k}$ is created by randomly drawing one training example $(x_{i},y_{i})$ from each of the $k$ classes. The validation set $Q$ is constructed similarly, but with more examples from each class. The lower-level objective $g_{S}(w,\lambda)$ is

where $l(\cdot,\cdot)$ is the $k$-way cross-entropy loss, and $nn(\cdot;w,\lambda)$ is a deep neural network parametrized by $w=\{w_{1},\dots,w_{V}\}$ and optionally hyperparameter $\lambda$. To prevent overfitting in the lower-level optimization, we regularize each parameter $w_{j}$ to be close to center $c_{j}$ with weight $\rho_{j}>0$. Both $c_{j}$ and $\rho_{j}$ are hyperparameters, as well as the inner learning rate $\gamma$. The upper-level objective is the loss of the trained network on the sampled validation set $Q$. In contrast to other experiments, this is a stochastic optimization problem. Also, $\mathcal{A}_{\lambda}(S)(x_{i})=nn(x_{i};\hat{w}^{*},\lambda)$ depends directly on the hyperparameter $\lambda$, in addition to the indirect dependence through $\hat{w}^{*}$ (i.e. $\nabla_{\lambda}f\neq 0$).

We use the Omniglot dataset and a similar neural network as used in with small modifications. Please refer to Appendix G.3 for more details about the model and the data splits. We set $T=50$ and optimize over the hyperparameter $\lambda=\{\lambda_{l_{1}},\lambda_{l_{2}},c,\rho,\gamma\}$. The average accuracy of each model is evaluated over $120$ randomly sampled training and validation sets from the meta-testing dataset. For comparison, we also try using full RMD with a very short horizon $T=1$, which is common in recent work on few-shot learning .

The statistics are shown in Table 4 and the learning curves in Figure 7. In addition to saving memory, all truncated methods are faster than full RMD, sometimes even five times faster. These results suggest that running few-step back-propagation with more hyper-iterations can be more efficient than the full RMD. To support this hypothesis, we also ran $1$-RMD and $10$-RMD for an especially large number of hyper-iterations ($15$k). Even with this many hyper-iterations, the total runtime is less than full RMD with $5000$ iterations, and the results are significantly improved. We also find that while using a short horizon ($T=1$) is faster, it achieves a lower accuracy at the same number of iterations.

CONCLUSION

We analyze $K$-RMD, a first-order heuristic for solving bilevel optimization problems when the lower-level optimization is itself approximated in an iterative way. We show that $K$-RMD is a valid alternative to full RMD from both theoretical and empirical standpoints. Theoretically, we identify sufficient conditions for which the hyperparameters converge to an approximate or exact stationary point of the upper-level objective. The key observation is that when $\hat{w}^{*}$ is near a strict local minimum of the lower-level objective, gradient approximation error decays exponentially with reverse depth. Empirically, we explore the properties of this optimization method with four proof-of-concept experiments. We find that although exact convergence appears to be uncommon in practice, the performance of $K$-RMD is close to full RMD in terms of application-specific metrics (such as generalization error). It is also roughly twice as fast. These results suggest that in hyperparameter optimization or meta learning applications with memory constraints, truncated back-propagation is a reasonable choice.

Our experiments use a modest number of parameters $M$, hyperparameters $N$, and horizon length $T$. This is because we need to be able to calculate both $K$-RMD and full RMD in order to compare their performance. One promising direction for future research is to use $K$-RMD for bilevel optimization problems that require powerful function approximators at both levels of optimization. Truncated RMD makes this approach feasible and enables comparing bilevel optimization to other meta-learning methods on difficult benchmarks.

Appendix

Appendix A Proof of Proposition 3.1

Let $\text{d}_{\lambda}f-h_{T-K}=e_{K}$. By definition of $h_{T-K}$,

Therefore, when $g$ is locally $\alpha$-strongly convex with respect to $w$ in the neighborhood of $\{w_{T-K-1},\dots,w_{T}\}$,

Suppose $g$ is $\beta$-smooth but nonconvex. In the worst case, if the smallest eigenvalue of $\nabla_{w,w}g(w_{t-1},\lambda)$ is $-\beta$, then $\|A_{t}\|=1+\gamma\beta\leq 2$ for $t=0,\dots,T-K$. This gives the bound in (9). However, if $g$ is globally strongly convex, then

The bound (10) uses the fact that $\sum_{t=0}^{T-K}(1-\gamma\alpha)^{t}\leq\sum_{t=0}^{\infty}(1-\gamma\alpha)^{t}=\frac{1}{\gamma\alpha}$ ∎

Appendix B Proof of Lemma 3.2

To illustrate the idea, here we prove the case where $K=1$. For $K>1$, similar steps can be applied. To prove the statement, we first expand the inner product by definition

where we recall $h_{T-1}=B_{T}\nabla_{\hat{w}^{*}}f$ as $\nabla_{\lambda}f=0$ by assumption.

Next we show a technical lemma, which provides a critical tool to bound the second term above; its proof is given in the next section.

Let $g$ be $\alpha$-strongly convex and $\beta$-smooth. Assume $B_{t}$ and $A_{t}$ are Lipschitz continuous in $w$, and assume $B_{T}$ has full column rank. For $\gamma\leq\frac{1}{\beta}$,

and $B_{T}^{\top}B_{T}$ is non-singular by assumption,

for some $c>0$, when $T$ is large enough and $\gamma$ is small enough. The implication holds because $\|\text{d}_{\lambda}f\|\leq O(\|\nabla_{\hat{w}^{*}}f\|)$. ∎

Let $C_{A}$ and $C_{B}$ be the Lipschitz constant of $A_{t}$ and $B_{t}$. First, we see that the inner product can be lower bounded by the following terms

The above lower bounds can be shown by the following inequalities:

Next we upper bound the error terms: $\Delta_{1}$, $\Delta_{2}$, and $\Delta_{3}$. We will use the fact that gradient descent converges linearly when optimizing a strongly convex and smooth function .

Let $w_{0}$ be the initial condition. Running gradient descent to optimize an $\alpha$-strongly convex and $\beta$-smooth function $g$, with step size $0<\gamma\leq\frac{1}{\beta}$, generates a sequence $\{w_{t}\}$ satisfying

where $D=\|w_{0}-w^{*}\|$ and $w^{*}=\operatorname*{arg\,min}g(w)$.

Lemma B.2 implies for $T\geq t$, $\|w_{T}-w_{t}\|\leq 2De^{-\alpha\gamma t}$.

Now we proceed to bound the errors $\Delta_{1}$, $\Delta_{2}$, and $\Delta_{3}$.

Using the bounds on $\Delta_{1}$, $\Delta_{2}$, and $\Delta_{3}$, we prove the final result.

Appendix C Proof of Theorem 3.3

The proof of this theorem is a standard proof of non-convex optimization with biased gradient estimates. Here we include it for completeness, as part of it will be used later in the proof of Theorem 3.4.

Let $\lambda_{\tau}$ be the $\tau$th iterate. For short hand, we write $\text{d}_{\lambda}f_{(\tau)}=\text{d}_{\lambda}f(\lambda_{\tau})$, and $h_{T-K,(\tau)}=h_{T-K}(\lambda_{\tau})$. Assume $F$ is $L$-smooth and $\|\text{d}_{\lambda}f_{(\tau)}\|\leq G$ and $\|h_{T-K,(\tau)}\|\leq G$ almost surely for all $\tau$. Then by $L$-smoothness, it satisfies

Let $e_{\tau}=\text{d}_{\lambda}f_{(\tau)}-h_{T-K,(\tau)}$ be the error in the gradient estimate. Substitute the recursive update $\lambda_{\tau+1}=\lambda_{\tau}-\eta_{t}h_{T-K,(\tau)}$ to the above inequality. Conditioned on $\lambda_{\tau}$, it satisfies

Performing telescoping sum with the above inequality, we have

Dividing both sides by $\sum_{\tau=1}^{R}\eta_{\tau}$ and using the facts that $\eta_{\tau}=O(\frac{1}{\sqrt{\tau}})$ and that

Appendix D Proof of Theorem 3.4

First we consider the special case when $S$ is deterministic. Let $H\geq K$. We decompose the full gradients into four parts

We assume that $w_{t}$ enters a locally strongly convex region for $t\geq H$. This implies, by Proposition 3.1, that $\|e\|\leq O(e^{-\alpha\gamma H}\|\nabla_{\hat{w}^{*}}f\|)$.

To prove the theorem, we first verify two conditions:

By Lemma 3.2, the assumption $\nabla_{\lambda}f^{\top}(\text{d}_{\lambda}f+h_{T-K}-\nabla_{\lambda}f)\geq\Omega(\|\nabla_{\lambda}f\|^{2})$, and $\|e\|\leq O(e^{-\alpha\gamma H}\|\nabla_{\hat{w}^{*}}f\|)$:

Therefore, for $H$ large enough, it holds that

By definition of $h_{T-K}=\nabla_{\lambda}f+q$, it holds that

Let $f$ be a lower-bound and $L$-smooth function. Consider the iterative update rule

where $g_{t}$ satisfies $g_{t}^{\top}\nabla f(x_{t})\geq c_{1}h_{t}^{2}$ and $\|g_{t}\|^{2}\leq c_{2}h_{t}^{2}$, for some constant $c_{1},c_{2}>0$ and scalar $h_{t}$. Suppose $f$ is lower-bounded and $\eta$ is chosen such that $\left(-c_{1}\eta+\frac{Lc_{2}\eta^{2}}{2}\right)\leq 0$. Then $\lim\limits_{t\to\infty}h_{t}=0$.

By telescoping sum, we can show $\sum_{t=0}^{\infty}\left(c\eta-\frac{L\eta^{2}}{2}\right)h_{t}^{2}<\infty$, which implies $\lim_{t\to\infty}h_{t}=0$. ∎

Finally, we prove the main theorem by applying Lemma D.1. Consider a deterministic problem. Take $h_{t}^{2}=\|\nabla_{\lambda}f(\lambda_{t})\|^{2}+\|\nabla_{\hat{w}^{*}}f(\lambda_{t})\|^{2}$. Because of (15) and (16), by Lemma D.1, it satisfies that

As $\|\text{d}_{\lambda}f\|\leq O(\|\nabla_{\lambda}f\|+\|\nabla_{\hat{w}^{*}}f\|)$, it shows $\|\text{d}_{\lambda}f\|$ converges to zero in the limit.

Appendix E Proof of Theorem 3.5

We prove the non-convergence using the following strategy. First we show that, when assumption 3 in Theorem 3.4, i.e.

does not hold, there is some problem such that $h_{T-k}\neq 0$ for all stationary points (i.e. $\lambda$ such that $\text{d}_{\lambda}f=0$). Then we show that, for such a problem, optimizing $\lambda$ with $h_{T-k}$ cannot converge to any of the stationary points.

To construct the counterexample, we consider a scalar deterministic bilevel optimization problem of the form

in which $\phi$ is some perturbation function that we will later define, and $\hat{w}^{*}$ is computed by performing $T>1$ steps of gradient descent in the lower-level optimization problem with some constant initial condition $w_{0}$ and constant step size $0<\gamma<1$, i.e.

We can observe this problem satisfies almost all the assumptions in Theorem 3.4:

The lower-level objective $g$ is smooth and strongly convex. (Proposition 3.1)

The upper-level objective $F$ is smooth. (Theorem 3.3)

The lower-level objective $g$ is second-order continuously differentiable (assumption 1 in Theorem 3.4)

The Jacobian if full rank, i.e. $B_{t}=\gamma>0$ (assumption 2 in Theorem 3.4)

The upper-level objective function is deterministic, i.e. $F=f$ (assumption 4 in Theorem 3.4)

But we will show that properly setting $\phi$ can break the non-interfering assumption in (17) (i.e. assumption 3 in Theorem 3.4) and then creates a problem such that optimizing $\lambda$ with $K$-RMD does not converge to an exact stationary point.

We follow the two-step strategy mentioned above.

Without loss of generality, let us consider optimizing $\lambda$ with $1$-RMD. In this case we can write the approximate and the exact gradients in closed form as

which are given by (5) and (8). We will show that by properly choosing $\phi$, we can define $f(\lambda)=\frac{1}{2}(\hat{w}^{*})^{2}+\phi(\lambda)$ such that, at any of the stationary points of $f$, the approximate gradient of $1$-RMD does not vanish. That is, we show when $\text{d}_{\lambda}f=0$, $h_{T-1}\neq 0$.

Before proceeding, let us define $u=w^{*}\gamma$ and $v=w^{*}\gamma\sum_{t=0}^{T}(1-\gamma)^{T-t}$ for convenience. To show how to construct $\phi$, let us consider the stationary points in the caseNote in this special case, assumption 3 in Theorem 3.4 holds trivially when $\phi(\lambda)=0$ (i.e. $\nabla_{\lambda}f=0$) and optimizing $\lambda$ with $K$-RMD converges to an exact stationary point. when $\phi=0$. Let $P_{0}$ denote the set of these stationary points, i.e. $P_{0}=\{\lambda:v=0\}$. Since $f$ is smooth and lower-bounded, we know that $P_{0}$ is non-empty, and from the construction of our counterexample we know that $P_{0}$ contains exactly the $\lambda$s such that $w^{*}=0$.

We use this fact to pick an adversarial $\phi$. Consider any smooth, lower-bounded $\phi$ whose stationary points are not in $P_{0}$, e.g. $\phi(\lambda)=\frac{1}{2}(\lambda-\lambda_{0})^{2}$ and $\lambda_{0}\notin P_{0}$. Then $f(\lambda)=\frac{1}{2}(\hat{w}^{*})^{2}+\phi(\lambda)$ has a non-empty set of stationary points $P_{\phi}$ such that $P_{\phi}\cap P_{0}=\emptyset$. We see that, for such $\phi$, the non-interfering assumption (assumption 3 in Theorem 3.4) is violated in $P_{\phi}$:

And we show for any $\lambda\in P_{\phi}$ it holds that $h_{T-1}\neq 0$. This can be seen from the definition

where the last inequality is because $w^{*}\neq 0$ for $\lambda\in P_{\phi}$.

We have shown that there is a problem which satisfies all the assumptions but assumption 3 of Theorem 3.4, and at any of its stationary points (i.e. when $\text{d}_{\lambda}f=0$) we have $h_{T-K}\neq 0$. Now we show this property implies failure to converge to the stationary points for the general problems considered in Theorem 3.5 (i.e. we do not rely on the form made in Step 1 anymore).

We prove this by contradiction. Let $\lambda^{*}$ be one of the stationary points. We choose $\delta_{0}>0$ such that, for some $\epsilon>0$, $\|h_{T-K}\|>\epsilon/\gamma$ for all $\lambda$ inside the neighborhood $\{\lambda:\|\lambda-\lambda^{*}\|<\frac{\delta_{0}}{2}\}$, where we recall $\gamma$ is the step size of the lower-level optimization problem. A non-zero $\delta_{0}$ exists because $h_{T-1}$ is continuous by our assumption and $h_{T-K}\neq 0$ at $\lambda^{*}$.

We are ready to show the contradiction. Let $\delta=\min\{\delta_{0},\epsilon\}$. Suppose there is a sequence $\{\lambda_{\tau}\}$ that converges to the stationary point $\lambda^{*}$. This means that there is $0<M<\infty$ such that, $\forall\tau\geq M$, $\|\lambda_{\tau}-\lambda^{*}\|<\frac{\delta}{2}$, which implies that $\forall\tau\geq M$, $\|\lambda_{\tau+1}-\lambda_{\tau}\|<\delta$. However, by our choice of $\delta_{0}$, $\|\lambda_{\tau+1}-\lambda_{\tau}\|=\gamma\|h_{T-K}\|>\epsilon\geq\delta$, leading to a contradiction.

Thus, no sequence $\{\lambda_{\tau}\}$ converges to any of the stationary points. This concludes our proof. ∎

Appendix F Proof of Proposition 3.6

Recall our shorthand that $\nabla_{\lambda,w}g$ and $\nabla_{w,w}g$ are evaluated at $(w^{*},\lambda)$. In the limit, it holds that

To prove the equality (12), we use Lemma (F.1).

For a matrix $A$ with $\|A\|<1$, it satisfies that

Since $\gamma\leq\frac{1}{\beta}$, we have $\gamma\alpha I\preceq\gamma\nabla_{w,w}g\preceq I$, so $\|I-\gamma\nabla_{w,w}g\|<1$. By Lemma F.1,

Appendix G Detailed experimental setup

In this appendix, we provide more details about the settings we used in each experiment. We use Adam to optimize the upper-level objective and vanilla gradient descent for the lower objective. We denote by $\hat{w}^{*}$ the results of running $T$ steps of gradient descent with step size $\gamma$.

In this appendix, we provide more details about the data hypercleaning experiment on MNIST from Section 4.2.1.

Both the training and the validation sets consist of 5000 class-balanced examples from the MNIST dataset. The test set consists of the remaining examples. For each training example, with probability $\frac{1}{2}$, we replaced the label with a uniformly random one.

For various $K$, we performed $K$-RMD for 1000 hyperiterations. Like in the toy experiment (Section 4.1) we adjusted the initial meta-learning rate $\eta_{0}$ for each $K$ so that the norm of the initial update was roughly the same for each $K$.

We asserted earlier that the reported F1 scores are not sensitive to our choice of threshold $\lambda_{i}<-3$. To validate this assertion, we repeated the experiment for various thresholds. F1 scores are reported in the table below.

We only ran these experiments for $150$ hyperiterations, because the F1 score has essentially converged by that point. Indeed, the plot below shows identification of corrupted labels for $K=1$, with cutoff $\lambda_{i}<-4$. The X axis is in units of 1000 hyperiterations. We see that $1$-RMD rapidly identifies most of the mislabeled examples, with a few false positives.

G.2 Task interaction

We use $T=100$ iterations of gradient descent with learning rate $0.1$ in the lower objective which yields $\hat{w}_{S}^{*}$. To ensure that $C$ is symmetric, and that $C_{ij}$ and $\rho$ are nonnegative, we re-parametrize them as $\rho=\text{softplus}(\nu)$ and $C=A+A^{\top}$, where $A_{ij}=\text{softplus}(B_{ij})$ and $B$ is a hyperparameter matrix. Thus, the hyperparameters to be optimized are $\lambda=\{B,\nu\}$.

Rather than using raw pixels, we extract image features from the output of the average pooling layer in Resnet-$18$ which is trained on ImageNet . We use the same data pre-processing that is used for training Resnet architecture.

When reporting test accuracy, we run 10 independent trials. In each trial, we sample the training and validation datasets with a balanced set of $m$ examples each ($m=50$ for CIFAR-10 and $m=300$ for CIFAR-100) and use the rest of the dataset for testing. To avoid over-fitting, we use early stopping when the testing error does not improve for $500$ hyper-iterations.

Although we are using a similar setting as Franceschi et al. , our results on full back-propagation are quite different from theirs. We believe it is because we are using a different network architecture and pre-processing method for feature extraction.

G.3 One-shot classification

The Omniglot dataset , a popular benchmark for few-shot learning, is used in this experiment. We consider $5$-way classification with $1$ training and $15$ validation examples for each of the five classes. To evaluate the generalization performance, we restrict the meta-training dataset to a random subset of $1200$ of the $1623$ Omniglot characters. The meta-validation dataset consists of $100$ other characters, and meta-testing dataset has the remaining $323$ characters. We use the meta-validation dataset for tuning the upper-level optimization parameters and report the performance of the algorithm on the meta-testing dataset. Note that no data augmentation method is used in the training.

The overall neural network architecture is shown in Figure 8. Our architecture inherits the hyper-representation model of Franceschi et al. with some modifications. The first two convolutional layers, parametrized by hyperparameter $\lambda=\{\lambda_{l_{1}},\lambda_{l_{2}}\}$, transform the input image into a “hyper-representation” space. The last three layers, parametrized by $w=\{w_{l_{3}},w_{l_{4}},w_{l_{5}}\}$ are fine-tuned in the lower-level optimization. Additionally, we have regularization hyperparameters $\lambda_{r}=\{\rho_{i}\}_{i=1}^{3}\cup\{c_{j}\}_{j=1}^{3}$. The overall setup corresponds essentially to meta-learning the two bottom layers of a CNN; for each task, the weights in the first two layers are frozen, and the $k$-way classifier of the last three layers is fine tuned. Overall, the model has $\approx 110$k hyperparameters and $\approx 75$k parameters.

We use a meta-batch-size of $4$ in each hyper-iteration. To limit the training time, we stop all the algorithms after $5000$ hyper-iterations. Needless to say, these results could be further improved by using data augmentation, higher meta-batch size, and running more hyper-iterations. However, our current setup is selected so that all the experiments can be run in a reasonable amount of time, while sharing a similar setting used in practical one-shot learning.