# Staged Training for Transformer Language Models

Sheng Shen† Pete Walsh⋄ Kurt Keutzer† Jesse Dodge⋄ Matthew Peters⋄ Iz Beltagy⋄
† University of California, Berkeley
⋄ Allen Institute for AI
<EMAIL_ADDRESS><EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

The current standard approach to scaling transformer language models trains
each model size from a different random initialization. As an alternative, we
consider a staged training setup that begins with a small model and
incrementally increases the amount of compute used for training by applying a
“growth operator” to increase the model depth and width. By initializing each
stage with the output of the previous one, the training process effectively
re-uses the compute from prior stages and becomes more efficient. Our growth
operators each take as input the entire training state (including model
parameters, optimizer state, learning rate schedule, etc.) and output a new
training state from which training continues. We identify two important
properties of these growth operators, namely that they preserve both the loss
and the “training dynamics” after applying the operator. While the loss-
preserving property has been discussed previously, to the best of our
knowledge this work is the first to identify the importance of preserving the
training dynamics (the rate of decrease of the loss during training). To find
the optimal schedule for stages, we use the scaling laws from Kaplan et al.
(2020) to find a precise schedule that gives the most compute saving by
starting a new stage when training efficiency starts decreasing. We
empirically validate our growth operators and staged training for
autoregressive language models, showing up to 22% compute savings compared to
a strong baseline trained from scratch. Our code is available at
https://github.com/allenai/staged-training.

††footnotetext: Corresponding author<EMAIL_ADDRESS>
##  1 Introduction

Language models form the backbone of many modern NLP systems, and these
language models have become progressively larger in recent years. Parameter
counts for these models have grown significantly from ELMo (94 M) (Peters et
al., 2018) to GPT-3 (175 B) (Brown et al., 2020). While larger models with
more learnable parameters perform better on a wide range of tasks, the
computational cost to train or even just to evaluate these models has become
prohibitively expensive (Schwartz et al., 2020). In this paper, we demonstrate
a method to reduce the compute cost of training transformer language models
Vaswani et al. (2017) through a staged training setup that iteratively builds
a large model from a smaller one.

Figure 1:  We train a GPT2${}_{\textsc{large}}$ (768M parameters) transformer
language model by first training a model 1/4 the size (orange line), then
increasing the model size by 4x by applying a growth operator to the entire
training state, and restarting training (green line). The result is a large
size model with comparable loss to one trained from scratch (blue line) but
with reduced compute cost illustrated initially by the dashed red arrow.

Most prior work on scaling language models initializes each model size from a
random initialization and trains to convergence Kaplan et al. (2020). This
work illustrated an intriguing property of model training shown in Figure 1,
namely that smaller models are initially more compute efficient then larger
models, but eventually the larger model will reach a lower loss. Our central
idea is to take advantage of this property by first training a smaller model
in the compute efficient region, applying a growth operator to the entire
training state, and restarting training with a larger sized model. We
introduce two operators to perform this growing operation along the model
depth and width dimensions. We also identify two important properties of the
operators: first, the loss before growing the model is preserved, and second,
the training dynamics of the grown model match that of an equivalent model
trained from scratch. To maintain training dynamics, growth operators must
take into the entire training state, including the optimizer state and
learning rate schedule, in addition to the model weights. As can be seen in
Figure 1 our growth operator is loss-preserving, and we show in subsequent
sections that it also preserves the training dynamics. These properties also
make applying the growth operators conceptually and algorithmically simple as
they won’t disrupt the training process, and prior results regarding the model
size needed to converge to a particular loss still hold.

While prior work Rusu et al. (2016); Wei et al. (2016); Gong et al. (2019);
Liu et al. (2019); Li et al. (2020b); Press et al. (2020); Gu et al. (2021);
Li et al. (2021); Rae et al. (2021); Evci et al. (2022) has examined some
aspects of staged training, our work is the first to address all aspects
including how to grow the entire training state and set the stage schedule. We
begin by describing the details of our growth operators for the model weights
and optimizer state. We then present a principled way to chose the stage
schedule, including how to chose the model sizes and number of gradient for
each stage. Intuitively, we should start a new stage when the training
efficiency decreases and the rate of loss decrease starts to slow down. To
formalize this intuition, we use the scaling laws from Kaplan et al. (2020) to
find the optimal schedule that gives the maximum compute saving. We then show
how to approximate the optimal schedule in the realistic scenario without
perfect knowledge of the scaling laws. We empirically validate our approach
with GPT2 style Radford et al. (2019) auto-regressive language model models
and demonstrate 5-30 % compute savings measured by validation loss, and zero-
shot perplexity using two benchmark datasets.

(a) Width growth with GPT2${}_{\textsc{large}}$ as target

(b) Depth growth with GPT2${}_{\textsc{large}}$ as target

Figure 2:  Our growth operators are loss-preserving and training-dynamics
preserving. Using (a) as an example, GPT2${}_{\text{{large}/4}}$ is the
original model which is 4x smaller than the target model
GPT2${}_{\textsc{large}}$. The model GPT2${}_{\text{{large}/4x4}}$ is the
grown model resulting from growing GPT2${}_{\text{{large}/4}}$ by 4x (doubling
model width). The pre-growth point is highlighted on the original model
GPT2${}_{\text{{large}/4}}$, and the post-growth point is highlighted on the
grown model GPT2${}_{\text{{large}/4x4}}$. The pre-growth and post-growth
points have the same loss, showing that the width growth operator is loss-
preserving. To demonstrate that it is also training dynamics-preserving, we
overlay the loss curve for the grown model over the target model and confirm
the rate of loss decrease with respect to the number of tokens is the same as
the target model trained from scratch. The x-axis is number of training tokens
since random initialization, or from the start of the training stage (for
GPT2${}_{\text{{large}/4x4}}$). A similar result is seen in (b) for the depth
growth operator.

##  2 Definitions and Properties of Growth Operators

###  2.1 Definitions

We begin by defining some terms which will be used throughout the paper. We
consider a model $\mathbf{y}=\phi(\mathbf{x},\theta)$ which takes input
$\mathbf{x}$, outputs $\mathbf{y}$, with parameters $\theta$. The model is
trained by minimizing a loss function
$\textit{loss}(\mathbf{y},\hat{\mathbf{y}})\in\rm\mathbb{R}$ through a
sequence of parameter updates, obtained by running an optimizer. We will also
write $\textit{loss}(\phi,\mathcal{D})$ for the total loss over a dataset
$\mathcal{D}=\\{\mathbf{x}_{i},\mathbf{y}_{i}\\}$ (or just
$\textit{loss}(\phi)$). The parameter updates for a particular mini-batch are
determined by both the mini-batch and the training state,
$\mathcal{T}=\\{\theta,(m,v),\lambda(t)\\}$, including the model parameters
$\theta$, optimizer state ($m,v$, here the first and second moments of the
Adam optimizer), and learning rate schedule $\lambda(t)$. Given a training
state, we apply a growth operator
$\mathbb{G}(\mathcal{T}_{\text{orig}})=\mathcal{T}_{\text{grow}}$ that takes
the original training state and outputs a grown training state where the model
size has increased, along with a corresponding compatible optimizer state and
learning rate schedule.

###  2.2 Desired properties

In this section, we define two key properties of growth operators. Building on
(Chen et al., 2016) we revisit the loss-preserving property, and we introduce
a more challenging property, the training-dynamic-preserving property.

####  2.2.1 Preserving Loss

A function-preserving growth operator is one that takes as input an original
model and returns a grown model that represents the same function as the
original model. If an operator is function-preserving then it is also loss-
preserving.

Mathematically, we can formulate loss-preserving as

$\textit{loss}(\mathbb{G}(\phi)(\mathbf{x}),\mathbf{y})=\textit{loss}(\phi(\mathbf{x}),\mathbf{y})$
(1)

for any $(\mathbf{x},\mathbf{y})$. A growth operator that is not loss-
preserving wastes time and compute initially until it recovers the same
performance of the original model. Figure 2 (and the more detailed Figure 4)
show examples of the proposed width and depth growth operators being loss-
preserving.

####  2.2.2 Preserving Training Dynamics

We define the training dynamics as the rate of decrease of loss relative to
the amount of compute, and a training-dynamics-preserving growth operator is
one that allows the grown model’s loss curve to match that of a target model
(a model of the same size as the grown model but trained from scratch).

Formally, let $\phi_{k+1}$ be the resulting model after applying the optimizer
update to model $\phi_{k}$ with training state $\mathcal{T}_{k}$. Applying the
update requires some amount of compute $C_{k}$. The training process produces
a loss curve that associates the loss with the total amount of compute used
for training

$\mathcal{L}(\phi,C)=\\{(\overline{C}_{k},\mathit{loss}(\phi_{k},\mathcal{D}_{k})),\hskip
4.30554ptk=1,2,\ldots\\}$

where $\overline{C}_{k}=\sum_{i<=k}C_{i}$ is the total compute used at step
$k$. The training dynamics is the compute efficiency of training, the expected
decrease in the loss relative to the amount of compute:

$\frac{\partial}{\partial
C}\mathcal{L}(\phi,C)=\mathbb{E}_{D}[\frac{\mathit{loss}(\phi_{k},\mathcal{D})-\mathit{loss}(\phi_{k+1},\mathcal{D})}{C_{k}}]$

which we denote by $\frac{\partial\mathcal{L}}{\partial C}$ as an abuse of
notation. Practically it is easy to estimate during training by monitoring the
model’s loss.

We can now define a training-dynamics-preserving growth operator $\mathbb{G}$
at a point on the loss curve with loss $L_{\mathbb{G}}$ as one that preserves
the efficiency of training of the grown model vs. a target model trained from
scratch:

$\frac{\partial}{\partial
C}\mathcal{L}(\mathbb{G}(\phi_{orig}),C)=\frac{\partial}{\partial
C}\mathcal{L}(\phi_{target},C)$ (2)

where efficiency is evaluated at $L_{\mathbb{G}}$.

Notice that while loss-preserving is a property comparing the original and
grown models, training-dynamics-preserving is a property comparing the grown
and target models. This property makes it possible for the grown model to
“jump” from the loss curve of the smaller model to the large model and always
benefit from faster convergence. A growth operator that does not satisfy this
requirement could be creating a larger model but with limited capacity or one
that is more difficult to train. Figure 2 (and the more detailed Figure 4)
shows examples of the width and depth growth operators preserving the training
dynamics, where the grown model perfectly follows the loss curves of the
target model.

While the function-preserving property of a growth operator can be confirmed
based on the implementation itself, preserving the training dynamics goes
beyond just growing the model size; it must be empirically evaluated. To
preserve training dynamics, one must address the whole training state
including the learning rate and the optimizer state, which are hardly
discussed in prior work. We discuss this further in the next section.

##  3 Growth Operators

We introduce two growth operators below; the operators are described
generally, though of course the implementations are model-specific. Our
experiments are using the GPT2 transformer architecture Radford et al. (2019).

###  3.1 Width

Our width operator doubles the hidden dimension of the entire model, and
therefore increases the number of parameters by approximately 4x. This
operator applies to all weights in the network including embeddings, feed
forward layers, bias, and normalization layers, not just the feed forward
layers as in prior work (Gu et al., 2021). It grows each layer in slightly
different ways. Layer-normalization layers and bias terms with weights
$W\in\mathbb{R}^{d}$ are duplicated:

$\mathbb{G}(W)=[W,W].$

where $[\cdot,\cdot]$ represents concatenation and we have overloaded
$\mathbb{G}(\cdot)$ to apply to a weight matrix instead of the entire training
state. Embedding layers are handled in a similar manner. For the feed forward
weights $W\in R^{nxm}$, we design the growth operators as

$\mathbb{G}(W)=\begin{pmatrix}W&Z\\\ Z&W\\\ \end{pmatrix}$

where $Z$ is a zero matrix of size similar to $W$.

In this case, the input before the grown last feed forward layer that produces
the logits is two times wider and the final logits are two times larger. To
keep the whole network loss-preserving, we divide the grown weights of the
last feed forward layer by a factor of two.

###  3.2 Depth

Our depth operator doubles the number of layers, and therefore increases the
number of non-embedding parameters by 2x. Given a model with layers
$(\phi_{0},\phi_{1},\ldots)$, the depth operator adds an identity layer
$\phi_{\text{id}}$ after each layer in the original model, so that the grown
model has layers
$(\phi_{0},\phi_{\text{id}},\phi_{1},\phi_{\text{id}},\ldots)$. The identity
layer is a layer where the input matches the output
$\phi_{\text{id}}(\mathbf{x})=\mathbf{x}$.

To construct the identity layer, we start with the formulation of each layer
in GPT2 as two sublayers:

$\displaystyle{\mathbf{x}^{\prime}}$
$\displaystyle={\mathbf{x}}+\text{Attention}(\text{LN}({\mathbf{x}})),$ (3)
$\displaystyle{\mathbf{y}}$
$\displaystyle={\mathbf{x}^{\prime}}+\text{FFN}(\text{LN}({\mathbf{x}^{\prime}}))$

where $\bf x\in\rm R^{d}$ is the input, $\bf y\in\rm R^{d}$ is the output. LN,
Attention, and FFN stands for the layer normalization, multi-head attention
and feed-forward operations. We initialize both the scale and bias parameters
of each LN in the identity layers to zero, so that
$\text{LN}(\mathbf{x})=\mathbf{0}$. We also set the bias parameters of all
linear layers to zero, which combined with the LN initialization gives
$\text{Attention}(\text{LN}({\bf x}))=\mathbf{0}$ and
$\text{FFN}(\text{LN}({\bf x^{\prime}}))=\mathbf{0}$, and the entire layer
reduces to an identity layer at initialization. Overall, the resulting depth
growth operator is loss-preserving.

###  3.3 Growth operator’s impact on training state

In practice, to build a growth operator that preserves training dynamics, we
find it important that optimizer state should be grown in a similar way to the
model parameters; initial experiments indicated that it can take many training
steps to re-estimate the optimizer state, and the initial phase after growth
can be unstable. This is expected because training dynamics is the rate of
loss change $\frac{\partial\mathcal{L}}{\partial C}$. To match the rate of
loss change of the target model, the growth operator needs to reproduce the
same scale of model updates, which are controlled by the update rule of the
optimizer. Using the Adam Kingma and Ba (2014) optimizer as an example, the
update rules are

$\displaystyle m_{t}$ $\displaystyle=\beta_{1}*m_{t-1}+(1-\beta_{1})g_{t}$ (4)
$\displaystyle v_{t}$ $\displaystyle=\beta_{2}*v_{t-1}+(1-\beta_{2})g_{t}^{2}$
$\displaystyle\theta_{t+1}$
$\displaystyle=\theta_{t}-\frac{\lambda(t)}{\sqrt{v_{t}}+\epsilon}m_{t}$

where $g$ and $g^{2}$ are the first-order gradients and the element-wise
squared first-order gradients, $m$ is the first moment (average of $g$), $v$
is the second moment (average of $g^{2}$), $\beta_{1},\beta_{2},\in[0,1)$ are
the exponential decay rates for the moment estimates, $\epsilon$ is a small
constant, and $t$ is the time-step. For the grown model to be updated at a
rate similar to that of the target model, it needs to match its learning rate
$\lambda(t)$ which we discuss in the next section. It also needs to produce an
optimizer state $m,v$ that’s compatible with the gradients of the grown model,
$g(\mathbb{G}(\phi))$ and $g^{2}(\mathbb{G}(\phi))$. Given that $m,v$ are
averages of $g,g^{2}$, we argue that $m,v$ should be grown with growth
operators $\mathbb{G}_{m}$, $\mathbb{G}_{v}$ that satisfy the following
properties:††Abusing the notation; Eq. 5, 6 are using $g$, $g^{2}$ as
functions to compute gradients of a model, not the gradients themselves.

$g(\mathbb{G}_{m}(\phi))={\mathbb{G}}_{m}(g({\phi}))$ (5)
${{g^{2}}(\mathbb{G}_{v}(\phi))}={\mathbb{G}_{v}(g^{2}(\phi))}$ (6)

The first condition states that the “gradients of the grown model” should
match “growing the gradients of the original model”. The second condition is
similar but for the squared gradients. To satisfy Eq. 5 and 6, the
implementations of $\mathbb{G}_{m}$, $\mathbb{G}_{v}$ are slightly different
from $\mathbb{G}$. For the width growth operator, some of the weights need to
be scaled by 0.5x or 0.25x to account for the 2x scaling in the forward pass
(see Section 3.1). For the depth growth operator, we copy $m$ and $v$ for the
original model layers and set $m$ and $v$ to zero for the identity layers.

Along with the loss-preserving property, the training dynamics preserving
property ensures that the new optimizer state is compatible with the grown
model weights.

##### Learning rate

To match training dynamics, the learning rate schedule of the grown model must
match that of the target model. The intuition is that our growth operators
allow the model state to “jump” from the loss curve of the original model
$\mathcal{L}(\phi_{orig},C)$ to $\mathcal{L}(\phi_{target},C)$ at a point with
loss $L_{\mathbb{G}}$. Because our growth operator is loss preserving, the
loss $L_{\mathbb{G}}$ defines two points on the loss curves, pre-growth on
$\mathcal{L}(\phi_{orig},C)$ and growth-target on
$\mathcal{L}(\phi_{target},C)$. To match the training dynamics of
$\mathcal{L}(\phi_{target},C)$, we start training the grown model with a
learning rate schedule that matches the target model but starts from growth-
target. We discuss finding that matched growth-target point in Section 4 and
5.

##  4 Optimal Schedule

Prior work Gu et al. (2021); Gong et al. (2019) used heuristics to determine
the training schedule. In contrast, our goal is to find the optimal training
schedule. An optimal training schedule is one that, given a target model size,
specifies the optimal sequence of growth operators, intermediate model sizes,
and number of training steps in each stage leading to the most compute saving.
This section will explain our intuition behind our optimal schedule, then
explains how to mathematically find it.

##### Training to optimality

We start from the scaling laws Kaplan et al. (2020), which showed that the
training of transformer language models is initially efficient with fast loss
reduction, then the compute-efficient regime ends and the rate of the loss
reduction slows down. In addition, the initial compute-efficient regime is
longer for larger models. These ideas are illustrated in Figure 1where
$\frac{\partial\mathcal{L}}{\partial C}$ is initially large then it slows
down. As shown in Kaplan et al. (2020) and Li et al. (2020b), the optimal
compute allocation should favor a large model size and stop the training by
the end of the initial compute-efficient regime when
$\frac{\partial\mathcal{L}}{\partial C}=\tau_{opt}$, where $\tau_{opt}$ is
some threshold. We call this training to “Optimality” as opposed to training
to “Completion” or to convergence. We discuss later this section how to find
the point of optimality using constrained optimization in an idealistic
scenario, then later in Section 5 using a more practical method that estimates
$\tau_{opt}$.

##### Intermediate Stages

We next discuss where in training to grow a model. Intuitively, the optimal
schedule is one where the original small model is trained until its compute-
efficient regime ends, then grown to a larger model to continue its compute-
efficient regime. Figure 1 highlights one such potential schedule. Notice that
there’s a specific point on the loss curve of the original model that leads to
the most compute saving; growing the model earlier means a wasted opportunity
skipping some of the fast loss reduction stage of the original model, and
growing the model later means wasting compute by continuing to train after the
loss of the original model begins to plateau.

##### Schedule which minimizes compute

Next, we describe how to mathematically find this optimal schedule. For that,
we use the scaling laws Kaplan et al. (2020) which derived empirical fits for
the language model loss $L$ as it relates to compute $C$, number of non-
embedding parameters $N$, number of gradient update steps $S$, and batch size
$B$. The total compute and the loss are given by

$\displaystyle C$ $\displaystyle\approx 6NBS,$ (7) $\displaystyle L(N,S)$
$\displaystyle=\left(\frac{N_{c}}{N}\right)^{\alpha_{N}}+\left(\frac{S_{c}}{S}\right)^{\alpha_{S}}$

where $\alpha_{N},\alpha_{S},\alpha_{B},N_{c},S_{c},B_{*}$ are all model-
specific constants. Thus, finding the the optimal schedule can be formulated
as a constrained optimization problem. The output is the intermediate model
sizes, and the amount of compute for each stage. We discuss the details in
Appendix C.

##  5 Practical Schedule

While the general scaling laws are known, re-estimating their constants
($\alpha_{N},\alpha_{S},\alpha_{B},N_{c},S_{c},B_{*}$ ) for our setup is
challenging because it requires running a large number of models of different
sizes. Instead of estimating all the constants, we make the observation that
we only need to find three key points:

  * •

pre-growth $\in\mathcal{L}(\phi_{orig},C)$ at which we grow the original model

  * •

growth-target $\in\mathcal{L}(\phi_{target},C)$ that the model is grown
towards to

  * •

optimality $\in\mathcal{L}(\phi_{grown},C)$ at which we stop training the
grown model.

Next we discuss the mathematical definition of each point, how to find them,
and how to use them in the actual training procedure.

#####  pre-growth and optimality points

We define pre-growth and optimality using the slope of the loss curve, as they
depend on the rate of change of the loss. Formally,

pre-growth $\displaystyle:\frac{\partial}{\partial
C}\mathcal{L}(\phi_{orig},C)=\tau_{\mathbb{G}}$ (8) optimality
$\displaystyle:\frac{\partial}{\partial
C}\mathcal{L}(\phi_{grown},C)=\tau_{opt}$

where $\tau_{\mathbb{G}}$ and $\tau_{opt}$ are empirically estimated
thresholds.

Importantly, both thresholds are independent of the model size, and this
independence can be derived from Eq. 7. We also empirically confirmed the
model-size independence by training many models of different sizes;
reconfirming results from Kaplan et al. (2020), the shape of the loss curves
for different sized models were similar, just shifted and scaled.
Additionally, while both thresholds are model-size independent,
$\tau_{\mathbb{G}}$ is a function of the growth operator.

#####  growth-target point

The importance of the growth-target point is that it specifics the learning
rate schedule of the grown model (Section 3.3). We will specify growth-target
using the number of training steps ${S_{\text{{growth-target}}}}$. We found
that it can be simply defined using

$\displaystyle{S_{\text{{growth-target}}}}=\rho\times S_{\text{{pre-growth}}}$
(9)

where $\rho$ is an empirically estimated constant, and $S_{\text{{pre-
growth}}}$ is number of training steps at pre-growth.

As above, and knowing that the loss curves of models of different sized models
are scaled and shifted versions of each other, we can use Eq. 7 to show that
$\rho$ is independent of the model size, and it is only a function of the
growth operator. We also verified this empirically using different model
sizes.

##### Estimating $\tau$ and $\rho$

Given that equations 8 and 9 are independent of model sizes, it is enough to
estimate the values of $\tau$ and $\rho$ using small models. To estimate them,
we first identify the three necessary points. Specifically, for a single
growth operator, we train an original model and a target model from scratch
then follow the intuition discussed in Section 4 to choose a pre-growth point
(on $\mathcal{L}(\phi_{orig},C)$) and an optimality point (on
$\mathcal{L}(\phi_{target},C))$. The growth-target point is simply the point
on $\mathcal{L}(\phi_{target},C))$ with the same loss at pre-growth. A plot
like Figure 1 (and the more detailed Figure 5) make it easy to manually
identify the three points, and we leave it to future work to automatically
find these points.

Notice that this method required training a target model from scratch, but in
practice we can estimate the constants once for smaller model sizes and apply
them to larger sizes. Using the identified points, we use equations 8 and 9 to
estimate values of $\tau_{\mathbb{G}}$ and $\rho$ for each growth operator,
and the value of $\tau_{opt}$. The empirical values we estimated are in
Appendix D. Notice that estimating these constants is much simpler than
estimating all the constants of the scaling laws
($\alpha_{N},\alpha_{S},\alpha_{B},N_{c},S_{c},B_{*}$) making this procedure
simpler to apply to a new setup.

Algorithm 1 Staged training for transformer LMs

0:

$\phi$: original model

$\lambda(t)$: learning rate schedule

$M$: number of stages

$(\mathbb{G}_{i},\tau_{{}_{\mathbb{G},i}},\rho_{{}_{\mathbb{G},i}})$:
$($growth op, $\tau$, $\rho)$ for stage $i$

$\mathbb{G}_{1}$: first stage operator assumed to be identity operator

$\tau_{opt}$: last stage’s $\tau$

$t\leftarrow 0$ // number of training steps

for $i=1$ to $M$ do

if $i=M$ then

$\tau\leftarrow\tau_{opt}$ // last stage, stop at optimality

else

$\tau\leftarrow\tau_{{}_{\mathbb{G},i}}$

end if

while $\frac{\partial}{\partial C}\mathcal{L}(\phi,C)\leq\tau$ do

Run training step and update $\phi$

$t\leftarrow t+1$ // update learning rate using $\lambda(t)$

end while

$\phi\leftarrow\mathbb{G}(\phi)$

$t\leftarrow t\times\rho_{{}_{\mathbb{G},i}}$ // set learning rate for next
stage

end for

return $\phi$

##### Training Procedure

Algorithm 1 summarizes the staged training procedure. Notice that it extends
it to the $M$ stage case. The algorithm starts with the model $\phi$, trains
it from scratch until pre-growth, grows the model, sets number of steps for
the grown model, and repeats until no more stages, then continues training to
optimality.

##  6 Experiments

In this section we present our main empirical results, focusing on the amount
of compute saving. We show results on in-domain data (validation loss) and in
the zero-shot transfer setting. We also compare our work to prior work in
growing transformer language models, establishing that previous methods fall
short in one or more areas.

###  6.1 Experimental setting

We experiment with GPT2 from Radford et al. (2019) (in base and large sizes)
using the public C4 Raffel et al. (2020) dataset. We follow the learning rate
schedule in Kaplan et al. (2020) for all model sizes, where the warmup period
is set to 3,000 steps, the batch size is set to 512 and the sequence length is
1,024. For the zero-shot transfer learning, we experiment on two tasks:
Wikitext-103 Merity et al. (2017) and LAMBADA Paperno et al. (2016), similar
to Li et al. (2021). We report the compute saving for in-domain validation
loss and zero-shot performance. We compare our practical schedule in Section 5
with the manual schedule that directly matches the loss of the pre-growth
model with the target model to select the growth-target point. We choose and
report different thresholds to decide the optimality of the training in each
stage.

###  6.2 Main results

Figure 1 shows the compute saving for growing GPT2${}_{\textsc{large/4}}$
original model to a target model of GPT2${}_{\textsc{large}}$ (we show the
results of other combinations of growth operators and the results of growing
to GPT2${}_{\textsc{base}}$ in Appendix A in Figure 5). It is clear that our
grown models are reaching the same loss as the target models but with less
compute. It is important to note that as both models train longer, the amount
of compute saving drops. This illustrates a key design in our schedule where
we stop training of the grown model at optimality when the compute-efficient
regime ends. Figure 1 also demonstrates that small models achieve better loss
trade-off when the total compute is limited, but saturate at higher loss when
more compute is added, highlighting the advantages of our schedule in applying
the growth operator to them before convergence.

|  | GPT2${}_{\textsc{Large}}$ | GPT2${}_{\textsc{Base}}$
---|---|---|---
|  | 5k | 10k | 14k | 3.75k | 7k | 11k
Baseline | (loss) | 3.21 | 3.03 | 2.97 | 3.61 | 3.45 | 3.38
Compute savings (percent saved vs. baseline)
1 stage${}_{\text{ manual}}$ | 2xW | 19.3 | 5.6 | 4.0 | 23.8 | 14.5 | 13.8
2xD | 33.5 | 7.8 | 6.3 | 24.0 | 23.6 | 21.8
1 stage${}_{\text{ practical}}$ | 2xW | 22.5 | 7.3 | 5.2 | 24.3 | 20.2 | 19.7
4xW | 18.0 | 5.3 | 3.8 | 16.0 | 8.6 | 5.5
2xD | 37.0 | 11.0 | 6.1 | 24.7 | 20.4 | 19.8
4xD | 22.5 | 7.3 | 5.2 | 18.8 | 10.1 | 6.4
2xDxW | 28.8 | 5.4 | 3.8 | 19.0 | 9.5 | 6.83
2 stage${}_{\text{ practical}}$ | 2x2xW | 26.8 | 10.9 | 7.8 | 26.8 | 17.9 | 11.4
2x2xD | 30.0 | 14.5 | 10.4 | 33.3 | 21.4 | 15.9

Table 1: Percentage compute savings for GPT2${}_{\textsc{Large}}$ and
GPT2${}_{\textsc{Base}}$ on in-domain validation loss on C4, “W” is width and
“D” is depth growth operators. Percent savings is how much less compute our
approach takes than the baseline to to train a model to equal or better loss
than the baseline. Significant savings can be found using our both width and
depth growth operators, and two growth stages can lead to even more savings
than one growth stage. The derivative threshold at 5k steps is -0.1, at 10k
steps it’s -0.05, and at 14k steps it’s -0.04; thus, 5k is undertrained, 10k
is approximately at the optimality threshold of -0.052, and 14k is trained
beyond optimality. Similar for GPT2${}_{\textsc{Base}}$.

|  | Zero-shot Wikitext-103 (PPL) | Zero-shot LAMBADA (accuracy)
---|---|---|---
|  | GPT2${}_{\textsc{Large}}$ | GPT2${}_{\textsc{Base}}$ | GPT2${}_{\textsc{Large}}$ | GPT2${}_{\textsc{Base}}$
|  | 5k | 10k | 14k | 3.75k | 7k | 11k | 5k | 10k | 14k | 3.75k | 7k | 11k
Baseline |  | 41.0 | 32.3 | 30.3 | 68.5 | 57.1 | 50.0 | 39.6 | 43.2 | 44.7 | 31.1 | 33.0 | 34.7
Compute savings (percent saved vs. baseline, negative means more compute than
baseline)
1 stage${}_{\text{ manual}}$ | 2xwidth | -18.5 | 6.2 | 5.8 | -19.7 | 0.3 | 2.4 | 3.2 | 6.7 | 8.3 | -8.3 | 0.2 | 13.6
2xdepth | 28.5 | 11.8 | 12.0 | 24.0 | 28.0 | 17.3 | 33.5 | 16.8 | 13.8 | 24.0 | 22.9 | 18.2
1 stage${}_{\text{ practical}}$ | 2xwidth | -20.5 | -0.25 | 8.7 | -15.7 | 5.9 | 3.8 | -15.5 | 12.3 | 14.8 | -15.7 | 3.3 | 10.6
4xwidth | -13.4 | 1.3 | 7.4 | -12.0 | 1.1 | 6.4 | -11.0 | 18.0 | 18.2 | -12.0 | 10.1 | 8.6
2xdepth | 32.0 | 13.5 | 11.4 | 31.3 | 33.5 | 17.5 | 32.0 | 23.5 | 21.4 | 24.7 | 16.8 | 13.9
4xdepth | 21.0 | 17.5 | 9.5 | 14.6 | 7.9 | 3.1 | 21.0 | 20.5 | 14.6 | 14.6 | 9.4 | 7.8
2xdepthxwidth | -10.5 | -0.3 | 1.6 | -12.0 | 3.6 | 4.5 | -95.0 | -37.5 | 0.7 | 5.4 | 6.4 | 6.4
2 stage${}_{\text{ practical}}$ | 2x2xwidth | -2.5 | 6.3 | 9.3 | 3.3 | 5.4 | 8.0 | -3.3 | 13.4 | 20.3 | -3.3 | 1.8 | 11.8
2x2xdepth | 30.0 | 12.5 | 12.8 | 33.3 | 14.3 | 11.8 | 19.0 | 14.5 | 23.6 | 19.9 | 10.6 | 15.0

Table 2: Percentage compute savings for GPT2${}_{\textsc{Large}}$ and
GPT2${}_{\textsc{Base}}$ on out-of-domain Wikitext-103 (perplexity) and
LAMBADA (accuracy). Percent savings is how much less compute our approach
takes than the baseline to to train a model to equal or better perplexity or
accuracy than the baseline; negative numbers mean our approach took more
compute than the baseline to achieve equal or better performance. Results are
mixed in the early stages of training, but our approach leads to compute
savings for all experiments later in training (14k for large, 11k for base).
The derivative threshold at 5k steps is -0.1, at 10k steps it’s -0.05, and at
14k steps it’s -0.04; thus, 5k is undertrained, 10k is approximately at the
optimality threshold of -0.052, and 14k is trained beyond optimality. Similar
for GPT2${}_{\textsc{Base}}$.

Table 1 shows the amount of compute saving of our growth operators for two
different model sizes. The table shows that our grown model reaches the same
performance of the baseline (target model) trained from scratch while having
considerable compute saving ranging from 30% to 5% for different thresholds.
The first row in Table 1 denotes the number of steps we used to train the
baseline model. Given that our growth operator is loss-preserving and
training-dynamic preserving, we can always reach the same loss of the target
model with less compute and the compute saving becomes larger when we decide
to stop the target model earlier. Also, given the same growth ratio (4x
growth), the depth growth operator is preferable versus width concerning the
compute saving.

We also evaluate our pretrained models on other language modeling tasks in the
zero-shot setting to verify that the grown models maintain their transfer
learning capabilities. Table 2 shows results on Wikitext-103 and LAMBADA. It
can be seen that using our practical schedule, we achieve comparable and
sometimes better performance versus using the manual schedule for both in-
domain loss and zero-shot transfer learning. In some cases we have negative
compute saving with the width operator or when combing the width and depth
operator on zero-shot transfer learning tasks early in training (indicating
the grown model used more compute to get to the same (or better) performance),
but the compute saving is always positive when training for longer. We assume
that this is due to instabilities in optimization right after applying the
growth operator. When the training proceeds, the zero-shot performance will
shortly recover and the grown models will have better positive compute saving
at optimality for the two model sizes. It is also less of an issue for the
base model size as growth-target. Moreover, the depth growth operator leads to
better performance compute saving trade-off compared to the width operator
under the same growth ratio. Finally, we show that applying the growth
operator twice (in two stages) lead to the best-performing compute saving and
performance. See Figure 7 for detailed evaluation plots.

###  6.3 Ablations and prior work comparison

We experiment with prior work and conduct ablation studies for our method. We
show that prior works fall short in one or more of the key components of our
proposed method; preserving loss, preserving training dynamics, or following
an optimal schedule.

##### Prior work comparison.

We evaluate against two growth operators from previous work (Gu et al., 2021;
Gong et al., 2019); one uses weight sharing to make a feed-forward network
module wider, and the other makes an entire network deeper. Given the nature
of these growth operators, the grown models do not represent the same function
as the original model; as shown in Figure 3, neither retain the same loss as
the original model, and thus they do not satisfy our loss-preserving property.

Prior work also mostly ignored the optimizer state. In Figure 3, we explore
this as an ablation by simply setting the optimizer state to zero for our
width and depth growth operators. Though the loss can be retained at the
starting pointing for the grown model, the training dynamic becomes extreme
unstable after applying the growth operators, and thus such a growth operator
will not have the training-dynamics-preserving property.

(a)  Width op. and optimizer

(b) Depth op. and optimizer

(c) LR and optimal schedule

(d) LR and optimal schedule

Figure 3:  Comparing with three different baselines from prior work and
ablation studies. $\bullet$ The width growth operator of Gu et al. (2021) in
GPT2${{}_{\text{{base}/4x4width ffn}}}$ and the depth growth operator of Gong
et al. (2019) in GPT2${}_{\text{{base}/2x2depth copy}}$ are not loss
preserving (higher initial loss). Also, GPT2${{}_{\text{{base}/4x4width
ffn}}}$ is significantly underperforming the target model
GPT2${}_{\textsc{base}}$. $\bullet$ Resetting the optimizer state to zero
instead of growing it (the zero_opt runs) have large instabilities and not
preserving of the training dynamics. $\bullet$ 3(c) shows that restarting the
learning rate schedule as in Rae et al. (2021) is not training-dynamics-
preserving. $\bullet$ 3(c), 3(d) also show that not following our optimal
schedule and grow the model too early or too late is still loss-preserving and
training-dynamics-preserving but leads to lower compute saving

##### Ablation studies

We further perform two ablation studies concerning the learning rate schedule
and growth schedule in Figure 3. For the learning rate schedule, we compare
our setting of the learning rate to the growth-target point with restarting
the learning rate schedule as in Rae et al. (2021). It can be seen that
resetting the learning rate schedule leads to much unstable training dynamics
that are not aligned with the growth-target. For the growth schedule, we
experiment with growing the small model earlier or later than the growth-
target point computed with our proposed practical schedule. It shows that
growing the model too late/early is still loss-preserving and training-
dynamics-preserving but the grown models lose the compute saving advantages.

##### Training to completion vs. to optimality.

While training to optimality is not one of our contributions, it is an
important part of our training schedule. Here we compare our models in Table 2
with equivalent models trained to completion. In Radford et al. (2019) ,
GPT2${}_{\textsc{base}}$ trained to completion achieves 37.5 PPL on
Wikitext-103 and 45.9 accuracy on LAMBADA. Our best-performing 2x2xdepth grown
GPT2${}_{\textsc{large}}$ model achieves the same PPL on Wikitext-103 as this
model with only 15% of the compute and the same accuracy as the this model on
LAMBADA with 33% of the compute. Notice that we are comparing different model
sizes, a large model trained to optimality vs. a smaller model trained to
completion.

##  7 Related Work

Perhaps the most similar prior work is Gu et al. (2021). However, this work
did not provide a method to decide when to apply a growth operator and instead
evaluated the performance of their operators at 100/300/500/700K steps of a
small model, or applied a heuristic to equally distribute training steps among
different model sizes. They did not discuss the optimizer state, and they
reset the learning rate to the maximum value of 1e-4 at the beginning of each
state without warmup. Their work built on progressive stacking (discussed
below), and their proposed method to grow the width only grew the feed-forward
layers instead of the entire model width.

Gong et al. (2019) proposed “progressive stacking” which doubles the depth of
a BERT transformer model by copying layers; to construct a 2$L$-layer model
from a $L$ layer model it copies layer $i\leq L$ in the smaller model to layer
$(i+L)$ in the larger model. The optimizer state is reset at the beginning of
each stage, but the learning rate is kept the same as the prior stage. They
use heuristics to set the stacking schedule: 50K steps for 3-layer, 70K steps
for 6-layer model, 280K steps for 12-layer model. In an ablation study they
examined the sensitivity to the number of steps before applying their growth
operator and concluded that there is a threshold number of steps and for
switching times such that switching to the larger model before the threshold
led to compute savings, but switching after the threshold didn’t. This is
consistent with our results that showed the amount of compute saving is
closely related to the stage length.

Li et al. (2020a) proposed growing encoder-decoder transformers in the context
of training a machine translation system. Their depth growth operator is
identical to progressive stacking, although they explore operations that
increase the model by only copying some of the layers from the small to large
model (e.g. growing from 12 to 18 layers). They do not mention optimizer
state, and reset the learning rate to max value at each stage.

In contemporaneous work, Gopher Rae et al. (2021) introduced a method which
tiles the weights from a small model to a larger one. However, their growth
operator does not satisfy our two properties from Section 2, and they focus on
training their models to completion. Evci et al. (2022) proposes a way to
initialize the grown weights by maximizing the gradient norm of the new
weights for vision models. Their growth operator also requires specified
activation functions that can not be directly applied to transformers.
Finally, Li et al. (2021) proposed applying a curriculum learning strategy to
the sequence length to reduce the training cost when training large language
models. Their work is orthogonal to our method and our methods could be
combined; we leave this to future work.

##  8 Conclusion and future work

One direction of future work is to combine batch size warmup Brown et al.
(2020) and sequence length growth Li et al. (2021) with our depth and width
growth operators. Another applies our proposed methods to train a massive
transformer.

We presented a staged training method for large transformer-based language
models that grows the model size during training. We demonstrated the
importance of two properties of the growth operators (loss-preserving and
training-dynamics-preserving), and provided depth and width operators that
satisfy both requirements. Finally, we devised a principled approach to find
the optimal schedule and a simple method to apply the schedule in practice.
Empirical evaluations show up to 22% compute saving.

## References

  * Brown et al. (2020) Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020.  Language models are few-shot learners.  In _Neural Information Processing Systems (NeurIPS)_.
  * Chen et al. (2016) Tianqi Chen, Ian Goodfellow, and Jonathon Shlens. 2016.  Net2net: Accelerating learning via knowledge transfer.  In _International Conference on Learning Representations_.
  * Evci et al. (2022) Utku Evci, Max Vladymyrov, Thomas Unterthiner, Bart van Merriënboer, and Fabian Pedregosa. 2022.  Gradmax: Growing neural networks using gradient information.  _arXiv preprint arXiv:2201.05125_.
  * Gong et al. (2019) Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang, and Tieyan Liu. 2019.  Efficient training of bert by progressively stacking.  In _International Conference on Machine Learning_ , pages 2337–2346. PMLR.
  * Gu et al. (2021) Xiaotao Gu, Liyuan Liu, Hongkun Yu, Jing Li, Chen Chen, and Jiawei Han. 2021.  On the transformer growth for progressive bert training.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 5174–5180.
  * Kaplan et al. (2020) Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020.  Scaling laws for neural language models.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba. 2014.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_.
  * Li et al. (2020a) Bei Li, Ziyang Wang, Hui Liu, Yufan Jiang, Quan Du, Tong Xiao, Huizhen Wang, and Jingbo Zhu. 2020a.  Shallow-to-deep training for neural machine translation.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 995–1005, Online. Association for Computational Linguistics.
  * Li et al. (2021) Conglong Li, Minjia Zhang, and Yuxiong He. 2021.  Curriculum learning: A regularization method for efficient and stable billion-scale gpt model pre-training.  _arXiv preprint arXiv:2108.06084_.
  * Li et al. (2020b) Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gonzalez. 2020b.  Train big, then compress: Rethinking model size for efficient training and inference of transformers.  In _International Conference on Machine Learning_ , pages 5958–5968. PMLR.
  * Liu et al. (2019) Qiang Liu, Lemeng Wu, and Dilin Wang. 2019.  Splitting steepest descent for growing neural architectures.  _Advances in neural information processing systems_.
  * Merity et al. (2017) Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2017.  Pointer sentinel mixture models.  In _International Conference on Learning Representations_.
  * Paperno et al. (2016) Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. 2016.  The lambada dataset: Word prediction requiring a broad discourse context.  In _ACL_.
  * Peters et al. (2018) Matthew E. Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. 2018.  Deep contextualized word representations.  In _Proc. of NAACL_.
  * Press et al. (2020) Ofir Press, Noah A Smith, and Mike Lewis. 2020.  Shortformer: Better language modeling using shorter inputs.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing_.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9.
  * Rae et al. (2021) Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. 2021.  Scaling language models: Methods, analysis & insights from training gopher.  _arXiv preprint arXiv:2112.11446_.
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _Journal of Machine Learning Research (JMLR)_ , 21(140).
  * Rusu et al. (2016) Andrei A Rusu, Neil C Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. 2016.  Progressive neural networks.  _arXiv preprint arXiv:1606.04671_.
  * Schwartz et al. (2020) Roy Schwartz, Jesse Dodge, Noah A. Smith, and Oren Etzioni. 2020.  Green AI.  _Communications of the ACM (CACM)_ , 63(12):54–63.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , pages 6000–6010.
  * Wei et al. (2016) Tao Wei, Changhu Wang, Yong Rui, and Chang Wen Chen. 2016.  Network morphism.  In _International Conference on Machine Learning_ , pages 564–572. PMLR.

##  Appendix A Additional Plots and Results

In Figure 4 and 5, we show loss curve plot and compute plot for applying
growth operator to GPT2${}_{\textsc{base}}$ model. These suggest our growth
operator is loss-preserving, training dynamic preserving and saves compute to
train the target base size model.

In Figure 6, we demonstrate the loss curve and derivatives regarding the
validation loss and compute in log scale. It clearly shows the used practical
threshold gives us a good estimate of the optimality of the loss curve.

In Figure 7, we present the evaluation results for every checkpoints across
the steps for target and grown GPT2${}_{\textsc{base}}$ model. The performance
of the GPT2${}_{\text{{base}/2x2depth}}$ can always perform on par with or
better than the baseline while GPT2${}_{\text{{base}/4x4width}}$ underperforms
baseline in the initial phase after growing but can catch up quickly. This
also explains the potential negative compute we have in the main text in the
early phase of the evaluation and better compute saving at optimality.

(a) Width growth with GPT2${}_{\textsc{large}}$ as target

(b) Depth growth with GPT2${}_{\textsc{large}}$ as target

Figure 4: Similar to Figure 2, our width and depth growth operators are loss-
preserving and training dynamics preserving. $\bullet$
GPT2${}_{\text{{large}/16x16-{width}}}$ indicates starting from a 16x smaller
model then growing it 16x by doubling the width twice. $\bullet$
GPT2${}_{\text{{large}/16x4x4-{width}}}$ indicates growing the model 16x over
two stages, by doubling the width once, continue training the model, then
doubling the width again. $\bullet$ The same applies to the depth growth
operator.

(a) Width growth with GPT2${}_{\textsc{large}}$ as target

(b) Width growth with GPT2${}_{\textsc{base}}$ as target

(c) Depth growth with GPT2${}_{\textsc{large}}$ as target

(d) Depth growth with GPT2${}_{\textsc{base}}$ as target

Figure 5: Growth operator and total compute. Our grown models are saving
compute compared with target model.

(a) First derivative on the validation loss curve for GPT2${}_{\textsc{base}}$
and GPT2${}_{\textsc{large}}$. This is an approximation to the training
dynamics $\frac{\partial\mathcal{L}}{\partial C}$.

(b) Validation loss curve for GPT2${}_{\textsc{base}}$ and
GPT2${}_{\textsc{large}}$

Figure 6:  6(a) shows the first derivative of the validation loss curve for
GPT2${}_{\textsc{base}}$ and GPT2${}_{\textsc{large}}$. The horizontal dotted
lines shows the empirical value of the derivative threshold we use to identify
the point of optimality. 6(b) shows how the threshold value in 6(a) maps to a
point in the loss curve for both model sizes (the first vertical line for
GPT2${}_{\textsc{base}}$ and the second for GPT2${}_{\textsc{large}}$). Both
point show slowed training and the end of the compute-efficient regime.

(a) Wikitext-103 evaluation with GPT2${}_{\textsc{base}}$ as target

(b) LAMBADA evaluation with GPT2base as target.

Figure 7: The detailed evaluation plots on Wikitext-103 and LAMBADA for
GPT2${}_{\textsc{base}}$ trained from scratch and two grown model using depth
or width operator. The performance of the GPT2${}_{\text{{base}/2x2depth}}$
perfectly follows or even outperform the target model, while
GPT2${}_{\text{{base}/4x4width}}$ underperforms the target model initially
after growing but it catches up quickly. This explains the initial “negative”
compute saving in Table 2 for the width growth operator, followed by positive
compute saving as the model trains longer.

##  Appendix B Model Architecture

We use GPT2${}_{\textsc{base}}$ and GPT2${}_{\textsc{large}}$ models. For
GPT2${}_{\textsc{base}}$ model, it consists of 12 layers, 768 hidden
dimensions, 12 heads and 125M parameters. For GPT2${}_{\textsc{large}}$ model,
it consists of 24 layers, 1536 hidden dimensions, 16 heads and 760M
parameters.

##  Appendix C Optimal Stage Schedule

This appendix defines a constrained optimization problem that produces the
optimal staged-training schedule. The scaling laws of Kaplan et al. (2020)
derived empirical fits for the language model loss $L$ as it relates to the
total amount of compute $C$, number of non-embedding parameters $N$, number of
gradient update steps $S$, and batch size $B$. The total compute††This
neglects contributions proportional to the context length, $n_{ctx}$, and may
not be valid in regime of large $n_{ctx}$ where $n_{ctx}\geq 12d_{model}$ is
given by

$C\approx 6NBS,$ (10)

and the loss $L$ for any model size $N$ and number of steps $S$ is given by:

$L(N,S)=\left(\frac{N_{c}}{N}\right)^{\alpha_{N}}+\left(\frac{S_{c}}{S}\right)^{\alpha_{S}}$
(11)

when training at the critical batch size
$B_{crit}=\frac{B_{*}}{L^{1/\alpha_{B}}}$, and where
$\alpha_{N},\alpha_{S},\alpha_{B},N_{c},S_{c},B_{*}$ are all model-specific
constants.

Our goal is to minimize the total compute to train a model of a given size
$N_{target}$ to a given target loss $L_{target}$. We assume we have access to
perfect growth operators that are loss-preserving and training dynamics-
preserving, and that can grow from any model size to any other size††We can
restrict the model size increases by adding additional constraints.. We
consider a training regime consisting of a number of stages. In each stage $k$
we train a model with size $N_{k}$ for $S_{k}$ gradient steps, with the goal
of reaching size $N_{target}$ and achieving a target loss $L_{target}$ at the
end of the final stage. We assume that $N_{k}\geq N_{k-1}$, and that there
exists some way to initialize the model with size $N_{k}$ from one of size
$N_{k-1}$ without changing the loss. For simplicity, we neglect the batch size
contribution to compute, and assume training is always at the critical batch
size $B_{*}/L_{target}^{1/\alpha_{B}}$.

With these assumptions the total compute at the end of training for $M$ stages
is:

$C=\sum_{k=1}^{M}6N_{k}\frac{B_{*}}{L_{target}^{1/\alpha_{B}}}S_{k}$ (12)

We can compute the loss at the end of each stage in an iterative fashion. The
loss at the end of the first stage is given by $L_{1}(N_{1},S_{1})$ from Eqn.
11. Then for each subsequent stage, we assume the loss curves can be
translated and the loss at the end of stage $k$ is computed by starting with
the loss at the end of the prior stage and decreased for $S_{k}$ steps. To do
so, first compute the effective number of steps $S_{eff,k}$ needed to reach
initial loss for the stage $L_{k-1}$ with model size $N_{k}$, and then compute
the loss at the end of the stage by $L_{k}(N_{k},S_{eff,k}+S_{k})$. In
summary:

$\displaystyle L_{1}$ $\displaystyle=$
$\displaystyle\left(\frac{N_{c}}{N_{1}}\right)^{\alpha_{N}}+\left(\frac{S_{c}}{S_{1}}\right)^{\alpha_{S}}$
$\displaystyle L_{k}$ $\displaystyle=$
$\displaystyle\left(\frac{N_{c}}{N_{k}}\right)^{\alpha_{N}}+\left(\frac{S_{c}}{S_{eff,k}+S_{k}}\right)^{\alpha_{S}},k>1$
$\displaystyle S_{eff,k}$ $\displaystyle=$
$\displaystyle\frac{S_{c}}{(L_{k-1}-(N_{c}/N_{k})^{\alpha_{N}})^{1/\alpha_{S}}}.$

With this in hand, we can use a constrained optimizer to solve for the optimal
schedule by minimizing Eqn. 12, subject to the constraint that the final model
size is the target size, and loss at the end of training is the target loss.
Formally,

$\min_{\\{(N_{k},S_{k})\\}}\sum_{k=1}^{M}6N_{k}\frac{B_{*}}{L_{target}^{1/\alpha_{B}}}S_{k}$

subject to

$\displaystyle L_{M}$ $\displaystyle=$ $\displaystyle L_{target}$
$\displaystyle N_{M}$ $\displaystyle=$ $\displaystyle N_{target}$
$\displaystyle 0$ $\displaystyle<$ $\displaystyle N_{1}$ $\displaystyle
N_{k-1}$ $\displaystyle\leq$ $\displaystyle N_{k},k>1$ $\displaystyle 0$
$\displaystyle\leq$ $\displaystyle S_{k}$

Note that in the single stage case ($M=1$) with no $N_{target}$ condition,
this formulation reduces to the optimal calculation in Appendix C of Kaplan et
al. (2020) to find the optimal model size to reach a target loss. This matches
our training to optimality.

##### Measuring compute saving

To measure the compute saving from staged training to reach a certain
$L_{target}$, we use a single staged training ($M=1$) with no $N_{target}$
condition, and use our optimization algorithm to find the optimal compute and
model size; this becomes $N_{target}$. Next, we run the optimization problem
with $N_{target}$, $L_{target}$, and $M>1$ to get the optimal training
schedule and the expected compute, which we compare with $M=1$ to find the
expected compute saving.

Number of stages | Compute reduction factor
---|---
1 | 1.0
2 | 0.83
3 | 0.792
4 | 0.779
5 | 0.771
10 | 0.763
Table 3: Decrease in compute costs over an optimal single stage training regime. Stage | $N_{k}$ | $S_{k}$ | $L_{k}$ | $C_{k}$
---|---|---|---|---
1 | 2.7M | 15.4K | 4.09 | 0.0033
2 | 22.1M | 24.3K | 3.58 | 0.0154
3 | 71.9M | 30.8K | 3.31 | 0.0365
4 | 163M | 36.0K | 3.13 | 0.0665
5 | 306M | 40.5K | 3.00 | 0.1056
Table 4: Sample optimal schedule for a five stage regime to train to
$L_{target}=3$ showing the number of non-embedding parameters $N_{k}$, the
number of gradient steps $S_{k}$, the loss at the end of the stage $L_{k}$,
and the compute in the stage $C_{k}$.

##### Observations

An implementation of this optimization shows that the increase in compute
efficiency for using multiple stages is independent of the target loss
$L_{target}$, and quickly approaches $~{}0.76$ as the number of stages
increases (using the scaling parameters for autoregressive transformer
language models).

We also found that constraining the ratio between consecutive model sizes
($\frac{N_{k}}{N_{k-1}}$) to be 2, 4 or 8, leads to almost the same compute
savings. These constraints come from the practical constraints of our
implementation of the growth operators.

##  Appendix D Practical Stage Schedule

The empirical values we estimated for the constants are:

$\displaystyle\tau_{opt}$ $\displaystyle=-0.052$ $\displaystyle\tau_{depth}$
$\displaystyle=-0.0575$ $\displaystyle\tau_{width}$ $\displaystyle=-0.0475$
$\displaystyle\tau_{depth-width}$ $\displaystyle=-0.03$
$\displaystyle\rho_{depth}$ $\displaystyle=0.70$ $\displaystyle\rho_{width}$
$\displaystyle=0.55$ $\displaystyle\rho_{depth-width}$ $\displaystyle=0.40$


# Molecular Dynamics Simulations of Anisotropic Particles Accelerated by
Neural-Net Predicted Interactions

B. Ruşen Argun  Mechanical Engineering, Grainger College of Engineering,
University of Illinois,Urbana-Champaign, 61801, IL, United States  Yu Fu
Physics, Grainger College of Engineering, University of Illinois,Urbana-
Champaign, 61801, IL, United States  Antonia Statt<EMAIL_ADDRESS>Materials Science and Engineering, Grainger College of Engineering, University
of Illinois, Urbana-Champaign, 61801, IL, United States

###### Abstract

Rigid bodies, made of smaller composite beads, are commonly used to simulate
anisotropic particles with molecular dynamics or Monte Carlo methods. To
accurately represent the particle shape and to obtain smooth and realistic
effective pair interactions between two rigid bodies, each body may need to
contain hundreds of spherical beads. Given an interacting pair of particles,
traditional molecular dynamics methods calculate all the inter-body distances
between the beads of the rigid bodies within a certain distance. For a system
containing many anisotropic particles, these distance calculations are
computationally costly and limit the attainable system size and simulation
time. However, the effective interaction between two rigid particles should
only depend on the distance between their center of masses and their relative
orientation. Therefore, a function capable of directly mapping the center of
mass distance and orientation to the interaction energy between the two rigid
bodies, would completely bypass inter-bead distance calculations. It is
challenging to derive such a general function analytically for almost any non-
spherical rigid body. In this study, we have trained neural nets, powerful
tools to fit nonlinear functions to complex datasets, to achieve this task.
The pair configuration (center of mass distance and relative orientation) is
taken as input and the energy, forces and torques between two rigid particles
are predicted directly. We show that molecular dynamics simulations of cubes
and cylinders performed with forces and torques obtained from the gradients of
the energy neural-nets quantitatively match traditional simulations that uses
composite rigid bodies. Both structural quantities and dynamic measures are in
agreement, while achieving up to 23 times speed up over traditional molecular
dynamics, depending on hardware and system size. The method presented here
can, in principle, be applied to any irregular shape with any pair
interaction, provided that sufficient training data can be obtained.

##  I Introduction

Advances in chemistry have made it possible to synthesize many non-spherical
colloidal particles Sacanna  _et al._ (2013). Those anisotropic colloidal
building blocks can self-assemble into various targeted nano-structures with
desirable structural, optical or electronic properties Cai  _et al._ (2021);
Henzie  _et al._ (2012). In the environment, pollutants such as microplastics
or engineered nano-particles can occur in different shapes Argun and Statt
(2023); Wang  _et al._ (2021). In addition, particles with non-spherical
morphologies are central to several biological processes Roth  _et al._
(1996). Thus, capturing anisotropic features of particles is essential to
accurately simulate and understand the fundamentals of wide range of phenomena
in science and engineering.

Traditional particle based simulations such as molecular dynamics (MD) and
Monte Carlo (MC) simulations often treat particles as perfect spheres,
however, there are several methods available that can take non-spherical
geometric features of particles into account. For athermal systems, where
particles interact with hard repulsion only, calculating the pairwise
potential reduces to the task of finding whether the two particles or shapes
overlap or not. In such systems, where the self-assembly is driven by entropic
effects, powerful methods to calculate the interactions for many geometries
exist. Donev (2006) introduced an event-driven MD algorithm that can
efficiently calculate the overlap between non-spherical geometries. This
method has been used to find closed-packing structures and phase behaviour of
hard superballs Ni  _et al._ (2012); Donev  _et al._ (2004). Damasceno  _et
al._ (2012) and othersAnderson  _et al._ (2016); Geng  _et al._ (2019); van
Anders  _et al._ (2015) have simulated a wide range of polyhedra with hard-
particle MC and investigated their self-assembly and phase behaviour using
different methods to calculate the overlap or contact between the shapes
Gilbert  _et al._ (1988). However, for systems where the self-assembly is
driven by enthalpic forces (i.e. attractive interactions), different methods
are required. The Gay-Berne Berardi  _et al._ (1998) potential is commonly
used to simulate ellipsoids with pairwise attractive interactions, depending
on the orientations of the particles in addition to the distance between their
center of masses. It is convenient for MD simulations, as forces and torques
can be calculated analytically, however, it is limited to ellipsoid shapes
only. The anisotropic Lennard-Jones potential Ramasubramani  _et al._ (2020)
extends this to polyhedra.

By using a composite bead approach, any particle shape with attractive and
hard repulsive interactions can be simulated. Due to its flexibility in both
shape and interaction, this method is commonly used. Here, smaller, spherical
beads, often called composite beads, are held together with stiff harmonic
bonds and angle potentials to keep the geometry of the larger composite shape
the same throughout the simulation. Alternatively, rigid body constraints
between the composite beads can be applied. Popular MD software such as HOOMD-
blue Anderson  _et al._ (2020); Nguyen  _et al._ (2011); Glaser  _et al._
(2020) and LAMMPS Plimpton (1995) have implemented rigid body constraints to
simulate non-spherical particles or molecules made out of smaller beads Nguyen
_et al._ (2011); Nguyen and Plimpton (2019). The composite bead approach is
flexible, as almost any concave or convex shape can be simulated with it.
Self-assembly behaviour of nanoparticles can be affected by small
morphological details Kim  _et al._ (2017). To obtain a precise geometry
including sharp corners and edges, the shape often needs to contain many
beads. The interaction between two bodies or shapes is now given by the double
sum of all pairwise interactions between the composite beads that make up the
two bodies. Thus, rigid bodies containing many composite beads significantly
increase the computational cost in both coarse-grained and atomistic
simulations. A similar computational cost challenge arises for example to
calculate the total van der Waals (vdW) interaction between two non-spherical
particles. Most notably, Lee and Arya (2020) devised an analytical method to
approximate the vdW interactions between two cubic rigid bodies bypassing the
cumbersome double sum. Their method allows for MC simulations of cubes, but it
is limited in its application to MD since forces and torques cannot be
obtained with it. There are several similar efforts in the literature to
calculate the vdW interaction for non-spherical bodies Priye and Marlow
(2013); Jaiswal and Beaudoin (2012); Hallez (2012); Yang and Qian (2009);
Maeda and Maeda (2015). Recent work Lee  _et al._ (2023) has calculated vdW
forces and torques for cubes with a numerical method. Ramasubramani  _et al._
(2020) introduced a Lennard Jones-like potential that can be used for
different polyhedra.

Machine Learning (ML) techniques are already applied to spherically symmetric
particles, especially to bridge quantum mechanical calculations and classical
atomistic MD simulations Deringer  _et al._ (2019). Some possibilities that ML
offers for anisotropic interactions are recently discovered. Campos-Villalobos
_et al._ (2022) has developed symmetry functions that can be used to construct
ML potentials for cylindrically symmetric particles. Anisotropic and rigid
molecules like benzene were coarse-grained by matching forces and torques to
accelerate all-atom condensed-phase MD simulations using a linear model Nguyen
and Huang (2022) and neural-networks Wilson and Huang (2023). Recent work Bag
_et al._ (2023) is using several ML models to detect whether or not two
composite rigid bodies are overlapping, bypassing explicit distance
calculations between primary beads of the rigid bodies. However, this method
does not allow attractive interactions.

In this work, we are using fully-connected feed-forward neural nets to predict
the total interaction between two example geometries of composite rigid
bodies: cubes and cylinders. We focus on forces and torques, since these are
the necessary inputs for a MD simulation. Here, we describe our generic method
of data sampling and processing and test the neural-net predictions for simple
example particle shapes. We show that the MD simulations performed with forces
and torques from the trained neural nets accurately reproduce structural and
dynamic properties of traditional rigid body simulations which were performed
with explicit distance calculations. Depending on the system simulated (number
of particles, temperature etc.) and hardware available, neural-net assisted
simulations are faster than traditional MD simulations. Our method is entirely
agnostic to the particle shape and can be, in principle, applied to any
irregular composite particle geometry with a wide range of both attractive and
repulsive pair interactions.

##  II Methods and Motivation

To calculate the total potential $U$ between a pair of composite-rigid bodies
that are made out of smaller beads (as in Fig. 1) that interact with pairwise
potentials $u_{bead}(r)$, the distances between each pair of beads of need to
be calculated.

$U=\sum_{i=1}^{N}\sum_{j=1}^{N}u_{bead}(r_{ij})\quad,$

Figure 1: A pair of cylindrical rigid-bodies made out of smaller beads. The
cylinder is approximately $2.75\sigma$ in diameter and $5.3\sigma$ long.

For two rigid bodies, each consisting of $N$ beads, calculating $N^{2}$
distances $r_{ij}$, can be computationally costly for high values of $N$. The
number of beads $N$ used to represent a rigid-body depends on several factors,
including the geometry of the particle and how accurate one would like to
reproduce that geometry. For example, to perform a systematic study of
particle shapes that range from a perfect cube to a sphere with several super-
ball shapes in between, edges and corners of the cube must be smoothed out
gradually as the shape approaches to a sphere, requiring a high resolution of
beads. Simulation methods that are used to incorporate a background fluid and
hydrodynamic interactions like multi-particle collision dynamics Malevanets
and Kapral (1999) may also require a high resolution of shape surface for
accurate coupling between the colloid and solvent particles Poblete  _et al._
(2014); Wani  _et al._ (2022). Another factor is the type of interaction
between the shapes. For hard-core repulsion interactions or WCA-like Weeks
_et al._ (1971) interactions, where there is no attraction, one can place
beads at the surface of the shape only and leave the core empty, thereby
reducing $N$ and saving computational time. However, this approach may not
work for attractive interactions. By only having surface beads, or shapes with
a small number of beads, it may not be possible to reproduce the desired
effective interaction between the two shapes in such cases. $N$ may need to be
on the order of $10^{2}$ to accurately represent the shape and reproduce the
desired effective interaction. We refer to the SI for a detailed illustration
of this issue.

In the following, we outline an alternative route to reduce computational
costs associated with the composite bead method by bypassing explicit distance
calculations, while keeping the full flexibility of the different shapes that
can be used. Any pair configuration of $3D$ bodies can be expressed as a
function of six numbers. Position and orientation of the first body 1 define a
coordinate frame. The position of the second body 2
$\mathbf{x}(x_{1},x_{2},x_{3})$ and its orientation
$\bm{\Omega}(\Omega_{1},\Omega_{2},\Omega_{3})$ in the relative coordinate
frame of the first body 1 are then sufficient to uniquely define any pair
configuration. Thus, the total interaction between the two composite-rigid
bodies is essentially a function of six numbers,

$U=g(\mathbf{x},\bm{\Omega})=g(x_{1},x_{2},x_{3},\Omega_{1},\Omega_{2},\Omega_{3})\quad.$

However, this function is quite complex and its general analytical solution is
not available even for simple shapes like cubes or non-aligned cylinders Lee
and Arya (2020). We follow a data-driven approach and use neural-nets
($g_{NN}$) to approximate the function $g$ relating the pair configuration
directly to the interaction energy, force or torque,

$U=\sum_{i=1}^{N}\sum_{j=1}^{N}u_{bead}(r_{ij})=g(\mathbf{x},\bm{\Omega})\approx
g_{NN_{U}}(\mathbf{x},\bm{\Omega})\quad,$

thus allowing us to skip the costly distance calculations at the expense of
the inference and training cost of a neural-net.

Training and testing data for the neural net was sampled from traditional
composite-rigid body MD simulations that were performed at a wide range of
temperatures, starting from $T=0.2\epsilon/k_{B}$ up to $2.0\epsilon/k_{B}$.
Both energetically favorable (i.e., pairs that are strongly attracting, for
example two cubes with their faces aligned) and unfavorable (i.e., pairs that
are slightly overlapping or touching, physically less likely to occur) pair
configurations can be efficiently sampled at low and high temperatures,
respectively. Sampled body pairs and their calculated energy, force and torque
vectors are then subjected to geometric transformations and operations in
order to express each pair as a function of $\bm{\Omega}$ and $\mathbf{x}$ and
to make use of the geometrical symmetries of the shapes. This process is
explained in detail in the SI Section 1 and the code is publicly available at
https://github.com/stattlab/NNAssistedRigidBodyMD. Here, we only mention that
unlike other $3D$ bodies, a pair of cylinders can be expressed as a function
of four numbers only
$g(\mathbf{x}(x_{1},x_{2}),\bm{\Omega}(\Omega_{1},\Omega_{2}))$, thanks to
their cylindrical symmetry.

Training a single neural-net that can predict the energy would be sufficient
to perform Monte Carlo simulations. For MD simulations, force
$\mathbf{f}=(f_{1},f_{2},f_{3})$ and torque
$\bm{\tau}=(\tau_{1},\tau_{2},\tau_{3})$ vectors are required to integrate the
equations of motion forward in time. We have used two different methods to
obtain $\mathbf{f}$ and $\bm{\tau}$. First, for cylinders, five separate
neural nets were trained for each component of
$\mathbf{f}=(f_{1},f_{2},f_{3})$ and $(\tau_{1},\tau_{2})$ and used directly
in the simulation. The torque around the axis of the cylinder, $\tau_{3}$ is
not needed. This leads to the following approximations

$f_{i}\approx g_{NN_{f_{i}}},\quad\tau_{i}\approx
g_{NN_{\tau_{i}}}\quad\text{with}\quad i={1,2,3}\quad.$

Second, for cubes of side length approx. $3.2\sigma$, the forces and torques
were calculated from the gradients of the energy neural-net as explained in SI
Section 6. The second method can provide the same accuracy as the first one
and it is more practical to implement. In addition, using one neural net for
energy also helps to ensure that the components of the forces and torques obey
energy conversation. Simple fully connected feed-forward neural nets are used
for these regression tasks.

##  III Results and Discussion

Generally, accurate force-torque prediction from our model should result in
correct molecular dynamics simulations. Before comparing the traditional and
neural-net assisted simulation results, we are presenting a more direct
measure of ML accuracy, namely parity plots obtained from the test data. In
Fig. 2, the predicted values of the $x$-component of the net force ($f_{x}$)
vs. the true value between two cylinders and cubes are plotted. The other
force components follow the same trends. For both shapes, data points outside
the 10% error lines were frequent, especially for values close to zero. We
note that using separate neural networks to predict $f_{i}$ for cubes, instead
of taking the gradient of the $g_{NN_{U}}$ results in a much lower rate of
error, as we show in SI Section 5. However, since the accuracy of the MD
simulation results are similar for both methods, the results obtained from
cube simulations assisted by $g_{NN_{f_{i}}}$ and $g_{NN_{\tau_{i}}}$ are not
presented here.

Figure 2: Parity plot of true vs. predicted $f_{x}$ values for the cylinder
(left) and cube (right) system. For cylinder $f_{x}$ is the direct output of
the force neural net $g_{NN_{f_{i}}}$, for cube it is obtained from the
gradient of the energy neural net $g_{NN_{U}}$. Dashed lines indicate the 10%
error line.

In Fig. 3, traditionally calculated and neural-net predicted energies of cubes
are compared. Here, we show the energies of three different relative
orientations as a function of their center of mass distances. Overall, there
is good agreement over a wide range of center of mass distances, with
noticeable deviations for small distances and high interaction energies. These
short distances are less likely to occur in a molecular dynamics simulation,
but might play a significant role for simulations at highly dense conditions
or higher temperatures above the temperatures used to generate the training
data set.

Figure 3: Interaction energy $U(r)$ as function of center of mass distance $r$
for three different cube configurations, as shown next to the curves. Each
cube had an approximate side length of $3.2\sigma$.

The next step is to compare results of neural-net assisted simulations with
traditional molecular dynamics simulations. For this, we performed two sets of
simulations, one using HOOMD-blue 2.9.4 with traditional molecular dynamics
simulations of rigid bodies in the NVT ensemble, and the second set with C++
code using neural-net forces and torques for the molecular dynamics
simulation. We are using half-power half-cosine as a base function,

$\mathrm{U}(r)=\begin{cases}-A\cos(c(r-b))&r>b\\\ (b-r)^{p}-2A&r\leq
b\end{cases}\quad,$ (1)

for the bead-to-bead potential as detailed in the SI Section 4 with the
following parameters, $A=0.00035\epsilon$, $b=1.46375\sigma$,
$c=3.3833\sigma^{-1}$, $p=2.5$ for both cylinders and cubes. When summed up
over the beads of the two interacting cubes, we effectively obtain a Lennard
Jones-like interaction between the two cubes as illustrated in Fig. 3.

The temperature was kept constant using a Nosé–Hoover thermostat with a
timestep of 0.005 $\tau$. We equilibrated the system for $2500\tau$, where
$\tau=\sigma\sqrt{m/\epsilon}$ is the unit of time. The particle volume
fraction $\phi=N*V_{\mathrm{particle}}/V$ of all simulations was set to
$\phi=0.16$. We used the size of the smaller composite beads $\sigma$, as unit
of length. Cubes have an approximate side length of $3.2\sigma$ and the
cylinders are approx. $2.75\sigma$ in diameter and $5.3\sigma$ long.

To assess the accuracy of our simulations a measure for the structure of the
system was used, the pair correlation function $g(r)$ of the center of mass of
the simulated shapes. In Fig. 4 $g(r)$ of cylinders was plotted for both nn-
assisted and traditional MD simulations. While qualitative agreement between
the two curves can be observed, and nn-assisted simulations reproduced the
peaks at correct positions, the height of the first peak is underestimated.
The first peak in $g(r)$ corresponds to the pair configuration where two
cylinders are aligned along their main axis and are close to each other. We
attribute the mismatch at the first peak to the assumption in nn-assisted
simulations that the cylinders are symmetric around their main axis. Even
though the cylinders were comprised of many composite beads for the
traditional MD simulations, the surface was still rough and rotations around
the main axis might at close distances result in different interactions. This
roughness is effectively coarse-grained out in the neural-net assisted
simulations, since the neural net cannot account for rotations around the main
axis by design of the input variables. This effect is the strongest when two
cylinders are closest to each other, almost touching, and decays as the
distance increased, as shown in Fig. SI-S6.

Figure 4: Pair correlation function $g(r)$ at $kT=1.0$ for a cylinder system.
The traditional simulation result is shown as dashed line, and the nn-assisted
simulation result is shown as filled symbols.

Next, we tested the structural accuracy of the cube simulations. In Fig. 5,
the $g(r)$ curves for the center of mass of cubes are plotted for simulations
performed at temperatures of $T=0.5\,\epsilon/k_{B}$, $0.6\,\epsilon/k_{B}$,
and $0.7\,\epsilon/k_{B}$. There is quantitative match at all temperatures
simulated. The nn-assisted simulations accurately capture the significant
structural changes observed in this temperature window, ranging from a non-
aggregated suspension of cubes at high temperatures $T=0.7\,\epsilon/k_{B}$,
down to a aggregated face-centered cubic structure at low temperatures
$T=0.5\,\epsilon/k_{B}$. Unlike in the case of cylinders, where we ignore
rotations around the main axis, rotations in the cube configurations are
represented explicitly and we speculate that this is part of the reason for
the improved match of the first peak in $g(r)$.

Figure 5: Pair correlation functions $g(r)$ for three different temperatures:
$0.5\,\epsilon/k_{B}$ (top), $0.6\,\epsilon/k_{B}$ (middle) and
$0.7\,\epsilon/k_{B}$ (bottom) for cubes. Dashed lines show the result from
traditional simulations, filled symbols show the result of the nn-assisted
simulations. The insets are representative snapshots of the corresponding
structures at these temperatures.

We also assessed the accuracy of dynamics in the neural-net assisted
simulations. For consistency, the mass of the cube rigid-body is set to be the
same in HOOMD and neural-net assisted simulations. In Fig. 6, mean square
displacements (MSD) of cubes at different temperatures for both neural-net
assisted and traditional simulations are shown. At all temperatures, good
agreement can be found. We did not determine long-term diffusion constants.
Due to the matching in relative masses, we expect that the long term center of
mass motion of the cubes will be similar in both simulation techniques. In
addition to MSD, mean square rotation (MSR) is also a fundamental measure of
the kinetics for non-spherical particles. MSR is the orientational-angular
analogue of the MSD, and its calculation is explained in detail in the SI
Section 5. In Fig. 6 the MSRs of the cubes at different temperatures are
plotted. We observe a reasonable match, except at lower temperatures (0.5 and
0.3). Cubes in the nn-assisted simulations are rotating slightly faster
compared to traditional MD simulations. It is not clear why we observe this
behaviour, and future investigations will be needed.

Figure 6: Mean squared displacement (top) and mean squared rotation (bottom)
of cubes at different temperatures. The solid lines represent the traditional
simulations, whereas the dashed lines show the nn-assisted results.

Finally, the speed, or performance, of the neural-net assisted simulations
were compared with HOOMD-blue simulation performance. Depending on the
hardware used and on the number $N$ of cubes simulated, nn-assisted
simulations can significantly increase the performance of simulations, as
shown in Table 1. We performed all simulations on a single RTX 2070 Super GPU
with 8GB memory and a single A100 GPU with 80GB memory and report the
timesteps per second. An acceleration factor between the two simulation
methods on each GPU was computed. Aside from the smallest system ($N=500$) on
the A100, all nn-assisted simulations are faster than their corresponding
HOOMD-blue simulations. In addition, on the RTX 2070, HOOMD rigid-body
simulations run out of the 8GB memory for systems larger than $N\approx 1300$.

Table 1: Timesteps per second and acceleration factor of nn-assisted simulations compared to HOOMD-blue for various system sizes $N$. A dash (—) indicates that GPU memory was not sufficient to run the simulation. GPU - Method N | 500 | 1000 | 1500 | 2000 | 3000 | 4000
---|---|---|---|---|---|---
RTX-2070 - HOOMD | 9.5 | 4.6 | — | — | — | —
RTX-2070 - NN-Assisted | 143 | 109 | 90 | 75 | 64 | 45
RTX-2070 - Speed up | x15 | x23 | — | — | — | —
A100 - HOOMD | 190 | 90 | 60 | 45 | 30 | 23
A100 - NN-Assisted | 153 | 125 | 108 | 91 | 78 | 60
A100 - Speed up | x0.8 | x1.4 | x1.8 | x2.0 | x2.6 | x2.6

##  IV Conclusions

In this work, we present a flexible neural-net assisted method to accelerate
molecular dynamics simulations of non-spherical composite rigid bodies of any
irregular shape. Here, we have used cubes and cylinders as examples, however,
the method has no restrictions on the shape of the particles as long as
sufficient training data can be obtained. Our method also allows any pair
interactions, including repulsive and attractive forces. We observed excellent
agreement in structural and dynamic properties between the nn-assisted and
traditional rigid-body composite simulations. We tested the performance of the
current method and observed a speedup of x15 to x23 on a single RTX 2070 Super
Nvidia GPU, and a speedup of x0.8 to x2.6 on a A100 Nvidia GPU. Depending on
the system of interest and available hardware, this speedup can be
significant.

Our method has several shortcomings that prevent further speed-up, which can
be addressed in future work. Ideally, the descriptors (for our model these are
the input parameters of the neural-net) should be easy to calculate from the
raw representation of the system, i.e., positions and quaternions Campos-
Villalobos  _et al._ (2022). However, the geometric operations that we perform
to consider the symmetries of a pair of cubes or cylinders are computationally
quite costly, reducing the performance. A second disadvantage is the use of a
neural-net, where inference is associated with significant computational cost,
especially considering inference is needed at every single timestep.

We note that cylinders have an advantage where any pair configuration can be
expressed with four numbers instead of six due to the cylindrical symmetry.
This symmetry greatly improved the performance of the neural-nets, or
equivalently, reduces the difficulty of fitting. However, this symmetry
assumption led to a slight loss in accuracy, particularly of parallel cylinder
configurations at short distances. Depending on the desired density and level
of surface roughness (i.e. how coarse-grained the cylinder is) these
inaccurate interaction predictions might be significant.

Our method, as presented here, is completely flexible in the input training
data. Consequently, it can be easily extended to small nanoparticle systems
that include ligands, charges, atomistic interactions, and other effects.
Generally, as particle sizes get smaller, the error in effective interactions
between particles increases, due to classical approximations and coarse-
graining Silvera Batista  _et al._ (2015). We speculate that our approach
could potentially attain the level of accuracy comparable to atomistic
simulations by training the neural nets using nanoparticles made out of atoms.
Similarly, for small nanoparticles, surface ligands play a significant role in
self-assembly, and neural-nets can be trained with effective forces and
torques between nanoparticles with explicit ligandsGiunta  _et al._ (2023);
Campos-Villalobos  _et al._ (2022). Self-assembly pathways of charged colloids
are also highly dependent on morphology Rosenberg  _et al._ (2020). It would
be possible to include coulombic interactions by adding charged beads to the
rigid bodies, which would be learned by the neural nets as a part of their
effective interactions. Depletion interactions of nanoparticles depend
strongly on their shape, and this method could provide an efficient way of
developing effective depletion pair interactions by using traditional
simulations with explicit depletants.

In this paper, the effective interaction we used was relatively short-ranged.
It decayed to zero around $1.4a$, where $a$ is the size of the shape, which is
shorter than a typical Lennard-Jones interaction that would decay to zero
around $2.5$ to $3a$. Especially for nanoparticles, relatively long-ranged
interactions can be relevant in self-assembly Kim  _et al._ (2017). Longer-
range shape interactions would significantly increase the potential speedup of
the nn-assisted method over traditional simulations. Unfortunately, more data
sampling and possibly longer training times would be required. Similarly, more
irregular shapes may require more composite beads for an accurate
representation, which would lead to a larger potential benefit for the nn-
assisted method. However, the regression task would be more complicated as
well. Future research and methods development will push the limits of the
method shown here, and we hope to include different physical interactions,
longer-ranged interactions, and more complicated geometries.

## Acknowledgements

This work was supported by funding from the Molecule Maker Lab Institute
(MMLI): An AI Research Institutes program supported by NSF under award No.
2019897. This work made use of the Illinois Campus Cluster, a computing
resource that was operated by the Illinois Campus Cluster Program (ICCP) in
conjunction with the National Center for Supercomputing Applications (NCSA)
and which was supported by funds from the University of Illinois at Urbana-
Champaign.

## Supplemental Information

The supplementary information contains the details of the geometric operations
performed to take the symmetries of pairs of cubes and cylinders into account.
A discussion of the cylindrical symmetry assumption for cylinders is included.
An example to demonstrate why many beads can be necessary for a rigid body is
also included. We show how one can obtain torques and forces from the energy
neural net. The code for running the nn-assisted simulations is available at
https://github.com/stattlab/NNAssistedRigidBodyMD.

## References

  * Sacanna  _et al._ (2013) S. Sacanna, M. Korpics, K. Rodriguez, L. Colón-Meléndez, S.-H. Kim, D. J. Pine, and G.-R. Yi, Nature communications 4, 1688 (2013).
  * Cai  _et al._ (2021) Z. Cai, Z. Li, S. Ravaine, M. He, Y. Song, Y. Yin, H. Zheng, J. Teng, and A. Zhang, Chemical Society Reviews 50, 5898 (2021).
  * Henzie  _et al._ (2012) J. Henzie, M. Grünwald, A. Widmer-Cooper, P. L. Geissler, and P. Yang, Nature materials 11, 131 (2012).
  * Argun and Statt (2023) B. R. Argun and A. Statt, Soft Matter 19, 8081 (2023).
  * Wang  _et al._ (2021) X. Wang, N. Bolan, D. C. Tsang, B. Sarkar, L. Bradney, and Y. Li, Journal of Hazardous Materials 402, 123496 (2021).
  * Roth  _et al._ (1996) C. M. Roth, B. L. Neal, and A. M. Lenhoff, Biophysical journal 70, 977 (1996).
  * Donev (2006) A. Donev,  _Jammed packings of hard particles_ (Princeton University, 2006).
  * Ni  _et al._ (2012) R. Ni, A. P. Gantapara, J. De Graaf, R. Van Roij, and M. Dijkstra, Soft Matter 8, 8826 (2012).
  * Donev  _et al._ (2004) A. Donev, I. Cisse, D. Sachs, E. A. Variano, F. H. Stillinger, R. Connelly, S. Torquato, and P. M. Chaikin, Science 303, 990 (2004).
  * Damasceno  _et al._ (2012) P. F. Damasceno, M. Engel, and S. C. Glotzer, Science 337, 453 (2012).
  * Anderson  _et al._ (2016) J. A. Anderson, M. Eric Irrgang, and S. C. Glotzer, Computer Physics Communications 204, 21 (2016).
  * Geng  _et al._ (2019) Y. Geng, G. van Anders, P. M. Dodd, J. Dshemuchadse, and S. C. Glotzer, Science advances 5, eaaw0514 (2019).
  * van Anders  _et al._ (2015) G. van Anders, D. Klotsa, A. S. Karas, P. M. Dodd, and S. C. Glotzer, Acs Nano 9, 9542 (2015).
  * Gilbert  _et al._ (1988) E. G. Gilbert, D. W. Johnson, and S. S. Keerthi, IEEE Journal on Robotics and Automation 4, 193 (1988).
  * Berardi  _et al._ (1998) R. Berardi, C. Fava, and C. Zannoni, Chemical physics letters 297, 8 (1998).
  * Ramasubramani  _et al._ (2020) V. Ramasubramani, T. Vo, J. A. Anderson, and S. C. Glotzer, The Journal of Chemical Physics 153 (2020).
  * Anderson  _et al._ (2020) J. A. Anderson, J. Glaser, and S. C. Glotzer, Computational Materials Science 173, 109363 (2020).
  * Nguyen  _et al._ (2011) T. D. Nguyen, C. L. Phillips, J. A. Anderson, and S. C. Glotzer, Computer Physics Communications 182, 2307 (2011).
  * Glaser  _et al._ (2020) J. Glaser, X. Zha, J. A. Anderson, S. C. Glotzer, and A. Travesset, Computational Materials Science 173, 109430 (2020).
  * Plimpton (1995) S. Plimpton, Journal of computational physics 117, 1 (1995).
  * Nguyen and Plimpton (2019) T. D. Nguyen and S. J. Plimpton, Computer Physics Communications 243, 12 (2019).
  * Kim  _et al._ (2017) J. Kim, Z. Ou, M. R. Jones, X. Song, and Q. Chen, Nature communications 8, 761 (2017).
  * Lee and Arya (2020) B. H.-j. Lee and G. Arya, Nanoscale Horizons 5, 1628 (2020).
  * Priye and Marlow (2013) A. Priye and W. H. Marlow, Journal of Physics D: Applied Physics 46, 425306 (2013).
  * Jaiswal and Beaudoin (2012) R. P. Jaiswal and S. P. Beaudoin, Langmuir 28, 8359 (2012).
  * Hallez (2012) Y. Hallez, Colloids and Surfaces A: Physicochemical and Engineering Aspects 414, 466 (2012).
  * Yang and Qian (2009) P. Yang and X. Qian, Journal of colloid and interface science 337, 594 (2009).
  * Maeda and Maeda (2015) H. Maeda and Y. Maeda, Langmuir 31, 7251 (2015).
  * Lee  _et al._ (2023) J. Lee, E. Nakouzi, J. Heo, B. A. Legg, G. K. Schenter, D. Li, C. Park, H. Ma, and J. Chun, Journal of Colloid and Interface Science 652, 1974 (2023).
  * Deringer  _et al._ (2019) V. L. Deringer, M. A. Caro, and G. Csányi, Advanced Materials 31, 1902765 (2019).
  * Campos-Villalobos  _et al._ (2022) G. Campos-Villalobos, G. Giunta, S. Marín-Aguilar, and M. Dijkstra, The Journal of Chemical Physics 157 (2022).
  * Nguyen and Huang (2022) H. T. Nguyen and D. M. Huang, The Journal of Chemical Physics 156 (2022).
  * Wilson and Huang (2023) M. O. Wilson and D. M. Huang, arXiv preprint arXiv:2301.10881 (2023).
  * Bag  _et al._ (2023) S. Bag, A. Jha, and F. Müller-Plathe, Advanced Theory and Simulations , 2300520 (2023).
  * Malevanets and Kapral (1999) A. Malevanets and R. Kapral, The Journal of chemical physics 110, 8605 (1999).
  * Poblete  _et al._ (2014) S. Poblete, A. Wysocki, G. Gompper, and R. G. Winkler, Physical Review E 90, 033314 (2014).
  * Wani  _et al._ (2022) Y. M. Wani, P. G. Kovakas, A. Nikoubashman, and M. P. Howard, The Journal of Chemical Physics 156, 024901 (2022).
  * Weeks  _et al._ (1971) J. D. Weeks, D. Chandler, and H. C. Andersen, The Journal of chemical physics 54, 5237 (1971).
  * Silvera Batista  _et al._ (2015) C. A. Silvera Batista, R. G. Larson, and N. A. Kotov, Science 350, 1242477 (2015).
  * Giunta  _et al._ (2023) G. Giunta, G. Campos-Villalobos, and M. Dijkstra, ACS nano (2023).
  * Rosenberg  _et al._ (2020) M. Rosenberg, F. Dekker, J. G. Donaldson, A. P. Philipse, and S. S. Kantorovich, Soft Matter 16, 4451 (2020).