On the Iteration Complexity of Hypergradient Computation Riccardo - - PowerPoint PPT Presentation

On the Iteration Complexity of Hypergradient Computation

Riccardo Grazzi

Computational Statistics and Machine Learning, Istituto Italiano di Tecnologia. Department of Computer Science, University College London. riccardo.grazzi@iit.it

Joint work with Luca Franceschi, Massimiliano Pontil and Saverio Salzo.

1

Bilevel Optimization Problem

min

𝜇∈Λ⊆ℝ𝑜 𝑔(𝜇) ∶= 𝐹(𝑥(𝜇), 𝜇)

(upper-level) 𝑥(𝜇) ∶= Φ(𝑥(𝜇), 𝜇) (lower-level)

• Hyperparameter optimization, meta-learning.
• Graph and recurrent neural networks.

How can we solve this optimization problem?

• Black-box methods (random/grid search, Bayesian optimization, ...).
• Gradient-based methods exploiting the hypergradient ∇𝑔(𝜇).

2

Bilevel Optimization Problem

min

𝜇∈Λ⊆ℝ𝑜 𝑔(𝜇) ∶= 𝐹(𝑥(𝜇), 𝜇)

(upper-level) 𝑥(𝜇) ∶= Φ(𝑥(𝜇), 𝜇) (lower-level)

• Hyperparameter optimization, meta-learning.
• Graph and recurrent neural networks.

How can we solve this optimization problem?

• Black-box methods (random/grid search, Bayesian optimization, ...).
• Gradient-based methods exploiting the hypergradient ∇𝑔(𝜇).

2

Computing the Hypergradient ∇𝑔(𝜇)

Can be really expensive or even impossible to compute Exactly. Two common approximation strategies are

• 1. Iterative Difgerentiation (ITD).
• 2. Approximate Implicit Difgerentiation (AID).

Which one is the best?

• Previous works provide mostly qualitative and empirical results.

Can we have quantitative results on the approximation error?

• Yes! If the fixed point map Φ(⋅, 𝜇) is a contraction.

3

Computing the Hypergradient ∇𝑔(𝜇)

Can be really expensive or even impossible to compute Exactly. Two common approximation strategies are

• 1. Iterative Difgerentiation (ITD).
• 2. Approximate Implicit Difgerentiation (AID).

Which one is the best?

• Previous works provide mostly qualitative and empirical results.

Can we have quantitative results on the approximation error?

• Yes! If the fixed point map Φ(⋅, 𝜇) is a contraction.

3

Our Contributions

Upper bounds on the approximation error for both ITD and AID

• Both methods achieve non-asymptotic linear convergence rates.
• We prove that ITD is generally worse than AID in terms of upper bounds.

Extensive experimental comparison among difgerent AID strategies and ITD

• If Φ(⋅, 𝜇) is a contraction, the results confirm the theory.
• If Φ(⋅, 𝜇) is NOT a contraction, ITD can be still a reliable strategy.

250 500 750 1000 1250 t 10

10

10

10

10

||∇f(λ) − g(λ)|| Logistic Regression ITD FP k = t FP k = 10 CG k = t CG k = 10 50 100 150 200 t 10

10

10

10

Kernel Ridge Regression 25 50 75 100 125 150 t 10

10 10

10

10

Biased Regularization 100 200 300 400 500 t 10

10

10

10

10

10

Hyper Representation

4

Motivation

Source: S.Ravi, H. Larochelle (2016).

• Hyperparameter optimization

(learn the kernel/regularization, ...).

• Meta-learning (MAML, L2LOpt, ...).

Source: snap.stanford.edu/proj/embeddings-www

• Graph Neural Networks.
• Some Recurrent Models.
• Deep Equilibrium Models.

All can be cast into the same bilevel framework where at the lower-level we seek for the solution to a parametric fixed point equation.

5

Motivation

Source: S.Ravi, H. Larochelle (2016).

• Hyperparameter optimization

(learn the kernel/regularization, ...).

• Meta-learning (MAML, L2LOpt, ...).

Source: snap.stanford.edu/proj/embeddings-www

• Graph Neural Networks.
• Some Recurrent Models.
• Deep Equilibrium Models.

All can be cast into the same bilevel framework where at the lower-level we seek for the solution to a parametric fixed point equation.

5

Motivation

Source: S.Ravi, H. Larochelle (2016).

• Hyperparameter optimization

(learn the kernel/regularization, ...).

• Meta-learning (MAML, L2LOpt, ...).

Source: snap.stanford.edu/proj/embeddings-www

• Graph Neural Networks.
• Some Recurrent Models.
• Deep Equilibrium Models.

All can be cast into the same bilevel framework where at the lower-level we seek for the solution to a parametric fixed point equation.

5

Example: Optimizing the Regularization Hyperparameter in Ridge Regression

min

𝜇∈(0,∞)

1 2‖𝑌val𝑥(𝜇) − 𝑧val‖2

2

𝑥(𝜇) = arg min

𝑥∈ℝ𝑒

{ℓ(𝑥, 𝜇) ∶= 1 2‖𝑌𝑥 − 𝑧‖2

2 + 𝜇

2 ‖𝑥‖2

2}

𝑥(𝜇) is the unique fixed point of the one step GD map Φ(𝑥, 𝜇) = 𝑥 − 𝛽∇1ℓ(𝑥, 𝜇) If the step size 𝛽 is suffjciently small, Φ(⋅, 𝜇) is also a contraction.

6

The Bilevel Framework

min

𝜇∈Λ⊆ℝ𝑜 𝑔(𝜇) ∶= 𝐹(𝑥(𝜇), 𝜇)

(upper-level) 𝑥(𝜇) ∶= Φ(𝑥(𝜇), 𝜇) (lower-level)

• 𝑥(𝜇) ∈ ℝ𝑒 is oħten not available in closed form.
• 𝑔 is usually non convex and expensive or impossible to evaluate exactly.
• ∇𝑔 is even harder to evaluate.

7

The Bilevel Framework

min

𝜇∈Λ⊆ℝ𝑜 𝑔(𝜇) ∶= 𝐹(𝑥(𝜇), 𝜇)

(upper-level) 𝑥(𝜇) ∶= Φ(𝑥(𝜇), 𝜇) (lower-level)

• 𝑥(𝜇) ∈ ℝ𝑒 is oħten not available in closed form.
• 𝑔 is usually non convex and expensive or impossible to evaluate exactly.
• ∇𝑔 is even harder to evaluate.

7

How to Compute the Hypergradient ∇𝑔(𝜇)?

Iterative Difgerentiaton (ITD)

• 1. Set 𝑥0(𝜇) = 0 and compute,

for 𝑗 = 1, 2, … 𝑢 ⌊ 𝑥𝑗(𝜇) = Φ(𝑥𝑗−1(𝜇), 𝜇).

• 2. Compute 𝑔𝑢(𝜇) = 𝐹(𝑥𝑢(𝜇), 𝜇).
• 3. Compute ∇𝑔𝑢(𝜇) effjciently using reverse

(RMAD) or forward (FMAD) mode automatic difgerentiation.

Approximate Implicit Difgerentiation (AID)

• 1. Get 𝑥𝑢(𝜇) with 𝑢 steps of a lower-level solver.
• 2. Compute 𝑤𝑢,𝑙(𝜇) with 𝑙 steps of a solver for

the linear system (𝐽 − 𝜖1Φ(𝑥𝑢(𝜇), 𝜇)⊤)𝑤 = ∇1𝐹(𝑥𝑢(𝜇), 𝜇).

• 3. Compute the approximate gradient as

̂ ∇𝑔(𝜇) ∶=∇2𝐹(𝑥𝑢(𝜇), 𝜇) + 𝜖2Φ(𝑥𝑢(𝜇), 𝜇)⊤𝑤𝑢,𝑙(𝜇).

Which one is the best?

8

How to Compute the Hypergradient ∇𝑔(𝜇)?

Iterative Difgerentiaton (ITD)

• 1. Set 𝑥0(𝜇) = 0 and compute,

for 𝑗 = 1, 2, … 𝑢 ⌊ 𝑥𝑗(𝜇) = Φ(𝑥𝑗−1(𝜇), 𝜇).

• 2. Compute 𝑔𝑢(𝜇) = 𝐹(𝑥𝑢(𝜇), 𝜇).
• 3. Compute ∇𝑔𝑢(𝜇) effjciently using reverse

(RMAD) or forward (FMAD) mode automatic difgerentiation.

Approximate Implicit Difgerentiation (AID)

• 1. Get 𝑥𝑢(𝜇) with 𝑢 steps of a lower-level solver.
• 2. Compute 𝑤𝑢,𝑙(𝜇) with 𝑙 steps of a solver for

the linear system (𝐽 − 𝜖1Φ(𝑥𝑢(𝜇), 𝜇)⊤)𝑤 = ∇1𝐹(𝑥𝑢(𝜇), 𝜇).

• 3. Compute the approximate gradient as

̂ ∇𝑔(𝜇) ∶=∇2𝐹(𝑥𝑢(𝜇), 𝜇) + 𝜖2Φ(𝑥𝑢(𝜇), 𝜇)⊤𝑤𝑢,𝑙(𝜇).

Which one is the best?

8

How to Compute the Hypergradient ∇𝑔(𝜇)?

Iterative Difgerentiaton (ITD)

• 1. Set 𝑥0(𝜇) = 0 and compute,

for 𝑗 = 1, 2, … 𝑢 ⌊ 𝑥𝑗(𝜇) = Φ(𝑥𝑗−1(𝜇), 𝜇).

• 2. Compute 𝑔𝑢(𝜇) = 𝐹(𝑥𝑢(𝜇), 𝜇).
• 3. Compute ∇𝑔𝑢(𝜇) effjciently using reverse

(RMAD) or forward (FMAD) mode automatic difgerentiation.

Approximate Implicit Difgerentiation (AID)

• 1. Get 𝑥𝑢(𝜇) with 𝑢 steps of a lower-level solver.
• 2. Compute 𝑤𝑢,𝑙(𝜇) with 𝑙 steps of a solver for

the linear system (𝐽 − 𝜖1Φ(𝑥𝑢(𝜇), 𝜇)⊤)𝑤 = ∇1𝐹(𝑥𝑢(𝜇), 𝜇).

• 3. Compute the approximate gradient as

̂ ∇𝑔(𝜇) ∶=∇2𝐹(𝑥𝑢(𝜇), 𝜇) + 𝜖2Φ(𝑥𝑢(𝜇), 𝜇)⊤𝑤𝑢,𝑙(𝜇).

Which one is the best?

8

A First Comparison

ITD

• Ignores the bilevel structure.
• Cost in time (RMAD): 𝑃(Cost(𝑔𝑢(𝜇)))
• Cost in memory (RMAD): 𝑃(𝑢𝑒).
• Can we control ‖∇𝑔𝑢(𝜇) − ∇𝑔(𝜇)‖?

AID

• Can use any lower-level solver.
• Cost in time (𝑙 = 𝑢): 𝑃(Cost(𝑔𝑢(𝜇))).
• Cost in memory: 𝑃(𝑒).
• Can we control ‖ ̂

∇𝑔(𝜇) − ∇𝑔(𝜇)‖?

𝑔𝑢(𝜇) = 𝐹(𝑥𝑢(𝜇), 𝜇).

9

Previous Work on the Approximation Error

ITD

• arg min 𝑔𝑢 −

− − →

𝑢→∞ arg min 𝑔

(Franceschi et al., 2018).

• We provide non-asymptotic upper

bounds on ‖∇𝑔𝑢(𝜇) − ∇𝑔(𝜇)‖. AID

• ‖ ̂

∇𝑔(𝜇) − ∇𝑔(𝜇)‖ − − − − →

𝑢,𝑙→∞ 0

(Pedregosa, 2016).

• ‖ ̂

∇𝑔(𝜇) − ∇𝑔(𝜇)‖ − − − − →

𝑢,𝑙→∞ 0 at a

linear rate in 𝑢 and 𝑙 for meta-learning with biased regularization (Rajeswaran et al., 2019).

• We provide non-asymptotic upper

bounds on ‖ ̂ ∇𝑔(𝜇) − ∇𝑔(𝜇)‖.

𝑔𝑢(𝜇) = 𝐹(𝑥𝑢(𝜇), 𝜇).

10

Previous Work on the Approximation Error

ITD

• arg min 𝑔𝑢 −

− − →

𝑢→∞ arg min 𝑔

(Franceschi et al., 2018).

• We provide non-asymptotic upper

bounds on ‖∇𝑔𝑢(𝜇) − ∇𝑔(𝜇)‖. AID

• ‖ ̂

∇𝑔(𝜇) − ∇𝑔(𝜇)‖ − − − − →

𝑢,𝑙→∞ 0

(Pedregosa, 2016).

• ‖ ̂

∇𝑔(𝜇) − ∇𝑔(𝜇)‖ − − − − →

𝑢,𝑙→∞ 0 at a

linear rate in 𝑢 and 𝑙 for meta-learning with biased regularization (Rajeswaran et al., 2019).

• We provide non-asymptotic upper

bounds on ‖ ̂ ∇𝑔(𝜇) − ∇𝑔(𝜇)‖.

𝑔𝑢(𝜇) = 𝐹(𝑥𝑢(𝜇), 𝜇).

10

Preliminaries

Assumptions

• Φ(⋅, 𝜇) is a contraction with constant 𝑟𝜇 < 1.
• 𝜖1Φ(⋅, 𝜇), 𝜖2Φ(⋅, 𝜇), ∇1𝐹(⋅, 𝜇) and ∇2𝐹(⋅, 𝜇) are Lipschitz continuous

⟹ 𝑔 difgerentiable and 𝑥′(𝜇) ∶= (𝐽 − 𝜖1Φ(𝑥(𝜇), 𝜇))−1𝜖2Φ(𝑥(𝜇), 𝜇) ∇𝑔(𝜇) = ∇2𝐹(𝑥(𝜇), 𝜇) + 𝑥′(𝜇)⊤∇1𝐹(𝑥(𝜇), 𝜇).

11

Main Contribution

Theorem (ITD error upper bound) ‖∇𝑔𝑢(𝜇) − ∇𝑔(𝜇)‖ ≤ (𝑑1(𝜇) + 𝑑2(𝜇) 𝑟𝜇 𝑢 + 𝑑3(𝜇))𝑟𝑢

𝜇,

Theorem (AID error upper bound) Let 𝑤𝑢(𝜇) ∶= (𝐽 − 𝜖1Φ(𝑥𝑢(𝜇), 𝜇)⊤)−1∇1𝐹(𝑥𝑢(𝜇), 𝜇) and assume that

• ‖𝑥𝑢(𝜇) − 𝑥(𝜇)‖ ≤ 𝜍𝜇(𝑢)‖𝑥(𝜇)‖,
• ‖𝑤𝑢,𝑙(𝜇) − 𝑤𝑢(𝜇)‖ ≤ 𝜏𝜇(𝑙)‖𝑤𝑢(𝜇)‖.

Then, ‖ ̂ ∇𝑔(𝜇) − ∇𝑔(𝜇)‖ ≤ (𝑑1(𝜇) + 𝑑2(𝜇) 1 − 𝑟𝜇 )𝜍𝜇(𝑢) + 𝑑3(𝜇)𝜏𝜇(𝑙).

12

Main Contribution (Part 2)

Effjcient solvers for the linear system in AID:

• fixed point method (FP)
• conjugate gradient (CG)

Theorem (CG and FP error upper bounds) Assume that the lower-level problem is solved as in ITD. Then (FP) ‖ ̂ ∇𝑔(𝜇) − ∇𝑔(𝜇)‖ ≤ (𝑑1(𝜇) + 𝑑2(𝜇)1 − 𝑟𝑙

𝜇

1 − 𝑟𝜇 )𝑟𝑢

𝜇 + 𝑑3(𝜇)𝑟𝑙 𝜇.

Moreover, when 𝜖1Φ(𝑥𝑢(𝜇), 𝜇) is symmetric, (CG) ‖ ̂ ∇𝑔(𝜇) − ∇𝑔(𝜇)‖ ≤ (𝑑1(𝜇) + 𝑑2(𝜇) 1 − 𝑟𝜇 )𝑟𝑢

𝜇 + 𝑑3(𝜇) ̂

𝑑(𝜇)𝑞𝑙

𝜇,

where 𝑞𝜇 < 𝑟𝜇.

13

So... Which method has the best approximation error?

From our analysis:

• ITD, CG and FP converge linearly (in 𝑢 and 𝑙) to ∇𝑔(𝜇).
• FP bound < ITD bound

for every 𝑢, when 𝑙 = 𝑢.

• CG bound < FP bound

for 𝑙 big enough when 𝜖1Φ(𝑥𝑢(𝜇), 𝜇) is symmetric.

Is this true also for the actual error in practice? What happens when Φ(⋅, 𝜇) is not a contraction?

14

So... Which method has the best approximation error?

From our analysis:

• ITD, CG and FP converge linearly (in 𝑢 and 𝑙) to ∇𝑔(𝜇).
• FP bound < ITD bound

for every 𝑢, when 𝑙 = 𝑢.

• CG bound < FP bound

for 𝑙 big enough when 𝜖1Φ(𝑥𝑢(𝜇), 𝜇) is symmetric.

Is this true also for the actual error in practice? What happens when Φ(⋅, 𝜇) is not a contraction?

14

Hypergradient Approximation on Synthetic Data

250 500 750 1000 1250 t 10

−6

10

−5

10

−4

10

−3

10

−2

||∇f(λ) − g(λ)|| Logistic Regression ITD FP k = t FP k = 10 CG k = t CG k = 10 50 100 150 200 t 10

1

10

3

10

5

10

7

Kernel Ridge Regression 25 50 75 100 125 150 t 10

−1

10 10

1

10

2

10

3

Biased Regularization 100 200 300 400 500 t 10

4

10

5

10

6

10

7

10

8

10

9

Hyper Representation

Hypergradient approximation errors (mean/std on randomly drawn values of 𝜇). 𝑕(𝜇) is equal to ∇𝑔𝑢(𝜇) for ITD and to ̂ ∇𝑔(𝜇) for CG and FP. In all settings Φ(⋅, 𝜇) is a contraction and 𝜖1Φ(𝑥, 𝜇) is symmetric.

• Aħter a while the error decreases linearly for all methods.
• Methods with lower error bounds have lower error on average.

15

Equilibrium Models 1 on MNIST (Proof of Concept)

min

𝛿=(𝐵,𝐶,𝑑),𝜄 𝑜

∑

𝑗=1

CE(𝑥𝑗(𝛿)⊤𝜄, 𝑧𝑗), 𝑥𝑗(𝛿) = 𝜚𝑗(𝑥𝑗(𝛿), 𝛿) = tanh(𝐵𝑥𝑗(𝛿) + 𝐶𝑦𝑗 + 𝑑)

200 400 600 Time (s) 10−5 10−4 10−3 10−2

Objective

CG † CG FP † FP ITD † ITD 10 20 30 40 50 Iterations ×100 0.925 0.930 0.935 0.940 0.945

Test accuracy

10 20 30 40 50 Iterations ×100 10−5 10−4 10−3 10−2 10−1

Hypergradient norm ||g(λ)||

10−3 10−2 10−1 100 Learning rate 0.88 0.89 0.90 0.91 0.92 0.93 0.94 0.95

Test accuracy vs learning rate

𝜚𝑗(⋅, 𝛿) NOT a contraction for † methods.

• When 𝜚𝑗(⋅, 𝛿) is a contraction all the methods perform similarly.
• ITD is the most stable when the contraction assumption is not satisfied.

1Shaojie Bai, J Zico Kolter, and Vladlen Koltun. “Deep equilibrium models”. In Advances in Neural Information Processing Systems. 2019, pp. 688–699.

16

Conclusions

We studied the iteration complexity of two strategies used to approximate the hypergradient in bilevel problems: iterative difgerentiation (ITD) and approximate implicit difgerentiation (AID). We proved non-asymptotic upper bounds on the approximation error

• CG, FP and ITD converge linearly to the exact hypergradient.
• ITD is generally worse than AID in terms of upper bounds.

We conducted experiments comparing ITD and AID

• If Φ(⋅, 𝜇) is a contraction, the results confirm the theory.
• If Φ(⋅, 𝜇) is NOT a contraction, ITD can be still a reliable strategy.

17

Thank you for the attention

CODE (PyTorch): https://github.com/prolearner/hypertorch

18

References

Franceschi, L., Frasconi, P., Salzo, S., Grazzi, R., and Pontil, M. (2018). Bilevel programming for hyperparameter optimization and meta-learning. In International Conference on Machine Learning, pages 1563–1572. Pedregosa, F. (2016). Hyperparameter optimization with approximate gradient. In International Conference on Machine Learning, pages 737–746. Rajeswaran, A., Finn, C., Kakade, S. M., and Levine, S. (2019). Meta-learning with implicit gradients. In Advances in Neural Information Processing Systems, pages 113–124.

19