# A Novel Framework for Neural Architecture Search in the Hill Climbing Domain

Mudit Verma, Pradyumna Sinha, Karan Goyal, Apoorva Verma and Seba Susan
Department of Information Technology
Delhi Technological University
Bawana Road, Shahbad Daulatpur, Rohini, Delhi, 110042
Email: {mudit.verma2014, pradyumna014, karangoyal2k96<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Neural networks have now long been used for solving complex problems of image
domain, yet designing the same needs manual expertise. Furthermore, techniques
for automatically generating a suitable deep learning architecture for a given
dataset have frequently made use of reinforcement learning and evolutionary
methods which take extensive computational resources and time. We propose a
new framework for neural architecture search based on a hill-climbing
procedure using morphism operators that makes use of a novel gradient update
scheme. The update is based on the aging of neural network layers and results
in the reduction in the overall training time. This technique can search in a
broader search space which subsequently yields competitive results. We achieve
a 4.96% error rate on the CIFAR-10 dataset in 19.4 hours of a single GPU
training.

###### Index Terms:

Neural Architecture Search, Hill-Climbing, Optimization

##  I Introduction

Deep Neural Networks have demonstrated impressive performance on many machine-
learning tasks such as computer vision, speech recognition and natural
language processing [1][2][3][4]. The high performance of deep neural networks
is strengthened by the cost of large-scale engineering and checks to find the
best architecture for a given problem. High-level design decisions such as
depth, units per layer, and layer connectivity are not always obvious and use
hit and trial approach which is exhausting and time-consuming. This led to a
growing interest in the automatic designing of neural networks. Since the
architecture search space is discrete, traditional optimization algorithms
such as gradient descent[20] are not applicable, hence approaches based on
reinforcement learning algorithms[13] and evolutionary algorithms[21] are used
for automated designing of neural networks. But these approaches are either
very costly (requires hundreds and thousands of GPU days) or yield non-
competitive results. In this work, we aim to reduce the computational cost by
using a simple hill-climbing approach for the optimization of search space and
our variant of network morphism[18] for making the network deeper and complex.
We have also added gradient stopping, linear morphism and several
implementational improvements for reducing the training time. Our work uses a
single GPU, and we validate our method on the CIFAR-10 dataset. It takes our
framework 19.4 hours to reach 4.96% error. When we use this same architecture
to train and validate on MNIST dataset, we obtain an error percentage of 0.28.
We have compared two ways of weight initialization, with and without Gradient
Stopping.

In the following sections, we will discuss the previous attempts at Neural
Architecture Search (NAS). Further, we present our NASGraph, explain its
various nodes and morphism operations, and propose our methodology. We
describe our implementation of morphism operations, gradient stopping
technique, linear morphism and finally give a flow chart for our algorithm. In
the end, we present our experiments and showcase our results in comparison
with handcrafted architectures on CIFAR-10 as well as other Neural
Architecture Search schemes, and also gain insights about using our generated
deep neural network on different datasets. We then conclude our work and
present a few pointers for possible future research opportunities based on our
work.

##  II Related Work

Research on neural architecture search started in the 1980s. Initially,
evolutionary algorithm based approaches were proposed which promises to find
out both the architecture and weights of the neural network in [5]. However,
they have been unable to match the performance of the hand-crafted network.
More recent work [6][7][8][9], use an evolutionary algorithm only to find out
the structure and stochastic gradient descent to find out the parameters of
the network. [9] and [10] use evolutionary algorithms to generate a complex
architecture from a smaller neural-network iteratively, but [9] used a large
number of resources (250 GPUs and ten days) whereas [10] is limited to
relatively smaller networks due to handling a population of networks.

RL-based approach. [12] proposed Neural Architecture Search (NAS) with the
help of reinforce algorithms[13] to learn a network architecture. It generates
a sequence of actions representing the architecture of a CNN. NASNet[14] uses
proximal policy-optimization (PPO) in-place of reinforce algorithms[13] which
further improves neural architecture search. It searches the architecture of a
small unit which they referred to as “block.” It repeatedly concatenated
itself to form a complete model. Other works in the field use approach like
policy gradient in[14], Q-learning in[15] for finding the model architecture.
Unfortunately, training a reinforcement learning agent for searching
architecture is extremely expensive: both [16] and [17] required over 10.000
fully trained networks, requiring hundreds to thousands of GPU days.

Network morphism/ transformation : Using the concept of transfer learning,
network morphism was introduced and first published in 2015 by Chen [18]. To
make the network more complex, the author described function preserving
operation such as ’deeper’ or ’wider’ with the goal of speeding up training
and exploring network architectures. [18] further proposed additional
operations, e.g., for handling non-idempotent activation functions or altering
the kernel size and introduced the term network morphism. [19] presents a
neural architecture search by hill climbing (NASH) and used morphism
operations for making the network more complex. Out of all the methods, [19]
resembles the most to our techniques of neural architecture search by hill
climbing (NASH) with network morphism. We extend their work by adding gradient
stopping, linear morphism, several implementational improvements like the
selection of nodes upon which morphism is to be applied, the addition of more
number of operators for growth of performance using less resources, time (1
GPU and less than one training day) and at the same time yielding higher
accuracy.

##  III NAS Graph

NASGraph represents an instance of neural-architecture in our state space of
neural network architectures for the dataset that is given to the hill
climbing algorithm. A specific set of attributes define a NASGraph. It is
implemented as a Graph data structure where each node can either be a
convolutional block (ConvLayer-BatchNorm-ReLU), a max pool node, a Merge node
or a linear node. It has an adjacency list, the matrix that defines the
connections between the nodes. It also stores a list of node identifiers and a
mapping between the identifiers and the actual instances. There is another
list that maintains the node identifiers of its parents and children. Morphism
operations have been defined on the NASGraph, and during the hill climbing
process, they are called to ”morph” the parent network into a child network.
At an abstract level, the graph is initialized with image (from the input
dataset) dimensions. A Convolution node is added with random parameters
initially, and linear layers are stacked at the end for classification.
Finally, the NASGraph is transformed into a network, and its parameters are
provided to the SGDR[23] optimizer.

Topological Ordering Condition: According to this condition, it is a must for
the NASGraph to remain acyclic at all times. Therefore any connections made
from node $x$ to $y$ in the NASGraph must be such that $x$ lies before $y$ in
the topologically sorted order of the NASGraph.

###  III-A NASGraph Nodes

NASGraph Nodes are the fundamental building blocks of all possible neural
architectures through our framework. These nodes have dual behavior; they are
the usual graph nodes and therefore participate in lists maintained for graph
representation. However, they are also network layers; this limits the
position of edges in the graph. Also, to implement the framework, algorithms
such as topological sort, cycle detection, etc. have been performed on these
which are used later to maintain the topological ordering condition.

There are a total of 4 NASGraph Nodes as follows.

  1. 1.

Convolution : This node is composed of a 2-D convolution layer followed by a
Batch Normalization[24] layer and ReLU activation.

  2. 2.

Maxpool : This node is composed of a max pool layer.

  3. 3.

Merge : This node may be an Add layer, which takes up multiple node outputs as
it’s input, where all inputs must be of the same dimension and output their
matrix addition. Additionally, this node may be a Merge layer, where different
inputs are stacked one below another along channel dimension and fed as one
input to the next layer.

  4. 4.

Linear : This is a fully connected neural network layer.

The first node of the NASGraph is supposed to be randomly created, and after
this, the iterative hill climbing begins. The first node is a convolution node
with kernel sampled uniformly from set $\\{3,5\\}$, padding, uniformly from
set $\\{True,False\\}$, $stride=1$, and channels from set $\\{8,16,32\\}$.
Beyond this point, the only way to modify the properties of an existing
NASGraph node is via morphism operations.

###  III-B NASGraph Morph Operations

Figure 1: Some of the NASGraph Operations

Our morphism technique differs from that of [19] in that we do consider layers
without padding. Convolution layers without padding have proven to be useful
as seen in popular architectures such as VGG16[25]. Therefore, this factor
must be included in the search space.

We allow our NASGraph to perform one of the following six operations for the
generation of the child node in the hill climbing algorithm :

####  III-B1 Skip Operation

Skip Operation allows the NASGraph to have skip connections in the neural
network. Skip connections help traverse information in neural networks. As the
NASGraph is found to generate very deep neural architectures, it is possible
for the gradient information to be lost since it passes through so many
layers, thus skip connections enable us to pass feature information to lower
layers to help classify directly. As we also provide the provision for max-
pooling layers, max pooling operation itself causes some amount of spatial
information to be lost, skip connections help in this as well.

The idea for performing Skip is first to find all the node pairs in which the
first node’s output is the same as the second node’s input. Such a situation
may cause the second node to have multiple parents each, providing this node
with same dimension inputs. Therefore, we use an Add node, described below, to
add all such parents and then supply the combined result to this node in Fig.
1. We can see an example of a skip connection being applied between the input
and output layer shown by the dashed line connection.

####  III-B2 Deepen

Deepen, as the name suggests, is a layer to deepen the existing network. This
layer, in our case, is a Convolutional Block composed of a Convolutional
Neural Layer, a Batch Normalization Layer, followed by a ReLU activation
function. The layer may also be a Linear Layer, that is a fully connected
layer.

To perform the deepen operation between nodes $x$ and $y$, let us assume,
layer $x$ has $a$ neurons, and layer $y$ has $b$ neurons. Then, the weight
matrix is of shape (a,b). The new layer $t$ must have input dimension of (b,b)
so that the composition of layers $x$ and $t$ gives the final output (a,b),
same as before.

Selection of Nodes for Deepen is given in section IV-B. Figure 1 shows a new
node $t$ being added in between nodes $x$ and $y$. The addition of a layer
adds another level of abstraction that cannot be as simply contained within
shallow layer architecture of the same number of parameters if they can be
contained at all.

####  III-B3 Widen

Widen operation increases the channels of a convolution node by a factor of 2
or 4. Widening requires changing the parent node weight matrix as well as the
subsequent child node weight matrix to adjust for the widening changes. To
simplify the process, we, therefore, create a set of those node pairs, in
their topological ordering, $(x,y)$ where has only one child $y$, and $y$ has
only one parent $x$. Not using this condition would require us to perform
changes in the architecture recursively, i.e., if the parent weight changes,
then all its children weight matrices have to change recursively.

Widen Operation has also been used for fully connected layer widen morphism
which is described in detail in section IV-E.

####  III-B4 Add

When a node in the graph has more than one parent, one of the options to
combine the inputs is to do an element-wise sum of the incoming tensor.
However, this is only possible if the incoming tensors are of the same shape.

####  III-B5 Merge

The other option to combine inputs from more than one parent is to concatenate
the channels of the image tensors. For this method to be used, it is required
that the height and width of the incoming image tensors be the same.

####  III-B6 Maxpool

The max pool property is essential to building CNN and is used extensively in
object detection tasks like in [25]. [19] does not allow for dynamic addition
of new max-pool nodes in the architecture, i.e., max pool layer must be
originally present in the seed architecture for it to exist in the final
architecture. We ease this restriction and allow for dynamic addition of max
pool nodes in our NASGraph. Similar to the addition of a Convolution Node, we
first find all node pairs, in topological order, which can accommodate a max
pool layer in between them and maintain all previous connections. Then, we
uniformly sample one such node pair and add a max pool.

##  IV Proposed Work

###  IV-A Morphism

Morphism has been used in[19] before for Neural Architecture Search; however,
in approaches like NASH[19], the use of Morphism restricted the network state
space, i.e., all Convolution Layers are required to be padded. They treat each
convolution layer as a function, which they split into two functions using the
morphism operation, thereby the final input and output dimensions remain the
same as before the morphism was applied. We, however, allow layers to be added
without padding, which dynamically changes the output dimension of the layer.

Morphism functioning for different operations has been explained in Section
III-B. The morphism has been applied in two ways, first with default Pytorch
version 0.4.1 [28] initialization of layer weights (Xavier Initialization for
Convolutional Layers), and the second, 0/1 initialization as in NASH. In 0/1
initialization when we use widen operation, the added neurons are set to value
0, and for deepening process, the added layer has all values initialized to 1.
We show our results with both types of initialization methods and find that
they do not affect the total time to get a trained network, and default
initialization performs better than 0/1 initialization only marginally.

###  IV-B Selection of Nodes For Morphism

NASH[19] states the morphism technique and the operations; however, the
selection of the nodes upon which the morphism operation would take place is
described as ”random.” This random selection works for the NASH procedure
because they begin with a seed architecture and each added layer has padding
enabled, thereby keeping the output always the same. In our case, however,
compatibility of nodes between which an operation is to be performed is an
issue. For example, in the Skip operation, we need to pick two nodes A and B
such that, $output(A)=input(B)$. Moreover, any such additional connection must
meet the topological order condition of NASGraph as in section III.

We always keep a topological ordering of the NASGraph nodes, and for any
request of connecting two nodes via Skip, compatible node pairs ordered
topologically are found, and one such pair is selected uniformly randomly.
This method of ’randomly’ picking nodes for the morphism operation can be
termed as fairly random for all practical purposes because it is essentially
selecting nodes for the operation randomly from the set of all compatible
nodes.

For adding a Convolution Layer with a given kernel shape and other parameters,
say c, we find a set of all topologically ordered node pairs (x,y) such that
$c(x.output).output=y.input$. Now we can randomly pick any element of this set
for our morphism operation.

###  IV-C Random Selections

Many parameters of the NASGraph nodes are sampled randomly and we describe all
such parameters below.

We need to select various properties of the NASGraph nodes, for example, the
kernel, stride, etc. of the Convolution Node. As in [19], we choose these
properties randomly from a predefined set of values. And finally, since we
allow some of the layers to exist without padding, we mark padding True or
False, with a probability of $0.5$ each.

For deepen operation which is the addition of a Convolution Node, we select
the kernel from the set $\\{3,5\\}$. The number of channels is chosen to be
equal to the number of channels of the closest preceding convolution.

For widen operation which is the expansion of a Convolution Node, the number
of channels increases by a widening factor, which is sampled randomly from the
set $\\{2,4\\}$.

The parameters for Maxpool operation are also required to be chosen uniformly.
Kernel for the same is selected uniformly from a set $\\{2,3\\}$.

The morphism operation applied is also chosen randomly from the set of all
possible morphisms. We describe morphism operations in the section III-B.
Finally, if a certain operation cannot be performed at some point of time,
then we temporarily remove that operation from the set and again sample
uniformly from this new operations set.

###  IV-D Gradient Stopping

Figure 2: Illustrating the nodes involved in gradient stopping for distance =
3

Repeated application of NASGraph morphism operations for a large number of
times can result in very deep neural networks. The problem with the
conventional approach to handling these dynamically growing neural
architectures is overtraining of specific ’old’ layers. Additionally, we try
to analyze the addition of a layer via a water drop ripple analogy. A drop of
water falling in a pool of water creates water ripples around the drop
location, thus leaving far away locations unaffected. Hence, the addition of a
convolution layer should effectively cause major changes in the weights of
near-by layers only, and training of the far-away layers may be stopped.
However, defining a set of nodes as ’far away’ is difficult, so we use a
Normal Distribution Curve to solve this.

By iterative generation of complete Neural Networks, we found the average
depth of the graph is around 36. We arbitrarily choose to let the training of
nodes continue which lie within $1/3$ of the average length of nodes which
comes to be approximately 12. Thus after adding a convolution node, we find
all the nodes connected up till distance = 6 (which would give at least 12
nodes as the total from either side). An example of distance = 3 is presented
in Fig. 2 where all the edges at a distance less than or equal to three have
been marked red. Now, all the nodes connected via these marked red edges can
be denoted via their distances.

To determine the learning rate of a node, a value from the Normal distribution
equation $f(distance)=\frac{6}{\sqrt{2\pi}\cdot
2.4}e^{\frac{-\left(x\right)^{2}}{2s^{2}}}$ corresponding to the distance
assigned to the node is obtained, say $alpha_{n}$. The distribution is set so
that this $alpha_{n}$ lies between the range $(0,1]$ and tends to zero as the
distance value reaches $6$. Beyond $distance=6$ we set the value of
$alpha_{n}=0$. Now, we obtain the learning rate value $alpha$ according to the
SGDR[23] update rule and get the new $alpha=alpha*alpha_{n}$. Note that in the
next iteration, $alpha$ is again obtained from SGDR rule, as it would be
without Gradient Stopping, and the changes are performed on that value. All
the nodes beyond distance 6 do not update their gradient until the next non-
layer-addition morphism operation.

Fig. 3 gives a graphical representation of gradient stopping, where the
darkest circle represents the newly added node, for which the learning rate is
unchanged, as $alpha_{n}=1$. As we move away from this node, the circle shade
fades off representing reduction in their learning rates, and beyond a
specific number of nodes (6 in this case), the training completely stops.
Gradient Stopping is a considerably good example of layer-wise training of a
neural network with dynamic changes to the architecture at the run time.

Figure 3: Gradient Stopping, using the normal distribution to update new
learning rates for neural network nodes.

###  IV-E Linear Morphisms

As mentioned previously, our framework supports convolution nodes to be added
without padding, as well as the addition of max pool layers. The output
dimensions of these layers are strictly less than the input dimensions;
therefore, the final output shape of tensor exiting the ’NasGraph’ (graph of
Convolution, MaxPool, Add, Merge Nodes) bubble reduces. For example, if the
previous output dimension is (28,28), now with the addition of Convolution
Node at the very end with kernel size (5,5) without padding, the output
dimension becomes (24,24). Hence the size of the flattened layer changes as in
Figure 4. Due to this change, we must make the immediate following Fully
Connected layer (FC layer) have a reduced number of neurons.

In our framework, we use two Fully Connected Layers for classification. The
first one, Variable FC Layer, is connected to the flattened NASGraph matrix
and is morphed to reduce in size dynamically. The number of neurons in this
Variable FC Layer is equal to the number of neurons in the Flattened Layer. We
can do this because the image size of the CIFAR-10 dataset is (32,32). Thus,
the maximum number of neurons in the Flatten layer can be $32*32=1024$. So now
Variable FC also has $1024$ neurons. Therefore the size of the weight matrix
becomes $1024*1024$ which is around 1 Million parameters, feasible to handle
by today’s computing standards. The next and the final layer has ten neurons
corresponding to the ten classes in the CIFAR-10 dataset. Therefore, any
changes in the NASGraph will be held until the Variable FC layer and we can
have a constant size final Fully Connected layer for Classification.

Figure 4: Linear Widening Morphism

###  IV-F Process Flow

The flowchart in Fig. 5 shows the steps of our algorithm in a nutshell. The
graphical architecture of the model is referred to as a NASGraph. First, a
random architecture is generated to be the starting point of the state space
search, called the NASGraph seed. The Hill climbing operation is carried for
$n_{steps}$. In each step, $n_{neigh}$ neighbors are generated for the current
best model. To create a neighbor, $n_{NM}$ morphism operations are applied as
described in section IV-A, IV-B and IV-C. The resultant neighbors imbibe the
properties of the parent model along with having new attributes. These
neighbors are then trained for $epoch_{neigh}$. Finally, the current best
model is updated by comparing with the results of the neighbors on the
validation dataset. After $n_{steps}$, the best model is trained for
$epoch_{final}$ and the trained model is reported as the output of the
algorithm.

Figure 5: Flowchart of our Iterative Hill Climbing technique for Neural
Architecture Search

##  V EXPERIMENTAL RESULTS

###  V-A Dataset Overview

For neural architecture search, we use CIFAR-10 dataset [29]. Convolution
Neural Networks have shown to be very effective in working on the image
problem domains. Thus, we showcase our results also on image datasets. To test
how well does the architecture generated by our method on CIFAR dataset
performs on images of a different domain, we use the handwritten digits MNIST
dataset. Baseline results on Neural Architecture Search are usually on CIFAR /
MNIST datasets. Therefore, a comparison of our method along with other methods
is suitable for the same.

###  V-B Discussion on Results

TABLE I: Results Comparison with different weight initialization techniques.  Learning Rate | Xavier Initialization | 0/1 Initialization
---|---|---
| Time (hrs) / Error (%) | Time (hrs) / Error (%)
| SGDR +
---
Gradient Stopping
19.4 / 4.96 | 19.6 / 5.2
SGDR | 23.4 / 5.8 | 23.3 / 5.8

For similarities in our methods, we choose the NASH [19] as our baseline.

Parameters for NASH used to compare the results are as provided in [19]. The
set values are given in Table II. NASH requires us to feed a simple seed
architecture which they mentioned as : Conv-Max pool-Conv-Max pool-Conv-FC-
Softmax where Conv = Conv + BatchNorm + ReLU. This architecture has already
been trained for 20 epochs. Using these parameters, the best possible error
rate achieved by NASH is 5.2% in one day of GPU training.

Table II also mentions the value of the parameters we have used, which is the
same for all our architectures presented in this paper. We explore the same
number of neighbors at any depth of the hill climbing as that of NASH.
However, the exploration depth is higher in our case which in turn enables us
to apply a higher number of morphisms on our network, and reach more depth,
without requiring much time.

We first explain the parameters used as follows.

  1. 1.

$n_{steps}$ is the depth of the hill climbing tree.

  2. 2.

$n_{NM}$ is the number of times the morphism operation is applied on a hill
climbing node.

  3. 3.

$n_{neigh}$ is the number of children generated in the hill climbing tree.

  4. 4.

$epoch_{neigh}$ is the number of epochs for which each child is trained for.

  5. 5.

$epoch_{final}$ is the number of epochs for which the final $model_{best}$ is
trained.

  6. 6.

$\lambda_{start}$ and $\lambda_{end}$ are the parameters required for SGDR.

TABLE II: Parameters for NASH and Our Method. Variable | NASH-2 | Ours
---|---|---
$n_{steps}$ | 8 | 10
$n_{NM}$ | 5 | 5
$n_{neigh}$ | 8 | 8
$epoch_{neigh}$ | 17 | 16
$epoch_{final}$ | 100 | 64
$\lambda_{start}$ | 0.05 | 0.1
$\lambda_{end}$ | 0.0 | 0.0
TABLE III: Comparison of Neural Architecture Search on CIFAR-10 on different models.  Model | Resource | Error(%)
---|---|---
Shake-Shake [11] | 2 days, 2 GPU | 2.9
RL-NAS [17] | ? days, 800 GPUs | 3.65
NASH-1 | 0.5 days, 1 GPU | 5.7
NASH-2 | 1 day, 1 GPU | 5.2
Ours | 19.4 hours, 1 GPU | 4.96

Since the objective of this paper is to propose a newer faster baseline, we
must focus our scrutiny on three areas primarily, amount of resources used,
the time required for the complete training process, and the final error
percentage. In all the three fronts we either match the NASH or are better.
Table III gives a comparison of our framework (with Linear Morphisms, Updated
Operations & Gradient Stopping with SGDR) with the two NASH models proposed in
[19], the Shake-Shake [11] model, and a reinforcement learning based Neural
Architecture Search Scheme [17]. NASH-1 has $n_{steps}=5$, while NASH-2
parameters are given in Table II. Shake-Shake [11] is a manually crafted
state-of-the-art model in CIFAR-10 classification which uses 2 GPUs and gives
2.9% error in 2 days. However, on comparing with a Neural Architecture Search
scheme like that of [17], that although gives a low error of 3.65%, but uses a
mammoth 800 GPU system, is quite unrealistic for practical purposes. Our
baseline, NASH scheme, uses 1 GPU and one day of training time to give a 5.2%
error, against our training time of 19.4 hours to produce better results at
4.96% error. The reported error for our model is the average of three full
runs on CIFAR-10 classification which yielded 5.1%, 4.9%, 4.9% errors
respectively averaging 19.4 hours of training time. We also note that our
method searches for comparatively more of the neural architecture space in
less time and produces better accuracy.

###  V-C On a Different Dataset

We also performed classification using the architectures produced with
CIFAR-10, on the MNIST dataset. It was done to gain insights regarding how the
architecture performs on other image datasets. Low error percentages by
architectures generated from one type of dataset, upon another dataset, would
mean either similarity in the dataset, or the robustness of the architecture
itself. Since MNIST and CIFAR-10 are known to be very different datasets,
where MNIST consists of handwritten digits, and CIFAR-10 consists of images of
various real-life objects, we chose MNIST for this purpose. Upon using the
same architecture, for which the results on CIFAR-10 are mentioned in Tables I
and III, we obtain 0.28% error on MNIST dataset. The current state-of-the-art
on MNIST is Rmdl: Random multimodel deep learning for classification [22]
which boasts an error percentage of 0.18. However, our architecture would
still stand in the top 10111https://paperswithcode.com/sota/image-
classification-mnist-handwritten-d error percentages reported until June 2018.

##  VI Conclusion & Future Work

We introduced a Neural Architecture Search in Hill Climbing domain using
morphism operations that use our novel Gradient Update scheme. Our method is
not just fast and uncomplicated, but it also outperforms many existing
automated architecture search methods with lesser GPU days and resources. Our
experiments performed on the CIFAR-10 and MNIST datasets and yielded
satisfactory results.

Our approach can quite comfortably deliver as a basis for the development of
more sophisticated techniques that yield further improvements in the existing
performance. Certain operations are being executed randomly in our method,
such as the selection of a morphism operation or selection of a node for
performing the morphism steps. One can make these selections intelligently by
using reinforcement learning or genetic algorithms. Essential criteria for
selection of nodes would be to pick convolution blocks from those regions of
the network which have been having significant changes in their weights in the
recent iterations. Again, there are different ways to initialize the weights
in the morphism operations. We experimented with two of many possibilities,
and expect further research in weight initialization would yield better
results.

Apart from improvements in the existing work, we find research in Learning
Curve Extrapolation for dynamic neural networks in its initial stages, but
would immensely speed up the overall hill climbing iterative algorithm.
However, it is currently there for static [26, 27] neural network graphs.

## References

  * [1] He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 770-778).
  * [2] Karpathy, A., & Fei-Fei, L. (2015). Deep visual-semantic alignments for generating image descriptions. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 3128-3137).
  * [3] Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., … & Dieleman, S. (2016). Mastering the game of Go with deep neural networks and tree search. nature, 529(7587), 484.
  * [4] Sutskever, I., Vinyals, O., & Le, Q. V. (2014). Sequence to sequence learning with neural networks. In Advances in neural information processing systems (pp. 3104-3112).4
  * [5] Schaffer, J. D., Whitley, D., & Eshelman, L. J. (1992, June). Combinations of genetic algorithms and neural networks: A survey of the state of the art. In [Proceedings] COGANN-92: International Workshop on Combinations of Genetic Algorithms and Neural Networks (pp. 1-37). IEEE.
  * [6] Liu, H., Simonyan, K., Vinyals, O., Fernando, C., & Kavukcuoglu, K. (2017). Hierarchical representations for efficient architecture search. arXiv preprint arXiv:1711.00436.
  * [7] Miikkulainen, R., Liang, J., Meyerson, E., Rawal, A., Fink, D., Francon, O., … & Hodjat, B. (2019). Evolving deep neural networks. In Artificial Intelligence in the Age of Neural Networks and Brain Computing (pp. 293-312). Academic Press.
  * [8] Real, E., Aggarwal, A., Huang, Y., & Le, Q. V. (2018). Regularized evolution for image classifier architecture search. arXiv preprint arXiv:1802.01548.
  * [9] Real, E., Moore, S., Selle, A., Saxena, S., Suematsu, Y. L., Tan, J., … & Kurakin, A. (2017, August). Large-scale evolution of image classifiers. In Proceedings of the 34th International Conference on Machine Learning-Volume 70 (pp. 2902-2911). JMLR. org.
  * [10] Suganuma, M., Shirakawa, S., & Nagao, T. (2017, July). A genetic programming approach to designing convolutional neural network architectures. In Proceedings of the Genetic and Evolutionary Computation Conference (pp. 497-504). ACM.
  * [11] Gastaldi, X. (2017). Shake-shake regularization. arXiv preprint arXiv:1705.07485.
  * [12] Zoph, B., & Le, Q. V. (2016). Neural architecture search with reinforcement learning. arXiv preprint arXiv:1611.01578.
  * [13] Williams, R.J.: Simple statistical gradient-following algorithms for connectionist reinforcement learning. Mach. Learn. 8(3-4) (May 1992) 229–256
  * [14] Zoph, B., Vasudevan, V., Shlens, J., & Le, Q. V. (2018). Learning transferable architectures for scalable image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 8697-8710).
  * [15] Cai, H., Chen, T., Zhang, W., Yu, Y., & Wang, J. (2018, April). Efficient architecture search by network transformation. In Thirty-Second AAAI Conference on Artificial Intelligence.
  * [16] Baker, B., Gupta, O., Naik, N., & Raskar, R. (2016). Designing neural network architectures using reinforcement learning. arXiv preprint arXiv:1611.02167.
  * [17] Bello, I., Zoph, B., Vasudevan, V., & Le, Q. V. (2017, August). Neural optimizer search with reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning-Volume 70 (pp. 459-468). JMLR. org.
  * [18] Wei, T., Wang, C., Rui, Y., & Chen, C. W. (2016, June). Network morphism. In International Conference on Machine Learning (pp. 564-572).
  * [19] Elsken, T., Metzen, J. H., & Hutter, F. (2017). Simple and efficient architecture search for convolutional neural networks. arXiv preprint arXiv:1711.04528.
  * [20] Ruder, S. (2016). An overview of gradient descent optimization algorithms. arXiv preprint arXiv:1609.04747.
  * [21] Miller, G. F., Todd, P. M., & Hegde, S. U. (1989, June). Designing Neural Networks using Genetic Algorithms. In ICGA (Vol. 89, pp. 379-384).
  * [22] Kowsari, K., Heidarysafa, M., Brown, D. E., Meimandi, K. J., & Barnes, L. E. (2018, April). Rmdl: Random multimodel deep learning for classification. In Proceedings of the 2nd International Conference on Information System and Data Mining (pp. 19-28). ACM.
  * [23] Loshchilov, I., & Hutter, F. (2016). Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983.
  * [24] Ioffe, S., & Szegedy, C. (2015). Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03167.
  * [25] Simonyan, K., & Zisserman, A. (2014). Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556.
  * [26] Klein, A., Falkner, S., Springenberg, J. T., & Hutter, F. (2016). Learning curve prediction with Bayesian neural networks.
  * [27] Domhan, T., Springenberg, J. T., & Hutter, F. (2015, June). Speeding up automatic hyperparameter optimization of deep neural networks by extrapolation of learning curves. In Twenty-Fourth International Joint Conference on Artificial Intelligence.
  * [28] Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., … & Lerer, A. (2017). Automatic differentiation in pytorch.
  * [29] Krizhevsky, A., & Hinton, G. (2009). Learning multiple layers of features from tiny images (Vol. 1, No. 4, p. 7). Technical report, University of Toronto.

[Sample Output of NASGraph]

Fig. 6 shows a sample graph produced by the NASGraph using the parameters
mentioned in Table II. The generated network can be seen to have convolution
blocks as represented by red bubbles, and Merge layers (which can be either
the Merge or the Add function) represented by the blue ’add’ bubbles. The
graph clearly holds the Topological Ordering condition as mentioned a must for
the NASGraph. In addition, it is seen to have Skip connections for example
node conv0 is connected to add27 providing input to conv1, along with conv13,
conv19, conv7, conv17, conv11. The NASGraph can generate quite complicated
neural networks and provide good accuracy rates on CIFAR-10.

Figure 6: Sample output of NASGraph on CIFAR-10 dataset


# Thermodynamic Natural Gradient Descent

Kaelan Donatella ,†
Normal Computing
&Samuel Duffield†
Normal Computing
Maxwell Aifer
Normal Computing
&Denis Melanson
Normal Computing
&Gavin Crooks
Normal Computing &Patrick J. Coles
Normal Computing Correspondence to<EMAIL_ADDRESS>
###### Abstract

Second-order training methods have better convergence properties than gradient
descent but are rarely used in practice for large-scale training due to their
computational overhead. This can be viewed as a hardware limitation (imposed
by digital computers). Here we show that natural gradient descent (NGD), a
second-order method, can have a similar computational complexity per iteration
to a first-order method, when employing appropriate hardware. We present a new
hybrid digital-analog algorithm for training neural networks that is
equivalent to NGD in a certain parameter regime but avoids prohibitively
costly linear system solves. Our algorithm exploits the thermodynamic
properties of an analog system at equilibrium, and hence requires an analog
thermodynamic computer. The training occurs in a hybrid digital-analog loop,
where the gradient and Fisher information matrix (or any other positive semi-
definite curvature matrix) are calculated at given time intervals while the
analog dynamics take place. We numerically demonstrate the superiority of this
approach over state-of-the-art digital first- and second-order training
methods on classification tasks and language model fine-tuning tasks.

$\dagger$$\dagger$footnotetext: These authors contributed equally to this
work.

##  1 Introduction

With the rise of more sophisticated AI models, the cost of training them is
exploding, as world-leading models now cost hundreds of millions of dollars to
train. This issue is compounded by the ending of both Moore’s Law and
Dennard’s Law for digital hardware [20], which impacts both the runtime and
energy efficiency of such hardware. This highlights a need and an opportunity
for specialized, unconventional hardware targeted at improving the efficiency
of training AI models.

Moreover, conventional digital hardware can be viewed as limiting the range of
training algorithms that a user may consider. Researchers are missing an
opportunity to co-design novel optimizers to exploit novel hardware
developments. Instead, relatively simplistic optimizers, such as stochastic
gradient descent (SGD), Adam [22], and their variants [27], are among the most
popular methods for training deep neural networks (DNNs) and other large AI
models. More sophisticated optimizers are rarely used due to the associated
computational overhead on digital hardware.

A clear example of this is second-order methods, which capture curvature
information of the loss landscape. These methods, while theoretically more
powerful in terms of convergence properties, remain computationally expensive
and harder to use, blocking their adoption. For example, natural gradient
descent (NGD) [4, 30] involves calculating estimates of second-order
quantities such as the Fisher information matrix and performing a costly
linear system solve at every epoch. Some approximations to NGD, such as the
Kronecker-factored approximate curvature (K-FAC) [31], have shown promise, and
K-FAC has shown superior performance to Adam [25, 14]. However, applying such
methods to arbitrary neural network architectures remains difficult [36].

In this article, we present thermodynamic natural gradient descent (TNGD), a
new method to perform second-order optimization. This method involves a hybrid
digital-analog loop, where a GPU communicates with an analog thermodynamic
computer. A nice feature of this paradigm is flexibility: the user provides
their model architecture and the analog computer serves only to accelerate the
training process. This is in contrast to many proposals to accelerate the
inference workload of AI models with analog computing, where the model is
hardwired into the hardware, and users are unable to change the model
architecture as they seamlessly would by using their preferred software tools
[21, 6, 11, 1].

The analog computer in TNGD uses thermodynamic processes as a computational
resource. Such thermodynamic devices have previously been proposed [10, 18,
15, 9, 26], have been theorized to exhibit runtime and energy efficiency gains
[2, 13], and have been successfully prototyped [34, 3]. Our TNGD algorithm
represents an instance of algorithmic co-design, where we propose a novel
optimizer to take advantage of a novel hardware paradigm. TNGD exploits a
physical Ornstein–Uhlenbeck process to implement the parameter update rule in
NGD. It has a runtime per iteration scaling linearly in the number of
parameters, and when properly parallelized it can be close to the runtime of
first-order optimizers such as Adam and SGD. Hence, it is theoretically
possible to achieve the computational efficiency of a first-order training
method while still accounting for the curvature of the loss landscape with a
second-order method. Moreover, our numerics show the competitiveness of TNGD
with first-order methods for classification and extractive question-answering
tasks.

##  2 Related work

There is a large body of theoretical research on natural gradient descent [4,
30, 7] arguing that NGD requires fewer iterations than SGD to converge to the
same value of the loss in specific settings. While less is known about the
theoretical convergence rate of Adam, there exists a large body of empirical
evidence that NGD can converge in fewer iterations than Adam [33, 31, 32, 14,
38, 16].

However, a single iteration of NGD is generally more computationally expensive
than that of SGD or Adam, which have a per-iteration cost scaling linearly in
the number of parameters $N$. NGD typically has a superlinear (assuming the
condition number scales as $\kappa=N^{\alpha},\alpha>0$ for NGD-CG) complexity
in the number of parameters (although this may be reduced to linear scaling at
the expense of higher-order scaling with batch size and output dimension, see
Section 3). K-FAC [31] aims to reduce this complexity and invokes a block-wise
approximation of the curvature matrix, which may not always hold. While first
introduced for multi-layer perceptrons, K-FAC has been applied to more complex
architectures, such as recurrent neural networks [32] and transformers [14],
where additional approximations have to be made and where the associated
computational overhead can vary.

There has been significant effort and progress towards reducing the time- and
space- complexity of operations used in the inference workload of AI models,
e.g., a variety of “linear attention" blocks have been proposed [40, 19, 42].
However, there has been less focus on reducing the complexity of training
methods. While various approaches are taken to accelerating training using
novel hardware, these efforts typically aim at reducing the constant
coefficients appearing in the time cost of computation. Especially relevant to
our work, analog computing devices have been proposed to achieve reduced time
and energy costs of training relative to available digital technology [21, 6,
11, 1]. These devices are generally limited to training a neural network that
has a specific architecture (corresponding to the structure of the analog
device). To our knowledge, there has not yet been a proposal that leverages
analog hardware to reduce the complexity of training algorithms such as NGD.

Given the existing results implying that fewer iterations are needed for NGD
relative to other commonly used optimizers, we focus on reducing the per-
iteration computational cost of NGD using a hybrid analog-digital algorithm to
perform each parameter update. Our algorithm therefore demonstrates that
complexity can be improved in training (not only in inference), and moreover
that the per-iteration complexity of NGD can be made similar to that of a
first-order training method.

##  3 Natural gradient descent

Let us consider a supervised learning setting, where the goal is to minimize
an objective function defined as:

$\ell(\theta)=\frac{1}{|\mathcal{D}|}\sum_{(x,y)\in\mathcal{D}}L(y,f_{\theta}(x)),$
(1)

where $L(y,f_{\theta}(x))\in\mathbb{R}$ is a loss function, $f_{\theta}(x)$ is
the forward function that is parametrized by $\theta\in\mathbb{R}^{N}$. These
functions depend on input data and labels $(x,y)\in\mathcal{D}$, with
$\mathcal{D}$ a given training dataset. Viewed through the lens of statistics,
minimizing the objective function is analogous to minimizing the Kullback-
Leibler (KL) divergence from the target joint distribution $q(x,y)$ to the
learned distribution $p(x,y|\theta)$ [30]. A straightforward way to optimize
$\ell(\theta)$ is to follow the direction of steepest descent, defined by the
negative gradient $-\nabla\ell$, defined as:

$\displaystyle\frac{-\nabla\ell}{||\nabla\ell||}=\underset{\epsilon\rightarrow
0}{\mathrm{lim}}\;\frac{1}{\epsilon}\,\underset{d:||d||\leq\epsilon}{\mathrm{arg\,min}}\,\ell(\theta+d),$
(2)

with $||\cdot||$ the Euclidean norm. The natural gradient, on the other hand
can be defined as the direction of steepest descent with respect to the KL
divergence defined as:

$\displaystyle\text{KL}(p(x,y|\theta+d)||p(x,y|\theta))=\iint p(x,y|\theta+d)\
\log\left(\frac{p(x,y|\theta+d)}{p(x,y|\theta)}\right)\
\mathrm{d}x\mathrm{d}y$ (3)

[5]. One may then Taylor-expand this divergence as

$\displaystyle\text{KL}(p(x,y|\theta+d)||p(x,y|\theta))=\frac{1}{2}d^{\top}Fd+O(d^{3}),$
(4)

where $F$ is the Fisher information matrix [30] (or the Fisher), defined as:

$\displaystyle F=\mathbb{E}_{p(x,y|\theta)}[\nabla\log p(x,y|\theta)\nabla\log
p(x,y|\theta)^{\top}].$ (5)

The natural gradient is then simply defined as

$\displaystyle\tilde{g}=F^{-1}\nabla\ell(\theta).$ (6)

For the NGD optimizer, the update rule is then given by:

$\theta_{k+1}=\theta_{k}-\eta F^{-1}\nabla\ell,$ (7)

with $\eta$ a learning rate. In practice, computing the Fisher information is
not always feasible because one must have access to the density
$p(x,y|\theta)$. A quantity that is always possible (and relatively cheap) to
compute thanks to auto-differentiation is the empirical Fisher information
matrix, defined as:

$\displaystyle\bar{F}=JJ^{\top}=\frac{1}{b}\sum_{(x,y)\in\mathcal{S}}\nabla\log
p(y|x,\theta)\nabla\log p(y|x,\theta)^{\top},$ (8)

where $\log p(y\mid x,\theta)=-L(y,f_{\theta}(x))$, $|\mathcal{S}|=b$ is the
batch size and $\mathcal{S}\subset\mathcal{D}$. The Jacobian matrix $J$ is
defined as

$J~{}=\frac{1}{\sqrt{b}}[\nabla\log p(y_{1}|x_{1},\theta),\nabla\log
p(y_{2}|x_{2},\theta),\ldots,\nabla\log p(y_{b}|x_{b},\theta)].$

Note that the squared gradient appearing in the second moment estimate of the
Adam optimizer [22] is the diagonal of the empirical Fisher matrix. Another
approximation to the Fisher matrix is the generalized Gauss-Newton (GGN)
matrix, defined as:

Time$\displaystyle t_{k-1}$$\displaystyle t_{k}$$\displaystyle t_{k+1}$ SPU
SPU SPU GPU GPU GPU Dynamical
Evolution Dynamical
Evolution
$\tilde{g}_{k-1}$$\nabla\ell_{k-2}$$F_{k-2}$$\tilde{g}_{k}$$\nabla\ell_{k-1}$$F_{k-1}$$\tilde{g}_{k+1}$$\nabla\ell_{k}$$F_{k}$
Automatic Differentiation Automatic Differentiation Output Parameters
$\displaystyle\theta$

Figure 1: Overview of Thermodynamic Natural Gradient Descent (TNGD). A GPU
that stores the model architecture and provides the gradient $\nabla\ell_{k}$
and Fisher matrix $F_{k}$ (through its representation given by the Jacobian
$J_{f}$ and Hessian $H_{L}$ matrices given by Eq. (9)) at step $k$ is
connected to a thermodynamic computer, called the stochastic processing unit
(SPU). At times $t_{k}$, the estimate of the natural gradient $\tilde{g}_{k}$
is sent to the GPU, which updates the parameters of the model and calculates
gradients and curvature matrices for some new data batch $(x_{k},y_{k})$.
During digital auto-differentiation, the SPU undergoes dynamical evolution,
either continuing to approach its steady-state or remaining in it. After some
time, gradient $\nabla\ell_{k}$ and Fisher matrix $F_{k}$ are sent to the SPU
through a DAC and digital controllers. This modifies the dynamics of the SPU,
and after some time interval, a new natural gradient estimate
$\tilde{g}_{k+1}$ is sent back to the GPU. Note that the time between two
measurements $t_{k+1}-t_{k}$ need not be greater than the time between two
auto-differentiation calls. The hybrid digital-thermodynamic process may be
used asynchronously as shown in the diagram (where the time of measurement of
$\tilde{g}$ and upload of the gradient and Fisher matrix are not the same).

$G=J_{f}H_{L}J_{f}^{\top}=\frac{1}{b}\sum_{(x,y)\in\mathcal{S}}J_{f}^{(x,y)}H_{L}^{(x,y)}J_{f}^{{(x,y)}\top},$
(9)

where $J_{f}^{(x,y)}$ is the Jacobian of $f_{\theta}(x)$ with respect to
$\theta$ and $H_{L}^{(x,y)}$ is the Hessian of $L(y,z)$ with respect to
$\theta$ evaluated at $z=f_{\theta}(x)$. $J_{f}$ is a $bd_{z}\times N$ matrix,
and $H_{L}$ is a $bd_{z}\times bd_{z}$ matrix, where $d_{z}$ is the output
dimension of $z=f_{\theta}(x)$ and $N$ is the number of parameters ($N$ also
depends on $d_{z}$, where for deep networks it is a weak dependence).

For loss functions of the exponential family (with natural parameter $z$), the
GGN matches the true Fisher matrix [30]. In addition, we have observed better
convergence with the GGN than with the empirical Fisher (as in other works
such as Refs. [33, 23], where better convergence than with the Hessian is also
observed). Therefore, we will consider the GGN in what follows. Note that the
methods we introduce in this work apply to any second-order optimization
algorithm with a positive semi-definite curvature matrix (by curvature matrix,
we mean any matrix capturing information about the loss landscape). In
particular, it applies most efficiently to matrices constructed as outer
products of rectangular matrices (such as the empirical Fisher and the GGN) as
explained below.

###  3.1 Fast matrix vector products

The linear system appearing in Eq. (6) can be solved using the conjugate
gradient (CG) method [33], which will be referred to as NGD-CG in what
follows. In fact, when $\ell$ is parametrized by a neural network, the GGN-
vector product $Gv$ involved in the conjugate gradient algorithm may be
evaluated in runtime $O(bN)$ thanks to fast Jacobian-vector products [8]
(JVPs). This approach also enables one to not explicitly construct the Fisher
matrix, thus also avoiding a $O(bd_{z}N^{2})$ runtime cost in computing it and
a $O(N^{2})$ memory cost in storing it. The efficiency of this approach
depends on the number of CG iterations required to obtain good performance.
Importantly, convergence in $\sqrt{\kappa}$ steps, with $\kappa$ the condition
number of $F$, is not required to obtain competitive performance [31, 16].
Crucially, due to the sequential nature of the algorithm, the CG iterations
cannot be parallelized.

In practice, since reaching convergence is computationally expensive, one
generally stops the CG algorithm after a set number of iterations. Because of
the way the step size is adapted in CG, we have observed that the solution
after $k$ steps $x_{k}$ is not necessarily closer to the true solution than
the initial guess $x_{0}$, in particular for ill-conditioned problems, which
can make NGD-CG difficult to use.

###  3.2 NGD with the Woodbury identity

In the machine learning setting, it is often the case that $b\ll N$ (and
$d_{z}\ll N$). This means that the curvature matrix is low-rank and the linear
system to solve is underdetermined. To mitigate this issue, the Fisher matrix
may be dampened as $F+\lambda\mathbb{I}$. In that case, the Woodbury identity
may be used to obtain the inverse Fisher vector-product $F^{-1}v$ appearing in
the NGD update. We have:

$\displaystyle F$ $\displaystyle=UV+\lambda\mathbb{I},\text{ with
}U=J_{f},V=H_{L}J_{f}^{\top}$ (10) $\displaystyle F^{-1}$
$\displaystyle=\lambda^{-1}\mathbb{I}-\lambda^{-2}U(\mathbb{I}+\lambda^{-1}VU)^{-1}V\qquad(\mathrm{Woodbury})$
(11) $\displaystyle F^{-1}v$
$\displaystyle=\lambda^{-1}\mathbb{I}-\lambda^{-2}U(\mathbb{I}+\lambda^{-1}VU)^{-1}Vv$
(12)

This is included in Ref. [38], and can be competitive with NGD-CG when the
batch size $b$ and output dimension $d_{z}$ are much smaller than the number
of trainable parameters $N$. Here one must construct the $V$ matrix, which has
runtime $O(d_{z}^{2}b^{2}N)$, and invert $(\mathbb{I}+\lambda^{-1}VU)$ which
is $O(b^{3}d_{z}^{3})$. While the batch size typically remains small, the
value of $d_{z}$ can make this inversion intractable. For example, in many
language-model tasks, $d_{z}\sim O(10^{4})$ is the vocabulary size.

##  4 Thermodynamic NGD

At a high level, TNGD combines the strength of GPUs (through auto-
differentiation) with the strength of thermodynamic devices at solving linear
systems. Regarding the latter, Ref. [2] showed that a thermodynamic device,
called a stochastic processing unit (SPU), can solve a linear system $Ax=b$
with reduced computational complexity relative to standard digital hardware.
The solution to the linear system is found by letting the SPU evolve under an
Ornstein–Uhlenbeck (OU) process given by the following stochastic differential
equation (SDE):

$dx=-(Ax-b)dt+\mathcal{N}\left[0,2\beta^{-1}\,dt\right],$ (13)

where $A$ is a positive matrix and $\beta$ is a positive scalar (which can be
seen as the inverse temperature of the noise). Operationally, one lets the SPU
settle to its equilibrium state under the dynamics of Eq. (13), at which point
$x$ is distributed according to the Boltzmann distribution given by:

$x\sim\mathcal{N}[A^{-1}b,\beta^{-1}A^{-1}].$ (14)

One can see that the first moment of this distribution is the solution to the
linear system $Ax=b$. Exploiting this approach, TNGD involves a subroutine
that estimates the solution to the linear system in Eq. (6). For this
particular linear system, the SDE in Eq. (13) becomes the following:

$\displaystyle d\tilde{g}_{k,t}$
$\displaystyle=-(F_{k-1}\tilde{g}_{k,t}-\nabla\ell_{k-1})dt+\mathcal{N}\left[0,2\kappa_{0}dt\right]$
(15)
$\displaystyle=-(J_{f,k-1}^{\top}H_{L,k-1}J_{f,k-1}\tilde{g}_{k,t}-\nabla\ell_{k-1})dt+\mathcal{N}\left[0,2\kappa_{0}dt\right]$
(16)

with $\tilde{g}_{k,t}$ the value of the natural gradient estimate at time $t$
and $\kappa_{0}$ the variance of the noise. Comparing Eqs. (13) and (15), we
see that in the equilibrium state (i.e. for large $t$), the mean of
$\tilde{g}_{k,t}$ provides an estimate of the natural gradient, in other
words:

$\displaystyle\tilde{g}_{k}:=\underset{t\rightarrow\infty}{\mathrm{lim}}\langle\tilde{g}_{k,t}\rangle=F_{k-1}^{-1}\nabla\ell_{k-1}.$
(17)

Optimizer | Runtime | Memory | Model calls
---|---|---|---
SGD/Adam | $O(bN)$ | $O(N)$ | $1$
NGD | $O(N^{3}+bd_{z}N^{2})$ | $O(N^{2})$ | $bd_{z}$
NGD-CG | $O(cbN)$ | $O(N)$ | $2c$
NGD-Woodbury | $O(b^{2}d_{z}^{2}N+b^{3}d_{z}^{3})$ | $O(bd_{z}N+b^{2}d_{z}^{2})$ | $bd_{z}$
Thermodynamic NGD | $O(bd_{z}N+t)$ | $O(bd_{z}N+b^{2}d_{z}^{2})$ | $bd_{z}$
Table 1: Runtime and memory complexity of optimizers considered in this paper.
All operations are per iteration. The first line corresponds to first-order
optimizers that evaluate the gradient only, and apply diagonal rescalings and
$O(N)$ operations to it only. Vanilla NGD (second line) includes the explicit
storage and inversion of the GGN matrix as well as its construction,
dominating the runtime and memory cost. NGD-CG (third line) can be performed
by running $c$ iterations, each dominated by GGN-vector products and has the
same memory cost as first-order methods. NGD-Woodbury can be performed by
constructing the matrix $VU$, and using the formula given by Eq. (12). This
results in a runtime cost dominated by constructing $VU$ and inverting it,
which also requires its storage.

The overall TNGD algorithm is illustrated in Fig. 1. Using the current
parameter estimates $\theta_{k}$, the GPU computes the matrices $J_{f}$ and
$H_{L}$, and the vector $\nabla\ell$, which can be accomplished efficiently
using auto-differentiation. The matrices $J_{f}$, $J_{f}^{\top}$, and $H_{L}$,
as well as the vector $\nabla\ell$, are uploaded to component values (see
Appendix) on the SPU, which is then allowed to equilibrate under the dynamics
of Eq. (15). Next, samples are taken of $\tilde{g}_{t,k}$, and are sent from
the SPU to the GPU, where samples are averaged to yield an estimate of
$\tilde{g}_{k}$. Finally, the parameters are updated using the equation

$\theta_{k+1}=\theta_{k}-\eta\tilde{g}_{k},$ (18)

and this process may be repeated until sufficient convergence is achieved
(other update equations may also be employed, see Section 5).

While Eq. (17) involves a long time limit, numerical evidence (see section 5)
shows that samples may be taken even before equilibrium has been reached
without harming performance significantly. Thus, the analog dynamics time $t$
is an important hyperparameter of TNGD. Furthermore, another hyperparameter
arises from the delay time $t_{d}$, defined as the time between a measurement
of $\theta_{k}$ and the update of the gradient and GGN on the device. As
discussed in Section 5, a non-zero delay time is not necessarily detrimental
to performance and can in fact improve it.

In addition to the advantage in time- and energy-efficiency, TNGD has another
advantage over NGD-CG in terms of stability. For some pathological linear
systems, CG fails to converge and instead diverges. However, the thermodynamic
algorithm is guaranteed to converge (on average) for any positive definite
matrix. To see this, note that the mean of $\tilde{g}_{k,t}$ evolves according
to

$\langle\tilde{g}_{k,t}\rangle=\exp{(-F_{k-1}t)}(\tilde{g}_{k,0}-F_{k-1}^{-1}\nabla\ell_{k-1})+F_{k-1}^{-1}\nabla\ell_{k-1}.$
(19)

There is still variance associated with the estimator of
$\langle\tilde{g}_{k,t}\rangle$ (the sample mean), but the sample mean
converges to the solution with high probability in all cases. We also note
that if we choose $\tilde{g}_{k,0}=\nabla\ell_{k-1}$, we obtain a smooth
interpolation between SGD ($t=0$) and NGD ($t=\infty$).

Figure 2: Runtime per iteration of second-order optimizers considered in this
paper. (a) The runtimes per iteration are compared for NGD, NGD-CG, NGD-
Woodbury, and TNGD (estimated) for various $N$. Here the convolutional network
we applied to MNIST is used and the dimension of the hidden layer is varied to
vary $N$ for fixed $d_{z}=20$. (b) The same comparison is shown for various
values of $d_{z}$. The same network is used and $d_{z}$ is varied (this also
has the effect of varying the $N$). Error bars are displayed as shaded area
but are smaller than the data markers.

###  4.1 Computational complexity and performance

The runtime complexity of TNGD and other second-order optimization (that do
not make assumptions on the structure of $G$, hence excluding K-FAC)
algorithms is reported in Table 1. As explained, Thermodynamic NGD (TNGD) has
a runtime and memory cost dominated by the construction and storage (before
sending them off to the analog hardware) of the Jacobian of $f_{\theta}(x)$
and the Hessian of the loss. The $t$ factor denotes the analog runtime, and
may be interpreted similarly to $c$ for NGD-CG as a parameter controlling the
approximation. For each optimizer the number of model calls is reported. For
all optimizers except NGD-CG these calls can be easily parallelized thanks to
vectorizing maps in PyTorch.

In Fig. 2 a comparison of the runtime per iteration of the four second-order
optimizers considered is shown. Fig. 2(a) shows the runtime as a function of
the number of parameters $N$. The scaling of NGD as $N^{3}$ can be observed,
and the NGD-CG data is close to flat, meaning the model calls parallelize well
for the range of parameter count considered. The linear scaling of NGD-
Woodbury and TNGD is also shown, although with a different overall behaviour
due to parallelization and a much shorter runtime per iteration for TNGD. This
shows that for the given range of $N$ at $d_{z}=20$, we can expect a
100$\times$ speedup over second-order optimizers. Fig. 2(b) shows the
dependence of runtime on the output dimension $d_{z}$ for the second-order
optimizers. These results indicate that TNGD is most competitive for
intermediate values of $d_{z}$. Finally we note that with better hardware, the
scaling with both $N$ and $d_{z}$ would be better, as the operations to
construct the Hessian and Jacobian can be more efficiently parallelized for
larger values.

##  5 Experiments

###  5.1 MNIST classification

We first consider the task of MNIST classification [24]. For our experiments,
we use a simple convolutional neural network consisting of a convolutional
layer followed by two feedforward layers, and we digitally simulate the TNGD
algorithm (see App. D). The goal of these experiments is twofold: (1) to
compare the estimated performance per runtime of TNGD against popular first-
order optimizers such as Adam, and (2) to provide some insights on other
features of TNGD, such as its performance as a function of the analog runtime
$t$ as well as its asynchronous execution as a function of the delay time
$t_{d}$.

Figure 3: Performance comparison of Adam and TNGD (simulated) on MNIST
classification. (a) Training (dashed lines) and test loss (solid lines) for
Adam (darker colors) and TNGD (lighter colors) are plotted against runtime
(measured for Adam, and estimated for TNGD from the timing model described in
Section 4.1). Shaded areas are standard deviations over five random seeds.
Note that Adam includes adaptive averaging of first and second moment
estimates with $(\beta_{1},\beta_{2})=(0.9,0.999)$, while TNGD does not. (b)
$1-\mathrm{Accuracy}$ for training and test sets.

In Fig. 3(a), the training and test losses as a function of runtime for both
Adam (measured) and TNGD (estimated) are presented. To estimate the TNGD
runtime, we took into account results for its runtime per iteration as
presented in the previous section, finding an overall $2\times$ runtime per
iteration with respect to Adam for this problem on an A100 GPU. One can see
from the figure that even while taking into account the difference in runtime
per iteration, TNGD still outperforms Adam, especially at the initial stages
of the optimization. Interestingly, it also generalizes better for the
considered experimental setup. In Fig.3(b), the training and test accuracies
are shown. We again see TNGD largely outperforming Adam, reaching the same
training accuracy orders of magnitude faster, while also displaying a better
test accuracy. These results are reminiscent of prior work on NGD [33],
however here the batch size is smaller than in other works, indicating that
even a noisy GGN matrix improves the optimization.

As mentioned previously, the continuous-time nature of TNGD allows one to
interpolate smoothly between first- ($t=0$) and second- ($t=\infty$) order
optimization, with a given optimizer choice (whether the optimizer update rule
is that of SGD or that of Adam as described in Alg. 1). In Fig. 4(a), the
training loss vs. iterations is shown for various analog dynamics times. These
results clearly demonstrate the effect mentioned above, where increasing the
analog runtime improves performances continuously until it approaches that of
exact NGD for $t\sim 50\tau$. In Fig. 4(b), the same quantity is shown for a
fixed analog dynamics time $t$, and varying delay times $t_{d}$. This leads to
a quadratic approximation of the objective function that is inaccurate (since
the GGN and gradients are calculated for parameters different than the value
around which the objective function is approximated). However, this results in
an improved performance, even for a small delay time. A likely explanation of
this result is that the state of the device retains information about the
curvature of the previous quadratic approximation, while being subject to the
updated quadratic approximation. This effect propagates across iterations
which is reminiscent of momentum.

Figure 4: Training loss vs. iterations for varying analog dynamics times. (a)
The training loss is shown for NGD (dashed line) and for TNGD with various
analog dynamics times $t$ (solid lines). (b) The training loss is shown for
NGD (dashed line) and for TNGD with fixed analog dynamics time $t=5\tau$ and
varying delay times $t_{d}$ (solid lines). The delay appears to have a
momentum effect, which can even lead to TNGD outperforming exact NGD for
certain analog dynamics and delay times. Shaded areas are standard deviations
over five random seeds.

###  5.2 Language model fine-tuning

In this section we show how thermodynamic NGD may be applied to language
modeling tasks, in more practically relevant settings than MNIST
classification. We consider the DistilBert model [39] which we fine-tune on
the Stanford question-answering dataset (SQuaD) [37], a common dataset to
evaluate model comprehension of technical domains through extractive question-
answering. As is commonly performed when fine-tuning, we apply a low-rank
adaptation [17] to the model, which reduces its trainable parameters (details
about this procedure are in App. E) to a manageable amount ($75k$ here) for
limited compute resources.

Figure 5(a) displays a comparison of the training loss for different
optimizers. The bare TNGD (as used in the previous section) shows a worse
performance than Adam in this setting. However, a hybrid approach, TNGD-Adam,
where the natural gradient estimate is used in conjunction with the Adam
update rule gives the best performance (this is explained in App. B). One
possible explanation for this result is that there are two pre-conditionings
of the gradient for TNGD-Adam: the first comes from the natural gradient,
which incorporates curvature information, and the second comes from the Adam
update rule, which acts as a signal-noise ratio as explained in Ref. [22],
which further adjusts the natural gradient values. In Fig. 5(b), we show that
the same results as in the previous section apply to TNGD-Adam, where
increasing the analog runtime boosts performance. Therefore, the analog
runtime in TNGD may be viewed as a resource in this sense, that is
computationally very cheap (as time constants can be engineered to be very
small).

Figure 5: Training loss vs. iterations for QA fine-tuning. (a) Comparison of
the performance per iteration of TNGD, Adam, and TNGD-Adam, where the latter
uses the natural gradient estimate in conjunction with the Adam update rule
with $(\beta_{1},\beta_{2})=(0,0)$. (b) Performance of the TNGD-Adam optimizer
for various analog dynamics times. Similar to Fig. 4, the performance improves
as $t$ grows.

##  6 Limitations

The practical impact of our work relies on the future availability of analog
thermodynamic computers, such as a scaled up version of the system in Ref.
[34]. We provide a circuit diagram of a potential thermodynamic computer in
the Appendix. Such computers can employ standard electrical components and
leverage CMOS-based fabrication infrastructure, and hence are likely
straightforward to scale up, although that remains to be demonstrated.

Analog computers, in general, tend to face precision issues, whereby the
solution accuracy is limited by the precision of the electrical components.
For analog thermodynamic computers, it is possible to mitigate this issue
through an averaging technique [3], and the method proposed in Ref. [3] can be
directly applied to the TNGD algorithm to improve solution accuracy.
Nevertheless, we suspect that training-based applications will have a
significant tolerance to precision-based errors, although a detailed study is
needed to confirm that hypothesis. We note that there is a growing body of
work on very low-precision inference [28] and training [41] which indicates
that high numerical precision is not crucial for good performance in machine
learning. We also remark that thermodynamic computers are predicted to be
robust to stochastic noise sources since stochasticity is a key component of
such computers [9], as is shown in Fig. 7 in the Appendix.

We have numerically tested TNGD for a small subset of potential tasks such as
MNIST classification and DistilBert fine-tuning on the SQuaD dataset, for a
small number of epochs. Hence, seeing if the advantage we observe for TNGD
also holds for other applications is an important direction.

##  7 Conclusion

This work introduced Thermodynamic Natural Gradient Descent (TNGD), a hybrid
digital-analog algorithm that leverages the thermodynamic properties of an
analog system to efficiently perform second-order optimization. TNGD greatly
reduces the computational overhead typically associated with second-order
methods for arbitrary model architectures. Our numerical results on MNIST
classification and language model fine-tuning tasks demonstrate that TNGD
outperforms state-of-the-art first-order methods, such as Adam, and provide
large speedups over other second-order optimizers. This suggests a promising
future for second-order methods when integrated with specialized hardware.

Looking forward, our research stimulates further investigation into TNGD,
particularly with enhancements such as averaging techniques and moving
averages. Extensions to approximate second-order methods such as K-FAC may
also be possible. Moreover, the principles of thermodynamic computing could
inspire new algorithms for Bayesian filtering. While the current impact of our
work relies on the development and availability of large-scale analog
thermodynamic computers, the theoretical and empirical advantages presented
here underscore the potential of co-designing algorithms and hardware to
overcome the limitations of conventional digital approaches.

## References

  * Aguirre et al. [2024] F. Aguirre, A. Sebastian, M. Le Gallo, W. Song, T. Wang, J. J. Yang, W. Lu, M.-F. Chang, D. Ielmini, Y. Yang, et al.  Hardware implementation of memristor-based artificial neural networks.  _Nature Communications_ , 15(1):1974, 2024.  URL https://www.nature.com/articles/s41467-024-45670-9.
  * Aifer et al. [2023] M. Aifer, K. Donatella, M. H. Gordon, T. Ahle, D. Simpson, G. E. Crooks, and P. J. Coles.  Thermodynamic linear algebra.  _arXiv preprint arXiv:2308.05660_ , 2023.  URL https://arxiv.org/abs/2308.05660.
  * Aifer et al. [2024] M. Aifer, D. Melanson, K. Donatella, G. Crooks, T. Ahle, and P. J. Coles.  Error mitigation for thermodynamic computing, 2024.  URL https://arxiv.org/abs/2401.16231.
  * Amari [1998] S.-I. Amari.  Natural gradient works efficiently in learning.  _Neural computation_ , 10(2):251–276, 1998.  URL http://cognet.mit.edu/journal/10.1162/089976698300017746.
  * Amari and Nagaoka [2000] S.-i. Amari and H. Nagaoka.  _Methods of information geometry_ , volume 191.  American Mathematical Soc., 2000.  URL https://bookstore.ams.org/mmono-191.
  * Ambrogio et al. [2018] S. Ambrogio, P. Narayanan, H. Tsai, R. M. Shelby, I. Boybat, C. Di Nolfo, S. Sidler, M. Giordano, M. Bodini, N. C. Farinha, et al.  Equivalent-accuracy accelerated neural-network training using analogue memory.  _Nature_ , 558(7708):60–67, 2018.  URL https://www.nature.com/articles/s41586-018-0180-5.
  * Bottou et al. [2018] L. Bottou, F. E. Curtis, and J. Nocedal.  Optimization methods for large-scale machine learning.  _SIAM review_ , 60(2):223–311, 2018.  URL https://epubs.siam.org/doi/10.1137/16M1080173.
  * Bradbury et al. [2018] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang.  JAX: composable transformations of Python+NumPy programs, 2018.  URL http://github.com/google/jax.
  * Coles et al. [2023] P. J. Coles, C. Szczepanski, D. Melanson, K. Donatella, A. J. Martinez, and F. Sbahi.  Thermodynamic AI and the fluctuation frontier, 2023.  URL https://arxiv.org/abs/2302.06584.
  * Conte et al. [2019] T. Conte, E. DeBenedictis, N. Ganesh, T. Hylton, J. P. Strachan, R. S. Williams, A. Alemi, L. Altenberg, G. E. Crooks, J. Crutchfield, et al.  Thermodynamic computing.  _arXiv preprint arXiv:1911.01968_ , 2019.  URL https://arxiv.org/abs/1911.01968.
  * Cristiano et al. [2018] G. Cristiano, M. Giordano, S. Ambrogio, L. P. Romero, C. Cheng, P. Narayanan, H. Tsai, R. M. Shelby, and G. W. Burr.  Perspective on training fully connected networks with resistive memories: Device requirements for multiple conductances of varying significance.  _Journal of Applied Physics_ , 124(15), 2018.  URL https://research.ibm.com/publications/perspective-on-training-fully-connected-networks-with-resistive-memories-device-requirements-for-multiple-conductances-of-varying-significance.
  * Duffield [2024] S. Duffield.  Posteriors.  https://github.com/normal-computing/posteriors, 2024.
  * Duffield et al. [2023] S. Duffield, M. Aifer, G. Crooks, T. Ahle, and P. J. Coles.  Thermodynamic matrix exponentials and thermodynamic parallelism.  _arXiv preprint arXiv:2311.12759_ , 2023.  URL https://arxiv.org/abs/2311.12759.
  * Eschenhagen et al. [2023] R. Eschenhagen, A. Immer, R. Turner, F. Schneider, and P. Hennig.  Kronecker-factored approximate curvature for modern neural network architectures.  In _Advances in Neural Information Processing Systems_ , volume 36, 2023.  URL https://proceedings.neurips.cc/paper_files/paper/2023/file/6a6679e3d5b9f7d5f09cdb79a5fc3fd8-Paper-Conference.pdf.
  * Ganesh [2017] N. Ganesh.  A thermodynamic treatment of intelligent systems.  In _2017 IEEE International Conference on Rebooting Computing (ICRC)_ , pages 1–4, 2017.  doi: 10.1109/ICRC.2017.8123676.  URL https://ieeexplore.ieee.org/abstract/document/8123676?casa_token=wcQ38wymLlIAAAAA:Z91ud4DdK-SN8ymrUrlYCog_MM4GmdUIifsoVqpTY-kskCW0qYKlwyWTw3e_eUe19MxG5osz.
  * Gargiani et al. [2020] M. Gargiani, A. Zanelli, M. Diehl, and F. Hutter.  On the promise of the stochastic generalized gauss-newton method for training dnns.  _arXiv preprint arXiv:2006.02409_ , 2020.  URL https://arxiv.org/abs/2006.02409.
  * Hu et al. [2021] E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen.  Lora: Low-rank adaptation of large language models.  _arXiv preprint arXiv:2106.09685_ , 2021.  URL https://arxiv.org/abs/2106.09685.
  * Hylton [2020] T. Hylton.  Thermodynamic neural network.  _Entropy_ , 22(3):256, 2020.  doi: 10.3390/e22030256.  URL https://www.mdpi.com/1099-4300/22/3/256.
  * Katharopoulos et al. [2020] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret.  Transformers are rnns: Fast autoregressive transformers with linear attention.  In _International conference on machine learning_ , pages 5156–5165. PMLR, 2020.  URL https://proceedings.mlr.press/v119/katharopoulos20a.html.
  * Khan et al. [2018] H. N. Khan, D. A. Hounshell, and E. R. Fuchs.  Science and research policy at the end of moore’s law.  _Nature Electronics_ , 1(1):14–21, 2018.  URL https://www.nature.com/articles/s41928-017-0005-9.
  * Kim et al. [2017] S. Kim, T. Gokmen, H.-M. Lee, and W. E. Haensch.  Analog cmos-based resistive processing unit for deep neural network training.  In _2017 IEEE 60th International Midwest Symposium on Circuits and Systems (MWSCAS)_ , pages 422–425. IEEE, 2017.  URL https://ieeexplore.ieee.org/abstract/document/8052950.
  * Kingma and Ba [2015] D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization.  In _Proceedings of the 3rd International Conference on Learning Representations (ICLR)_ , 2015.  URL http://arxiv.org/abs/1412.6980.
  * Kunstner et al. [2019] F. Kunstner, P. Hennig, and L. Balles.  Limitations of the empirical fisher approximation for natural gradient descent.  _Advances in neural information processing systems_ , 32, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/46a558d97954d0692411c861cf78ef79-Abstract.html.
  * LeCun [1998] Y. LeCun.  The mnist database of handwritten digits.  _http://yann. lecun. com/exdb/mnist/_ , 1998.  URL http://yann.lecun.com/exdb/mnist/.
  * Lin et al. [2023] W. Lin, F. Dangel, R. Eschenhagen, K. Neklyudov, A. Kristiadi, R. E. Turner, and A. Makhzani.  Structured inverse-free natural gradient: Memory-efficient & numerically-stable kfac for large neural nets.  _arXiv preprint arXiv:2312.05705_ , 2023.  URL https://arxiv.org/abs/2312.05705.
  * Lipka-Bartosik et al. [2023] P. Lipka-Bartosik, M. Perarnau-Llobet, and N. Brunner.  Thermodynamic computing via autonomous quantum thermal machines.  _arXiv preprint arXiv:2308.15905_ , 2023.  URL https://arxiv.org/abs/2308.15905.
  * Loshchilov and Hutter [2017] I. Loshchilov and F. Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.  URL https://arxiv.org/abs/1711.05101.
  * Ma et al. [2024] S. Ma, H. Wang, L. Ma, L. Wang, W. Wang, S. Huang, L. Dong, R. Wang, J. Xue, and F. Wei.  The era of 1-bit llms: All large language models are in 1.58 bits.  _arXiv preprint arXiv:2402.17764_ , 2024.  URL https://arxiv.org/abs/2402.17764.
  * Mangrulkar et al. [2022] S. Mangrulkar, S. Gugger, L. Debut, Y. Belkada, S. Paul, and B. Bossan.  Peft: State-of-the-art parameter-efficient fine-tuning methods.  https://github.com/huggingface/peft, 2022.
  * Martens [2020] J. Martens.  New insights and perspectives on the natural gradient method.  _The Journal of Machine Learning Research_ , 21(1):5776–5851, 2020.  URL https://arxiv.org/abs/1412.1193.
  * Martens and Grosse [2015] J. Martens and R. Grosse.  Optimizing neural networks with kronecker-factored approximate curvature.  In _International conference on machine learning_ , pages 2408–2417. PMLR, 2015.  URL https://proceedings.mlr.press/v37/martens15.html.
  * Martens et al. [2018] J. Martens, J. Ba, and M. Johnson.  Kronecker-factored curvature approximations for recurrent neural networks.  In _International Conference on Learning Representations_ , 2018.  URL https://openreview.net/pdf?id=HyMTkQZAb.
  * Martens et al. [2010] J. Martens et al.  Deep learning via hessian-free optimization.  In _ICML_ , volume 27, pages 735–742, 2010.  URL https://www.cs.toronto.edu/~asamir/cifar/HFO_James.pdf.
  * Melanson et al. [2023] D. Melanson, M. A. Khater, M. Aifer, K. Donatella, M. H. Gordon, T. Ahle, G. Crooks, A. J. Martinez, F. Sbahi, and P. J. Coles.  Thermodynamic computing system for AI applications.  _arXiv preprint arXiv:2312.04836_ , 2023.  URL https://arxiv.org/abs/2312.04836.
  * Paszke et al. [2019] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ , 32, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/bdbca288fee7f92f2bfa9f7012727740-Abstract.html.
  * Pauloski et al. [2020] J. G. Pauloski, Z. Zhang, L. Huang, W. Xu, and I. T. Foster.  Convolutional neural network training with distributed k-fac.  In _SC20: International Conference for High Performance Computing, Networking, Storage and Analysis_ , pages 1–12. IEEE, 2020.  URL https://ieeexplore.ieee.org/abstract/document/9355234.
  * Rajpurkar et al. [2016] P. Rajpurkar, J. Zhang, K. Lopyrev, and P. Liang.  Squad: 100,000+ questions for machine comprehension of text.  _arXiv preprint arXiv:1606.05250_ , 2016.  URL https://arxiv.org/abs/1606.05250.
  * Ren and Goldfarb [2019] Y. Ren and D. Goldfarb.  Efficient subsampled gauss-newton and natural gradient methods for training neural networks.  _arXiv preprint arXiv:1906.02353_ , 2019.  URL https://arxiv.org/abs/1906.02353.
  * Sanh et al. [2019] V. Sanh, L. Debut, J. Chaumond, and T. Wolf.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.  _arXiv preprint arXiv:1910.01108_ , 2019.  URL https://arxiv.org/abs/1910.01108.
  * Shen et al. [2021] Z. Shen, M. Zhang, H. Zhao, S. Yi, and H. Li.  Efficient attention: Attention with linear complexities.  In _Proceedings of the IEEE/CVF winter conference on applications of computer vision_ , pages 3531–3539, 2021.  URL https://ieeexplore.ieee.org/abstract/document/9423033.
  * Sun et al. [2020] X. Sun, N. Wang, C.-Y. Chen, J. Ni, A. Agrawal, X. Cui, S. Venkataramani, K. El Maghraoui, V. V. Srinivasan, and K. Gopalakrishnan.  Ultra-low precision 4-bit training of deep neural networks.  _Advances in Neural Information Processing Systems_ , 33:1796–1807, 2020.  URL https://proceedings.neurips.cc/paper/2020/hash/13b919438259814cd5be8cb45877d577-Abstract.html.
  * Wang et al. [2020] S. Wang, B. Z. Li, M. Khabsa, H. Fang, and H. Ma.  Linformer: Self-attention with linear complexity.  _arXiv preprint arXiv:2006.04768_ , 2020.  URL https://arxiv.org/abs/2006.04768.

##  Appendix A Levenberg-Marquardt regularization schedule

Poor-conditioning and singularity of the curvature matrix can greatly decrease
the performance of NGD, which is dealt with by adding a term
$\lambda\mathbb{I}$ to the curvature matrix. Following [33], it is possible to
use a simple method to adapt the value of $\lambda$ at each time step, known
as a Levenberg-Marquardt schedule. This involves computing the reduction ratio
$\rho$, defined as:

$\displaystyle\rho=\frac{\ell(\theta_{k+1})-\ell(\theta_{k})}{q_{\theta_{k}}(\theta_{k+1}-\theta_{k})-q_{\theta_{k}}(0)}$
(20)

with $q_{\theta_{k}}(p)$ the quadratic approximation to $\ell$ around
$\theta_{k}$ defined as:

$\displaystyle
q_{\theta_{k}}(p)=\ell(\theta_{k})+\nabla\ell(\theta_{k})^{\top}p+\frac{1}{2}p^{\top}G_{k}p.$
(21)

If $\rho>a$, $\lambda\leftarrow\alpha\lambda$, and if
$\rho<1-a,\lambda\leftarrow\lambda/\alpha$. This can be interpreted as
distrusting the quadratic model when $\rho$ is small, hence increasing
$\lambda$ for the next iteration. This procedure may be used for TNGD, however
this adds a supplementary digital cost of a GGN-vector product $G_{k}p$ (which
has a similar cost to two JVPs). For our experiments we did not find it to
significantly boost performance although it may be considered for future work.

##  Appendix B TNGD algorithm

In Alg. 1 we provide the steps for the TNGD algorithm. This algorithm may be
used in conjunction with various digital optimizers (such as SGD or Adam). The
thermodynamic linear solver (TLS) is performed by an analog thermodynamic
computer whose physical implementation is described in appendix C. The TLS
takes as inputs the Jacobian $J_{f,k}$, the Hessian $H_{L}$, the gradient
$g_{k}$ and an initial point $x_{0}$ (that can be reset at each iteration, or
not, in which case $t_{d}>0$).

Algorithm 1 Thermodynamic Natural Gradient Descent

$n>0$

Initialize $\theta_{0}$

$\tilde{g_{0}}\leftarrow\nabla\ell(\theta_{0})$

$\texttt{optimizer}\leftarrow\texttt{SGD}(\eta,\beta)\textbf{ or
}\texttt{Adam}(\eta,\beta_{1},\beta_{2})$

while $k\neq n$ do

$x_{k},y_{k}\leftarrow$ next batch

$g_{k}\leftarrow\nabla\ell(\theta_{k},x_{k},y_{k})$

$\tilde{g}_{k}\leftarrow\texttt{TLS}(J_{f},H_{L},b=g_{k},x_{0}=\tilde{g}_{k-1})$

$\texttt{optimizer}.\texttt{update}(\theta_{k},\tilde{g}_{k})$

$k\leftarrow k+1$

end while

##  Appendix C Hardware implementation

The thermodynamic NGD algorithm can be implemented in similar hardware to what
is presented in Ref. [2, 34]. However, this requires one to construct the full
curvature matrix, which is quadratic in the number of parameters, and then
send it to the analog hardware. Therefore, an alternative hardware
implementation that is described by the same electronic circuit equations is
preferred. This alternative implementation is comprised of three arrays of
resistors of size $(bd_{z},N),(bd_{z},bd_{z})$ and $(N,bd_{z})$ for storing
$J_{f},H_{L}$ and $J_{f}^{\top}$, respectively (hence two of these are
rectangular resistor arrays). These three arrays of resistors enable one to
implement the following differential equation in hardware:

$dV=-(J_{f}H_{L}J_{f}^{\top}+\lambda\mathbb{I})Vdt-
bdt+\mathcal{N}(0,2\kappa_{0}\sqrt{dt})$ (22)

where $\kappa_{0}$ is the noise variance and $V=(V_{1},V_{2},\ldots,V_{N})$ is
a vector of voltages.

To understand this let us consider a single resistor resistor array shown in
Fig. 6. By Kirchhoff’s current law, we obtain the equation of motion for the
vector of voltages $V=(V_{1},V_{2},V_{3})$ as:

$C\dot{V}=-\mathcal{G}V+R^{-1}V_{\mathrm{in}}+I_{\mathrm{n}},$ (23)

with $V_{\mathrm{in}}=(V_{\mathrm{in}1},V_{\mathrm{in}2},V_{\mathrm{in}3})$,
$R=\mathrm{diag}(R_{1},R_{2},R_{3})$, $C=\mathrm{diag}(C_{1},C_{2},C_{3})$,
and
$I_{\mathrm{n}}=(I_{\mathrm{n}_{1}},I_{\mathrm{n}_{2}},I_{\mathrm{n}_{3}})$,
which is the current noise vector and is assumed to be Gaussian with zero mean
and variance $2\kappa_{0}$. In this case we have

$\mathcal{G}=\begin{pmatrix}\frac{1}{R_{11}}&\frac{1}{R_{21}}&\frac{1}{R_{31}}\\\
\frac{1}{R_{12}}&\frac{1}{R_{22}}&\frac{1}{R_{32}}\\\
\frac{1}{R_{13}}&\frac{1}{R_{23}}&\frac{1}{R_{33}}\end{pmatrix}.$ (24)

These resistor values need to be tunable such that different Jacobians and
Hessians may be given as inputs throughout the training process (and for the
hardware to be used for different problems).

At steady state ($\dot{V}=0$), the average voltage vector is $\langle
V\rangle=\mathcal{G}^{-1}R^{-1}V_{\mathrm{in}}$, which corresponds to the
solution of the linear system $Ax=b$ with $A=\mathcal{G}$, $x=V$,
$b=R^{-1}V_{\mathrm{in}}$.

One may compose three arrays of resistors to obtain Eq. (22), giving a total
resistor count of $bd_{z}(bd_{z}+2N)$ (with two of them having no capacitors,
which enablies one to obtain a first-order SDE). The damping term proportional
$\lambda$ may be obtained by adding resistors, which has the physical effect
of stabilizing the electrical system (as it does numerically).

Figure 6: Circuit diagram for the thermodynamic device in the case of a single
resistor array implementation for a 3-dimensional problem. The device is
comprised of three voltage sources, each of which is connected to three
resistors $\\{R_{i}\\}$ . Each of these resistors is connected to a line that
goes into the negative pin of a different operational amplifier. A capacitor
connects the negative pin of the operational amplifier to the operational
amplifier’s output port, where the voltage is denoted by $\\{V_{i}\\}$. The
output ports of the operational amplifiers are connected to an array of
$N\times N$ resistors $\\{R_{ij}\\}$ (nine here), each of which connects to
the line going from the resistors $\\{R_{i}\\}$ to the negative pins of the
operational amplifiers. This feedback loop enables the circuit to run a
differential equation whose steady-state corresponds to the solution of a
linear system of equations.

One may run the thermodynamic linear solver by setting the voltage values
$V_{\mathrm{in}}$ to the the gradient $\nabla\ell$ with a digital-to-analog
converter, and set the values of the programmable resistors thanks to a
digital controller.

To obtain the comparisons to other digital methods, we considered the
following procedure to run the TLS on electrical hardware. In such a device
the $\tau=RC$ time constant sets the characteristic timescale of the
thermodynamic device.

  1. 1.

Digital-to-analog (DAC) conversion of the the gradient vector with a given
bit-precision.

  2. 2.

Set the configuration of the programmable resistors ($bd_{z}(bd_{z}+2N)$)
values with a given bit precision to set).

  3. 3.

Let the dynamics run for $t$ (the analog dynamic time). Note that for
simulations this time was chosen heuristically by exploring convergence in the
solutions of the problem of interest.

  4. 4.

Analog-to-digital (ADC) conversion of the solution measured at nodes $V_{i}$
to the digital device.

The runtime estimated are based on the following assumptions:

  * •

16 bits of precision.

  * •

A digital transfer speed of $50Gb/s$.

  * •

$R=10^{3}\,\Omega$, $C=1\,\text{nF}$, which means $RC=1\mu\text{s}$ is the
characteristic timescale of the system.

Finally, note that in all cases that were investigated, the dominant
contribution to the total runtime of TNGD was the digital steps to compute the
gradients, Jacobian and Hessian matrices. Hence some assumptions about the
DAC/ADC may be relaxed and the total TNGD runtime would be similar. The $RC$
time constant may also be reduced to make the algorithm faster, although this
is found to be easily sub-dominant with respect to input operations (setting
the configuration of the device).

##  Appendix D Simulating TNGD

The results reported in this paper require simulating the thermodynamic
device. To do so, we employ a Euler-Maruyama discretization of Eq. 15, where
the update equation is:

$\tilde{g}^{(k+1)}=\tilde{g}^{(k)}+\delta
t(G\tilde{g}^{(k)}-\nabla\ell)+z\sqrt{2\kappa_{0}\delta t}$ (25)

where $\delta t$ is a step size, $z\sim\mathcal{N}(0,1)$ and the GGN-vector
product $G\tilde{g}^{(k)}$ is evaluated in linear time, with no need to
construct $G$ as explained in Section 3. One may consider higher-order
schemes, which in general will cost $d$ GGN-vector products (hence $2d+1$
model calls, accounting for the gradient) for each step of an order $d$
solver.

With an Euler-Maruyama scheme, one therefore requires 3 model calls per time
step, which results in long simulation times for the larger $t$ values we
report.

##  Appendix E Experimental details

All experiments except the one reported in Fig. 2 were carried out on a Nvidia
A100 GPU with 80 Gb of RAM. The experiment corresponding to Fig. 2 was carried
out on a AMD EPYC 7763 64-Core CPU with 32 GB of RAM (the results on the GPU
had too much variance even for a large number of repetitions). For Fig. 2,
$b=32$, $c=200$, and the results were obtained by repeating over 5 manual
random seeds, with the standard deviation over runs being shown as shaded
areas. Modifying $c$ has the simple effect of shifting the curve on the scale.

All experiments are written in PyTorch [35], and we have used the posteriors
library [12] which supports GGN-vector products, constructing GGN matrices
(for exact NGD and NGD-Woodbury) and the conjugate gradient solver in PyTorch.

###  E.1 MNIST

For the MNIST experiments, we train on 10,000 images and test on 10,000 other
images, with $b=64$ for 10 epochs. Below is a table with hyperparameters for
each figure in the main text:

Figure – Optimizer | Optimizer parameters
---|---
Fig. 3 – Adam | $\eta=0.001,\beta_{1}=0.9,\beta=0.999,\epsilon=1e-8$
Fig. 3 – TNGD | $\eta=0.01,\beta=0,t=50\tau,t_{d}=0,\lambda=0.01,\delta t=0.1\tau$
Fig. 2(a) – TNGD | $\eta=0.01,\beta=0,\lambda=0.01,\delta t=0.1\tau$

For these experiments, the plotted data is the mean (and standard deviation)
of the moving average over $200$ points for the five same manual random seeds.
The Adam experiments took $\sim 5$ minutes to run, while the longest TNGD
experiments took $\sim 14$ hours (due to many time steps being required, see
Section D. We performed sweeps over the damping and learning rate value of
TNGD which we do not report in the paper that took $\sim 10$ days of
accumulated total runtime.

####  E.1.1 Noisy simulation

One key feature of TNGD is that it is noise-resilient. Indeed, because the
solution of the linear system of Eq. (6) is encoded in the first moment of the
equilibrium distribution, any noise that is approximately Gaussian will not
affect much the quality of the results for reasonable noise levels. In Fig. 7,
the loss vs. iterations are shown for varying noise levels, which are defined
by the value of the noise variance $\kappa_{0}$. For $\kappa_{0}<0.01$, the
noise essentially does not affect the performance of TNGD (as the influence of
noise of performance starts to saturate at this value), even for a very small
analog dynamics time (here, $t=\tau$). For a realistic electrical device where
$\theta_{t}$ are voltages, the contribution of thermal noise to the noise
level would be $\kappa_{0}\sim 10^{-6}V$. Other noise sources may contribute
to the noise level, but because of the nature of the TNGD algorithm, it
exhibits a high noise-resilience. Note that it is also possible to collect
more samples from the device to reduce influence of the noise if the noise
level is large.

Figure 7: Training loss vs. iterations for varying noise levels. The noise
level is defined by the noise variance $\kappa_{0}$ entering Eq. (15). Here
$t=\tau$.

###  E.2 Extractive question-answering

For the QA experiments, we train on 800 articles and test on 200 other
articles of the SQuaD dataset, with $b=32$ for 5 epochs. Below is a table with
hyperparameters for Fig. 5:

Figure – Optimizer | Optimizer parameters
---|---
Fig. 5(a) – TNGD | $\eta=0.01,\beta=0,t=0.4\tau,t_{d}=0,\lambda=1$
Fig. 5(a) – Adam | $\eta=0.001,\beta_{1}=0.9,\beta_{2}=0.999,\epsilon=1e-8$
Fig. 5(b) – TNGD-Adam | $\eta=0.001,\beta_{1}=0,\beta_{2}=0,\epsilon=1e-8,t_{d}=0,\lambda=1,\delta t=0.02\tau$

We apply low-rank adaptation (LoRA) to the $Q,K,V$ modules and output
projection matrices of the attention layers with parameters $r=2,\alpha=32$
and a dropout of 0.1. LoRA consists in replacing the pre-trained weight
matrices of the targeted layers $W_{0}$ by:

$\tilde{W}=W_{0}+AB,$ (26)

where $A$ and $B$ are two rectangular matrices with their smaller dimension
being $\alpha$ (hence $AB$ is low-rank). We used the peft package [29], which
interfaces smoothly with PyTorch and posteriors.

For these experiments we only report a single fixed random seed due to long
simulation times. The Adam experiments took $\sim 20$ minutes, while the
longest TNGD experiments took $\sim 2$ days (due to many time steps being
required, see Section D. We performed sweeps over the damping and learning
rate value of TNGD and Adam which we do not report in the paper that took
$\sim 10$ days of accumulated total runtime.