Gathering and Exploiting Higher-Order Information when Training Large Structured Models

## Abstract

Abstract We consider a gradient-based optimization method
applied to a function ℒ ℒ \mathcal{L} of a
vector of variables 𝜽 𝜽 \boldsymbol{\theta} , in the case where 𝜽 𝜽 \boldsymbol{\theta} is represented as a tuple of tensors ( 𝐓 1 , ⋯ , 𝐓 S ) subscript 𝐓 1 ⋯ subscript 𝐓 𝑆 (\mathbf{T}_{1},\cdots,\mathbf{T}_{S}) .
This framework encompasses many common use-cases,
such as training neural networks by gradient descent. First, we propose a computationally inexpensive technique
providing higher-order information on ℒ ℒ \mathcal{L} ,
especially about the interactions between the tensors 𝐓 s subscript 𝐓 𝑠 \mathbf{T}_{s} ,
based on automatic differentiation and computational tricks.
Second, we use this technique at order 2 to build a second-order
optimization method which is suitable, among other things,
for training deep neural networks of various architectures. This second-order method leverages the partition structure of 𝜽 𝜽 \boldsymbol{\theta} into tensors ( 𝐓 1 , ⋯ , 𝐓 S ) subscript 𝐓 1 ⋯ subscript 𝐓 𝑆 (\mathbf{T}_{1},\cdots,\mathbf{T}_{S}) , in such a way that it requires neither the computation
of the Hessian of ℒ ℒ \mathcal{L} according to 𝜽 𝜽 \boldsymbol{\theta} ,
nor any approximation of it. The key part
consists in computing a smaller matrix interpretable as
a “Hessian according to the partition”, which can be computed exactly and efficiently.
In contrast to many existing practical second-order methods used in neural networks,
which perform a diagonal or block-diagonal approximation of the Hessian or its inverse,
the method we propose does not neglect interactions between layers. Our code is available on GitHub: https://github.com/p-wol/GroupedNewton .

## 1 Introduction

The appealing theoretical properties of Newton’s method
have led to numerous attempts to adapt it to
neural network optimization.
Therefore, the study of the Hessian
of a loss according to many parameters
has become an area of research in itself,
leading to a large number of methods to approximate it
accurately with a small computational cost.

#### Newton’s method applied to neural network optimization.

When it comes to neural networks, Newton’s method suffers from several
problems. Some are *technical*, such as building an accurate and
computationally efficient method to estimate the Hessian. But some of them
are *essential*, in the sense that they cannot be solved solely by a perfect
knowledge of the full Hessian. For instance, several works (Sagun et al.,, )
have shown that many eigenvalues of the Hessian are close to zero
when training neural networks, making it
impossible to use Newton’s method in practice, even if the Hessian is perfectly known.

Therefore, in this work, we do not aim to build a technique to estimate
the Hessian accurately and efficiently. We propose to take a step aside
and focus on two related goals.
First, we aim to access efficiently to higher-order information in order to use it for optimization.
To do so, we leverage the available computational tools, that is, automatic differentiation
and partition of the set of parameters into tensors stored on a GPU.
Second, we aim to build a second-order optimization method with properties similar
to the properties of Newton’s method (but inevitably weaker).
Notably, we want a method invariant by layer-wise affine reparameterizations of the model,
providing a Hessian-inspired matrix showing the interactions between each pair of tensors,
along with layer-wise step sizes.
Naturally, the computational cost should remain reasonable, and
the problem of close-to-zero eigenvalues of Hessian has to be dealt with.

#### First contribution: extracting higher-order information.

Formally, we study a loss $\mathcal{L}$ to minimize according to
a vector of parameters $\boldsymbol{\theta}\in\mathbb{R}^{P}$,
which can be represented as a tuple of tensors $(\mathbf{T}_{1},\cdots,\mathbf{T}_{S})$.
In a multilayer perceptron with $L$ layers, the $S=2L$ tensors
$(\mathbf{T}_{s})_{1\leq s\leq 2L}$ are the tensors of weights
and the vectors of biases of each layer. In that case, $S\ll P$.
Within this framework, we propose a technique *summarizing* the order-$d$ derivative of
the loss, which is a tensor belonging to $\mathbb{R}^{P^{d}}$, into
a tensor belonging to $\mathbb{R}^{S^{d}}$, which is significantly
smaller and easier to compute.

#### Second contribution: a scalable second-order
optimization method.

Then, we make use of the preceding technique
at orders $2$ and $3$ to build a second-order optimization method.
Formally, the method presented here and Newton’s method look alike:
in both cases, a linear system $\mathbf{H}_{0}\mathbf{x}=\mathbf{g}_{0}$ has to be solved
(according to $\mathbf{x}$), where $\mathbf{g}_{0}$ and $\mathbf{H}_{0}$ contain
respectively first-order and second-order information about $\mathcal{L}$.
Despite this formal resemblance, the difference is enormous:
with Newton’s method, $\mathbf{H}_{0}$ is equal to the Hessian $\mathbf{H}$ of $\mathcal{L}$ of size $P\times P$,
while with ours, $\mathbf{H}_{0}$ is equal to a matrix $\bar{\mathbf{H}}$ of size $S\times S$.
Thus, $\bar{\mathbf{H}}$ is undoubtedly smaller and easier to compute than $\mathbf{H}$
when $S\ll P$. Nevertheless, since $\bar{\mathbf{H}}$ is a dense matrix,
it still contains information about the interactions between the tensors $\mathbf{T}_{s}$
when they are used in $\mathcal{L}$.
This point is crucial, since most second-order optimization methods
applied in neural networks use a simplified version of the Hessian
(or its inverse), usually a diagonal or block-diagonal approximation,
which ignores the interactions between layers.
Finally, we propose an anisotropic version of Nesterov’s cubic regularization
(Nesterov and Polyak,, ), which uses order-3 information to regularize $\bar{\mathbf{H}}$ and avoid
instabilities when computing $\bar{\mathbf{H}}^{-1}\bar{\mathbf{g}}$.

As a preview of the experimental results, we report in Figure [1](#S1.F1)
the *summary of the Hessian* $\bar{\mathbf{H}}$ and its inverse $\bar{\mathbf{H}}^{-1}$
at two stages of training of LeNet-5:
interactions between different layers exist and are not negligible.
And, to achieve the proof-of-concept, we also provide in Section [5](#S5)
a short series of experiments showing
that this method is able to achieve training losses comparable to those of Adam and K-FAC.

Figure: Figure 1: LeNet-5 trained by SGD on CIFAR-10. Submatrices of $\bar{\mathbf{H}}$ (first row) and $\bar{\mathbf{H}}^{-1}$ (second row), where focus is on interactions: weight-weight, weight-bias, bias-bias of the different layers. We show it at initialization and before the 73rd epoch (best valid. loss). Legend: white: close-to-zero value; dark red: large positive value; dark blue: negative value, large in norm.
Refer to caption: /html/2312.03885/assets/x1.png

#### Structure of the paper.

First, we show the context and motivation of our work in Section [2](#S2).
Then, we provide in Sections [3](#S3) and
[4](#S4) standalone presentations of the contributions,
respectively the higher-order information extraction technique
and the second-order optimization method.
In Section [5](#S5), we present experimental results
showing that the developed methods are usable in practice.
Finally, we discuss the results in Section [6](#S6).

## 2 Context and motivation

### 2.1 Higher-order information

Extracting the maximum higher-order information about a loss
at a minimal computational cost and using it to improve optimization is not a novel idea.
This is typically what is done by Dangel, (),
despite it does not go beyond the second-order derivative.
In that research direction, the *Hessian-vector product*
technique (Pearlmutter,, )
is decisive, since it allows to compute at a very small cost
the projection of higher-order derivatives in given directions (see Appendix [A](#A1)).
Regarding the derivatives of order $3$ and beyond,
Nesterov’s cubic regularization of Newton’s method
(Nesterov and Polyak,, ) uses some information of order $3$
to avoid too large training steps. Incidentally,
we develop an anisotropic variation of it in Section [4.2](#S4.SS2).

Overall, the current need for information contained in higher-order derivatives
is still limited.
Anyway, we hope that making it available in large models would
lead to the emergence of new uses.

### 2.2 Second-order methods

The Hessian $\mathbf{H}$ of the loss $\mathcal{L}$ according to the vector
of parameters $\boldsymbol{\theta}$ is known to contain useful information about $\mathcal{L}$.
Above all, the Hessian is used to develop second-order optimization
algorithms. Let us denote by $\boldsymbol{\theta}_{t}$ the value of $\boldsymbol{\theta}$ at time step $t$,
$\mathbf{g}_{t}\in\mathbb{R}^{P}$ the gradient of $\mathcal{L}$ at step $t$ and $\mathbf{H}_{t}$
its Hessian at step $t$.
One of the most widely known second-order optimization method is Newton’s method,
whose step is (Nocedal and Wright,, , Chap. 3.3):

$$ $\displaystyle\boldsymbol{\theta}_{t+1}:=\boldsymbol{\theta}_{t}-\mathbf{H}_{t}^{-1}\mathbf{g}_{t}.$ $$

Under some conditions, including strong convexity of $\mathcal{L}$,
Newton’s method is known to have a quadratic convergence rate
(Nocedal and Wright,, , Th. 3.7), which makes this method
very appealing.
Besides, there exist other methods making use of second-order information,
without requiring the full computation of the Hessian.
For instance, Cauchy’s steepest descent (Cauchy,, ) is
a variation of the usual gradient descent, where the step size is
tuned by extracting very little information from the Hessian:

$$ $\displaystyle\boldsymbol{\theta}_{t+1}:=\boldsymbol{\theta}_{t}-\eta_{t}^{*}\mathbf{g}_{t},\quad\text{where}\quad\eta_{t}^{*}:=\frac{\mathbf{g}_{t}^{T}\mathbf{g}_{t}}{\mathbf{g}_{t}^{T}\mathbf{H}_{t}\mathbf{g}_{t}},$ $$

where the value of $\mathbf{g}_{t}^{T}\mathbf{H}_{t}\mathbf{g}_{t}$ can be obtained with little computational cost
(see Appendix [A](#A1)).
However, when optimizing a quadratic function $f$ with Cauchy’s steepest descent,
$f(\theta_{t})$ is known to decrease at a rate $(\frac{\lambda_{\mathrm{max}}-\lambda_{\mathrm{min}}}{\lambda_{\mathrm{max}}+\lambda_{\mathrm{min}}})^{2}$, where $\lambda_{\mathrm{max}}$ and $\lambda_{\mathrm{min}}$
are respectively the largest and the smallest eigenvalues of the Hessian of $f$
(Luenberger and Ye,, , Chap. 8.2, Th. 2).
If the Hessian of $f$ is
strongly anisotropic, then this rate is close to one and optimizatoin is slow.
For a comparison of both methods,
one may refer to (Gill et al.,, ; Luenberger and Ye,, ; Nocedal and Wright,, ).

Finally, one can guess that there is some space between Newton’s method,
which requires to know the full Hessian,
and Cauchy’s steepest descent, which requires only minimal and computationally cheap information
about the Hessian.
The optimization method presented in Section [4](#S4)
explores precisely this in-between space.

#### Quasi-Newton methods.

In many high-dimensional situations, the computation of the Hessian $\mathbf{H}_{t}$,
as well as the inversion of the linear system $\mathbf{g}_{t}=\mathbf{H}_{t}\mathbf{x}$,
is computationally intensive.
Quasi-Newton methods are designed to avoid any direct computation of the Hessian,
and make an extensive use of the gradients and of finite difference methods
to approximate the direction of $\mathbf{H}_{t}^{-1}\mathbf{g}_{t}$.
A list of common quasi-Newton methods may be found in
(Nocedal and Wright,, , Chap. 8).

As explained in Nocedal and Wright, (), since it is easy to compute the Hessian
by using Automatic Differentiation (AutoDiff), quasi-Newton methods tend to lose their interest.
Nevertheless, they should remain useful in situations where such computation
is too difficult.
Besides, the availability of a tool as powerful as AutoDiff
should push us to invent methods avoiding the computation of the full Hessian.

#### Block-diagonal approximations of the Hessian.

In the development of neural network training algorithms,
several attempts have been made to use the Hessian to build
second-order optimization methods with good properties.
A first rough attempt has been made by Wang and Lin, ().
In this work, the Hessian matrix is divided into blocks,
following the division of the network into layers,
and its off-diagonal blocks are removed.
From another perspective, Ollivier, ()
has kept this block-diagonal structure, but performed an additional
approximation on the remaining blocks. In both cases,
the interactions between the different layers are neglected.

#### Kronecker-Factored Approximate Curvature (K-FAC).

This technique of approximation of the Hessian has been proposed
in (Martens and Grosse,, ) in the context of neural network training.
K-FAC leverages the specific architecture of neural networks
to output an approximation of the true Hessian, which is
usually too costly to compute.
However, the resulting Hessian suffers from several issues.
First, the main approximation is quite rough, since
“[it assumes] statistical independence between products […] of unit activities
and products […] of unit input derivatives” (Martens and Grosse,, , Sec. 3.1).
Second, even with an approximation of the full Hessian, one
still has to invert it, which is computationally intensive
even for small networks.
To overcome this difficulty, Martens and Grosse, ()
make a block-diagonal or block-tridiagonal approximation of the inverse of
the Hessian, which kills many of the interactions between the parameters
or the layers. However, this approximation is less rough than
a direct block-diagonal approximation of the Hessian.

#### Summarizing the Hessian.

One should also mention Lu et al., (), which proposes
to approximate the Hessian with a matrix composed of
blocks in which all the coefficient are identical.
Thus, the Hessian can be compressed into a smaller matrix,
which looks like the summary of Hessian matrix $\bar{\mathbf{H}}$ used in Section [4](#S4).
In a completely different setup, Yuan et al., ()
proposes a “Sketched Newton-Raphson”, which
is driven by the same spirit as the method presented in Section [4](#S4):
instead of dealing with a complicated
large matrix, one should “project” it on spaces of lower dimension.

#### Optimization methods invariant by affine reparameterization.

Several optimization methods based on the Hessian,
such as Newton’s method, have an optimization step invariant by
affine reparameterization of $\boldsymbol{\theta}$
(Amari,, ) (Nesterov,, , Chap. 4.1.2).
Specifically, when using Newton’s method,
it is equivalent to optimize $\mathcal{L}$ according to $\boldsymbol{\theta}$ and according to
$\tilde{\boldsymbol{\theta}}=\mathbf{A}\boldsymbol{\theta}+\mathbf{B}$ ($\mathbf{A}\in\mathbb{R}^{P\times P}$
invertible, $\mathbf{B}\in\mathbb{R}^{P}$).
This affine-invariance property holds also when the function $\mathcal{L}$ to minimize
is a negative $\log$-likelihood, and one chooses to minimize $\boldsymbol{\theta}$
with the *natural gradient* method (Amari,, ).
This method also requires, at some point, the computation of the Hessian of $\mathcal{L}$.

### 2.3 Motivation

#### What are we really looking for?

The methods presented in Section [2.2](#S2.SS2), which all attempt to construct the Hessian matrix $\mathbf{H}$
or its inverse $\mathbf{H}^{-1}$, do the same implicit assumption:
we want to use a Newton-like optimization method, so we must have access
to $\mathbf{H}$ or $\mathbf{H}^{-1}$.
This assumption is certainly correct when the loss to optimize is strongly convex.
But, when the loss is non convex and very complicated, for instance when training a
neural network, this assumption lacks justification.
Worse, it has been shown empirically that, at the end of the training of a neural network,
the eigenvalues of the Hessian are concentrated around zero
(Sagun et al.,, ), with only a few large positive eigenvalues.

Therefore, Newton’s method itself does not seem to be recommended in neural network
training, so we may not need to compute the full Hessian at all,
which would relieve us from a tedious, if not impossible, task.(^1^11“No one is bound to do the impossible.”)

#### Importance of the interactions between layers.

Also, we would like to show the importance of keeping the off-diagonal
blocks in the Hessian or its inverse.
Several empirical works have shown that the role and the behavior
of each layer must be considered along its interactions between the other layers.
We give two examples. First, Zhang et al., () has shown that,
at the end of their training, many networks exhibit a strange
feature: some (but not all) layers can be reinitialized to their initial value
with little damage to the performance.
Second, Kornblith et al., () has compared the similarity between
the representations of the data after each layer. It has been remarked
that the similarity matrix of the layers may change qualitatively
when changing the number of layers (Kornblith et al.,, , Fig. 3).
Among all, these empirical results motivate our search for
mathematical objects showing how layers interact.

#### Per-layer scaling of the learning rates.

A whole line of research is concerned with building a well-grounded method
to find a good scaling of the distribution of initialization of the parameters,
along with a good scaling for the learning rates, which may be chosen layer-wise.
One may refer to the paper introducing the Neural Tangent
Kernels (Jacot et al.,, ), in which a layer-wise scaling for the weights
has been proposed and theoretically grounded.
Also, in the “feature learning” line of works,
(Yang and Hu,, ) proposes a general relation between several
scalings related to weight initialization and training.
Therefore, there is an interest in finding a scalable and theoretically grounded method to
build per-layer learning rates.

#### Unleashing the power of AutoDiff.

Nowadays, several libraries provide easy-to-use automatic
differentiation packages, allowing the user to compute
numerically the gradient of a function, and even higher-order derivatives.(^2^22With PyTorch: torch.autograd.grad.)
Disregarding the computational cost, the full Hessian could theoretically
be computed numerically without approximation.
In order to make this computation feasible in practice, one should
aim for an easier goal: instead of computing the full Hessian, one may
consider computing a smaller matrix, consisting of projections of the Hessian.

Moreover, we may hope that such projections would “squeeze” the close-to-zero
eigenvalues of the Hessian, in such a way that the eigenvalues of the projected matrix
would be mostly not too close to zero.

## 3 Summarizing higher-order information

Let us consider the minimization of a loss function $\mathcal{L}:\mathbb{R}^{P}\rightarrow\mathbb{R}$
according to a variable $\boldsymbol{\theta}\in\mathbb{R}^{P}$.

#### Full computation of the derivatives.

The object containing all the information of order $d$ about $\mathcal{L}$ is
a $d$-linear form, which can be represented as a tensor of order $d$
as follows:

$$ $\displaystyle\frac{\partial^{d}\mathcal{L}}{\partial\boldsymbol{\theta}^{d}}(\boldsymbol{\theta}):\mathbb{R}^{P}\times\cdots\times\mathbb{R}^{P}$ $\displaystyle\rightarrow\mathbb{R}$ $\displaystyle(\mathbf{u}^{1},\cdots,\mathbf{u}^{d})$ $\displaystyle\mapsto\frac{\partial^{d}\mathcal{L}}{\partial\boldsymbol{\theta}^{d}}(\boldsymbol{\theta})[\mathbf{u}^{1},\cdots,\mathbf{u}^{d}].$ $$

This tensor belongs to $\mathbb{R}^{P^{d}}$ and contains $P^{d}$ scalars.
Even when considering the symmetries,
it is computationally too demanding to compute it exactly for $d\geq 2$ in most cases.
For instance, it is not even possible to compute numerically the full Hessian of $\mathcal{L}$
according to the parameters of a small neural network, i.e.,
with $P=10^{5}$ and $d=2$, the Hessian contains $P^{d}=10^{10}$ scalars.

#### Terms of the Taylor expansion.

At the opposite, one can obtain cheap higher-order information
about $\mathcal{L}$ at $\boldsymbol{\theta}$ by considering a specific direction $\mathbf{u}\in\mathbb{R}^{P}$.
The Taylor expansion of $\mathcal{L}(\boldsymbol{\theta}+\mathbf{u})$ gives:

$$ $\displaystyle\mathcal{L}(\boldsymbol{\theta}+\mathbf{u})$ $\displaystyle=\mathcal{L}(\boldsymbol{\theta})+\sum_{d=1}^{D}\frac{1}{d!}\frac{\partial^{d}\mathcal{L}}{\partial\boldsymbol{\theta}^{d}}(\boldsymbol{\theta})[\mathbf{u},\cdots,\mathbf{u}]$ $\displaystyle\hphantom{\;=\;}+o(\|\mathbf{u}\|^{D}).$ $$

The terms of the Taylor expansion contain higher-order information
about $\mathcal{L}$ in the direction $\mathbf{u}$. Notably, they can be used to
predict how $\mathcal{L}(\boldsymbol{\theta})$ would change when translating $\boldsymbol{\theta}$
in the direction of $\mathbf{u}$.
Additionally, computing the first $D$ terms has a complexity of order
$D\times P$, which is manageable even for large models.
The trick allowing for such a low complexity, the *Hessian-vector product*,
has been proposed by
Pearlmutter, () and is recalled in Appendix [A](#A1).

#### An intermediate solution.

Now, let us assume that, in the practical implementation of a gradient-based
method of optimization of $\mathcal{L}(\boldsymbol{\theta})$,
$\boldsymbol{\theta}$ is represented by a tuple of tensors $(\mathbf{T}_{1},\cdots,\mathbf{T}_{S})$.
So, each Taylor term can be expressed as:

$$ $\displaystyle\;\hphantom{=}\;\frac{\partial^{d}\mathcal{L}}{\partial\boldsymbol{\theta}^{d}}(\boldsymbol{\theta})[\mathbf{u},\cdots,\mathbf{u}]$ $\displaystyle=\sum_{s_{1}=1}^{S}\cdots\sum_{s_{d}=1}^{S}\frac{\partial^{d}\mathcal{L}}{\partial\mathbf{T}_{s_{1}}\cdots\partial\mathbf{T}_{s_{d}}}(\boldsymbol{\theta})[\mathbf{U}^{s_{1}},\cdots,\mathbf{U}^{s_{d}}]$ $\displaystyle=\mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})[\mathds{1}_{S},\cdots,\mathds{1}_{S}],$ $$

where $\mathds{1}_{S}\in\mathbb{R}^{S}$ is a vector full of ones,
the tuple of tensors $(\mathbf{U}^{1},\cdots,\mathbf{U}^{S})$ represents $\mathbf{u}$,(^3^33$(\mathbf{U}^{1},\cdots,\mathbf{U}^{S})$ is to $\mathbf{u}$ as $(\mathbf{T}_{1},\cdots,\mathbf{T}_{S})$ is to $\boldsymbol{\theta}$.)
and $\mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})\in\mathbb{R}^{S^{d}}$ is a
tensor of order $d$ with size $S$ in every dimension with values:

$$ $\displaystyle(\mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u}))_{s_{1},\cdots,s_{d}}=\frac{\partial^{d}\mathcal{L}}{\partial\mathbf{T}_{s_{1}}\cdots\partial\mathbf{T}_{s_{d}}}(\boldsymbol{\theta})[\mathbf{U}^{s_{1}},\cdots,\mathbf{U}^{s_{d}}].$ $$

Thus, $\mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})$ is equivalently a $d$-linear form on $\mathbb{R}^{S}$
and a tensor of order $d$ and size $S$ in every dimension.
Moreover, the trick of Pearlmutter, () applies also to
the computation of $\mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})$, which is then much less costly to
compute than the Hessian.

#### Properties of 𝐃 𝜽 ( d ) ​ ( 𝐮 ) superscript subscript 𝐃 𝜽 𝑑 𝐮 \mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u}) .

We show a comparison between the three techniques in Table [1](#S3.T1).
If $S$ is small enough, computing $\mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})$ becomes
feasible for $d\geq 2$.
For usual multilayer perceptrons with $L$ layers,
there is one tensor of weights and one vector of biases per layer,
so $S=2L$.
This allows to compute $\mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})$ in practice
for $d=2$ even when $L\approx 20$.

**Table 1: Comparison between three techniques extracting higher-order information about $\mathcal{L}$: size of the result and complexity of the computation.**
|  | size | complexity |
| --- | --- | --- |
| Full derivative $\frac{\partial^{d}\mathcal{L}}{\partial\boldsymbol{\theta}^{d}}(\boldsymbol{\theta})$ | $P^{d}$ | $P^{d}$ |
| Taylor term $\mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})[\mathds{1}_{S},\cdots]$ | $1$ | $d\times P$ |
| Tensor $\mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})$ | $S^{d}$ | $S^{d-1}\times P$ |

Therefore, the tensors $\mathbf{D}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})$ extract
more information than the naive Taylor terms, while keeping a reasonable computational cost.
Moreover, through their off-diagonal elements, they give access to information about the interactions between the
tensors $(\mathbf{T}_{1},\cdots,\mathbf{T}_{S})$ when there are processed in the function $\mathcal{L}$.

## 4 A scalable second-order optimization method

### 4.1 Presentation of the method

The method presented here consists in partitioning the set of indices of parameters
$\{1,\cdots,P\}$ into $S$ subsets $(\mathcal{I}_{s})_{1\leq s\leq S}$,
attribute for all $1\leq s\leq S$ the same learning rate $\eta_{s}$ to
the parameters $(\theta_{p})_{p\in\mathcal{I}_{s}}$,
and find the vector of learning rates
$\boldsymbol{\eta}=(\eta_{1},\cdots,\eta_{S})$
optimizing the decrease of the loss $\mathcal{L}$
for the current training step $t$, by using its order-2 Taylor approximation.(^4^44With the notation of Section [3](#S3), $\mathcal{I}_{s}$ is the set of indices $p$ of the parameters $\theta_{p}$ belonging to the tensor $\mathbf{T}_{s}$, so the scalars $(\theta_{p})_{p\in\mathcal{I}_{s}}$ correspond exactly to the scalars belonging to $\mathbf{T}_{s}$. Consequently, everything is as if a specific learning rate $\eta_{s}$ is assigned to each tensor $\mathbf{T}_{s}$.)
Formally, given a direction $\mathbf{u}_{t}\in\mathbb{R}^{P}$ in the
parameter space (typically, $\mathbf{u}_{t}=\mathbf{g}_{t}$, the gradient)
and $\mathbf{U}_{t}:=\mathrm{Diag}(\mathbf{u}_{t})\in\mathbb{R}^{P\times P}$,
we consider the training step:

$$ $\displaystyle\boldsymbol{\theta}_{t+1}:=\boldsymbol{\theta}_{t}-\mathbf{U}_{t}\mathbf{I}_{P:S}\boldsymbol{\eta}_{t},$ $$

which is a training step in a direction based on $\mathbf{u}_{t}$,
distorted by a subset-wise step size $\boldsymbol{\eta}_{t}$. Then, we
minimize the order-2 Taylor approximation $\boldsymbol{\Delta}_{2}$ of
$\mathcal{L}(\boldsymbol{\theta}_{t+1})-\mathcal{L}(\boldsymbol{\theta}_{t})$:

$$ $\displaystyle\boldsymbol{\Delta}_{2}(\boldsymbol{\eta}_{t}):=-\mathbf{g}_{t}^{T}\mathbf{U}_{t}\mathbf{I}_{P:S}\boldsymbol{\eta}_{t}+\frac{1}{2}\boldsymbol{\eta}_{t}^{T}\mathbf{I}_{S:P}\mathbf{U}_{t}\mathbf{H}_{t}\mathbf{U}_{t}\mathbf{I}_{P:S}\boldsymbol{\eta}_{t},$ $$

which gives:

$$ $\displaystyle\boldsymbol{\theta}_{t+1}$ $\displaystyle=\boldsymbol{\theta}_{t}-\mathbf{U}_{t}\mathbf{I}_{P:S}\boldsymbol{\eta}_{t}^{*},$ (1) $\displaystyle\boldsymbol{\eta}_{t}^{*}$ $\displaystyle:=(\mathbf{I}_{S:P}\mathbf{U}_{t}\mathbf{H}_{t}\mathbf{U}_{t}\mathbf{I}_{P:S})^{-1}\mathbf{I}_{S:P}\mathbf{U}_{t}\mathbf{g}_{t},$ $$

where $\mathbf{I}_{S:P}\in\mathbb{R}^{S\times P}$ is the *partition matrix*,
verifying $(\mathbf{I}_{S:P})_{sp}=1$ if $p\in\mathcal{I}_{s}$ and
0 0 otherwise, and $\mathbf{I}_{P:S}:=\mathbf{I}_{S:P}^{T}$. Alternatively, $\boldsymbol{\eta}_{t}^{*}$
can be written:

$$ $\displaystyle\boldsymbol{\eta}_{t}^{*}=\bar{\mathbf{H}}_{t}^{-1}\bar{\mathbf{g}}_{t},\quad\bar{\mathbf{H}}_{t}$ $\displaystyle:=\mathbf{I}_{S:P}\mathbf{U}_{t}\mathbf{H}_{t}\mathbf{U}_{t}\mathbf{I}_{P:S}$ $\displaystyle\in\mathbb{R}^{S\times S},$ $\displaystyle\bar{\mathbf{g}}_{t}$ $\displaystyle:=\mathbf{I}_{S:P}\mathbf{U}_{t}\mathbf{g}_{t}$ $\displaystyle\in\mathbb{R}^{S}.$ $$

Details are provided in Appendix [B](#A2).

With the notation of Section [3](#S3),
$\bar{\mathbf{H}}_{t}=\mathbf{D}_{\boldsymbol{\theta}_{t}}^{(2)}(\mathbf{u}_{t})$
and $\bar{\mathbf{g}}_{t}=\mathbf{D}_{\boldsymbol{\theta}_{t}}^{(1)}(\mathbf{u}_{t})$.
Incidentally, computing $\bar{\mathbf{H}}$ is of complexity
$SP$, and solving the system $\bar{\mathbf{H}}\mathbf{x}=\bar{\mathbf{g}}$
is of complexity $S^{2}$.

### 4.2 Regularizing 𝐇 ¯ ¯ 𝐇 \bar{\mathbf{H}} by using order-3 information

The method proposed in Section [4.1](#S4.SS1)
requires to compute $\boldsymbol{\eta}^{*}=\bar{\mathbf{H}}^{-1}\bar{\mathbf{g}}$.
Usually, inverting such a linear system at every step is considered as hazardous
and unstable. That is why, when using Newton’s method, instead of computing
a direction of descent $\mathbf{u}:=\mathbf{H}^{-1}\mathbf{g}$,
it is very common to add a regularization term:
$\mathbf{u}_{\lambda}:=\left(\mathbf{H}+\lambda\mathbf{I}\right)^{-1}\mathbf{g}$
(Nocedal and Wright,, , Chap. 6.3).

However, this regularization technique suffers from several drawbacks.
The most obvious is computational: finding the best hyperpameter $\lambda$
requires either many runs of the full optimization process,
either one run with a technique adjusting $\lambda$ on the fly. Anyway,
there is an inevitable extra computational cost.

Besides of that, the theoretical ground of such a regularization technique
is not fully satisfactory.
Basically, the main problem is not having a matrix $\bar{\mathbf{H}}$ with close-to-zero eigenvalues:
after all, if the loss landscape is very flat in a specific direction,
it is better to make a large training step.
The problem lies in the order-2 approximation of the loss
made in the training step ([1](#S4.E1)), as well as in Newton’s method:
instead of optimizing the true decrease of the loss, we optimize the decrease
of its order-2 approximation.
Thus, the practical question is: does this approximation model faithfully
the loss at the current point $\boldsymbol{\theta}_{t}$, in a region encompassing also
the next point $\boldsymbol{\theta}_{t+1}$?

To answer that question, one has to take into account order-3 information,
and regularize $\bar{\mathbf{H}}$ in such a way that the resulting update remains
in an area around $\boldsymbol{\theta}_{t}$ where the cubic term of the Taylor approximation
is negligible.
In practice, we propose an anisotropic version of Nesterov’s cubic regularization
Nesterov and Polyak, ().

#### Anisotropic Nesterov cubic regularization.

By using the technique presented in Section [3](#S3),
the diagonal coefficients $(D_{1},\cdots,D_{S})$
of $\mathbf{D}_{\boldsymbol{\theta}}^{(3)}(\mathbf{u})\in\mathbb{R}^{S\times S\times S}$
are available with little computational cost.
Let:

$$ $\displaystyle\mathbf{D}:=\mathrm{Diag}(|D_{1}|^{1/3},\cdots,|D_{S}|^{1/3})\in\mathbb{R}^{S}.$ $$

We modify the method of Nesterov and Polyak, () by
integrating an anisotropic factor $\mathbf{D}$ into the cubic term.
Thus, our goal is to minimize according to $\boldsymbol{\eta}$ the function $T$:

$$ $\displaystyle T(\boldsymbol{\eta}):=-\boldsymbol{\eta}^{T}\bar{\mathbf{g}}+\frac{1}{2}\boldsymbol{\eta}\bar{\mathbf{H}}\boldsymbol{\eta}+\frac{\lambda_{\mathrm{int}}}{6}\|\mathbf{D}\boldsymbol{\eta}\|^{3},$ $$

where $\lambda_{\mathrm{int}}$ is the *internal damping*
coefficient, which can be used to tune the strength of the cubic
regularization. Under conditions detailed in Appendix [D](#A4),
this minimization problem is equivalent to finding a solution $\boldsymbol{\eta}_{*}$ such that:

$$ $\displaystyle\boldsymbol{\eta}_{*}=\left(\bar{\mathbf{H}}+\frac{\lambda_{\mathrm{int}}}{2}\|\mathbf{D}\boldsymbol{\eta}_{*}\|\mathbf{D}^{2}\right)^{-1}\bar{\mathbf{g}},$ (2) $$

which is no more than a regularized version of ([1](#S4.E1)).
Finally, this multi-dimensional minimization
problem boils down to a scalar root finding problem (see Appendix [D](#A4)).

### 4.3 Properties

The final method is a combination of the training step ([1](#S4.E1))
with regularization ([2](#S4.E2)):

###### Method 4.1 .

Training step $\boldsymbol{\theta}_{t+1}=\boldsymbol{\theta}_{t}-\mathbf{U}_{t}\mathbf{I}_{P:S}\boldsymbol{\eta}_{t}^{*}$,
where $\eta_{*}$ is the solution with the largest norm $\|\mathbf{D}\boldsymbol{\eta}\|$ of the equation:

$$ $\displaystyle\boldsymbol{\eta}=\left(\bar{\mathbf{H}}+\frac{\lambda_{\mathrm{int}}}{2}\|\mathbf{D}\boldsymbol{\eta}\|\mathbf{D}^{2}\right)^{-1}\bar{\mathbf{g}}.$ $$

Method [4.1](#S4.Thmtheorem1) has several interesting properties.

###### Method 4.1 .

#### Encompassing Newton’s method and Cauchy’s steepest descent.

Without the cubic regularization ($\lambda_{\mathrm{int}}=0$),
Newton’s method is recovered when using the *discrete partition*, that
is, $S=P$ with $\mathcal{I}_{s}=\{s\}$ for all $s$,
and Cauchy’s steepest descent is recovered when using
the *trivial partition*, that is, $S=1$ with
$\mathcal{I}_{1}=\{1,\cdots,P\}$. See Appendix [C](#A3)
for more details.

#### No need to compute or approximate the full Hessian.

The full computation of the Hessian $\mathbf{H}_{t}\in\mathbb{R}^{P\times P}$ is not required.
Instead, one only needs to compute the $S\times S$ matrix
$\bar{\mathbf{H}}_{t}:=\mathbf{I}_{S:P}\mathbf{U}_{t}\mathbf{H}_{t}\mathbf{U}_{t}\mathbf{I}_{P:S}$,
which can be done efficiently by computing
$\mathbf{u}^{T}\mathbf{H}_{t}\mathbf{v}$ for a number $S\times S$ of pairs of well-chosen directions
$(\mathbf{u},\mathbf{v})\in\mathbb{R}^{P}\times\mathbb{R}^{P}$.
This property is useful especially when $S\ll P$.
When optimizing a neural network with $L=10$
layers and $P=10^{6}$ parameters, one can naturally partition
the set of parameters into $S=2L$ subsets,
each one containing either all the weights or all the biases of each of the
$L$ layers. In this situation, one has to solve a linear system
of size $2L=20$ at each step, which is much more reasonable
than solving a linear system with $P=10^{6}$ equations.
We call this natural partition of the parameters of a neural network
the *canonical partition*.

#### No need to solve a very large linear system.

Using Equations ([1](#S4.E1)) or ([2](#S4.E2)) only requires solving
a linear system of $S$ equations, instead of $P$ in Newton’s method.
With the cubic regularization, only a constant term is added to the complexity,
since it is a matter of scalar root finding.

#### The interactions between different tensors are not neglected.

The matrix $\bar{\mathbf{H}}_{t}$, which simulates the Hessian $\mathbf{H}_{t}$,
is basically dense, which means that
it does not exhibit a diagonal or block-diagonal structure.
So, the interactions between subsets of parameters are taken into account
when performing optimization steps.
In the framework of neural networks with the canonical partition,
it means that interactions between layers are taken into account
during optimization, including even when the layers are far from
each other. This is a great advantage compared to
many existing approximations of the Hessian or its inverse,
which are diagonal or block-diagonal.

#### Invariance by subset-wise affine reparameterization.

As showed in Appendix [E](#A5),
under a condition on the directions $\mathbf{u}_{t}$,
the trajectory of optimization of a model trained by Method [4.1](#S4.Thmtheorem1)
is invariant by affine reparameterization
of the sub-vectors of parameters $\boldsymbol{\theta}_{\mathcal{I}_{s}}:=\mathrm{vec}(\{\theta_{p}:p\in\mathcal{I}_{s}\})$.(^5^55This holds typically if $\mathbf{u}_{t}$ is the gradient or a moving average of the gradients (momentum).)
More precisely, let $(\alpha_{s})_{1\leq s\leq S}$ and
$(\beta_{s})_{1\leq s\leq S}$ be respectively a sequence of nonzero scalings
and a sequence of offsets, and $\tilde{\boldsymbol{\theta}}$ such that,
for all $1\leq s\leq S$, $\tilde{\boldsymbol{\theta}}_{\mathcal{I}_{s}}=\alpha_{s}\boldsymbol{\theta}_{\mathcal{I}_{s}}+\beta_{s}$.
Then, the training trajectory of the model is the same with both parameterizations
$\boldsymbol{\theta}$ and $\tilde{\boldsymbol{\theta}}$.
This property is interesting in the case of neural networks,
where one can either use the usual parameterization or
the NTK parameterization, which consists in a layer-wise scaling of
the parameters.

Compared to the standard regularization $\bar{\mathbf{H}}+\lambda\mathbf{I}$ and
Nesterov’s cubic regularization, the anisotropic Nesterov regularization
does not break the property of invariance by subset-wise scaling of the parameters
of ([1](#S4.E1)).
Among all, this is due to our choice to keep only the diagonal coefficients of
$\mathbf{D}_{\boldsymbol{\theta}}^{(3)}(\mathbf{u})$ while discarding the others.
Notably, the off-diagonal coefficients contain cross-derivatives,
which would be difficult to include in an invariant training step.

## 5 Experiments

### 5.1 Empirical computation of 𝐇 ¯ ¯ 𝐇 \bar{\mathbf{H}} and 𝜼 𝜼 \boldsymbol{\eta}

As recalled in Section [2](#S2), many works
perform a diagonal, block-diagonal or block-tridiagional (Martens and Grosse,, )
approximation of the Hessian or its inverse.

Since a summary $\bar{\mathbf{H}}$ of the Hessian and its inverse $\bar{\mathbf{H}}^{-1}$ are available
and all their off-diagonal coefficients have been computed and kept, one can
to check if these coefficient are indeed negligible.

#### Setup.

We have trained LeNet-5 and VGG-11’(^6^66VGG-11’ is a variant of VGG-11 with only one fully-connected layer at the end, instead of 3.)
on CIFAR-10 by SGD with momentum.
Before each epoch, we compute the full-batch gradient, denoted by $\mathbf{u}$,
which we use as a direction to compute $\bar{\mathbf{H}}$, again in full-batch.
We report submatrices of $\bar{\mathbf{H}}$ and $\bar{\mathbf{H}}^{-1}$ at initialization
and at the epoch where the validation loss is the best in
Figure [1](#S1.F1) (LeNet) and Figure [2](#S5.F2) (VGG-11’).

For the sake of readability,
$\bar{\mathbf{H}}$ has been divided into blocks:
a weight-weight block $\bar{\mathbf{H}}_{\text{WW}}$,
a bias-bias block $\bar{\mathbf{H}}_{\text{BB}}$,
and a weight-bias block $\bar{\mathbf{H}}_{\text{WB}}$.
They represent the interactions between the layers:
for instance, $(\bar{\mathbf{H}}_{\text{WB}})_{l_{1}l_{2}}$
represents the interaction between the tensor of weights of layer $l_{1}$
and the vector of biases of layer $l_{2}$. Naturally, $l_{1}$ and $l_{2}$ are not
necessarily different.

#### Results on 𝐇 ¯ ¯ 𝐇 \bar{\mathbf{H}} .

First, the block-diagonal
approximation of the Hessian is indeed very rough,
while the block-diagonal approximation of the inverse Hessian
seems to be more reasonable (at least in these setups), which has already been
shown by Martens and Grosse, ().
Second, long-range interactions between layers seem to exist, both at initialization
and after several epochs. For instance, in LeNet, all the layers (except the first one)
seem to interact together at initialization (Fig. [1](#S1.F1)).
In te matrix $\bar{\mathbf{H}}^{-1}$ computed on VGG, the block of the last 3 layers
interact strongly and the block of the last 6 layers also interact, but a little less.

According to these observations, a neural networks should also be regarded as a whole,
in which layers can hardly be studied independently from each other.

As far as we know, this result is the first scalable illustration of
interactions between layers far from each other, based on
second-order information.

Figure: Figure 2: VGG-11’ trained by SGD on CIFAR-10. Submatrices of $\bar{\mathbf{H}}$ (first row) and $\bar{\mathbf{H}}^{-1}$ (second row), where focus is on interactions: weight-weight, weight-bias, bias-bias of the different layers, at initialization and before the 4th epoch (best validation loss).
Refer to caption: /html/2312.03885/assets/x2.png

#### Results on 𝜼 ∗ subscript 𝜼 \boldsymbol{\eta}_{*} .

The evolution of the learning rates $\eta_{*}$ computed according to ([2](#S4.E2))
in LeNet and VGG has been reported in Figure [3](#S5.F3).
First, the learning rates computed
for the biases are larger than for the weights.
Second, even when considering only the weights, the computed $\eta_{*}$
can differ by several order of magnitude.
Finally, the first two layers of LeNet (which are convolutional)
have smaller $\eta_{*}$ than the last three (which are fully-connected).
Conversely, in VGG, the weights of the last (convolutional) layers have
a smaller $\eta_{*}$ than in the first ones.

Figure: Figure 3: Setup: LeNet, VGG-11’ trained by SGD on CIFAR-10. Learning rates $\eta_{*}$ computed according to ([2](#S4.E2)), specific to each tensor of weights and tensor of biases of each layer. For each epoch $k\in\{10,30,50,70,90\}$, the reported value has been averaged over the epochs $[k-10,k+9]$ to remove the noise.
Refer to caption: /html/2312.03885/assets/x3.png

### 5.2 Training experiments

In this section, we show a proof-of-concept of the optimization method
[4.1](#S4.Thmtheorem1), on simple vision tasks
and with medium-sized neural networks.
All the implementation details are available in Appendix [F](#A6).
Notably, we have introduced a damping factor $\lambda_{1}$, which leads to the
following modification of the training step ([1](#S4.E1)):

$$ $\displaystyle\boldsymbol{\theta}_{t+1}$ $\displaystyle=\boldsymbol{\theta}_{t}-\lambda_{1}\mathbf{U}_{t}\mathbf{I}_{P:S}\boldsymbol{\eta}_{t}^{*}.$ $$

#### Setup.

We consider 4 image classification setups:

- •
MLP: multilayer perceptron trained on MNIST with
layers of sizes 1024, 200, 100, 10, with $\tanh$ activation function;
- •
LeNet: LeNet-5 (LeCun et al.,, ) model trained
on CIFAR-10 with 2 convolutional layers of sizes 6, 16, and
3 fully connected layers of sizes 120, 84, 10;
- •
VGG: VGG-11’ trained on CIFAR-10.
VGG-11’ is a variant of VGG-11 (Simonyan and Zisserman,, )
with only one fully-connected layer at the end, instead of 3,
with $\mathrm{ELU}$ activation function (Clevert et al.,, ), without batch-norm;
- •
BigMLP: multilayer perceptron trained on CIFAR-10,
with 20 layers of size 1024 and one classification layer of size 10,
with $\mathrm{ELU}$ activation function.

And we have tested 3 optimization methods:

- •
Adam: the best learning rate has been selected by grid-search;
- •
K-FAC: the best learning rate and damping have been selected by grid-search;
- •
NewtonSummary (ours): the best $\lambda_{1}$ and
$\lambda_{\mathrm{int}}$ have been selected by grid search.

#### Results.

The evolution of the training loss has been plotted in Figure [4](#S5.F4)
for each of the 3 optimization methods, for 5 different seeds.
In each series of experiments, training is successful,
while being slow or unstable at some points.
Anyway, the minimum training loss achieved by Method [4.1](#S4.Thmtheorem1) (NewtonSummary)
is comparable to the minimum training loss achieved by K-KAC or Adam
in all the series except with BigMLP, whose training is extremely slow.

Some runs have been stopped early due to failures caused by
a very large step size (large $\boldsymbol{\eta}_{*}$). Actually, we did not
use any safeguard, such as a regularization term $\lambda\mathbf{I}$ added to $\bar{\mathbf{H}}$,
or clipping the learning rates exceeding a given value.
This choice has been made in order to maintain the property of invariance
by reparameterization, and to avoid increasing the number of hyperparameters to tune.

Figure: Figure 4: Training curves in different setups. The reported loss is the negative log-likelihood computed on the training set.
Refer to caption: /html/2312.03885/assets/x4.png

#### Extension to very large models.

Since the matrix $\bar{\mathbf{H}}$ can be computed numerically
as long as $S$ remains relatively small,
then, for very large models, this method may become unpractical.
However, Method [4.1](#S4.Thmtheorem1) is flexible enough to
be adapted to such models: one can regroup tensors “of the same
kind” in order to build a coarser partition of the parameters,
and consequently obtain a small $S$, which is exactly what is needed to compute
$\bar{\mathbf{H}}$ and invert it.
The main difficulty would then be to find a good partition of the parameters,
by grouping all the tensors that “look alike”.
We provide an example in Appendix [G](#A7),
with a very deep multilayer perceptron.

## 6 Discussion

#### Convergence rate.

Method [4.1](#S4.Thmtheorem1) does not come with any convergence rate.
Given the convergence rates of Newton’s method and Cauchy’s steepest
descent, we may expect to find some in-between convergence rates.
Since Cauchy’s steepest method is vulnerable
to a highly anisotropic Hessian, it would be valuable to know how much
this weakness is overcome with our method.

#### Practicality.

Despite the interesting properties of Method [4.1](#S4.Thmtheorem1)
(scalabity, invariance by reparameterization, evaluation of long-range
interactions between layers),
we have proposed nothing more than a proof-of-concept.
As showed in the experimental section, this method
is subject to instabilities during training, which is
certainly not surprising for a second-order method,
but it is not acceptable for the end user.
Therefore, some additional tricks should be added
to improve the stability of training, which
is a common practice, but comes usually with additional hyperparameters to tune.
Besides, this method has to be improved to reduce the duration of each epoch,
which is longer than with K-FAC.

## 7 Acknowledgments

This work was granted access to the HPC resources of IDRIS under the allocation
2023-AD011013762R1 made by GENCI.

## Appendix A Fast computation of the terms of the Taylor expansion

In this appendix, we show how to use the trick of Pearlmutter, ()
to compute the terms of the Taylor expansion of $\mathcal{L}$.
We recall that we want to compute:

$$ $\displaystyle\tilde{\mathbf{D}}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u}):=\frac{\partial^{d}\mathcal{L}}{\partial\boldsymbol{\theta}^{d}}(\boldsymbol{\theta})[\mathbf{u},\cdots,\mathbf{u}]\in\mathbb{R}.$ $$

We use the following recursion formula:

$$ $\displaystyle\tilde{\mathbf{D}}_{\boldsymbol{\theta}}^{(d+1)}(\mathbf{u})=\left(\frac{\partial\tilde{\mathbf{D}}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})}{\partial\boldsymbol{\theta}}\right)^{T}\mathbf{u}.$ $$

Therefore, at each step $d$, we only have to compute the gradient of
a scalar $\tilde{\mathbf{D}}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})$ according to $\boldsymbol{\theta}\in\mathbb{R}^{P}$,
and compute a dot product in the space $\boldsymbol{\theta}\in\mathbb{R}^{P}$.
So, computing $\tilde{\mathbf{D}}_{\boldsymbol{\theta}}^{(d)}(\mathbf{u})$
has a complexity proportional to $d\times P$,
and does not require the computation of the full
tensor $\frac{\partial^{d}\mathcal{L}}{\partial\boldsymbol{\theta}^{d}}(\boldsymbol{\theta})\in\mathbb{R}^{P^{d}}$.

## Appendix B Derivation of the second-order method

We consider an update of $\boldsymbol{\theta}$ with one learning rate $\eta_{s}$
for each subset $\mathcal{I}_{s}$ of parameters.
Let $\mathbf{I}_{S:P}\in\mathbb{R}^{S\times P}$ be the *partition matrix*,
verifying $(\mathbf{I}_{S:P})_{sp}=1$ if $p\in\mathcal{I}_{s}$ and
0 0 otherwise, and $\mathbf{I}_{P:S}:=\mathbf{I}_{S:P}^{T}$.
We consider an update based on a given direction $\mathbf{u}_{t}$ and
we define $\mathbf{U}_{t}:=\mathrm{Diag}(\mathbf{u}_{t})$:

$$ $\displaystyle\boldsymbol{\theta}_{t+1}=\boldsymbol{\theta}_{t}-\mathbf{U}_{t}\mathbf{I}_{P:S}\boldsymbol{\eta},$ $$

where $\boldsymbol{\eta}=(\eta_{1},\cdots,\eta_{S})\in\mathbb{R}^{S}$.

The second-order approximation of $\mathcal{L}$ gives:

$$ $\displaystyle\mathcal{L}(\boldsymbol{\theta}_{t+1})$ $\displaystyle=\mathcal{L}(\boldsymbol{\theta}_{t}-\mathbf{U}_{t}\mathbf{I}_{P:S}\boldsymbol{\eta})$ $\displaystyle=\mathcal{L}(\boldsymbol{\theta}_{t})-\boldsymbol{\eta}^{T}\mathbf{I}_{S:P}\mathbf{U}_{t}\frac{\partial\mathcal{L}}{\partial\boldsymbol{\theta}}(\boldsymbol{\theta}_{t})$ $\displaystyle\hphantom{=}\;+\frac{1}{2}\boldsymbol{\eta}^{T}\mathbf{I}_{S:P}\mathbf{U}_{t}\frac{\partial^{2}\mathcal{L}}{\partial\boldsymbol{\theta}^{2}}(\boldsymbol{\theta}_{t})\mathbf{U}_{t}\mathbf{I}_{P:S}\boldsymbol{\eta}+o(\|\boldsymbol{\eta}\|^{2})$ $\displaystyle=\mathcal{L}(\boldsymbol{\theta}_{t})-\boldsymbol{\eta}^{T}\mathbf{I}_{S:P}\mathbf{U}_{t}\mathbf{g}_{t}$ $\displaystyle\hphantom{=}\;+\frac{1}{2}\boldsymbol{\eta}^{T}\mathbf{I}_{S:P}\mathbf{U}_{t}\mathbf{H}_{t}\mathbf{U}_{t}\mathbf{I}_{P:S}\boldsymbol{\eta}+o(\|\boldsymbol{\eta}\|^{2})$ $\displaystyle=\mathcal{L}(\boldsymbol{\theta}_{t})-\boldsymbol{\eta}^{T}\bar{\mathbf{g}}_{t}+\frac{1}{2}\boldsymbol{\eta}^{T}\bar{\mathbf{H}}_{t}\boldsymbol{\eta}+o(\|\boldsymbol{\eta}\|^{2}),$ $$

where:

$$ $\displaystyle\bar{\mathbf{g}}_{t}$ $\displaystyle:=\mathbf{I}_{S:P}\mathbf{U}_{t}\mathbf{g}_{t}$ $\displaystyle\in\mathbb{R}^{S},$ $\displaystyle\bar{\mathbf{H}}_{t}$ $\displaystyle:=\mathbf{I}_{S:P}\mathbf{U}_{t}\mathbf{H}_{t}\mathbf{U}_{t}\mathbf{I}_{P:S}$ $\displaystyle\in\mathbb{R}^{S\times S}.$ $$

Now, we omit the $o(\|\boldsymbol{\eta}\|^{2})$ term and we want to minimize
according to $\boldsymbol{\eta}$ the variation of the loss:

$$ $\displaystyle\mathcal{L}(\boldsymbol{\theta}_{t+1})-\mathcal{L}(\boldsymbol{\theta}_{t})$ $\displaystyle\approx\boldsymbol{\Delta}_{2}(\boldsymbol{\eta}):=\mathcal{L}(\boldsymbol{\theta}_{t})-\boldsymbol{\eta}^{T}\bar{\mathbf{g}}_{t}+\frac{1}{2}\boldsymbol{\eta}^{T}\bar{\mathbf{H}}_{t}\boldsymbol{\eta}.$ $$

We have:

$$ $\displaystyle\frac{\partial\boldsymbol{\Delta}_{2}}{\partial\boldsymbol{\eta}}=-\bar{\mathbf{g}}_{t}+\bar{\mathbf{H}}_{t}\boldsymbol{\eta},$ $$

which is zero if, and only if:

$$ $\displaystyle\bar{\mathbf{g}}_{t}=\bar{\mathbf{H}}_{t}\boldsymbol{\eta}.$ $$

If this linear system can be inverted, one may choose:

$$ $\displaystyle\boldsymbol{\eta}=\boldsymbol{\eta}_{t}^{*}:=\bar{\mathbf{H}}_{t}^{-1}\bar{\mathbf{g}}_{t}.$ $$

## Appendix C Link with Cauchy’s steepest descent and Newton’s method

#### Cauchy’s steepest descent.

Let us consider the trivial partition: $S=1$, $\mathcal{I}_{1}=\{1,\cdots,P\}$.
So, $\mathbf{I}_{S:P}=(1,\cdots,1)=\mathds{1}_{S}^{T}$. Therefore, the training step is:

$$ $\displaystyle\boldsymbol{\theta}_{t+1}$ $\displaystyle:=\boldsymbol{\theta}_{t}-\mathbf{G}_{t}\mathds{1}_{S}(\mathds{1}_{S}^{T}\mathbf{G}_{t}\mathbf{H}_{t}\mathbf{G}_{t}\mathds{1}_{S})^{-1}\mathds{1}_{S}^{T}\mathbf{G}_{t}\mathbf{g}_{t}$ $\displaystyle=\boldsymbol{\theta}_{t}-\mathbf{g}_{t}\frac{\mathbf{g}_{t}^{T}\mathbf{g}_{t}}{\mathbf{g}_{t}^{T}\mathbf{H}_{t}\mathbf{g}_{t}},$ $$

since $\mathbf{G}_{t}\mathds{1}_{S}=\mathbf{g}_{t}$.
We recover Cauchy’s steepest descent.

#### Newton’s method.

Since we aim to recover Newton’s method, we assume
that the Hessian $\mathbf{H}_{t}$ is positive definite.
Let us consider the discrete partition: $S=P$, $\mathcal{I}_{s}=\{s\}$.
So, $\mathbf{I}_{S:P}=\mathbf{I}_{P}$, the identity matrix of $\mathbb{R}^{P\times P}$.
Therefore, the training step is:

$$ $\displaystyle\boldsymbol{\theta}_{t+1}$ $\displaystyle:=\boldsymbol{\theta}_{t}-\mathbf{G}_{t}(\mathbf{G}_{t}\mathbf{H}_{t}\mathbf{G}_{t})^{-1}\mathbf{G}_{t}\mathbf{g}_{t}.$ $$

To perform the training step, we have to find $\mathbf{x}\in\mathbb{R}^{P}$
such that: $(\mathbf{G}_{t}\mathbf{H}_{t}\mathbf{G}_{t})^{-1}\mathbf{G}_{t}\mathbf{g}_{t}=\mathbf{x}$.
That is, solve the linear system
$\mathbf{G}_{t}\mathbf{H}_{t}\mathbf{G}_{t}\mathbf{x}=\mathbf{G}_{t}\mathbf{g}_{t}$.
In the case where all the coordinates of the gradient $\mathbf{g}_{t}$ are nonzero,
we can write:

$$ $\displaystyle\mathbf{x}=\mathbf{G}_{t}^{-1}\mathbf{H}_{t}^{-1}\mathbf{G}_{t}^{-1}\mathbf{G}_{t}\mathbf{g}_{t}=\mathbf{G}_{t}^{-1}\mathbf{H}_{t}^{-1}\mathbf{g}_{t},$ $$

so the training step becomes:

$$ $\displaystyle\boldsymbol{\theta}_{t+1}$ $\displaystyle:=\boldsymbol{\theta}_{t}-\mathbf{G}_{t}\mathbf{x}=\boldsymbol{\theta}_{t}-\mathbf{H}_{t}^{-1}\mathbf{g}_{t},$ $$

which corresponds to Newton’s method.

## Appendix D Anisotropic Nesterov cubic regularization

Let $\mathbf{D}$ be a diagonal matrix whose diagonal coefficients are all
strictly positive: $\mathbf{D}=\mathrm{Diag}(d_{1},\cdots,d_{S})$,
with $d_{i}>0$ for all $i$.

We want to minimize the function:

$$ $\displaystyle T(\boldsymbol{\eta}):=-\boldsymbol{\eta}^{T}\bar{\mathbf{g}}+\frac{1}{2}\boldsymbol{\eta}\bar{\mathbf{H}}\boldsymbol{\eta}+\frac{\lambda_{\mathrm{int}}}{6}\|\mathbf{D}\boldsymbol{\eta}\|^{3}.$ $$

The function $T$ is strictly convex if, and only if,
$\bar{\mathbf{H}}$ is positive definite.
Moreover, $T$ is differentiable twice and has at least one global minimum $\boldsymbol{\eta}_{*}$,
so $\frac{\partial T}{\partial\boldsymbol{\eta}}(\boldsymbol{\eta}_{*})=0$.
Therefore, we first look for the solutions of the equation
$\frac{\partial T}{\partial\boldsymbol{\eta}}(\boldsymbol{\eta})=0$.

We have:

$$ $\displaystyle\frac{\partial T}{\partial\boldsymbol{\eta}}(\boldsymbol{\eta})$ $\displaystyle=-\bar{\mathbf{g}}+\bar{\mathbf{H}}\boldsymbol{\eta}+\frac{\lambda_{\mathrm{int}}}{2}\|\mathbf{D}\boldsymbol{\eta}\|\mathbf{D}^{2}\boldsymbol{\eta}$ $\displaystyle=-\bar{\mathbf{g}}+\left(\bar{\mathbf{H}}+\frac{\lambda_{\mathrm{int}}}{2}\|\mathbf{D}\boldsymbol{\eta}\|\mathbf{D}^{2}\right)\boldsymbol{\eta},$ $$

which is equal to zero if, and only if:

$$ $\displaystyle\bar{\mathbf{g}}=\left(\bar{\mathbf{H}}+\frac{\lambda_{\mathrm{int}}}{2}\|\mathbf{D}\boldsymbol{\eta}\|\mathbf{D}^{2}\right)\boldsymbol{\eta}.$ (3) $$

Let $\boldsymbol{\eta}^{\prime}:=\mathbf{D}\boldsymbol{\eta}$. Eqn. ([3](#A4.E3)) is then equivalent to:

$$ $\displaystyle\bar{\mathbf{g}}$ $\displaystyle=\left(\bar{\mathbf{H}}\mathbf{D}^{-1}+\frac{\lambda_{\mathrm{int}}}{2}\|\boldsymbol{\eta}^{\prime}\|\mathbf{D}\right)\boldsymbol{\eta}^{\prime}.$ $\displaystyle=\frac{\lambda_{\mathrm{int}}}{2}\mathbf{D}\left(\frac{2}{\lambda_{\mathrm{int}}}\mathbf{D}^{-1}\bar{\mathbf{H}}\mathbf{D}^{-1}+\|\boldsymbol{\eta}^{\prime}\|\mathbf{I}\right)\boldsymbol{\eta}^{\prime}$ $$

Let $\mathbf{K}:=\frac{2}{\lambda_{\mathrm{int}}}\mathbf{D}^{-1}\bar{\mathbf{H}}\mathbf{D}^{-1}$.
We want to solve:

$$ $\displaystyle\bar{\mathbf{g}}$ $\displaystyle=\frac{\lambda_{\mathrm{int}}}{2}\mathbf{D}\left(\mathbf{K}+\|\boldsymbol{\eta}^{\prime}\|\mathbf{I}\right)\boldsymbol{\eta}^{\prime}$ (4) $$

Since $\mathbf{K}$ is positive definite if, and only if, $\bar{\mathbf{H}}$ is positive definite,
we consider the following cases.

#### Case 1: 𝐇 ¯ ¯ 𝐇 \bar{\mathbf{H}} is positive definite.

In this case, Eqn. ([4](#A4.E4)) is equivalent to:

$$ $\displaystyle\boldsymbol{\eta}^{\prime}$ $\displaystyle=\frac{2}{\lambda_{\mathrm{int}}}\left(\mathbf{K}+\|\boldsymbol{\eta}^{\prime}\|\mathbf{I}\right)^{-1}\mathbf{D}^{-1}\bar{\mathbf{g}}.$ $$

Now, let $r=\|\boldsymbol{\eta}^{\prime}\|$. We want to solve:

$$ $\displaystyle r=\frac{2}{\lambda_{\mathrm{int}}}\left\|\left(\mathbf{K}+r\mathbf{I}\right)^{-1}\mathbf{D}^{-1}\bar{\mathbf{g}}\right\|.$ (5) $$

Trivially: $\boldsymbol{\eta}$ solution of ([3](#A4.E3))
$\Rightarrow$ $\mathbf{D}\boldsymbol{\eta}$ solution of ([4](#A4.E4))
$\Rightarrow$ $\|\mathbf{D}\boldsymbol{\eta}\|$ solution of ([5](#A4.E5)).
Reciprocally: $r$ solution of ([5](#A4.E5))
$\Rightarrow$ $\boldsymbol{\eta}^{\prime}:=(\bar{\mathbf{H}}\mathbf{D}^{-1}+\frac{\lambda_{\mathrm{int}}}{2}r\mathbf{D})^{-1}\bar{\mathbf{g}}$ solution of ([4](#A4.E4))
$\Rightarrow$ $\mathbf{D}^{-1}\boldsymbol{\eta}^{\prime}$ solution of ([3](#A4.E3)).

Therefore, in order to find the unique global minimum of $T$,
it is sufficient to solve Eqn. ([5](#A4.E5)).
This is doable numerically.

#### Case 2: 𝐇 ¯ ¯ 𝐇 \bar{\mathbf{H}} is not positive definite.

We follow the procedure proposed in (Nesterov and Polyak,, , Section 5).
Let $\lambda_{\mathrm{min}}$ be the minimum eigenvalue of $\mathbf{K}$.
So, $\lambda_{\mathrm{min}}\leq 0$.
Following Nesterov and Polyak, (), we look for the unique $\boldsymbol{\eta}^{\prime}$
belonging to $\mathcal{C}:=\{\boldsymbol{\eta}^{\prime}\in\mathbb{R}^{S}:\|\boldsymbol{\eta}^{\prime}\|>|\lambda_{\mathrm{min}}|\}$,
which is also the solution of maximum norm of Eqn. ([4](#A4.E4)).
Conditionally to $\boldsymbol{\eta}^{\prime}\in\mathcal{C}$,
$(\mathbf{K}+\|\boldsymbol{\eta}^{\prime}\|\mathbf{I})$ is invertible. So
we only need to solve:

$$ $\displaystyle r>|\lambda_{\mathrm{min}}|:\quad r=\frac{2}{\lambda_{\mathrm{int}}}\left\|\left(\mathbf{K}+r\mathbf{I}\right)^{-1}\mathbf{D}^{-1}\bar{\mathbf{g}}\right\|,$ (6) $$

which has exactly one solution $r_{*}$. Then,
we compute $\boldsymbol{\eta}_{*}:=\mathbf{D}^{-1}(\bar{\mathbf{H}}\mathbf{D}^{-1}+\frac{\lambda_{\mathrm{int}}}{2}r_{*}\mathbf{D})^{-1}\bar{\mathbf{g}}$.

## Appendix E Approximate invariance by subset-wise affine reparameterization

We consider a parameter $\tilde{\boldsymbol{\theta}}$ such that:

$$ $\displaystyle\boldsymbol{\theta}=\varphi(\tilde{\boldsymbol{\theta}}),$ $$

where $\varphi$ is an invertible map, affine on each subset of parameters.
Therefore, its Jacobian is: $\mathbf{J}=\mathrm{Diag}(\alpha_{1},\cdots,\alpha_{p})$,
where, for all $1\leq s\leq S$ and $1\leq p_{1},p_{2}\leq P$,
we have:

$$ $\displaystyle p_{1},p_{2}\in\mathcal{I}_{s}\Rightarrow\alpha_{p_{1}}=\alpha_{p_{2}}=:a_{s}.$ $$

Also, let $\bar{\mathbf{J}}=\mathrm{Diag}(a_{1},\cdots,a_{S})$.

We want to compare the training trajectory of $\mathcal{L}(\boldsymbol{\theta})$
and $\mathcal{L}(\varphi(\tilde{\boldsymbol{\theta}}))$ when using Method [4.1](#S4.Thmtheorem1).
For any quantity $\mathbf{x}$ computed with the parameterization $\boldsymbol{\theta}$,
we denote by $\tilde{\mathbf{x}}$ its counterpart computed with the parameterization $\tilde{\boldsymbol{\theta}}$.

We compute $\tilde{\boldsymbol{\eta}}_{*}$. Equation ([2](#S4.E2)) gives:

$$ $\displaystyle\tilde{\boldsymbol{\eta}}_{*}=\left(\tilde{\bar{\mathbf{H}}}+\frac{\lambda_{\mathrm{int}}}{2}\|\tilde{\mathbf{D}}\tilde{\boldsymbol{\eta}}_{*}\|\tilde{\mathbf{D}}^{2}\right)^{-1}\tilde{\bar{\mathbf{g}}}.$ (7) $$

Besides:

$$ $\displaystyle\tilde{\bar{\mathbf{H}}}$ $\displaystyle:=\mathbf{I}_{S:P}\tilde{\mathbf{U}}\tilde{\mathbf{H}}\tilde{\mathbf{U}}\mathbf{I}_{P:S}$ $\displaystyle\tilde{\bar{\mathbf{g}}}$ $\displaystyle:=\mathbf{I}_{S:P}\tilde{\mathbf{U}}\tilde{\mathbf{g}}$ $$

To go further, we need to do an assumption about the direction $\mathbf{u}$.

###### Assumption E.1 .

We assume that $\mathbf{U}_{t}$ is computed in such a way that $\tilde{\mathbf{U}}_{t}=\mathbf{J}\mathbf{U}_{t}$
at every step.

This assumption holds typically when $\mathbf{u}_{t}$ is the gradient at time step $t$.
It holds also when $\mathbf{u}_{t}$ is a linear combination of the past gradients:

$$ $\displaystyle\mathbf{u}_{1}$ $\displaystyle:=\mathbf{g}_{1}$ $\displaystyle\mathbf{u}_{t+1}$ $\displaystyle:=\mu\mathbf{u}_{t}+\mu^{\prime}\mathbf{g}_{t+1},$ $$

which includes the momentum.

To summarize, we have:

$$ $\displaystyle\tilde{\mathbf{U}}=\mathbf{J}\mathbf{U},\quad\tilde{\mathbf{H}}=\mathbf{J}\mathbf{H}\mathbf{J},\quad\tilde{\mathbf{g}}=\mathbf{J}\mathbf{g},$ $$

So:

$$ $\displaystyle\tilde{\bar{\mathbf{H}}}$ $\displaystyle=\tilde{\mathbf{J}}^{2}\mathbf{I}_{S:P}\mathbf{U}\mathbf{H}\mathbf{U}\mathbf{I}_{P:S}\tilde{\mathbf{J}}^{2}=\tilde{\mathbf{J}}^{2}\bar{\mathbf{H}}\tilde{\mathbf{J}}^{2},$ $\displaystyle\tilde{\bar{\mathbf{g}}}$ $\displaystyle=\tilde{\mathbf{J}}^{2}\mathbf{I}_{S:P}\mathbf{U}\mathbf{g}=\tilde{\mathbf{J}}^{2}\bar{\mathbf{g}},$ $$

since $\mathbf{J}$ and $\mathbf{U}$ are diagonal.
And:

$$ $\displaystyle\text{since: }\quad\mathbf{D}_{ii}$ $\displaystyle=\left|(\mathbf{D}_{\boldsymbol{\theta}}^{(3)}(\mathbf{u}))_{iii}\right|^{1/3},$ $\displaystyle\text{then: }\quad\tilde{\mathbf{D}}_{ii}$ $\displaystyle=a_{i}^{2}\mathbf{D}_{ii},$ $\displaystyle\text{thus:\leavevmode\nobreak\ \leavevmode\nobreak\ }\,\quad\tilde{\mathbf{D}}$ $\displaystyle=\tilde{\mathbf{J}}^{2}\mathbf{D}.$ $$

Thus, Eqn. ([7](#A5.E7)) becomes:

$$ $\displaystyle\tilde{\boldsymbol{\eta}}_{*}=\left(\tilde{\mathbf{J}}^{2}\bar{\mathbf{H}}\tilde{\mathbf{J}}^{2}+\frac{\lambda_{\mathrm{int}}}{2}\|\tilde{\mathbf{J}}^{2}\mathbf{D}\tilde{\boldsymbol{\eta}}_{*}\|\tilde{\mathbf{J}}^{4}\mathbf{D}^{2}\right)^{-1}\tilde{\mathbf{J}}^{2}\bar{\mathbf{g}},$ $$

which can be rewritten (since $\tilde{\mathbf{J}}$ is invertible):

$$ $\displaystyle\tilde{\mathbf{J}}^{2}\tilde{\boldsymbol{\eta}}_{*}=\left(\bar{\mathbf{H}}+\frac{\lambda_{\mathrm{int}}}{2}\|\mathbf{D}\tilde{\mathbf{J}}^{2}\tilde{\boldsymbol{\eta}}_{*}\|\mathbf{D}^{2}\right)^{-1}\bar{\mathbf{g}}.$ $$

Therefore, $\tilde{\boldsymbol{\eta}}_{*}$ is a solution of Eqn. ([2](#S4.E2))
in the parameterization $\tilde{\boldsymbol{\theta}}$ if, and only if,
$\tilde{\mathbf{J}}^{2}\tilde{\boldsymbol{\eta}}_{*}$ is a solution in the parameterization
$\boldsymbol{\theta}$. Moreover, $\|\tilde{\mathbf{D}}\tilde{\boldsymbol{\eta}}_{*}\|=\|\mathbf{D}\mathbf{J}^{2}\tilde{\boldsymbol{\eta}}_{*}\|$,
so $\tilde{\boldsymbol{\eta}}_{*}$ is the solution of maximum norm $\|\tilde{\mathbf{D}}\tilde{\boldsymbol{\eta}}_{*}\|$
of ([2](#S4.E2)) with parameterization $\tilde{\boldsymbol{\theta}}$ iff
$\tilde{\mathbf{J}}^{2}\tilde{\boldsymbol{\eta}}_{*}$ is a the solution of maximum norm
$\|\mathbf{D}\mathbf{J}^{2}\tilde{\boldsymbol{\eta}}_{*}\|$ of ([2](#S4.E2)) with parameterization $\boldsymbol{\theta}$.

Thus, $\boldsymbol{\eta}_{*}=\tilde{\mathbf{J}}^{2}\tilde{\boldsymbol{\eta}}_{*}$,
and the update step in parameterization $\tilde{\boldsymbol{\theta}}$ is:

$$ $\displaystyle\tilde{\boldsymbol{\theta}}_{t+1}$ $\displaystyle=\tilde{\boldsymbol{\theta}}_{t}-\tilde{\mathbf{U}}_{t}\mathbf{I}_{P:S}\tilde{\boldsymbol{\eta}}_{*}$ $\displaystyle=\tilde{\boldsymbol{\theta}}_{t}-\tilde{\mathbf{U}}_{t}\mathbf{I}_{P:S}\tilde{\mathbf{J}}^{-2}\boldsymbol{\eta}_{*}$ $$

which can be rewritten:

$$ $\displaystyle\mathbf{J}^{-1}\boldsymbol{\theta}_{t+1}$ $\displaystyle=\mathbf{J}^{-1}\boldsymbol{\theta}_{t}-\mathbf{U}\mathbf{J}\mathbf{I}_{P:S}\tilde{\mathbf{J}}^{-2}\boldsymbol{\eta}_{*},$ (8) $$

since $\varphi$ is an affine function with factor $\mathbf{J}$.
Finally, En. ([8](#A5.E8)) boils down to:

$$ $\displaystyle\boldsymbol{\theta}_{t+1}$ $\displaystyle=\boldsymbol{\theta}_{t}-\mathbf{U}\mathbf{I}_{P:S}\boldsymbol{\eta}_{*},$ $$

which is exactly Method [4.1](#S4.Thmtheorem1) in parameterization $\boldsymbol{\theta}$.

###### Assumption E.1 .

## Appendix F Experimental details

#### Practical implementation.

To implement the method proposed in Section [4](#S4),
we propose Algorithm [1](#alg1). The key function is
compute_lr$(\lambda_{\mathrm{int}};\mathcal{L},\boldsymbol{\theta},\tilde{Z},\mathbf{u})$,
which returns a solution $\boldsymbol{\eta}_{*}$ of:

$$ $\displaystyle\boldsymbol{\eta}_{*}$ $\displaystyle=\left(\bar{\mathbf{H}}+\frac{\lambda_{\mathrm{int}}}{2}\|\mathbf{D}\boldsymbol{\eta}_{*}\|\mathbf{D}^{2}\right)^{-1}\bar{\mathbf{g}},$ $\displaystyle\text{with: }\quad\bar{\mathbf{H}}$ $\displaystyle:=\mathbf{I}_{S:P}\mathrm{Diag}(\mathbf{u})\frac{\partial^{2}\mathcal{L}}{\partial\boldsymbol{\theta}^{2}}(\boldsymbol{\theta},\tilde{Z})\mathrm{Diag}(\mathbf{u})\mathbf{I}_{P:S},$ $\displaystyle\quad\bar{\mathbf{g}}$ $\displaystyle:=\mathbf{I}_{S:P}\mathrm{Diag}(\mathbf{u})\mathbf{g}_{t},$ $\displaystyle\quad\mathbf{D}$ $\displaystyle:=\mathrm{Diag}\left(\left(\left|\mathbf{D}^{(3)}_{\boldsymbol{\theta}}(\mathbf{u})\right|_{iii}^{1/3}\right)_{1\leq i\leq S}\right).$ $$

“momentum$(\mu,\mathbf{x},\tilde{\mathbf{x}})$”
returns $\mathbf{x}$ if $\tilde{\mathbf{x}}$ is undefined, else
$\mu\tilde{\mathbf{x}}+(1-\mu)\mathbf{x}$. “schedule$(\tau_{\mathrm{sch}},p_{\mathrm{sch}},f_{\mathrm{sch}};\cdots)$” corresponds to torch.optim.lr_scheduler.ReduceLROnPlateau
called every $\tau_{\mathrm{sch}}$ with patience $p_{\mathrm{sch}}$ and factor $f_{\mathrm{sch}}$,
in order to reduce the damping $\lambda_{t}$ when the loss attains a plateau.(^7^77See torch.optim.lr_scheduler.ReduceLROnPlateau.)
The samplers $\mathcal{D}_{g}$ and $\mathcal{D}_{\mathrm{newt}}$ are respectively
used to compute the gradients $\mathbf{g}_{t}$ and $(\bar{\mathbf{H}},\bar{\mathbf{g}})$ used in “compute_lr”.

The hyperparameters are: the initial damping factor $\lambda_{1}$, the
momentum $\mu_{g}$ on the gradients $\mathbf{g}_{t}$, the minibatch size $B$ to sample
the $\tilde{Z}$ (used to compute $\bar{\mathbf{g}}$, $\bar{\mathbf{H}}$ and $\mathbf{D}$),
the number of steps $\tau$ between each call of compute_lr,
the momentum $\mu_{\eta}$ on the learning rates $\boldsymbol{\eta}_{t}$,
the internal damping $\lambda_{\mathrm{int}}$, and
the parameters of the scheduler $\tau_{\mathrm{sch}}$, $p_{\mathrm{sch}}$, $f_{\mathrm{sch}}$.

Figure: Algorithm 1 Complete implementation of the second-order optimization method described in Section [4](#S4). $\lambda_{1}$ and $\lambda_{\mathrm{int}}$ are the only hyperparameter to be tuned across the experiments, the others are fixed.

#### Explanation.

The “momentum” functions are used to deal with the stochastic part of the training process, since
our method has not been designed to be robust against noise. The period $\tau$
is usually strictly greater than $1$, in order to avoid calling “compute_lr” at every step,
which would be costly. The minibatch size $B$ should be large enough to reduce noise in the estimation of
$\boldsymbol{\eta}_{*}$. If we denote by $B_{g}$ the size of the minibatches in $\mathcal{D}_{g}$,
then we recommend the following setup: $\tau=\frac{B}{B_{g}}=\frac{1}{1-\mu_{g}}$. That way, we ensure that the training data are sampled from
$\mathcal{D}_{g}$ and $\mathcal{D}_{\mathrm{newt}}$ at the same rate, and that
$\tilde{\mathbf{g}}_{t}$ memorizes the preceding gradients $\mathbf{g}_{t}$ for $\tau$ steps.
Besides, we have to take the positive part $(\boldsymbol{\eta}_{t})_{+}$ of $\boldsymbol{\eta}_{t}$ in order to avoid negative learning
rates.

#### Experimental setup.

We provide in Table [2](#A6.T2) the hyperparameters
fixed for all the experiments. In Table [3](#A6.T3),
we report the results of the grid-search for the hyperparameters
of the 3 tested optimization methods.

**Table 2: Hyperparameters fixed in all the series of experiments. $N_{e}$ is the number of training steps per epoch.**
| $\mu_{g}$ | $B_{g}$ | $B$ | $\tau$ | $\mu_{\eta}$ | $\tau_{\mathrm{sch}}$ | $p_{\mathrm{sch}}$ | $f_{\mathrm{sch}}$ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| $0.9$ | $10^{2}$ | $10^{3}$ | $10$ | $0.5$ | $N_{e}$ | $5$ | $0.5$ |

**Table 3: Hyperparams tuned for each series of experiments. $\eta$: learning rate, $\lambda_{1}$: initial damping factor.**
|  | MLP | LeNet | VGG-11’ | BigMLP |
| --- | --- | --- | --- | --- |
| Adam: $\eta$ | $3\cdot 10^{-4}$ | $3\cdot 10^{-4}$ | $10^{-5}$ | $10^{-5}$ |
| KFAC: $\eta$ | $10^{-4}$ | $10^{-4}$ | $3\cdot 10^{-4}$ | $10^{-5}$ |
| KFAC: $\lambda_{1}$ | $10^{-2}$ | $3\cdot 10^{-2}$ | $3\cdot 10^{-2}$ | $10^{-2}$ |
| Ours: $\lambda_{1}$ | $10^{-1}$ | $3\cdot 10^{-1}$ | $3\cdot 10^{-1}$ | $10^{-1}$ |
| Ours: $\lambda_{\mathrm{int}}$ | $10$ | $3$ | $3$ | $10$ |

## Appendix G Very deep multilayer perceptron

#### Grouping the layers.

In addition to the neural networks considered in Section [5](#S5),
we have also tested “VBigMLP”, a very deep multilayer perceptron with
100 layers of size 1024 trained on CIFAR-10.
Instead of considering $S=2L=200$ groups of parameters,
we split the sequence of layers of VBigMLP into 5 chunks.
Then, each chunk is divided into 2 parts, one containing the weight tensors,
and the other the bias vectors.
Finally, we have $S=10$ subsets of parameters, grouped by role (weight/bias)
and by position inside the network.

#### Experimental results.

We show in Figure [5](#A7.F5) the matrices $\bar{\mathbf{H}}$
and $\bar{\mathbf{H}}^{-1}$ at different stages of training.
At initialization, even if the neural network is very deep,
we observe that all the chunks of the network interact together,
even the first one with the last one.
However, after several training steps, the long-range interactions seem
to disappear. Incidentally, the matrices become tridiagonal,
which ties in with the block-tridiagonal approximation of the inverse of the Hessian
done by Martens and Grosse, ().

Figure: Figure 5: VBigMLP trained by SGD on CIFAR-10. Submatrices of $\bar{\mathbf{H}}$ (first row) and $\bar{\mathbf{H}}^{-1}$ (second row), at initialization and before the 100th epoch.
Refer to caption: /html/2312.03885/assets/x5.png

In Figure [6](#A7.F6), we observe the evolution of the
learning rates $\eta_{*}$ computed according to ([2](#S4.E2)).
First, there are all increasing during training.
Second, the biases in the last layers of the network seem to need
larger learning rates than biases in the first layers.
Third, the learning rate computed for the weights of the first
chunk of layers is smaller than the others.

Figure: Figure 6: VBigMLP trained by SGD on CIFAR-10. Legend: solid lines: weights; dotted lines: biases. Learning rates $\eta_{*}$ computed according to ([2](#S4.E2)), specific to each subset of parameters. For each epoch $k\in\{25,75,125,175,225,275,325,375\}$, the reported value has been averaged over the epochs $[k-25,k+24]$ to remove the noise.
Refer to caption: /html/2312.03885/assets/x6.png

Finally, the training curves in Figure [7](#A7.F7)
indicate that our method can be used to train very deep networks, but quite slowly.
In this setup, it is far from being competitive with Adam.
Besides, we did not manage to tune the learning rate and th damping of K-FAC to
make it work in this setup.

Figure: Figure 7: Training curves with VBigMLP, CIFAR-10. The training of Adam has been stopped at 200 epochs, since no significant improvement was observed in the last 50 epochs.
Refer to caption: /html/2312.03885/assets/x7.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See above
soundness: 3
strengths: - The paper is in general well motivated with the need to capture the interactions between parameters in different layers, which is often ignored by block-diagonal methods. This is operationalized in a natural way by studying the Hessian with suitable contractions. 

- The approach of getting layerwise learning rates through their layerwise grouping is neat. This offers a principled extension to the Cauchy's steepest descent rule. 

- The method could also find utility in studying the behaviour of feature learning across layers, and thus be used more than just in optimization.
summary: This work considers building summaries of higher order loss derivatives, like the Hessian and the third-order Tensor, which bucket interactions at the level of layers or some arbitrary partitions, instead of each individual parameter. In particular, by considering a particular contraction (like with the all-ones vector or gradient direction), very compact higher-order summaries can be built (which scales polynomially in the number of layers, instead of parameters). As an application, they use it to derive a layerwise scaling of learning rates, which neatly interpolates between the two extremes of using Newton's method and Cauchy's steepest descent rule. The method is demonstrated on some very simple experimental setups.
weaknesses: - The experimental section is quite weak. I understand that the authors themselves pitch it as a proof-of-concept, but I am not so sure about even if you can call it a proof of concept. The experiments are on small datasets like CIFAR, even over there none of the methods get the typical 90% and above accuracy, the test accuracy of their method is much worse than K-FAC. 

- More fundamentally, it is unclear to me where lies the bigger problem: correcting the curvature across layers, or that within the layers. It is well known that the Hessian tends to have a significant energy on its block diagonals, and thus maybe correcting across the layers, may not possess significantly more information. 

- Also, even for their method, I am curious what amount of the performance can be explained by simply estimating the scales of the Hessian on the diagonal blocks. In particular, if they instead use diag(\bar{H}), and then use it to get layerwise learning rates, how does that perform? This would form a test bed to showcase how crucial is it to capture cross-layer information. 

- Besides, I think in the vision setting the Hessian tends to be more homogeneous across the networks as opposed to that in Transformers with language modelling [1]. Hence, I think their approach might be more suited to that setting, and would be interesting to see if it can outperform methods like Adam-Mini [2]. 

- There are very limited baselines considered by the authors. I would have liked to see the Cauchy step size, AdaHessian, and even a block-diagonal quasi Newton method. 

- The overall runtime cost can be quite large, as there are multiple Hessian vector products. Can the authors do a wall-clock comparison?

References:

[1] Ormaniec, et. al. (2024). What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis. arXiv preprint arXiv:2410.10986.

[2] Zhang, et. al. (2024). Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406.16793.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: See weaknesses section. Do the authors agree with this perspective?
soundness: 4
strengths: The authors present a very interesting idea, and thoroughly motivate the idea. It really is a big short-coming that many related papers on this topic neglect inter-layer interactions when designing optimization algorithms. And finding computationally tractable ways to "summarize" the Hessian is a very strong idea.

The discussion of related work is quite comprehensive and clear.

Generally the work feels quite thoughtful: considering what are the issues with Newton's method, and how to try to overcome them.
summary: The paper proposes to make second-order optimization computationally tractable by taking a coarse view of high-dimensional parameter spaces. The authors break the Hessian of a deep learning model into blocks, with one block for every pair of layers, and compute one summary scalar per block. This allows for modeling inter-layer interactions, unlike many other approaches to approximate second-order optimization that neglect these terms. The authors propose a cubic-regularized version of their algorithm, and present experimental evidence that inter-layer interactions in deep learning models are non-trivial. Ultimately, the experimental results of the authors' proposed method are mixed.
weaknesses: I think the method could potentially be introduced in a more straightforward way, and I want to suggest one way. I think the method could be viewed as a change of variables to a smaller set of local optimization variables. In particular, instead of viewing the loss as a function of general perturbations to all the weight tensors:

loss( W_1 + ∆W_1, W_2 + ∆W_2, ..., W_S + ∆W_S)

we can view the loss as as a function of scalar-parameterized perturbations to each layer:

loss( W_1 - η_1 * G_1 , W_2 - η_2 * G_2, ..., W_S - η_S * G_S)			(*)

where η_1, η_2, ..., η_S are a collection of scalars and G_1, G_2, ..., G_S are the gradients of the loss with respect to each weight tensor, evaluated at the point W_1, W_2, ..., W_S. We can consider (*) to be a loss function with S variables η_1, η_2, ..., η_S. It's then clear that Hessian of this "reduced-dimensionality" loss is an S x S matrix. And we can throw any of a wide range of optimization methods toward solving this local S-dimensional optimization problem.

I also think the title could possibly be improved. What about something like "Multi-Tensor Optimization via Second-Order Scalar Summaries"?

Algorithmically, there is a weakness in the method that it only searches in the gradient direction for each tensor. This means the method would miss Shampoo-style [1] changes to the gradient direction which have been found to speed up deep learning training empirically.

The experimental results are mixed and perhaps not terribly promising at the moment. But it's good that you are open and up front about this. I think it's unlikely the broader community would focus on the paper too closely without more thorough experimental results.

Some light proof-reading of the writing would be helpful in places. For example line 90--91 "This is typically what is done by Dangel (2023), despite it does not go beyond the second-order derivative."

[1] https://arxiv.org/abs/1802.09568

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1) L43 - I'm not sure this is correct, assuming unlimited compute, and a perfectly known Hessian, even if the Hessian is singular you can still invert it on the nonsingular subspace, using SVD and find the pseudo-inverse.
2) L71 - Why do the authors need to regularize instead of computing just the pseudo inverse? is it inefficient? this should be stated explicitly
3) L76 - "so its preserves "
4) L218-220 seems wrong, the d-derivative is a map from $\mathbb{R}$ to $\mathbb{R}^{p^d}$ and not the other way around, even if the intention was a map from weight space to the operator output it should be from $\mathbb{R}^p$ to $\mathbb{R}^{p^d}$, so this is unclear to me. Additionally, the second term is unclear, $u$ are not defined at this point and the brackets $[.]$ are undefined as well. In the appendix it seems clear that the intention is tensor contraction between the previous term and the brackets but this is not standard.
5) L258 - "Therefore, the tensors $D_\theta^d(u)$ extract more information than the naive Taylor terms, while keeping a reasonable computational cost. " why is this statement obvious "therefore"? is the statement that regular Taylor terms lose the layer-wise structure of the network? I understand that the equality between Taylor and this decomposition is only obtained after tracing out the $s_i$ indices, but what is the intuition? it would be useful to show explicitly for a low $d$ derivative with a fixed number of parameters to illustrate the difference.
6) the authors don't comment on the compute time of their method compared to single gradient based method (Which seem to be better so far), it would seem like $t*p*S$ vs $t*p$, making it substantially slower for deep networks. 
7) While the method provided in this text is not shown to be superior to standard algorithms, it might be interesting to consider the computation method in the context of sharpness aware minimization.
soundness: 2
strengths: The paper has several strengths:

1) The interest in efficient computational schemes of higher order information for deep neural networks is significant, and improved methods of estimating the Hessian, as well as higher order terms could help improve interpretability and shed light on what neural networks learn during optimization.
2) The partitioning scheme is the main contribution of this paper, and seems to be novel as far as I know, with possible applications to many interesting avenues.
summary: The authors suggest a layer-wise partitioning of order $d$ derivatives, that transforms the $R^{p^d}$ tensor to a $R^{S^d}$ tensor which is computationally tractable even in deep networks.  The authors then leverage this partitioning to first compute empirically the Hessian for some deep neural networks, and subsequently suggest a second order method based on partitioning for optimization.
weaknesses: In spite of the strengths, the paper has some clear drawbacks in my opinion:

1) The paper is not written well enough, with a substantial lack of a literature survey on properties of Hessians in deep networks, as well as works on second order methods from recent years. In terms of the writing itself, all of the equations on page 5 are unnumbered making them hard to refer to, and the chosen notation for tensor contraction $A[ u, u... u]$ is not standard and never explained. It must be understood from context in the main text or the appendices.
2) The main contribution - the fast computation of lower complexity tensor containing similar information as the original tensor is not given in detail in the main text, and even in the appendix should be explained explicitly.
3) In general, I believe the focus of the paper is completely misguided. It seems clear that the suggestion of the second order optimization method is not superior (at least in its current form) to other simple gradient based methods, and should then not be the focus of the paper. Instead, if the paper focused on the partitioning/computation methods and then applied this to real-world or even particular solvable examples, extending Sec. 5.1, the paper would be much stronger.

Meta-review: This paper proposes an optimization algorithm “NewtonSummary”, which can be viewed as a middle ground of Newton’s method and Cauchy’s steepest descent. By grouping parameters and calculating “summarized” high-order derivatives of the loss function, the algorithm implements a computationally inexpensive (but crude) version of Newton’s method with Nesterov cubic regularization.

The reviewers agreed that the algorithm has some novel elements such as a computationally tractable way to summarize 2nd and 3rd order derivatives of the loss function. Indeed, as pointed out by one of the reviewers during reviewer-AC discussion, the computational method for higher order information presented in the paper can be useful, in particular for interpretability and understanding the landscape of DNNs.

However, it was unanimously pointed out that the empirical demonstration of the proposed algorithm’s effectiveness leaves much to be desired. Although the authors claims to provide a preliminary “proof-of-concept”, the algorithm is evaluated in a limited range of benchmarks and models (although the authors indeed test it on a 100-layer model), is compared against only a few baselines, and yet does not show clear advantages over existing methods. This makes me question the motivation for using higher-order methods for neural network training—if gradient-based methods like adam are better, why use this method? Moreover, the paper does not provide any convergence guarantees, which is also a weakness.

Overall, my conclusion is that the paper is slightly below the acceptance threshold for ICLR.

Decision: Reject