The symbiotic relationship between optimization and deep learning

Pierre Ablin

CNRS - Université paris-dauphine

Joint works with :

T. Moreau, M. Massias, A.Gramfort, M. Sander, M.Blondel and G.peyré

optimization

Optimization: how to minimize a function ?

 

$$\min_{x\in\mathcal{X}} f(x)$$

Challenges:

Algorithm design

Find algorithms depending on assumptions on \(f\) (convex, smooth, ...) and on \(\mathcal{X}\) (convex, manifold, discrete...)

Theoretical guarantees

Convergence ? In which sense ? At which speed ?

Implementation

Numerical complexity, practical computational costs, hardware 

Deep learning

Design a parametrized transform $$ \phi_{\theta}: x \to y$$ by composition of simple, differentiable blocks

Challenges:

Network design

Design networks that work well depending on the task / input data / output data

Theoretical guarantees

What can we say about generalization of the network ? About the learned weights?

Implementation

Fast training and inference, deal with memory issues,...

The obvious link

Neural networks are usually trained by optimizing a function

 

Empirical risk minimization:

 

$$\min_{\theta} \frac1n\sum_{i=1}^n \ell_i(\phi_{\theta}(x_i))$$
 

Basic algorithm: stochastic gradient descent

 

$$\text{sample } i\sim [1, n]$$

$$\theta \leftarrow \theta - \rho \nabla_{\theta}[\ell_i(\phi_{\theta}(x_i))]$$

Today's talk :

other links

Learning to optimize: neural networks for optimization

Inverse problems

Latent process \(z \) generates observed outputs  \(x\):

 

\(z \to x \) 

The forward operation "\( \to\)" is generally known:

 

\(x = f(z) + \varepsilon \)

Goal of inverse problems: find a mapping 

 

\(x \to z\)

Example: MEG acquisition

\( z \) : current density in the brain

\( x \) : observed MEG signals

\(D\) : linear operator given by physics (Maxwell's equations)

\(  x \)

\(  D \)

\(  = \)

\(  z \)

Linear regression

Linear forward model : \(z \in \mathbb{R}^m\), \(x\in\mathbb{R}^n\), \(D \in \mathbb{R}^{n \times m} \)

 

\(x = Dz + \varepsilon \)

Problem: in some applications, \(m \gg n \), least-squares ill-posed

\(\to\) bet on sparsity : only a few coefficients in \(z^*\) are \( \neq 0 \)

 

\(z\) is sparse

Simple solution: least squares

\( z^* \in \arg\min \frac12 \|x - Dz\|^2\)

The Lasso

\( \lambda > 0 \) regularization parameter :

 

\(z^*\in\arg\min \frac12\|x - Dz\|^2 + \lambda \|z\|_1 = F_x(z)\) 

Enforces sparsity of the solution.

 

Easier to see on the equivalent problem: \(z^* \in \arg\min \frac12 \|x-Dz\|^2 \) s.t. \(\|z\|_1\leq C\)

Tibshirani, Regression shrinkage and selection via the lasso, 1996

Lasso induces sparsity

\(z^*\in \arg\min \frac12 \|x - Dz \|^2 \)s.t. \(\|z\|_1\leq C\)

\(z^*\in \arg\min \frac12 \|x - Dz \|^2 \)

Iterative shrinkage-thresholding algorithm

ISTA: simple algorithm to fit the Lasso.

 \(F_x(z) = \frac12\|x-Dz\|^2 + \lambda \|z\|_1\)

Idea: use proximal gradient descent

\(\to\) \(\frac12\|x - Dz\|^2\) is a smooth function

$$\nabla_z \left(\frac12\|x-Dz\|^2\right) = D^{\top}(Dz-x)$$

\(\to\) \(\lambda \|z\|_1\) has a simple proximal operator

Iterative shrinkage-thresholding algorithm

ISTA: simple algorithm to fit the Lasso.

 \(F_x(z) = \frac12\|x-Dz\|^2 + \lambda \|z\|_1\)

Daubechies et al.,  An iterative thresholding algorithm for linear inverse problems with a sparsity constraint. , 2004

ISTA: gradient descent step on the smooth part + proximal step

\(z^{(t+1)} = \text{st}(z^{(t)} - \frac1LD^{\top}(Dz^{(t)} - x), \frac{\lambda}{L})\)

ISTA as a Recurrent neural network

Solving the lasso many times

Assume that we want to solve the Lasso for many observations \(x_1, \cdots, x_N\) with a fixed dictionary \(D\) 

e.g.  MEG inverse problem:

\(D\) is fixed given by Maxwell's equations, \(x_i\) is one sample of the recording:

 

up to 100K samples !

This is an optimization problem:

Learning to optimize

We want to solve the Lasso many times with same \(D\); can we accelerate ISTA ? 

ISTA is a Neural Net

ISTA:

\(z^{(t+1)} = \text{st}(z^{(t)} - \frac1LD^{\top}(Dz^{(t)} - x), \frac{\lambda}{L})\)

 

Let \(W_1 = I_m - \frac1LD^{\top}D\) and \(W_2 = \frac1LD^{\top}\):

\(z^{(t+1)} = \text{st}(W_1z^{(t)} +W_2x, \frac{\lambda}{L})\)

 

 

3 iterations of ISTA = 3 layers NN

Learned-ISTA

Gregor, LeCun, Learning Fast Approximations of Sparse Coding, 2010

A \(T\)-layer Lista network is a function \(\Phi\) parametrized by \(T\) parameters \( \Theta = (W^t_1,W^t_2, \beta^t )_{t=0}^{T-1}\)

Learned-ISTA

A \(T\)-layer Lista network is a function \(\Phi\) parametrized by \(T\) parameters \( \Theta = (W^t_1,W^t_2, \beta^t )_{t=0}^{T-1}\) 

  • \(z^{(0)} = 0\)
  • \(z^{(t+1)} = st(W^t_1z^{(t)} + W^t_2x, \beta^t) \)
  • Return \(z^{(T)}  = \Phi_{\Theta}(x) \)

 

The parameters of the network are learned to get better results than ISTA

learning parameters

Unsupervised

Learn to solve the Lasso:

$$\Theta \in \argmin\mathcal{L}(\Theta) = \sum_{i=1}^N F_{x_i}(\Phi_{\Theta}(x_i))$$

If we see a new sample \(x\), we expect:

$$F_x(\Phi_{\Theta}(x)) \leq F_x(ISTA(x)) \enspace, $$

where ISTA is applied for \(T\) iterations.

Lista

Advantages:

  • Can handle large-scale datasets
  • GPU friendly

Drawbacks:

  • Learning problem is non-convex, non-differentiable: no practical guarantees, and generally hard to train
  • Can fail badly on unseen data

what does the network learn?

Consider a "deep" LISTA network: \(T\gg 1\), assume that we have solved the expected optimization problem:

$$\Theta \in \argmin \mathbb{E}_{x\sim p}\left[F_x(\Phi_{\Theta}(x))\right]$$

As \(T \to + \infty \), assume \( (W_1^t, W_2^t, \beta^t) \to (W_1^*,W_2^*, \beta^*)\). Call \(\alpha = \beta^* / \lambda\). Then:

 

$$ W_1^* = I_m - \alpha D^{\top}D$$

$$ W_2^* = \alpha D^{\top}$$

$$ \beta^* = \alpha \lambda$$

A., Moreau et al., Learning Step Sizes for Unfolded Sparse Coding, 2019

Corresponds to ISTA with step size \(\alpha\) instead of \(1/L\)

$$ W_1 = I_m - \frac1LD^{\top}D$$

$$ W_2 = \frac1L D^{\top}$$

$$ \beta = \frac{\lambda}{L}$$

what does the network learn?

As \(T \to + \infty \), assume \( (W_1^t, W_2^t, \beta^t) \to (W_1^*,W_2^*, \beta^*)\). Call \(\alpha = \beta^* / \lambda\). Then:

 

$$ W_1^* = I_m - \alpha D^{\top}D$$

$$ W_2^* = \alpha D^{\top}$$

$$ \beta^* = \alpha \lambda$$

A., Moreau et al., Learning Step Sizes for Unfolded Sparse Coding, 2019

Corresponds to ISTA with step size \(\alpha\) instead of \(1/L\)

$$ W_1 = I_m - \frac1LD^{\top}D$$

$$ W_2 = \frac1L D^{\top}$$

$$ \beta = \frac{\lambda}{L}$$

Optimization theory helps up characterize precisely what the network learns !

Network architecture guided by optimization

Residual networks

Classical networks :

 

\(x_{n+1} = f(x_n, \theta_n)\), e.g. \(x_{n+1} = \sigma(Wx_n + b)\)
 

Problem :

Stacking too many layers degrades perf

 

Residual networks

 

\(x_{n+1} = x_n +  f(x_n, \theta_n)\)
 

Easy to learn identity !

Can stack many layers :)

Residual networks

Allows to stack many layers (100s')

 

- State of the art on many problems for a long time

- Still widely used today

He et al., Deep residual learning for image recognition, 2015

Memory issues...

Forward pass: evaluate network


\(x_0 = x\)

\(x_{n+1}= x_n + f(x_n , \theta_n)\)

\(y = x_N\)

 

Backward pass: compute gradients


\(\frac{\partial \ell(y)}{\partial y} = \ell'(y)\)

\(\frac{\partial \ell(y)}{\partial x_n} = \frac{\partial \ell(y)}{\partial x_{n+1}}(I + \partial_{x}f(x_n, \theta_n))\)

\(\frac{\partial \ell(y)}{\partial \theta_n} =  \frac{\partial \ell(y)}{\partial x_{n+1}}\partial_{\theta}f(x_n, \theta_n)\)

Start from output : need to store \(x_n\)

Memory issues...

Need to store all activations \(x_n\) for backprop : huge memory cost !

On a classical image classification task : 

PARALLEL with optimization

Residual network

\(x_{n+1} = x_n + f(x_n, \theta) \)

Gradient descent

\(x_{n+1} = x_n -\rho \nabla g(x_n, \theta) \)

Equivalent if \(f = - \rho \nabla_xg\)

Momentum Gradient descent

\(v_{n+1}= \beta v_n - \rho \nabla g(x_n, \theta)\)

\(x_{n+1} = x_n +v_{n+1} \)

Momentum residual network

\(v_{n+1}= \beta v_n + f(x_n, \theta)\)

\(x_{n+1} = x_n +v_{n+1} \)

Sander, A., Blondel & Peyré, Momentum residual neural networks, 2021

Invertible layers

Momentum residual network

\(v_{n+1}= \beta v_n + f(x_n, \theta)\)

\(x_{n+1} = x_n +v_{n+1} \)

Inverted by:

\(x_n = v_{n+1} - x_n\)

\(v_n = \frac1\beta(v_{n+1} - f(x_n, \theta))\)

No need to store activations ! They can be recomputed on the fly during backprop

representation capacity

Momentum residual network

\(v_{n+1}= \beta v_n + f(x_n, \theta)\)

\(x_{n+1} = x_n +v_{n+1} \)

Continuous equivalent:

\(\varepsilon \ddot x +\dot x = f(x, \theta)\)

Residual network

\(x_{n+1} =x_n + f(x_n, \theta)\)

 

Continuous equivalent:

\(\dot x = f(x, \theta)\)

Open source code

>>> pip install momentumnet

Get python code : 

import torch
from torchvision.models import resnet101

from momentumnet import transform_to_momentumnet

resnet = resnet101(pretrained=True)
mresnet101 = transform_to_momentumnet(resnet)

Transform a resnet into a momentum net : 

Documentation: 

 

michaelsdr.github.io/momentumnet/  

Conclusion

 

 

- There are many links between deep learning and optimization

 

- Deep neural networks can accelerate optimization when training is possible

 

- Optimization can guide the design of deep networks and lead to intriguing properties

 

Thanks  ! 

Made with Slides.com