with $\mathbf{v}_{resid}$ (Figure 20). We found that the dot products of most
heads and MLPs was low, and that the S-inhibition heads were the only heads
whose dot product differed between different patterns ABB and BAB. This shows
that only the S-inhibition heads write to the $\mathbf{v}_{\text{resid}}$
direction (as one would hope). Importantly, this test separates
$\mathbf{v}_{resid}$ from the interpretability illusion $\mathbf{v}_{MLP}$.
While patching $\mathbf{v}_{\text{MLP8}}$ also writes to
$\mathbf{v}_{\text{resid8}}$ (i.e.
$\mathbf{v}_{\text{MLP8}}W_{out}\approx\mathbf{v}_{\text{resid8}}$), the MLP
layer does not write this subspace on the IOI task (see Figure 4). This
further supports the observation that the $\mathbf{v}_{\text{MLP}}$ patch
activates a dormant pathway in the model.

Generalization beyond the IOI distribution. We also investigate how the
subspace generalizes. We sample prompts from OpenWebText-10k and look at those
with particularly high and low activations in $\mathbf{v}_{\text{sinhib}}$.
Representative examples are shown in Figure 21 together with the name movers
attention at the position of interest, how the probability changes after
subspace ablation, and how the name movers attention changes.

Stability of found solution. Finally, we note that solutions found by DAS in
the residual stream are stable, including when trained on a subset of
S-inhibition heads (see Figure 18).

Figure 18: Cosine Similarity between learned position subspaces in the
S-inhibition heads is high even when using only a subset of S-inhibition heads
for training Figure 19: The IOI position subspace activates at words that
predict a repeated name. S-inhibition subspace activations for different IOI
prompts per position Figure 20: S-Inhibition heads but not MLP8 write to the
position subspace in the residual stream that is causally connected to the
name movers on the IOI task Figure 21: The IOI position subspace generalizes
to arbitrary OpenWebText prompts Figure 22:  Figure 23:

##  Appendix D Additional details for Section 6

###  D.1 Dataset construction and training details

We use the first 1000 examples from the CounterFact dataset (Meng et al.,
2022a). We filter the facts which GPT2-XL correctly recalls. Out of the
remaining facts, for each relation we form all pairs of distinct facts, and we
sample 5 such pairs from each relation with at least 5 facts. This results in
a collection of 40 fact pairs spanning 8 different relations. We then use
these facts as follows:

  * •

for the ROME experiments in Subsection 6.3, we define edits by requesting one
of the facts in each pair to be rewritten with the object of the other fact;

  * •

for the activation patching experiments in Subsection 6.1, we patch from the
last token of $s^{\prime}$ in $B$ to the last token of $s$ in $A$ (prior work
has shown that the fact is retrieved on $s$ (Geva et al., 2023)), and we again
use DAS Geiger et al. (2023b) to optimize for a direction that maximizes the
logit difference between $o^{\prime}$ and $o$.

###  D.2 Additional fact patching experiments

In figure 24, we show the distribution of the fractional logit difference
metric (see Subsection 4.2 for a definition) when patching between facts as
described in Subsection 6.1. Like in the related Figure 7, we observe that,
while patching along the directions found by DAS achieves strongly negative
values (indicating that the facts are very often successfully changed by the
patch), the interventions that replace the entire MLP layer or only the
causally relevant component of the DAS directions have no such effect.

Fractional logit differenceIntervention layer Figure 24: Fractional logit
difference distributions under three interventions: patching along the
direction found by DAS (blue), patching the component of the DAS direction in
the rowspace of $W_{out}$ (green), and patching the entire hidden MLP
activation (orange).

Next, we observe that the nullspace component of the patching direction is the
one similar to the variation in the inputs (difference of last-token
activations at the two subjects). Specifically, in Figure 25, we plot the
(absolute value of the) cosine similarity between the difference in
activations for the two last subject tokens, and the nullspace component of
the DAS direction. We note that this similarity is consistently significantly
high (note that it can be at most $1$, which would indicate perfect
alignment).

Absolute cosine similarityIntervention layer Figure 25: Distribution of the
absolute value of the cosine similarity between the nullspace component of the
DAS fact patching directions and the difference in activations of the last
tokens of the two subjects.

Finally, we observe that the nullspace component of the patching direction is
a non-trivial part of the direction in Figure 26, where we plot the
distribution of the $\ell_{2}$ norm of this component.

Norm of nullspace componentIntervention layer Figure 26: Distribution of the
norm of the nullspace component of the DAS direction across intervention
layers.

###  D.3 ROME implementation details

ROME takes as input a vector $\mathbf{k}\in\mathbb{R}^{d_{\text{MLP}}}$
representing the subject (e.g. an average of last-token representations of the
subject) and a vector $\mathbf{v}\in\mathbb{R}^{d_{\text{resid}}}$ which, when
output by the MLP layer, will cause the model to predict a new object for the
factual prompt, but at the same time won’t change other facts about the
subject. ROME modifies the MLP weight by setting
$W_{out}^{\prime}=W_{out}+\mathbf{a}\mathbf{b}^{\top}$, where
$\mathbf{a}\in\mathbb{R}^{d_{\text{resid}}},\mathbf{b}\in\mathbb{R}^{d_{\text{MLP}}}$
are chosen so that $W_{out}^{\prime}\mathbf{k}=\mathbf{v}$, and the MLP’s
output is otherwise minimally changed. Without loss of generality, the first
condition implies that $\mathbf{a}=\mathbf{v}-W_{out}\mathbf{k}$ and
$\mathbf{b}^{\top}\mathbf{k}=1$; the second condition is then modeled by
minimizing the variance of $\mathbf{b}^{\top}\mathbf{x}$ when
$\mathbf{x}\sim\mathcal{N}(0,\Sigma)$ for an empirical estimate
$\Sigma\in\mathbb{R}^{d_{\text{MLP}}\times d_{\text{MLP}}}$ of the covariance
of MLP activations (see Lemma D.1 in Appendix D for details and a proof). In
all our experiments involving ROME, we use GPT2-XL (Radford et al., 2019), and
we use the precomputed values of $\Sigma$ from Meng et al. (2022a) accessible
online here.

###  D.4 ROME as an Optimization Problem

We now review the ROME method from Meng et al. (2022a) and show how it can be
characterized as the solution of a simple optimization problem. Following the
terminology of 6.4, let us have an MLP layer with an output projection
$W_{out}$, a key vector $\mathbf{k}\in\mathbb{R}^{d_{\text{MLP}}}$ and a value
vector $\mathbf{v}\in\mathbb{R}^{d_{\text{resid}}}$.

In Meng et al. (2022a), equation 2, the formula for the rank-1 update to
$W_{out}$ is given by

$\displaystyle
W_{out}^{\prime}=W_{out}+(\mathbf{v}-W_{out}\mathbf{k})\frac{\mathbf{k}^{\top}\Sigma^{-1}}{\mathbf{k}^{\top}\Sigma^{-1}\mathbf{k}}$
(6)

where $\Sigma$ is an empirical estimate of the uncentered covariance of the
pre-$W_{out}$ activations. We derive the following equivalent characterization
of this solution (which may be of independent interest):

######  Lemma D.1.

Given a matrix $W_{out}\in\mathbb{R}^{d_{\text{resid}}\times d_{\text{MLP}}}$,
a key vector $\mathbf{k}\in\mathbb{R}^{d_{\text{MLP}}}$ and a value vector
$\mathbf{v}\in\mathbb{R}^{d_{\text{resid}}}$, let $\Sigma\succ
0,\Sigma\in\mathbb{R}^{d_{\text{MLP}}\times d_{\text{MLP}}}$ be a positive
definite matrix (specifically, the uncentered empirical covariance), and let
$\mathbf{x}\sim\mathcal{N}(0,\Sigma)$ be a normally distributed random vector
with mean $0$ and covariance $\Sigma$. Then, the ROME weight update is
$W_{out}^{\prime}=W_{out}+\mathbf{a}\mathbf{b}^{\top}$ where
$\mathbf{a}\in\mathbb{R}^{d_{\text{resid}}},\mathbf{b}\in\mathbb{R}^{d_{\text{MLP}}}$
solve the optimization problem

$\displaystyle\min_{\mathbf{a},\mathbf{b}}\operatorname{trace}(\operatorname{Cov}_{\mathbf{x}}\left[W_{out}^{\prime}\mathbf{x}-W_{out}\mathbf{x}\right])\quad\text{subject
to}\quad W_{out}^{\prime}\mathbf{k}=\mathbf{v}.$

In other words, the ROME update is the update that causes $W_{out}$ to output
$\mathbf{v}$ on input $\mathbf{k}$, and minimizes the total variance of the
extra contribution of the update in the output of the MLP layer under the
assumption that the pre-$W_{out}$ activations are normally distributed with
covariance $\Sigma$.

###### Proof.

We have
$W_{out}^{\prime}\mathbf{x}-W_{out}\mathbf{x}=\mathbf{a}\mathbf{b}^{\top}\mathbf{x}$.
Next, Using $\mathbb{E}_{\mathbf{x}}[\mathbf{x}\mathbf{x}^{\top}]=\Sigma$ and
the cyclic property of the trace, we see that

$\displaystyle\operatorname{trace}(\operatorname{Cov}_{\mathbf{x}}\left[W_{out}^{\prime}\mathbf{x}-W_{out}\mathbf{x}\right])=\|\mathbf{a}\|_{2}^{2}\mathbf{b}^{\top}\Sigma\mathbf{b}$

We must have
$\mathbf{a}\mathbf{b}^{\top}\mathbf{k}=\mathbf{v}-W_{out}\mathbf{k}$, so
without loss of generality we can rescale $\mathbf{a},\mathbf{b}$ so that
$\mathbf{a}=\mathbf{v}-W_{out}\mathbf{k}$. Then, we want to solve the problem

$\displaystyle\min_{\mathbf{b}}\mathbf{b}^{\top}\Sigma\mathbf{b}\quad\text{subject
to}\quad\mathbf{b}^{\top}\mathbf{k}=1$

which we can solve using Lagrange multipliers. The Lagrangian is

$\displaystyle\mathcal{L}(\mathbf{b},\lambda)=\frac{1}{2}\mathbf{b}^{\top}\Sigma\mathbf{b}-\lambda\mathbf{b}^{\top}\mathbf{k}$

and the derivative w.r.t. $\mathbf{b}$ is
$\Sigma\mathbf{b}-\lambda\mathbf{k}=0$, which tells us that $\mathbf{b}$ is in
the direction of $\Sigma^{-1}\mathbf{k}$. Then the constraint
$\mathbf{b}^{\top}\mathbf{k}=1$ forces the constant of proportionality, and we
arrive at
$\mathbf{b}=\frac{\mathbf{k}^{\top}\Sigma^{-1}}{\mathbf{k}^{\top}\Sigma^{-1}\mathbf{k}}$
∎

###  D.5 Connection between 1-dimensional activation patching and model
editing

######  Lemma D.2.

Given prompts A and B, two token positions $t_{A}$, $t_{B}$, and an MLP layer
with output projection weight $W_{out}\in\mathbb{R}^{d_{\text{resid}}\times
d_{\text{MLP}}}$, let $u_{A},u_{B}\in\mathbb{R}^{d_{\text{MLP}}}$ be the
respective (post-nonlinearity) activations at these token positions in this
layer. If $v$ is a direction in the activation space of the MLP layer, then
there exists a ROME edit $W_{out}^{\prime}=W_{out}+ab^{\top}$ such that the
activation patch from $u_{B}$ into $u_{A}$ along $v$ and the edit result in
equal outputs of the MLP layer at token $t_{A}$ when run on prompt A.
Moreover, the ROME edit is given by

$\displaystyle a=\left((u_{B}-u_{A})^{\top}v\right)W_{out}v\quad\text{ and any
$b$ that satisfies }\quad b^{\top}u_{A}=1.$

Choosing $b=\frac{\Sigma^{-1}u_{A}}{u_{A}^{T}\Sigma^{-1}u_{A}}$ minimizes the
change to the model (in the sense of Meng et al. (2022a)) over all such rank-1
edits.

###### Proof.

The activation after patching from B into A along $v$ is
$u_{A}^{\prime}=u_{A}+((u_{B}-u_{A})^{\top}v)v$, which means that the change
in the output of the MLP layer at this token will be

$\displaystyle
W_{out}u_{A}^{\prime}-W_{out}u_{A}=((u_{B}-u_{A})^{\top}v)W_{out}v$

The change introduced by a fact edit at this token is

$\displaystyle
W_{out}^{\prime}u_{A}-W_{out}u_{A}=ab^{\top}u_{A}=\left(b^{\top}u_{A}\right)\left((u_{B}-u_{A})^{\top}v\right)W_{out}v$

and the two are equal because $b^{\top}u_{A}=1$.

To find the $b$ that minimizes the change to the model, we minimize the
variance of $b^{\top}x$ when $x\sim\mathcal{N}(0,\Sigma)$ subject to
$b^{\top}u_{A}=1$. The variance is equal to $b^{\top}\Sigma b$, so we have a
constrained (convex) minimization problem

$\displaystyle\min\frac{1}{2}b^{\top}\Sigma b\quad\text{subject to}\quad
b^{\top}u_{A}=1$

The rest of the proof is the same as in Lemma D.1. Namely, we can solve this
optimization problem using Lagrange multiplies. The Lagrangian is

$\displaystyle\mathcal{L}(b,\lambda)=\frac{1}{2}b^{\top}\Sigma b-\lambda
b^{\top}u_{A}$

and the derivative w.r.t. $b$ is $\Sigma b-\lambda u_{A}=0$, which tells us
that $b$ is in the direction of $\Sigma^{-1}u_{A}$. Then the constraint
$b^{\top}u_{A}=1$ forces the constant of proportionality. ∎

###  D.6 Additional experiments comparing fact patching and rank-1 editing

In Figure 27, we plot the distributions of the logit difference between the
correct object for a fact and the object we are trying to substitute when
patching the 1-dimensional subspaces found by DAS, and performing the
equivalent rank-1 weight edit according to Lemma D.2. We observe that the two
metrics quite closely track each other, indicating that the additional effects
of using a weight edit (as opposed to only intervening at a single token) are
negligible.

Similarly, in Figure 28, we show the success rate of the the two methods in
terms of making the model output the object of the fact we are patching from.
Again, we observe that they quite closely track each other.

Logit differenceIntervention layer Figure 27: Comparison of logit difference
between 1-dimensional fact patches and their derived rank-1 model edits Fact
patch success rateIntervention layer Figure 28: Comparison of fact editing
success rate between 1-dimensional fact patches and their derived rank-1 model
edits

###  D.7 From Rank-1 Model Edits to Subspace Interventions

In this section, we describe how, given a rank-1 edit
$W_{out}^{\prime}=W_{out}+ab^{T}$, to obtain a direction
$v\in\mathbb{R}^{d_{MLP}}$ such that intervening on the model by setting the
projection on $v$ to some constant $c\in\mathbb{R}$ (at each token) is
approximately equivalent to intervening via the rank-1 edit.

Specifically, given an activation $x\in\mathbb{R}^{d_{MLP}}$, the patched
activation is $x^{\prime}=x+\left(c-v^{T}x\right)v$ and the extra contribution
of the subspace intervention to the output of the MLP layer will be

$\displaystyle\operatorname{contrib}_{\text{subspace}}(x)=W_{out}x^{\prime}-W_{out}x=\left(c-v^{T}x\right)Wv.$

Similarly, the extra contribution of the rank-1 edit to the output of the MLP
layer is

$\displaystyle\operatorname{contrib}_{\text{rank-1}}(x)=W_{out}^{\prime}x-W_{out}x=\left(b^{T}x\right)a.$

Recall (see Appendix D.4) that the ROME method (Meng et al., 2022a) implicitly
treats the activation $x$ as a random vector sampled from
$\mathcal{N}\left(0,\Sigma\right)$ where $\Sigma$ is an empirical estimate of
the covariance. In particular, this distribution is used to quantify the
amount to which a rank-1 edit changes the model.

Motivated by this, we formalize approximating the rank-1 edit by the subspace
intervention using the following criteria analogous to the ROME method:

  * •

$\mathbb{E}_{x\sim\mathcal{N}(0,\Sigma)}\left[\operatorname{contrib}_{\text{subspace}}(x)\right]=\mathbb{E}_{x\sim\mathcal{N}(0,\Sigma)}\left[\operatorname{contrib}_{\text{rank-1}}(x)\right]$,
i.e. the interventions have the same expectation;

  * •

$W_{out}v\parallel a$, i.e. the interventions point in the same direciton;

  * •

$\operatorname{trace}\left(\operatorname{Cov}_{x}\left[\operatorname{contrib}_{\text{subspace}}(x)-\operatorname{contrib}_{\text{rank-1}}(x)\right]\right)$
is minimized, i.e. the two interventions are maximally similar with respect to
the activation distribution modeled as $x\sim\mathcal{N}\left(0,\Sigma\right)$
(this is the criterion used by ROME; recall D.4).

The expectation of $\operatorname{contrib}_{\text{rank-1}}(x)$ is zero, while
the expectation of $\operatorname{contrib}_{\text{subspace}}(x)$ is
$cW_{out}v$, and since $W_{out}v=0$ would lead to a trivial intervention, we
must have

$\displaystyle c=0.$

Next, to ensure $W_{out}v\parallel a$, we have to pick $v=\alpha
W_{out}^{+}a+u$ where $u\in\ker W_{out}$. With this, the covariance
minimization can then be written as

$\displaystyle\min_{\alpha,v}\left\|a\right\|_{2}^{2}\left(b+\alpha
v\right)^{T}\Sigma\left(b+\alpha v\right)$

(this is a similar derivation to the one in Appendix D.4). After removing
constant terms and setting $w=u/\alpha$, we are left with

$\displaystyle\min_{\alpha,w}\left[\alpha^{4}\left(W_{out}^{+}a+w\right)^{T}\Sigma\left(W_{out}^{+}a+w\right)+2\alpha^{2}b^{T}\Sigma\left(W_{out}^{+}a+w\right)\right].$

subject to $W_{out}w=0$. The Lagrangian is

$\displaystyle\mathcal{L}\left(\alpha,w,\lambda\right)=\alpha^{4}\left(W_{out}^{+}a+w\right)^{T}\Sigma\left(W_{out}^{+}a+w\right)+2\alpha^{2}b^{T}\Sigma\left(W_{out}^{+}a+w\right)+\lambda^{T}W_{out}w$

with the first-order conditions

$\displaystyle\frac{\partial\mathcal{L}}{\partial
w}=2\alpha^{4}\Sigma\left(W_{out}^{+}a+w\right)+2\alpha^{2}\Sigma
b+W_{out}^{T}\lambda=0$

and $\partial\mathcal{L}/\partial\lambda=W_{out}w=0$. Multiplying the $w$
derivative with $W_{out}\Sigma^{-1}$ on the left gives us a linear system for
$\lambda$:

$\displaystyle
W_{out}\Sigma^{-1}W_{out}^{T}\lambda=-2\alpha^{2}W_{out}b-2\alpha^{4}a,$

which can be solved assuming we know $\alpha$, and then substituting $\lambda$
in $\frac{\partial\mathcal{L}}{\partial w}=0$ gives us $w$. In practice, we
guess several values for $\alpha$ (typically, $\alpha^{2}=0.05$ performs best)
and pick the one resulting in the best value for the objective.

##  Appendix E Additional Details for Section 7

###  E.1 Prevalence of Causal Directions in MLP Layers

Given an MLP activation $\mathbf{x}$ and a vector
$\mathbf{u}\in\mathbb{R}^{d_{MLP}}$, changing the projection of $\mathbf{x}$
on $\mathbf{u}$ means replacing $\mathbf{x}$ with the new activation
$\mathbf{x}^{\prime}=\mathbf{x}+\alpha\mathbf{u}$ for some
$\alpha\in\mathbb{R}$. This translates to the new output of the MLP layer
being

$\displaystyle
W_{out}\left(\mathbf{x}+\alpha\mathbf{u}\right)=W_{out}\mathbf{x}+\alpha
W_{out}\mathbf{u}.$

Under our assumptions, the direction $\mathbf{u}$ will be causally relevant if
the extra contribution to the residual stream $\alpha W_{out}\mathbf{u}$
points along $\mathbf{v}$; thus it suffices to find a $\mathbf{u}$ such that
$W_{out}\mathbf{u}\parallel\mathbf{v}$.

As it turns out, we can simply choose $\mathbf{u}=W_{out}^{+}\mathbf{v}$.
Indeed, we empirically observe that $W_{out}\in\mathbb{R}^{d_{resid}\times
d_{MLP}}$ is a full-rank matrix999This is also heuristically plausible: models
want to maximize their expressive capacity, and pre-training datasets are very
complex, so making $W_{out}$ low-rank would not be preferred by optimization.,
with almost all singular values bounded well away from $0$ (see Appendix E.5).
Since $d_{MLP}>d_{resid}$, it follows that
$W_{out}W_{out}^{+}\mathbf{v}=\mathbf{v}$. This establishes that $\mathbf{u}$
is a causal direction.

###  E.2 Prevalence of Directions Discriminating for $C$ in MLP layers

For a feature $\mathbf{u}\in\mathbb{R}^{d_{MLP}}$ to discriminate between
values of $C$, we need projections of the post-nonlinearity activations on
$\mathbf{u}$ to linearly separate examples according to the values of $C$. By
assumption, $\mathbf{v}$ is a good linear separator for the values of $C$ in
the residual stream. We can thus frame our goal as a more general question:

_If two sets of activations are linearly separable in the residual stream, are
their images after the non-linearity also (approximately) linearly separable?_

The transformation from residual vectors $\mathbf{x}\in\mathbb{R}^{d_{resid}}$
to post-nonlinearity activations is given by the steps

$\displaystyle\mathbf{x}\mapsto\operatorname{LayerNorm}\left(W_{in}\mathbf{x}\right)\mapsto\operatorname{gelu}\left(\operatorname{LayerNorm}\left(W_{in}\mathbf{x}\right)\right)$

The composition of $\operatorname{LayerNorm}$ and $W_{in}$ is approximately a
linear operation (Elhage et al., 2021), so the values of the concept $C$ are
also linearly separated in the pre-$\operatorname{gelu}$ activations. However,
it is not a priori clear if the $\operatorname{gelu}$ operation
(approximately) preserves linear separability.

We show ample empirical evidence in Appendix E.4 that this transformation
approximately preserves the Euclidean geometry of activations in a certain
restricted sense; then, we prove that this preservation implies that points
remain approximately linearly separable after this transformation. We further
argue this empirically in Appendix E.6, where we show that linear separability
is approximately preserved in MLP activations for random directions
$\mathbf{v}$ in the residual stream.

###  E.3 Empirical Analysis of Distortion Introduced by the Non-linearity

#### Methodology.

We use the first 10K texts of OpenWebText dataset (Gokaslan & Cohen, 2019).
Each of these texts contains 1,024 tokens; we pass each text through GPT-2
Small, and for each layer collect the pre-$\operatorname{gelu}$ activations
$\mathbf{x}_{i}$ of the MLP layer, as well as the values
$\mathbf{z}_{i}=\operatorname{proj}_{\ker
W_{out}}\left(\operatorname{gelu}\left(\mathbf{x}_{i}\right)\right)$. We
sample 250 quadruples of distinct $i,j,k,l$ per text, and compute the values

$\displaystyle
a_{ijkl}=\left(\mathbf{x}_{i}-\mathbf{x}_{j}\right)^{\top}\left(\mathbf{x}_{k}-\mathbf{x}_{l}\right)$
$\displaystyle
b_{ijkl}=\left(\mathbf{z}_{i}-\mathbf{z}_{j}\right)^{\top}\left(\mathbf{z}_{k}-\mathbf{z}_{l}\right)$

We collect these numbers across the first 1000 texts out of the first 10K,
resulting in 250K datapoints per layer, and perform linear regression of
$b_{i}$ against $a_{i}$.

We note that there is some inherent linearity in the quantities
$a_{ijkl},b_{ijkl}$ that could in principle skew the results of the linear
regression towards a higher $r^{2}$ statistic in the presence of enough
samples. In particular, there are linear dependencies of the form

$\displaystyle a_{ijkl}=a_{ijkp}+a_{ijpl}$

for any $p$, and similarly for the $b_{ijkl}$ quantities. This makes these
quantities potentially misleading targets for linear regression. However, in
our regime, we sample 250 4-element subsets from the set
$\\{1,\ldots,1024\\}$, and the probability of sampling quadruples that are
linearly related is quite small.

#### Results.

We find that the coefficients of determination $r^{2}$ for the linear
regression are consistently high ($\approx 0.8$ or higher) for all layers
except for layer 0, indicating a high degree of fit. The $r^{2}$ values are
given in Table TODO, and regression plots are shown in Figure TODO. We also
remark that the intercept coefficients are $\approx 0$ relative to the
standard deviation in the dependent variable.

###  E.4 Details for Subsection E.2

We will show the stronger property that activations remain approximately
linearly separable even after projecting on the kernel of $W_{out}$. Define
the function

$\displaystyle
f:\mathbf{x}^{\prime}\mapsto\operatorname{gelu}\left(\mathbf{x}^{\prime}\right)\mapsto\operatorname{proj}_{\ker
W_{out}}\left(\operatorname{gelu}\left(\mathbf{x}^{\prime}\right)\right)$

where $\operatorname{proj}_{\ker W_{out}}$ is orthogonal projection on the
kernel of $W_{out}$.

To overcome the non-linearity of $f$, we establish an empirical property of
$f$ on activations from the model’s pre-training distribution. Specifically,
we show that $f$ approximately preserves the Euclidean geometry of activations
in a certain restricted sense:

$\displaystyle\left(f(\mathbf{x}_{i})-f(\mathbf{x}_{j})\right)^{\top}\left(f(\mathbf{x}_{k})-f(\mathbf{x}_{l})\right)\approx\lambda\left(\mathbf{x}_{i}-\mathbf{x}_{j}\right)^{\top}\left(\mathbf{x}_{k}-\mathbf{x}_{l}\right)+\eta$
(7)

for $\lambda>0$ and $\eta\approx 0$ (relative to the standard deviation of the
expression on the left-hand side of the approximation). Specifically, we
perform linear regression of
$\left(f(\mathbf{x}_{i})-f(\mathbf{x}_{j})\right)^{\top}\left(f(\mathbf{x}_{k})-f(\mathbf{x}_{l})\right)$
using
$\left(\mathbf{x}_{i}-\mathbf{x}_{j}\right)^{\top}\left(\mathbf{x}_{k}-\mathbf{x}_{l}\right)$
as the predictor variable in Appendix E.3, and find very high coefficients of
determination ($r^{2}\approx 0.8$) in all layers except for layer $0$. To
avoid relying solely on the coefficient of determination, we also generate
regression plots for the data; see Figure 29 for regression lines over samples
of $10^{4}$ points from each layer of GPT2-Small.

Figure 29: Regression plots accompanying the experiments in Appendix E.4.

Finally, we prove that $f$ maintains linear separability if we assume that
Equation 7 holds exactly with $\eta=0$:

######  Lemma E.1.

Let $x_{i}\in\mathbb{R}^{d},1\leq i\leq n$ be linearly separable with respect
to binary labels $y_{i}\in\\{-1,1\\}$. Let $f:\mathbb{R}^{d}\to\mathbb{R}^{d}$
be a transformation with the property that

$\displaystyle\left(f(x_{i})-f(x_{j})\right)^{\top}\left(f(x_{k})-f(x_{l})\right)=\lambda\left(x_{i}-x_{j}\right)^{\top}\left(x_{k}-x_{l}\right)$

for all distinct $i,j,k,l$ and some $\lambda>0$. Then, $f(x_{i})$ are linearly
separable with respect to the labels $y_{i}$ as well.

###### Proof.

Consider the hard SVM objective for the points $(x_{i},y_{i})$,

$\displaystyle\min_{w,b}\frac{1}{2}\left\|w\right\|_{2}^{2}\quad\text{subject
to}\quad y_{i}\left(w^{\top}x_{i}+b\right)\geq 1.$

Since the examples are linearly separable, we know that the minimizer
$(w^{*},b^{*})$ exists and satisfies all constraints. Furthermore, from the
optimality conditions of the dual formulation of the objective we know that we
can write $w^{*}=\sum_{i}\alpha_{i}x_{i}$ where $\sum_{i}\alpha_{i}=0$ (see
for example Awad et al. (2015)). Let
$S=\\{s_{1},\ldots,s_{t}\\}\subset\\{1,\ldots,n\\}$ be the support of
$\alpha$, i.e. the indices $s_{j}$ such that $\alpha_{s_{j}}\neq 0$. Since
$\sum_{i}\alpha_{i}=0$, we can rewrite $w^{*}$ as

$\displaystyle w^{*}=\sum_{j}\beta_{j}\left(x_{s_{j}}-x_{s_{j+1}}\right)$

with indices modulo $\left|S\right|$. Since $(w^{*},b^{*})$ is a separating
hyperplane for $(x_{i},y_{i})$, we have

$\displaystyle\left(w^{*}\right)^{\top}x_{i}$ $\displaystyle\geq
1-b\quad\text{when}\quad y_{i}=1$
$\displaystyle\left(w^{*}\right)^{\top}x_{i}$
$\displaystyle\leq-1-b\quad\text{when}\quad y_{i}=-1$

and thus

$\displaystyle\left(w^{*}\right)^{\top}\left(x_{i}-x_{j}\right)\geq
2\quad\text{when}\quad y_{i}=1,y_{j}=-1.$

Using the expansion of $w^{*}$ as a linear combination of differences between
examples, this says

$\displaystyle\sum_{j}\beta_{j}\left(x_{s_{j}}-x_{s_{j+1}}\right)^{\top}\left(x_{i}-x_{j}\right)\geq
2\quad\text{when}\quad y_{i}=1,y_{j}=-1.$

and thus

$\displaystyle\sum_{j}\beta_{j}\left(f(x_{s_{j}})-f(x_{s_{j+1}})\right)^{\top}\left(f(x_{i})-f(x_{j})\right)\geq
2\lambda>0\quad\text{when}\quad y_{i}=1,y_{j}=-1.$

Now we claim that

$\displaystyle\widehat{w}=\sum_{j}\beta_{j}\left(f(x_{s_{j}})-f(x_{s_{j+1}})\right)$

is a linear separator for $(f(x_{i}),y_{i})$ for some bias to be determined
later. Indeed, let $M=\min_{y_{i}=1}\widehat{w}^{\top}f(x_{i})$ and
$m=\max_{y_{i}=-1}\widehat{w}^{\top}f(x_{i})$. Then we have $M-m\geq
2\lambda>0$. Choosing any $\widehat{b}\in(m,M)$, we have

$\displaystyle\widehat{w}^{\top}f\left(x_{i}\right)-\widehat{b}\geq
M-\widehat{b}>0\quad\text{when}\quad y_{i}=1$
$\displaystyle\widehat{w}^{\top}f\left(x_{i}\right)-\widehat{b}\leq
m-\widehat{b}<0\quad\text{when}\quad y_{i}=-1$

which shows that $(\widehat{w},\widehat{b})$ linearly separates the points
$(f(x_{i}),y_{i})$. ∎

###  E.5 MLP weights are full-rank matrices

In figure 30, we plot the 100 smallest singular values of the MLP weights in
GPT2-Small for all 12 layers. We observe that they the vast majority are
bounded well away from $0$. This confirms that both MLP weights are full-rank
transformations.

Singular valueIndex (decreasing)

Index (decreasing)

Figure 30: Smallest 100 singular values of the $W_{in}$ (left) and $W_{out}$
(right) MLP weights by layer in in GPT2-Small

###  E.6 Features in the residual stream propagate to hidden MLP activations

Intuition. Suppose we have two classes of examples that are linearly separable
in the residual stream. The transformation from the residual stream to the
hidden MLP activations is a linear map followed by a nonlinearity,
specifically $x\mapsto\operatorname{gelu}(W_{in}x)$. As we observed in E.5,
the $W_{in}$ matrix is full-rank, meaning that all the information linearly
present in $x$ will also be so in $W_{in}x$. Even better, since $W_{in}$ maps
$x$ from a $d_{\text{resid}}$-dimensional space to a
$d_{\text{MLP}}=4d_{\text{resid}}$-dimensional space, this should intuitively
make it much easier to linearly separate the points, because in a higher-
dimensional space there are many more linear separators. On the other hand,
the non-linearity has an opposite effect: by compressing the space of
activations, it makes it harder for points to be separable. So it is a priori
unclear which intuition is decisive.

Empirical validation. However, it turns out that empirically this is not such
a problem. To test this, we run the model GPT2-Small on random samples from
its data distribution (we used OpenWebText-10k), and extract 2000 activations
of an MLP-layer after the non-linearity. We train a linear regression with
$\ell_{2}$-regularization to recover the dot product of the residual stream
immediately before the MLP-layer of interest and a randomly chosen direction.
We repeat this experiment with different random vectors and for each layer. We
observe that all regressions are better than chance and explain a significant
amount of variance on the held-out test set ($R^{2}=0.71\pm
0.17,\text{MSE}=0.31\pm 0.18,p<0.005$). Results are shown in Figure 31 (right)
(every marker corresponds to one regression model using a different random
direction).

The position information in the IOI task is really a binary feature, so we are
also interested in whether _binary_ information in general is linearly
recoverable from the MLP activations. To test this, we sample activations from
the model run on randomly-sampled prompts. This time however, we add or
subtract a multiple of a random direction $v$ to the residual stream
activation $u$, and calculate the MLP activations using this new residual
stream vector $u^{\prime}$:

$\displaystyle u^{\prime}=u+y\times z\times\|u\|_{2}\times v$

where $y\in\\{-1,1\\}$ is uniformly random, $z$ is a scaling factor we
manipulate, and $v$ is a randomly chosen direction of unit norm. For each
classifier, we randomly sample a direction $v$ that we either add or subtract
(using $y$) from the residual stream. The classifier is trained to predict
$y$. We rescale v to match the average norm of a residual vector and then
scale it with a small scalar $z$.

Then, a logistic classifier is trained on 1600 samples. Again, we repeat this
experiment for different $v$ and $z$, and for each layer. We observe that the
classifier works quite well across layers even with very small values of $z$
(still, accuracy drops for $z=0.0001$). Results are shown in Figure 31
(right), and Table 2.

Table 2: Mean Accuracy for Different Values of $z$ $z$ | Mean Accuracy
---|---
0.0001 | 0.69
0.001 | 0.83
0.01 | 0.87
0.1 | 0.996

Figure 31: Recovering residual stream features linearly from hidden MLP
activations: classification (left) and regression (right).


2021

These authors contributed equally to this work. [1]Jun Yan These authors
contributed equally to this work.

1] School of Electronic and Information Engineering, Tongji University,
Caoangong Street, Shanghai, 201804, China

2]Mechanical Engineering, Michigan State University, Lansing, 48824, USA

# Multi-agent Reinforcement Learning for Cooperative Lane Changing of
Connected and Autonomous Vehicles in Mixed Traffic

Wei Zhou<EMAIL_ADDRESS>Dong Chen<EMAIL_ADDRESS><EMAIL_ADDRESS>Zhaojian Li<EMAIL_ADDRESS>Huilin Yin
<EMAIL_ADDRESS>Wanchen Ge<EMAIL_ADDRESS>[  [

###### Abstract

Autonomous driving has attracted significant research interests in the past
two decades as it offers many potential benefits, including releasing drivers
from exhausting driving and mitigating traffic congestion, among others.
Despite promising progress, lane-changing remains a great challenge for
autonomous vehicles (AV), especially in mixed and dynamic traffic scenarios.
Recently, reinforcement learning (RL), a powerful data-driven control method,
has been widely explored for lane-changing decision makings in AVs with
encouraging results demonstrated. However, the majority of those studies are
focused on a single-vehicle setting, and lane-changing in the context of
multiple AVs coexisting with human-driven vehicles (HDVs) have received scarce
attention. In this paper, we formulate the lane-changing decision making of
multiple AVs in a mixed-traffic highway environment as a multi-agent
reinforcement learning (MARL) problem, where each AV makes lane-changing
decisions based on the motions of both neighboring AVs and HDVs. Specifically,
a multi-agent advantage actor-critic network (MA2C) is developed with a novel
local reward design and a parameter sharing scheme. In particular, a multi-
objective reward function is proposed to incorporate fuel efficiency, driving
comfort, and safety of autonomous driving. Comprehensive experimental results,
conducted under three different traffic densities and various levels of human
driver aggressiveness, show that our proposed MARL framework consistently
outperforms several state-of-the-art benchmarks in terms of efficiency, safety
and driver comfort.

###### keywords:

Multi-agent deep reinforcement learning, lane-changing, connected autonomous
vehicles, mixed traffic

##  1 Introduction

Autonomous driving has received significant research interests in the past two
decades due to its many potential societal and economical benefits. Compared
to traditional vehicles, autonomous vehicles (AVs) not only promise fewer
emissions Paden  but are also expected to improve safety and efficiency.
Despite remarkable progress, high-level decision-making in AVs remains a big
challenge due to the complex and dynamic traffic environment, especially in
mixed traffic co-existing with other road users. Lane changing is one such
challenging high-level decision-making in AVs, which has significant
influences on traffic safety and efficiency desiraju2014minimizing ;
li2020cooperative , and is the focus of this paper.

The considered lane-changing scenario is illustrated in Fig. 1, where AVs and
HDVs coexist on a one-way highway with two lanes. The AVs aim to safely travel
through the traffic while making necessary lane changes to overtake slow-
moving vehicles for improved efficiency. Furthermore, in the presence of
multiple AVs, the AVs are expected to collaboratively learn a policy to adapt
to HDVs and enable safe and efficient lane changes. As HDVs bring
unknown/uncertain behaviors, planning and control in such mixed traffic to
realize safe and efficient maneuvers is a challenging task chen2020autonomous
.

Figure 1: Illustration of the considered lane-changing scenario (green: AVs,
blue: HDVs, arrow curve: a possible trajectory of the ego vehicle AV1 to make
the lane change).

Recently, reinforcement learning (RL) has emerged as a promising framework for
autonomous driving due to its online adaptation capabilities and the ability
to solve complex problems wang2019continuous ; xi2020efficient . Several
recent studies have explored the use of RL in AV lane-changing
wang2021harmonious ; du2020cooperative ; chen2020autonomous , which consider a
single AV setting where the ego vehicle learns a lane-changing behavior by
taking all other vehicles as part of the driving environment for decision
making. While completely scalable, this single-agent approach will lead to
unsatisfactory performance in the complex environment like multi-AV lane-
changing in mixed traffic that requires close collaboration and coordination
among AVs Hoel .

On the other hand, multi-agent reinforcement learning (MARL) has been greatly
advanced and successfully applied to a variety of complex multi-agent systems
such as games vinyals2019grandmaster , traffic light control Chu  and fleet
management lin2018efficient . The applications of MARL to autonomous driving
also exist shalev2016safe ; wangmulti ; ha2020leveraging ; palanisamy2020multi
, with the objective of accomplishing autonomous driving tasks cooperatively
and reacting timely to HDVs. In particular, previous works chen2021graph ;
ha2020leveraging  also apply the MARL to highway lane change tasks and show
promising and scalable performance, in which AVs learn cooperatively via
sharing the same objective (i.e., reward/cost function) that considers safety
and efficiency. However, those reward designs often ignore the passengers’
comfort, which may lead to sudden acceleration and deceleration that can cause
ride discomfort. In addition, they assume that the HDVs follow unchanged,
universal human driver behaviors, which is clearly oversimplified and
impractical in the real world as different human drivers tend to behave quite
differently. Learning algorithms should thus work with different human driving
behaviors, e.g., aggressive or mild.

To address the above issues, we develop a multi-agent reinforcement learning
algorithm by employing a multi-agent advantage actor-critic network (MA2C) for
multi-AV lane-changing decision making, featuring a novel local reward design
that incorporates the safety, efficiency and passenger comfort as well as a
parameter sharing scheme to foster inter-agent collaborations. The main
contributions and the technical advancements of this paper are summarized as
follows.

  1. 1.

We formulate the multi-AV highway lane changing in mixed traffic is modeled as
a decentralized cooperative MARL problem, where agents cooperatively learn a
safe and efficient driving policy.

  2. 2.

We develop a novel, efficient, and scalable MARL algorithm, multi-agent
advantage actor-critic network, by introducing a parameter-sharing mechanism
and effective reward function design.

  3. 3.

We conduct comprehensive experiments on three different traffic densities and
two levels of drivers’ behavior modes, and the results show that the proposed
approach consistently outperforms several state-of-the-art algorithms in terms
of driving safety, efficiency and driver comfort.

The rest of the paper is organized as follows. Section 2 reviews the state-of-
the-art dynamics-based and RL/MARL algorithms for autonomous driving tasks.
The preliminaries of RL and the proposed MARL algorithm are introduced in
Section 3. Experiments, results, and discussions are presented in Section 4.
Finally, we summarize the paper and discuss future work in Section 5.

##  2 Related Work

In this section, we survey the existing literature on decision-making tasks in
autonomous driving, which can be mainly classified into two categories: non-
data-driven and data-driven methods.

###  2.1 Non-data-driven Methods

Conventional rule-based or model-based approaches ho2009lane ; nilsson2016if ;
nilsson2016lane  rely on hard-coded rules or dynamical models to construct
predefined logic mechanisms to determine the behaviors of ego vehicles under
different situations. For instance, in ho2009lane , lane-changing guidance is
provided by establishing virtual trajectory references for every vehicle, and
a safe trajectory is then planned by considering the trajectories of other
vehicles. In nilsson2016if , a low-complexity lane-changing algorithm is
developed by following heuristic rules such as keeping appropriate inter-
vehicle traffic gaps and time instances to perform the maneuver. In addition,
an optimization-based lane change approach is proposed in nilsson2016lane ,
which formulates the trajectory planning problem as coupled longitudinal and
lateral predictive control problems and is then solved via Quadratic Programs
under specific system constraints. However, the rules and optimization
criteria for real-world driving problems, especially in mixed-traffic
scenarios with unknown and stochastic driver’s behaviors, may become too
complex to be explicitly formulated for all scenarios.

###  2.2 Data-Driven Methods

Recently, data-driven methods, such as reinforcement learning (RL), have
received great attention and been widely explored for autonomous driving
tasks. Particularly, a model-free RL approach based on deep deterministic
policy gradient (DDPG) is proposed in wang2019continuous  to learn a
continuous control policy efficient lane changing. In chen2020autonomous , a
safe RL framework is presented by integrating a lane-changing regret model
into a safety supervisor based on an extended double deep Q-network (DDQN). In
chen2019attention , a hierarchical RL algorithm is developed to learn lane-
changing behaviors in dense traffic by applying the designed temporal and
spatial attention strategies, and promising performance are demonstrated in
the TORCS simulator under various lane change scenarios. However, the
aforementioned methods are designed for the single-agent (i.e., one ego
vehicle) scenarios, treating all other vehicles as part of the environment,
which makes them implausible for the considered multi-agent lane-changing
setting where collaboration and coordination among AVs are required.

On the other hand, multi-agent reinforcement learning (MARL) has also been
explored for autonomous driving tasks shalev2016safe ; dong2020drl ; wangmulti
; chen2021deep . In particular, in shalev2016safe , a MARL algorithm with
hard-coded safety constraints is proposed to solve the double-merge problem.
Also, a hierarchical temporal abstraction method is applied to reduce the
effective horizon and the variance of the gradient estimation error. In
chen2021deep , a MARL algorithm is proposed to solve the on-ramp merging
problem with safety enhancement by a novel priority-based safety supervisor.
In addition, the authors in dong2020drl  propose a novel MARL approach
combining Graphic Convolution Neural Network (GCN) kipf2016semi  and Deep Q
Network (DQN) mnih2013playing  to better fuse the acquired information from
collaborative sensing, showing promising results on a 3-lane freeway
containing 2 off-ramps highway environment. While these MARL algorithms only
consider the efficiency and safety in their designed reward function, another
important factor, the passenger comfort, is not considered in their reward
function design. Furthermore, those approaches assume the HDVs follow a
constant, universal driving behavior, which has limited implications for real-
world applications as different human drivers may behave totally differently.

In this paper, we formulate the decision making of multiple AVs on highway
lane changing as a MARL problem, where a multi-objective reward function is
proposed to simultaneously promote safety, efficiency and passenger comfort. A
parameter-sharing scheme is exploited to foster inter-agent collaborations.
Experimental results on three different traffic densities with two levels of
driver aggressiveness show the proposed MARL performs well on different lane
change scenarios.

##  3 Problem Formulation

In this section, we review the preliminaries of RL and formulate the
considered highway lane-changing problem as a partially observable Markov
decision process (POMDP). Then we present the proposed multi-agent actor-
critic algorithm, featuring a parameter-sharing mechanism and efficient reward
function design, to solve the formulated POMDP.

###  3.1 Preliminary of RL

In the standard RL setting, the agent aims to learn an optimal policy
$\pi^{*}$ to maximize the accumulated future reward
$R_{t}=\sum_{k=0}^{T}\gamma^{k}r_{t+k}$ from the time step $t$ with discount
factor $\gamma\in(0,1]$ by continuous interacting with the environment.
Specially, at time step $t$, the agent receives a state
$s_{t}\in\mathcal{R}^{n}$ from the environment and selects an action
$a_{t}\in\mathcal{A}^{m}$ according to its policy
$\pi:\,\mathcal{S}\rightarrow\Pr(\mathcal{A})$. As a result, the agent
receives the next state $s_{t+1}$ and receives a scalar reward $r_{t}$. If the
agent can only observe a part of the state $s_{t}$, the underlying dynamics
becomes a POMDP spaan2012partially  and the goal is then to learn a policy
that maps from the partial observation to an appropriate action to maximize
the rewards.

The action-value function $Q^{\pi}(s,a)=E[R_{t}{\mid}{s=s_{t}},a]$ is defined
as the expected return obtained by selecting an action $a$ in state $s_{t}$
and following policy $\pi$ afterwards. The optimal Q-function is given by
$Q^{*}(s,a)=\max_{\pi}Q^{\pi}(s,a)$ for state $s$ and action $a$. Similarly,
the state-value function is defined as
$V^{\pi}(s_{t})=E_{\pi}{[R_{t}{\mid}{s=s_{t}}]}$ representing the expected
return for following the policy $\pi$ from state $s_{t}$. In model-free RL
methods, the policy is often represented by a neural network denoted as
$\pi_{\theta}(a_{t}{\mid}s_{t})$, where $\theta$ is the learnable parameters.
In actor-critic (A2C) algorithms mnih2016asynchronous , a critic network,
parameterized by $\omega$, learns the state-value function
$V_{\omega}^{\pi_{\theta}}(s_{t})$ and an actor network
$\pi_{\theta}(a_{t}{\mid}s_{t})$ parameterized by $\theta$ is applied to
update the policy distribution in the direction suggested by the critic
network as follows:

$\theta\leftarrow\theta+E_{\pi_{\theta}}\left[\Big{(}\nabla_{\theta}\log\pi_{\theta}(a_{t}{\mid}s_{t})\Big{)}A_{t}\right],$
(1)

where the advantage function
$A_{t}=Q^{\pi_{\theta}}(s,a)-V_{\omega}^{\pi_{\theta}}(s_{t})$
mnih2016asynchronous  is introduced to reduce the sample variance. The
parameters of the state-value function are then updated by minimizing the
following loss function:

$\min_{\omega}E_{\mathcal{B}}\Big{(}R_{t}+\gamma
V_{\omega^{\prime}}(s_{t+1})-V_{\omega}(s_{t})\Big{)}^{2},$ (2)

where $\mathcal{B}$ is the experience replay buffer that stores previously
encountered trajectories and $\omega^{\prime}$ denotes the parameters of the
target network mnih2013playing .

###  3.2 Lane Changing as MARL

In this subsection, we develop a decentralized, MARL-based approach for
highway lane-changing of multiple AVs. In particular, we model the mixed-
traffic lane-changing environment as a multi-agent network:
$\mathcal{G}=(\text{$\nu$},\text{$\varepsilon$})$, where each agent (i.e., ego
vehicle) $i\in\text{$\nu$}$ communicates with its neighbors $\mathcal{N}_{i}$
via the communication link $\varepsilon_{ij}\in\large\text{$\varepsilon$}$.
The corresponding POMDP is characterized as
$(\\{\mathcal{A}_{i},\mathcal{O}_{i},\mathcal{R}_{i}\\}_{i\subseteq\nu},\mathcal{T})$,
where $\mathcal{O}_{i}\in\mathcal{S}_{i}$ is the partial description of the
environment state as stated in chu2020multi . In a multi-agent POMDP, each
agent $i$ follows a decentralized policy
$\pi_{i}:\mathcal{O}_{i}\times\mathcal{S}_{i}\rightarrow[0,1]$ to choose the
action $a_{t}$ at time step $t$. The described POMDP is defined as:

  1. 1.

State Space: The state space $\mathcal{O}_{i}$ of Agent $i$ is defined as a
matrix $\mathcal{N}_{N_{i}}\times\mathcal{F}$, where $\mathcal{N}_{N_{i}}$ is
the number of detected vehicles, and $\mathcal{F}$ is the number of features,
which is used to represent the current state of vehicles. It includes the
longitudinal position $x$, the lateral position $y$ of the observed vehicle
relative to the ego vehicle, the longitudinal speed $v_{x}$, and the lateral
speed $v_{y}$ of the observed vehicle relative to the ego vehicle.

  2. 2.

Action Space: The action space $\mathcal{A}_{i}$ of agent $i$ is defined as a
set of high-level control decisions, including speed up, slow down, cruising,
turn left, and turn right. The action space combination for AVs is defined as
$\mathcal{A}=\mathcal{A}_{1}\times\mathcal{A}_{2}\times\cdot\cdot\cdot\times\mathcal{A}_{N}$,
where $N$ is the total number of vehicles in the scene.

  3. 3.

Reward Function: Multiple metrics including safety, traffic efficiency and
passenger’s comfort are considered in the reward function design:

     * $\bullet$

safety evaluation $r_{s}$: The vehicle should operate without collisions.

     * $\bullet$

headway evaluation $r_{d}$: The vehicle should maintain a safe distance from
the preceding vehicles during driving to avoid collisions.

     * $\bullet$

speed evaluation $r_{v}$: Under the premise of ensuring safety, the vehicle is
expected to drive at a high and stable speed.

     * $\bullet$

driving comfort $r_{c}$: Smooth acceleration and deceleration are expected to
ensure safety and comfort. In addition, frequent lane changes should be
avoided.

As such, a multi-objective reward $r_{i,t}$ at the time step $t$ is defined
as:

$r_{i,t}=\omega_{s}r_{s}+\omega_{d}r_{d}+\omega_{v}r_{v}-\omega_{c}r_{c},$ (3)

where $\omega_{s}$, $\omega_{d}$, $\omega_{v}$ and $\omega_{c}$ are the
weighting coefficients. We set the safety factor $\omega_{s}$ to a large
value, because safety is the most important criterion during driving. The
details of the four performance measurements are discussed next:

    1. (1)

If there is no collision, the collision evaluation $r_{s}$ is set to 0,
otherwise, $r_{s}$ is set as -1.

    2. (2)

The headway evaluation is defined as

$r_{d}=\log\frac{d_{headway}}{v_{t}t_{d}},$ (4)

where $d_{headway}$ is the distance to the preceding vehicle, and $v_{t}$ and
$t_{d}$ are the current vehicle speed and time headway threshold,
respectively.

    3. (3)

The speed evaluation $r_{v}$ is defined as

$r_{v}=\min\left\\{\frac{v_{t}-v_{min}}{v_{max}-v_{min}},1\right\\},$ (5)

where $v_{t}$, $v_{min}$ and $v_{max}$ are the current, minimum, and maximum
speeds of the ego vehicle, respectively. Within the specified speed range,
higher speed is preferred to improve the driving efficiency.

    4. (4)

The driving comfort $r_{c}$ is defined as

$r_{c}=r_{a}+r_{lc},$ (6)

where

$r_{a}=\left\\{\begin{array}[]{rcl}-1,&|a_{t}|\geq a_{th}\\\
0,&|a_{t}|<a_{th}\end{array}\right.$

is the penalty term of rapid acceleration and deceleration than a given
threshold $a_{th}$. Here $a_{t}$ presents the acceleration at time $t$.

$r_{lc}=\left\\{\begin{array}[]{rcl}-1,&\rm\normalsize change&\rm\normalsize
lane\\\ 0,&\rm\normalsize keep&\rm\normalsize lane\end{array}\right.$

is defined as the lane change penalty. Excessive lane changes can cause
discomfort and safety issues.

  4. 4.

Transition Probability: The transition probability $T(s^{,}\mid s,a)$
characterizes the transition from one state to another. Since our MARL
algorithm is a model-free design, we do not assume any prior knowledge about
transition probability.

###  3.3 MA2C for AVs

In this paper, we extend the actor-critic network mnih2016asynchronous  to the
multi-agent setting as a multi-agent actor-critic network (i.e., MA2C). MA2C
improves the stability and scalability of the learning process by allowing
certain communication among agents chu2020multi . To take the advantage of
homogeneous agents in the considered MARL setting, we assume all the agents
share the same network structure and parameters, while they are still able to
make different maneuvers according to different input states. The goal in
cooperative MARL setting is to maximize the global reward of all the agents.
To overcome the communication overhead and the credit assignment problem
sutton2018reinforcement , we adopt the local reward design chen2021deep  as
follows:

$r_{i,t}=\frac{1}{\mid{\nu_{i}}\mid}\sum_{j\in\nu_{i}}r_{j,t},$ (7)

where $\mid\nu_{i}\mid$ denotes the cardinality of a set containing the ego
vehicle and its close neighbors. Compared to the global reward design
previously used in kaushik2018parameter ; dong2020drl , the designed local
reward design mitigates the impact of remote agents.

The backbone of the proposed MA2C network is shown in Fig. 2, in which states
separated by physical units are first processed by separate 64-neuron fully
connected (FC) layers. Then all hidden units are combined and fed into the
128-neuron FC layer. Then the shared actor-critic network will update the
policy and value networks with the extracted features. As mentioned in
chen2021deep , the adopted parameter sharing scheme lin2018efficient  between
the actor and value networks can greatly improve the learning efficiency.

Figure 2: The architecture of the proposed MA2C network with shared actor-
critic network design, where $x$ and $y$ are the longitudinal and lateral
position of the observed vehicle relative to the ego vehicle, and $v_{x}$ and
$v_{y}$ are the longitudinal and lateral speed of the observed vehicle
relative to the ego vehicle.

The pseudo-code of the proposed MA2C algorithm is shown in Algorithm 1. The
hyperparameters include: the (time)-discount factor $\gamma$, the learning
rate $\eta$, the politeness coefficient $p$ and the epoch length $T$.
Specifically, the agent receives the observation $O_{i,t}$ from the
environment and updates the action by its policy (Line 3-6). After each
episode is completed, the network parameters are updated accordingly (Line
9-11). If an episode is completed or a collision occurs, the “DONE” signal is
released and the environment will be reset to its initial state to start a new
epoch (Line 13-14).

Algorithm 1 MARL for AVs.

1:$\gamma,\eta,p,T$.

2:$\theta$.

3:Initialize $o_{0},t\leftarrow 0$.

4:repeat

5: for $i\in V$ do

6: Observe $o_{i,t}$;

7: Update $a_{i,t}\sim\pi_{\theta_{i,t}}$;

8: end for

9: Update $t=t+1$;

10: if DONE then

11: for $i\in V$ do

12: Update
$\theta_{i}\leftarrow\theta_{i}+\eta\nabla_{\theta_{i}}{J(\theta_{i})}$;

13: end for

14: end if

15: if $t=T$ then

16: Initialize $o_{0},t\leftarrow 0$;

17: end if

18:until Stop condition is reached

##  4 Experiments and Discussion

In this section, we evaluate the performance of the proposed MARL algorithm in
terms of training efficiency, safety and driving comfort in the considered
highway lane changing scenario shown in Fig. 1.

###  4.1 HDV Models

In this experiment, we assume that the longitudinal control of HDVs follows
the Intelligent Driver Model (IDM) Treiber , which is a deterministic
continuous-time model describing the dynamics of the position and speed of
each vehicle. It takes into account the expected speed, distance between the
vehicles and the behavior of the acceleration/deceleration process caused by
the different driving habits. In addition, Minimize Overall Braking Induced By
Lane Change model (MOBIL) Kesting  is adopted for the lateral control. It
takes vehicle acceleration as the input variable of the model and can work
well with most car-following models. The acceleration expression is defined as
follows:

$\tilde{a}_{n}\geq-b_{safe},$ (8)

where $\tilde{a}_{n}$ is the acceleration of the new follower after the lane
change, and $b_{safe}$ is the maximum braking imposed to the new follower. If
the inequality in Eqn. 8 is satisfied, the ego vehicle is able to change
lanes. The incentive condition is defined as:

${\underbrace{\tilde{a}_{c}-a_{c}}_{\text{\rm
ego~{}vehicle}}}+p\big{(}{\underbrace{\tilde{a}_{n}-a_{n}}_{\text{\rm
new~{}follower}}}+{\underbrace{\tilde{a}_{o}-a_{o}}_{\text{\rm
old~{}follower}}}\big{)}\geq\Delta a_{th},$ (9)

where $a$ and $\tilde{a}$ are the acceleration of the ego vehicle before and
after the lane change, respectively, $\Delta a_{th}$ is the threshold that
determines whether to trigger the lane change or not, and $p$ is a politeness
coefficient that controls how much effect we want to take into account for the
followers, where $p=1$ represents the most considerate drivers whose decision
on change lanes may give way to the following blocked vehicles whereas $p=0$
characterizes the most aggressive drivers where the HDV makes selfish lane-
changing decisions by only considering their own speed gains and ignoring
other vehicles. The performance evaluation of different $p$ values is
discussed in Section 4.3.3.

###  4.2 Experimental Settings

The simulation environment is modified from the gym-based highway-env
simulator highway-env . We set the highway road length to $520~{}m$, and the
vehicles beyond the road are ignored. The vehicles are randomly spawned on the
highway with different initial speeds $25-30~{}m/s$ ($56~{}mph-67~{}mph$). The
vehicle control sampling frequency is set as the default value of $5~{}Hz$.
The motions of HDVs follow the IDM and MOBIL model, where the maximum
deceleration for safety purposes is limited by $b_{safe}=-9~{}m/s^{2}$,
politeness factor $p$ is 0, and the lane-changing threshold $\Delta a_{th}$ is
set as $0.1~{}m/s^{2}$. To evaluate the effectiveness of the proposed methods,
three traffic density levels are employed, which correspond to low, middle,
high levels of traffic congestion, respectively. The number of vehicles in
different traffic modes is shown in Table 1.

Table 1: Traffic density modes. Traffic density modes | AVs | HDVs | Explanation
---|---|---|---
1 | 1-3 | 1-3 | low level
2 | 2-4 | 2-4 | middle level
3 | 4-6 | 4-6 | high level

We train the MARL algorithms for 1 million steps (10,000 epochs) by applying
two different random seeds and the same random seed is shared among agents. We
evaluate each model 3 times every 200 training episodes. The parameters
$\gamma$ and learning rate $\eta$ are set as $0.99$ and $5\times 10^{-4}$,
respectively. The weighting coefficients in the reward function are set as
$\omega_{s}=200$, $\omega_{d}=4$, $\omega_{v}=1$ and $\omega_{c}=1$,
respectively. These experiments are conducted on a macOS server with a 2.7 GHz
Intel Core i5 processor and 8GB of memory.

###  4.3 Results & Analysis

####  4.3.1 Local v.s. Global Reward Designs

Fig. 3 shows the performance comparison between the proposed local reward and
the global reward design kaushik2018parameter ; dong2020drl  (with shared
actor-critic parameters). In all three traffic modes, the local reward design
consistently outperforms the global reward design in terms of larger
evaluation rewards and smaller variance. In addition, the performance gaps are
enlarged as the number of vehicles increases. This is due to the fact that the
global reward design is more likely to cause credit assignment issues as
mentioned in sutton2018reinforcement .

Figure 3: Performance comparisons between local and global reward designs. The
shaded region denotes the standard deviation over 2 random seeds.

####  4.3.2 Sharing v.s. Separate Actor-critic Network

Fig. 4 shows the performance comparison between strategies with or without
sharing the actor-critic network parameters during training. Obviously,
sharing an actor-critic network has better performance than without sharing.
Specifically, sharing actor-critic parameters in all three modes results in
higher rewards and lower variance. The reason is that, in separate actor-
critic networks, the critic network can only guide the actor network to the
correct training direction until the critic network is well-trained which may
take a long time to achieve. In contrast, the actor network can benefit from
the shared state representation via the critic network in a shared actor-
critic network chen2021deep ; graesser2019foundations .

Figure 4: Performance comparisons between with and without actor-critic
network sharing.

####  4.3.3 Verification of Driving Comfort

In this subsection, we evaluate the effectiveness of the proposed multi-
objective reward function with the driving comfort in Eqn. 3. Fig. 5 shows the
acceleration and deceleration of the AV with or without the comfort
measurement defined in Eqn. 6. It is clear that the proposed reward design
with the comfort measurement has a low variance (average deviation:
$0.237m/s^{2}$) and is more smooth than the reward design without comfort term
(average deviation: $0.582m/s^{2}$), which shows the proposed reward design
presents good driving comfort.

Figure 5: Performance comparisons of acceleration between the reward design
with or without comfort measurement.

####  4.3.4 Adaptability of the Proposed Method

In this subsection, we evaluate the proposed MA2C under different HDV
behaviors, which is controlled by the politeness coefficient $p$ denoted in
Eqn. 9, in which $p=0$ means the most aggressive behavior while $p=1$
represents the most polite behavior. Fig. 6 shows the training performance of
two different HDV models (i.e., aggressive or politeness) under different
traffic densities. It is clear that the proposed algorithm achieves scalable
and stable performance whenever the HDVs take aggressive or courteous
behaviors.

Figure 6: Performance comparisons on different politeness coefficients $p$
under different traffic densities.

####  4.3.5 Comparison with the state-of-the-art benchmarks

In order to demonstrate the performance of the proposed MARL approach, we
compared it with several state-of-the-art MARL methods:

  1. 1.

Multi-agent Deep Q-Network (MADQN) Ji : This is the multi-agent version of
Deep Q-Network (DQN) mnih2013playing , which is an off-policy RL method by
applying a deep neural network to approximate the value function and an
experience replay buffer to break the correlations between samples to
stabilize the training.

  2. 2.

Multi-agent actor-critic using Kronecker-Factored Trust Region (MAACKTR): This
is the multi-agent version of actor-critic using Kronecker-Factored Trust
Region (ACKTR) Yuhuaiwu , which is an on-policy RL algorithm by optimizing
both the actor and the critic using Kronecker-factored approximate curvature
(K-FAC) with trust region.

  3. 3.

Multi-agent Proximal Policy Optimization (MAPPO) Schulman : This is a multi-
agent version of Proximal Policy Optimization (PPO) schulman2017proximal ,
which improves the trust region policy optimization (TRPO) schulman2015trust
by using a clipped surrogate objective and adaptive KL penalty coefficient.

  4. 4.

The Proposed MA2C: This is our proposed method with the designed multi-
objective reward function design, parameter sharing and local reward design
schemes.

Table 2 shows the average return for the MARL algorithms during the
evaluation. Obviously, the proposed MA2C algorithm shows the best performance
under the density1 scenario than other MARL algorithms. It also shows
promising results on the density2 and density3 scenarios and outperforms
MAACKTR and MAPPO algorithms. Note that even though MADQN shows a better
average reward than the MA2C algorithm, it shows larger reward deviations
which may cause unstable training and safety issues. Similarly, the evaluation
curves during the training process are shown in Fig. 7. As expected, the
proposed MA2C algorithm outperforms other benchmarks in terms of evaluation
reward and reward standard deviations.

Figure 7: Performance comparisons on accumulated rewards in MADQN, MA2C, MAACKTR, and MAPPO. Table 2: Mean episode reward in different traffic flow scenario. Method | Density 1 | Density 2 | Density 3
---|---|---|---
MADQN | 47.451 | 51.568 | 48.509
| ($\pm 27.948$) | ($\pm 32.943$) | ($\pm 24.078$)
MA2C | 58.000 | 44.744 | 32.579
| ($\pm 9.308$) | ($\pm 10.895$) | ( $\pm 8.160$)
MAACKTR | 8.812 | 3.759 | 4.892
| ($\pm 6.217$) | ($\pm 10.858$) | ($\pm 10.986$)
MAPPO | 31.988 | 19.300 | 5.073
| ($\pm 6.567$) | ($\pm 16.097$) | ($\pm 19.762$)

####  4.3.6 Policy Interpretation

In this subsection, we attempt to interpret the learned AVs’ behavior. Fig. 8
shows the snapshots during testing at time steps 20, 28, and 40. As shown in
Fig. 8(a), ego vehicle ➅ attempts to make a lane change to achieve a higher
speed. To make a safe lane change, ego vehicle ➅ and ego vehicle ➆ are
expected to work cooperatively. Specially, the ego vehicle ➆ should slow down
to make space for the ego vehicle ➅ to avoid collisions, which is also
represented in Fig. 9, where the ego vehicle ➆ starts to slow down at about
20-time steps. Then the ego vehicle ➅ begins to speed up to make the lane
change as shown in Fig. 8(b) and Fig. 9. Meanwhile, the ego vehicle ➆
continues to slow down to ensure a safe headway distance with ego vehicle ➅ as
shown in Fig. 9. Fig. 8(c) shows the completed lane changes, at which time the
ego vehicle ➆ starts to speed up. This demonstration shows the proposed MARL
framework learns a reasonable and cooperative policy for ego vehicles.

(a) initial state

(b) changing lanes

(c) lane change completed

Figure 8: Lane change in simulation environment (vehicles ➀-➂: HDVs, vehicles
➅-➇: AVs). Figure 9: Speeds of the AVs ➅, ➆ and HDVs ➀.

##  5 Conclusion

In this paper, we formulated the highway lane-changing problem in mixed
traffic as an on-policy MARL problem, and extended the A2C into the multi-
agent setting, featuring a novel local reward design and parameter sharing
schemes. Specifically, a multi-objective reward function is proposed to
simultaneously promote the driving efficiency, comfort, and safety of
autonomous driving. Comprehensive experimental results, conducted on three
different traffic densities under different levels of HDV aggressiveness, show
that our proposed MARL framework consistently outperforms several state-of-
the-art benchmarks in terms of efficiency, safety and driver comfort.

Acknowledgments The authors are grateful for the efforts of our colleagues in
the Sino-German Center of Intelligent Systems, Tongji University. We are
grateful for the suggestions on our manuscript from Dr. Qi Deng.

Authors’ Contribution Wei Zhou, Dong Chen, Jun Yan, and Prof. Zhaojian Li
participated in the framework design and manuscript writing, and Wei Zhou
implemented experiments inspired by the encouragement and guidance of Dong
Chen. Prof. Huilin Yin helped revise the manuscript. Prof. Wancheng Ge is the
master supervisor of Wei Zhou in Tongji University which provides the
opportunity to complete the research work.

Funding Jun Yan is supported by the National Natural Science Foundation of
China under Grant No. 61701348 hosted by Pro. Huilin Yin. Jun Yan, Prof.
Huilin Yin, and Prof. Wancheng Ge are grateful to the generous support of TÜV
SÜD. Wei Zhou is supported by the DAAD scholarship for a dual-degree program
between Tongji University and Technical University of Munich.

Availability of data and materials Not applicable.

Code availability Not applicable.

Competing interests There are no conflicts of interest for this paper.

## References

  * (1) Brian Paden, Michal Cáp, Sze Zheng Yong, Dmitry S. Yershov, and Emilio Frazzoli.  A survey of motion planning and control techniques for self-driving urban vehicles.  IEEE Trans. Intell. Veh., 1(1):33–55, 2016.
  * (2) Divya Desiraju, Thidapat Chantem, and Kevin Heaslip.  Minimizing the disruption of traffic flow of automated vehicles during lane changes.  IEEE Transactions on Intelligent Transportation Systems, 16(3):1249–1258, 2014.
  * (3) Tingting Li, Jianping Wu, Ching-Yao Chan, Mingyu Liu, Chunli Zhu, Weixin Lu, and Kezhen Hu.  A cooperative lane change model for connected and automated vehicles.  IEEE Access, 8:54940–54951, 2020.
  * (4) Dong Chen, Longsheng Jiang, Yue Wang, and Zhaojian Li.  Autonomous driving using safe reinforcement learning by incorporating a regret-based human lane-changing decision model.  In 2020 American Control Conference (ACC), pages 4355–4361. IEEE, 2020.
  * (5) Pin Wang, Hanhan Li, and Ching-Yao Chan.  Continuous control for automated lane change behavior based on deep deterministic policy gradient algorithm.  In 2019 IEEE Intelligent Vehicles Symposium (IV), pages 1454–1460. IEEE, 2019.
  * (6) Chenyang Xi, Tianyu Shi, Yuankai Wu, and Lijun Sun.  Efficient motion planning for automated lane change based on imitation learning and mixed-integer optimization.  In 2020 IEEE 23rd International Conference on Intelligent Transportation Systems (ITSC), pages 1–6. IEEE, 2020.
  * (7) Guan Wang, Jianming Hu, Zhiheng Li, and Li Li.  Harmonious lane changing via deep reinforcement learning.  IEEE Transactions on Intelligent Transportation Systems, 2021.
  * (8) Runjia Du, Sikai Chen, Yujie Li, Jiqian Dong, Paul Young Joun Ha, and Samuel Labi.  A cooperative control framework for cav lane change in a mixed traffic environment.  arXiv preprint arXiv:2010.05439, 2020.
  * (9) Carl-Johan Hoel, Krister Wolff, and Leo Laine.  Automated speed and lane change decision making using deep reinforcement learning.  In Wei-Bin Zhang, Alexandre M. Bayen, Javier J. Sánchez Medina, and Matthew J. Barth, editors, 21st International Conference on Intelligent Transportation Systems, pages 2148–2155. IEEE, 2018.
  * (10) Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al.  Grandmaster level in starcraft ii using multi-agent reinforcement learning.  Nature, 575(7782):350–354, 2019.
  * (11) T. Chu, J. Wang, Lara Codeca, and Z. Li.  Multi-agent deep reinforcement learning for large-scale traffic signal control.  IEEE Transactions on Intelligent Transportation Systems, pages 1–10, 2019.
  * (12) Kaixiang Lin, Renyu Zhao, Zhe Xu, and Jiayu Zhou.  Efficient large-scale fleet management via multi-agent deep reinforcement learning.  In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 1774–1783, 2018.
  * (13) Shai Shalev-Shwartz, Shaked Shammah, and Amnon Shashua.  Safe, multi-agent, reinforcement learning for autonomous driving.  arXiv preprint arXiv:1610.03295, 2016.
  * (14) Jiawei Wang, Tianyu Shi, Yuankai Wu, Luis Miranda-Moreno, and Lijun Sun.  Multi-agent graph reinforcement learning for connected automated driving.
  * (15) Paul Young Joun Ha, Sikai Chen, Jiqian Dong, Runjia Du, Yujie Li, and Samuel Labi.  Leveraging the capabilities of connected and autonomous vehicles and multi-agent reinforcement learning to mitigate highway bottleneck congestion.  arXiv preprint arXiv:2010.05436, 2020.
  * (16) Praveen Palanisamy.  Multi-agent connected autonomous driving using deep reinforcement learning.  In 2020 International Joint Conference on Neural Networks (IJCNN), pages 1–7. IEEE, 2020.
  * (17) Sikai Chen, Jiqian Dong, Paul Ha, Yujie Li, and Samuel Labi.  Graph neural network and reinforcement learning for multi-agent cooperative control of connected autonomous vehicles.  Computer-Aided Civil and Infrastructure Engineering, 36(7):838–857, 2021.
  * (18) Man Lung Ho, Ping T Chan, and AB Rad.  Lane change algorithm for autonomous vehicles via virtual curvature method.  Journal of advanced Transportation, 43(1):47–70, 2009.
  * (19) Julia Nilsson, Jonatan Silvlin, Mattias Brannstrom, Erik Coelingh, and Jonas Fredriksson.  If, when, and how to perform lane change maneuvers on highways.  IEEE Intelligent Transportation Systems Magazine, 8(4):68–78, 2016\.
  * (20) Julia Nilsson, Mattias Brännström, Erik Coelingh, and Jonas Fredriksson.  Lane change maneuvers for automated vehicles.  IEEE Transactions on Intelligent Transportation Systems, 18(5):1087–1096, 2016.
  * (21) Yilun Chen, Chiyu Dong, Praveen Palanisamy, Priyantha Mudalige, Katharina Muelling, and John M Dolan.  Attention-based hierarchical deep reinforcement learning for lane change behaviors in autonomous driving.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pages 0–0, 2019.
  * (22) Jiqian Dong, Sikai Chen, Paul Young Joun Ha, Yujie Li, and Samuel Labi.  A drl-based multiagent cooperative control framework for cav networks: a graphic convolution q network.  arXiv preprint arXiv:2010.05437, 2020.
  * (23) Dong Chen, Zhaojian Li, Yongqiang Wang, Longsheng Jiang, and Yue Wang.  Deep multi-agent reinforcement learning for highway on-ramp merging in mixed traffic.  arXiv preprint arXiv:2105.05701, 2021.
  * (24) Thomas N Kipf and Max Welling.  Semi-supervised classification with graph convolutional networks.  arXiv preprint arXiv:1609.02907, 2016.
  * (25) Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller.  Playing atari with deep reinforcement learning.  CoRR, abs/1312.5602, 2013.
  * (26) Matthijs TJ Spaan.  Partially observable markov decision processes.  In Reinforcement Learning, pages 387–414. Springer, 2012.
  * (27) Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu.  Asynchronous methods for deep reinforcement learning.  In International conference on machine learning, pages 1928–1937. PMLR, 2016.
  * (28) Tianshu Chu, Sandeep Chinchali, and Sachin Katti.  Multi-agent reinforcement learning for networked system control.  arXiv preprint arXiv:2004.01339, 2020.
  * (29) Richard S Sutton and Andrew G Barto.  Reinforcement learning: An introduction.  MIT press, 2018.
  * (30) Meha Kaushik, K Madhava Krishna, et al.  Parameter sharing reinforcement learning architecture for multi agent driving behaviors.  arXiv preprint arXiv:1811.07214, 2018.
  * (31) M. Treiber, A. Hennecke, and D. Helbing.  Congested traffic states in empirical observations and microscopic simulations.  Physical Review E, 62:1805–1824, 2000.
  * (32) Arne Kesting, Martin Treiber, and Dirk Helbing.  Connectivity statistics of store-and-forward intervehicle communication.  IEEE Trans. Intell. Transp. Syst., 11(1):172–181, 2010.
  * (33) Edouard Leurent.  An environment for autonomous driving decision-making.  https://github.com/eleurent/highway-env, 2018.
  * (34) Laura Graesser and Wah Loon Keng.  Foundations of deep reinforcement learning: theory and practice in Python.  Addison-Wesley Professional, 2019.
  * (35) Guanglin Ji, Junyan Yan, Jingxin Du, Wanquan Yan, Jibiao Chen, Yongkang Lu, Juan Rojas, and Shing Shin Cheng.  Towards safe control of continuum manipulator using shielded multiagent reinforcement learning.  IEEE Robotics Autom. Lett., 6(4):7461–7468, 2021.
  * (36) Yuhuai Wu, Elman Mansimov, Roger B. Grosse, Shun Liao, and Jimmy Ba.  Scalable trust-region method for deep reinforcement learning using kronecker-factored approximation.  In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems, pages 5279–5288, 2017.
  * (37) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  CoRR, abs/1707.06347, 2017.
  * (38) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  arXiv preprint arXiv:1707.06347, 2017.
  * (39) John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz.  Trust region policy optimization.  In International conference on machine learning, pages 1889–1897. PMLR, 2015.


# First Step Advantage: Importance of Starting Right in Multi-Step Reasoning

Kushal Jain ♠ Kumar Shridhar ⋄

♠ UC San Diego ⋄ ETH Zurich

###### Abstract

Large Language Models (LLMs) can solve complex reasoning tasks by generating
rationales for their predictions. Distilling these capabilities into a
smaller, compact model can facilitate the creation of specialized, cost-
effective models tailored for specific tasks. However, smaller models often
face challenges in complex reasoning tasks and often deviate from the correct
reasoning path. We show that LLMs can guide smaller models and bring them back
to the correct reasoning path only if they intervene at the right time. We
show that smaller models fail to reason primarily due to their difficulty in
initiating the process, and that guiding them in the right direction can lead
to a performance gain of over 100%. We explore different model sizes and
evaluate the benefits of providing guidance to improve reasoning in smaller
models.

First Step Advantage: Importance of Starting Right in Multi-Step Reasoning

Kushal Jain ♠ Kumar Shridhar ⋄ ♠ UC San Diego ⋄ ETH Zurich

##  1 Introduction

Over the years, Large Language Models (LLMs) have improved their reasoning
skills by explaining their intermediate thoughts Wei et al. (2022). This
allows LLMs to transfer intermediate knowledge to student models 111We refer
to smaller models as student models and larger models as teachers. The
distinction is not based on the number of parameters, but rather on relative
size, with smaller models often referred to as students. to improve their
reasoning skills; often referred to as knowledge distillation Yuan et al.
(2023); Magister et al. (2023); Shridhar et al. (2023b); Hsieh et al. (2023).
While training student models can certainly improve their reasoning skills,
there are instances where teacher intervention remains essential to guide the
student model when it encounters uncertainty or confusion. This is similar to
the situation in a classroom, where a student can acquire knowledge
independently by learning from textbooks, but often benefits from the guidance
of a teacher Wood et al. (1976); Van de Pol et al. (2015).

Figure 1: The First Step Advantage: Figure demonstrating the effect of first
step guidance from a teacher on the student’s overall performance (right)
versus no guidance (left).

While teacher intervention can provide valuable assistance to students,
understanding _when_ and _how_ to provide guidance plays a critical role. In
general, more guidance typically leads to a consistent improvement in student
performance Saha et al. (2023), but the question arises: _should guidance at
different stages be given equal importance?_

Our observations, as shown in Figure 1, suggest that for multi-step reasoning
tasks, intervening at the first step yields the most significant benefits,
with the effects diminishing at subsequent steps. As expected, expert teachers
(models with superior performance) tend to provide more effective guidance,
resulting in better student performance.

On a mathematical dataset of multi-step word problems, we demonstrate the
effectiveness of _first step guidance_ on various combinations of teacher-
student pairs. A consistent improvement in student performance was observed
with first step guidance. We show that a smaller student model (LLaMA 13B
Touvron et al. (2023)), when correctly guided at its first step, can achieve
the same performance as a larger student model (LLaMA 70B) without any
guidance. Furthermore, our results show a continuous and upward improvement of
the student model’s performance with expert guidance, with GPT-4 OpenAI (2023)
as the teacher reaching a performance level similar to that of a human
instructor.

##  2 Related Work

Previous work has shown that it is possible to elicit reasoning abilities from
LLMs through in-context learning Wei et al. (2022); Zhou et al. (2022). A key
recipe in most methods is to spread the reasoning process over multiple
tokens, rather than expecting them to provide an immediate response token. One
way is to provide the model with the intermediate steps, or chain of thought
(CoT), that leads to the final answer Wei et al. (2022); Kojima et al. (2023);
Yang et al. (2023); Wang et al. (2023). In parallel, in-context learning has
been used to teach the model how to break a problem down into smaller, easier
sub-problems, and then solve those sub-problems that eventually lead to the
final answer Shridhar et al. (2022); Zhou et al. (2023).

However, if the problem is misinterpreted, it can lead to a cascade of errors
in subsequent steps. To counter this, several techniques have been proposed to
intervene and correct intermediate steps Welleck et al. (2022) or to provide
feedback on their own generations, essentially “self-correcting” their own
generations Madaan et al. (2023); Shridhar et al. (2023a). It is important to
note that complex reasoning and self-correcting capabilities only emerge in
very large language models. For smaller models, CoT follows a rather flat
scaling curve Wei et al. (2022). Our work presents an effective way to
transfer such reasoning capabilities into smaller models that does not require
large scale pre-training, making it more accessible for researchers with
limited compute.

While the LLM’s ability to revise its own generations may prove helpful in
many cases, it sometimes leads to worse outcomes upon refinement, requiring a
“rolling-back” to the previous output Shridhar et al. (2023a). The “rolling-
back” dilemma can be avoided if we can know when to intervene. This would be
especially beneficial in settings of teacher-student knowledge distillation,
where the teacher can generate his explanation in CoT style and distill this
knowledge into a smaller model Shridhar et al. (2023b); Magister et al.
(2023); Ho et al. (2023); Fu et al. (2023), and later guide the student
towards correct reasoning. However, such approaches do not motivate an
interactive learning environment where the student’s mistakes can be corrected
by the teacher in real-time. Saha et al. (2023) presented an approach based on
ToM Kosinski (2023); Kadavath et al. (2022), where a teacher model intervenes
in a student model only for harder questions by creating an implicit mental
model of the student’s understanding. In contrast, our work shows that it is
not necessary to intervene and help a student with the entire solution.
Rather, just starting correctly has a significant impact on the student’s
performance and avoids the need for backtracking and correcting mistakes,
saving time and effort.

Student | Teacher |
---|---|---
Model | Type | No | LLaMA 13B | LLaMA 70B | ChatGPT | GPT-4 | Human
LLaMA 7B | Pre-Trained | 10.53 | 14.86 | 19.48 | 21.00 | 23.27 | 22.74
LLaMA 7B | Fine-Tuned | 34.19 | 38.26 | 45.90 | 45.94 | 47.61 | 47.15
LLaMA 13B | Pre-Trained | 24.70 | - | 26.39 | 33.24 | 35.75 | 33.75
LLaMA 13B | Fine-Tuned | 46.24 | - | 55.34 | 59.28 | 60.50 | 61.86
LLaMA 70B | Pre-Trained | 58.90 | - | - | 63.53 | 67.40 | 67.55
LLaMA 70B | Fine-Tuned | 63.30 | - | - | 70.05 | 72.32 | 74.14
Mistral 7B | Pre-Trained | 40.25 | - | 46.17 | 48.82 | 49.50 | 50.34
MetaMath 7B | Pre-Trained | 62.69 | - | 61.48 | 66.94 | 69.52 | 65.95
MetaMath 13B | Pre-Trained | 67.85 | - | 66.79 | 71.41 | 75.51 | 73.00
Table 1: Accuracy comparison for different configurations of the student and
the teacher models. Best results are presented in bold.

##  3 Experimental Design

### Dataset

We examine our intervention on multi-step mathematical dataset : GSM8K Cobbe
et al. (2021). The dataset is a grade-school-level math word problem dataset
with a training set of 7473 samples and a test set of 1319 samples, each
requiring two to eight steps to solve.

### Setup

We used variants of LLaMA models (LLaMA 7B, LLaMA 13B and LLaMA 70B) Touvron
et al. (2023) and its variants (Mistral 7B Jiang et al. (2023) and MetaMath 7B
and 13B Yu et al. (2023))as student models. On the other hand, we used LLaMA
13B and LLaMA 70B as teachers alongside ChatGPT (gpt-3.5-turbo) Brown et al.
(2020) and GPT-4 (gpt-4) OpenAI (2023). For pre-trained models as students in
few-shot settings, 4-shot demonstrations were provided, which were chosen
randomly from the train set. Fine-tuned students were trained on the training
set with no modifications for 3 epochs. All models were evaluated in the
greedy approach (temp=0, top p=1). Fine-tuning was performed on 1 node of 8
A100 GPUs. We report the accuracy (maj@1) on the test set.

##  4 Results and Discussion

Figure 2: Accuracy of LLaMA 7B fine-tuned student model on GSM8K dataset with
correct intervention at different steps by humans. The baseline is the
represented by the dotted line with an accuracy of 34.19.

### Early Intervention is Key

Figure 2 compares the effect of intervention by humans as teachers for LLaMA
7B student model on the GSM8K dataset. Intervention at the first step (in
blue) proves to be most beneficial with maximum gains over baseline without
any interventions (dotted line). The gains go down when the intervention is
done at “step 2” (in orange) and intervening at later stages (“step 3” and
beyond) leads to diminishing returns. This is because students may have
already internalized incorrect concepts or approaches, making it harder to
correct and replace them with the correct ones. This underscores the necessity
of early and accurate guidance.

Figure 3: Accuracy comparison for No teacher vs GPT-4 as teacher for LLaMA 7B
fine-tuned model across different steps for GSM8K dataset. 2-8 represents the
number of steps needed to solve the problem. Question Gene is sewing a quilt
out of old souvenir t-shirts. He has one shirt from each vacation he has been
on. Every shirt is its own quilt block. Each row is made of blocks from a
different year of vacations. He goes on four vacations a year and has been
vacationing since he was 23 years old. He is now 34. How many quilt blocks
does he have in total?
---
No Guidance He has been on 4 vacations a year for 34 years so he has been on
4*34 = <<4*34=136>>136 vacations. Each vacation is a quilt block so he has 136
quilt blocks. Each quilt block is one shirt so he has 136 shirts. The answer
is 136. ✗
GPT4 Guidance Gene has been on 34 - 23 = <<34-23=11>>11 years of vacations, so
he has 11 rows of blocks. Each row is made of 4 blocks, so he has 11*4 =
<<11*4=44>>44 blocks. The answer is 44. ✓
Figure 4: Example demonstrating the usefulness of first step guidance. Figure
5: Comparison of curiculum style training for a 7B student model (first step
training followed by next steps training) vs fine-tuning (Fine-tuned CoT) on
GSM8K dataset.

### Starting right: first step to successful reasoning

Table 1 demonstrates the usefulness of a teacher guiding a student with first-
step guidance. LLaMA 7B pre-trained model’s performance improves by 40% (10.5
$\rightarrow$ 14.8) when a slightly larger model (LLaMA 13B) acts as a teacher
to provide first-step guidance. The performance increases monotonically with
larger and expert models (10.5 $\rightarrow$ 19.5 with LLaMA 70B as teacher
and to 21.0 with ChatGPT as teacher) and goes beyond 100% improvement with
GPT-4 (10.5 $\rightarrow$ 23.2) or Human teachers (10.5 $\rightarrow$ 22.7).
This demonstrates the importance of first-step guidance in reaching to the
correct answer. Similar results are observed for LLaMA 13B as a student where
the performance goes from 24.70 with no guidance to 35.75 with GPT-4 first
step guidance. The trend follows similarly for even a relatively larger
student model (LLaMA 70B) where the performance goes up from 58.9 without
guidance to 67.4 with GPT-4 guidance.

Next, we fine-tune all the student models on the training data to see if first
step guidance can be learned during training. Although fine-tuning improves
the baseline performance significantly, the guidance still proved useful in
boosting the performance for all models. LLaMA 7B fine-tuned model’s
performance improves by 40% (34.1 $\rightarrow$ 47.6) with expert guidance
from GPT-4 as the teacher. A qualitative example is presented in Figure 4
where if LLaMA 7B model is guided to start correctly, it will to the correct
answer. The trend is similar for LLaMA 13B fine-tuned version with a gain of
30% with GPT-4 guidance (46.2 $\rightarrow$ 60.5) and LLaMA 70B goes up by 15%
(63.3 $\rightarrow$ 72.3).

Finally, we test the applicability of first-step guidance across better
student models: Mistral 7B and MetaMath 7B and 13B. Mistral 7B achieves a
performance boost of 25% with GPT-4 as teacher (40.25 $\rightarrow$ 49.50)
while MetaMath 7B and 13B gain more than 10% each (62.69 $\rightarrow$ 69.52
and 67.85 $\rightarrow$ 75.51 respectively). In all the cases above, it is
worth noting that the guidance of GPT-4 is very close to human’s guidance and
in many cases surpasses it. This demonstrates the capabilities of GPT-4 as an
alternative to teachers in educational domains.

### Starting right helps even for longer reasoning chains

Figure 3 demonstrates the performance of LLaMA 7B model with and without first
step guidance from a teacher for different steps in GSM8K dataset. Across all
steps (2 to 8), guidance improves the performance and suggests that starting
with a solid foundation can help over longer context. However, the improvement
is higher for problems with 2 to 5 steps compared to 6 to 8, suggesting that
over a longer reasoning chain, the chances of making mistakes increase with
the increase in the number of steps required to solve it. Nonetheless,
starting right has a positive impact on longer reasoning chains too.

### Can smaller models be aligned to start better?

Each problem can be broken down into a first step and the next consecutive
steps, where the first step can serve as a guidance for the consecutive steps.
We train the student model to first learn the initial step and then fine-tune
it further to learn the next steps required to solve the problem. This two-
step training mechanism has similarity with curriculum learning Platanios et
al. (2019); Xu et al. (2020) where the simpler first step is learnt first,
followed by the subsequent more difficult steps. Figure 5 shows a drop of 40
points once the next steps are learned and overall performance gets worse than
learning all steps at once in a fine-tuning style. Since the first-step
accuracy is close to 60%, only 3/5 samples get the correct guidance during the
next steps training and we suspect this might be the reason for a lower
overall performance.

##  5 Conclusion

Distilling reasoning capabilities in smaller models is a challenging task due
to their limited abilities to learn complex reasoning strategies. To make
these skills more accessible to smaller models, we present an effective way of
first-step guidance, where LLMs can guide smaller models in the right
direction to solve a reasoning task. On a multi-step reasoning dataset, we
show the importance of starting right with a performance improvement of over
100%. Finally, our experiments reveal the quality of guidance which the
monotonically increases with the size of the expertise of the teacher model.

## References

  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D 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 Ziegler, Jeffrey Wu, Clemens Winter, Chris 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 _Advances in Neural Information Processing Systems_ , volume 33, pages 1877–1901. Curran Associates, Inc.
  * Cobbe et al. (2021) Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021.  Training verifiers to solve math word problems.
  * Fu et al. (2023) Yao Fu, Hao Peng, Litu Ou, Ashish Sabharwal, and Tushar Khot. 2023.  Specializing smaller language models towards multi-step reasoning.
  * Ho et al. (2023) Namgyu Ho, Laura Schmid, and Se-Young Yun. 2023.  Large language models are reasoning teachers.
  * Hsieh et al. (2023) Cheng-Yu Hsieh, Chun-Liang Li, Chih-Kuan Yeh, Hootan Nakhost, Yasuhisa Fujii, Alexander J. Ratner, Ranjay Krishna, Chen-Yu Lee, and Tomas Pfister. 2023.  Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes.  _ArXiv_ , abs/2305.02301.
  * Jiang et al. (2023) Albert Qiaochu Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, L’elio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. 2023.  Mistral 7b.  _ArXiv_ , abs/2310.06825.
  * Kadavath et al. (2022) Saurav Kadavath, Tom Conerly, Amanda Askell, Tom Henighan, Dawn Drain, Ethan Perez, Nicholas Schiefer, Zac Hatfield-Dodds, Nova DasSarma, Eli Tran-Johnson, Scott Johnston, Sheer El-Showk, Andy Jones, Nelson Elhage, Tristan Hume, Anna Chen, Yuntao Bai, Sam Bowman, Stanislav Fort, Deep Ganguli, Danny Hernandez, Josh Jacobson, Jackson Kernion, Shauna Kravec, Liane Lovitt, Kamal Ndousse, Catherine Olsson, Sam Ringer, Dario Amodei, Tom Brown, Jack Clark, Nicholas Joseph, Ben Mann, Sam McCandlish, Chris Olah, and Jared Kaplan. 2022.  Language models (mostly) know what they know.
  * Kojima et al. (2023) Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. 2023.  Large language models are zero-shot reasoners.
  * Kosinski (2023) Michal Kosinski. 2023.  Theory of mind might have spontaneously emerged in large language models.
  * Madaan et al. (2023) Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, Shashank Gupta, Bodhisattwa Prasad Majumder, Katherine Hermann, Sean Welleck, Amir Yazdanbakhsh, and Peter Clark. 2023.  Self-refine: Iterative refinement with self-feedback.
  * Magister et al. (2023) Lucie Charlotte Magister, Jonathan Mallinson, Jakub Adamek, Eric Malmi, and Aliaksei Severyn. 2023.  Teaching small language models to reason.
  * OpenAI (2023) OpenAI. 2023.  Gpt-4 technical report.
  * Platanios et al. (2019) Emmanouil Antonios Platanios, Otilia Stretcu, Graham Neubig, Barnabas Poczos, and Tom Mitchell. 2019.  Competence-based curriculum learning for neural machine translation.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 1162–1172, Minneapolis, Minnesota. Association for Computational Linguistics.
  * Saha et al. (2023) Swarnadeep Saha, Peter Hase, and Mohit Bansal. 2023.  Can language models teach? teacher explanations improve student performance via personalization.  In _Thirty-seventh Conference on Neural Information Processing Systems_.
  * Shridhar et al. (2023a) Kumar Shridhar, Harsh Jhamtani, Hao Fang, Benjamin Van Durme, Jason Eisner, and Patrick Xia. 2023a.  Screws: A modular framework for reasoning with revisions.
  * Shridhar et al. (2022) Kumar Shridhar, Jakub Macina, Mennatallah El-Assady, Tanmay Sinha, Manu Kapur, and Mrinmaya Sachan. 2022.  Automatic generation of socratic subquestions for teaching math word problems.
  * Shridhar et al. (2023b) Kumar Shridhar, Alessandro Stolfo, and Mrinmaya Sachan. 2023b.  Distilling reasoning capabilities into smaller language models.  In _Findings of the Association for Computational Linguistics: ACL 2023_ , pages 7059–7073, Toronto, Canada. Association for Computational Linguistics.
  * Touvron et al. (2023) Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Daniel M. Bikel, Lukas Blecher, Cristian Cantón Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony S. Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel M. Kloumann, A. V. Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, R. Subramanian, Xia Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zhengxu Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. 2023.  Llama 2: Open foundation and fine-tuned chat models.  _ArXiv_ , abs/2307.09288.
  * Van de Pol et al. (2015) Janneke Van de Pol, Monique Volman, Frans Oort, and Jos Beishuizen. 2015.  The effects of scaffolding in the classroom: support contingency and student independent working time in relation to student achievement, task effort and appreciation of support.  _Instructional Science_ , 43:615–641.
  * Wang et al. (2023) Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. 2023.  Self-consistency improves chain of thought reasoning in language models.
  * Wei et al. (2022) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoc V Le, and Denny Zhou. 2022.  Chain of thought prompting elicits reasoning in large language models.  In _Advances in Neural Information Processing Systems_.
  * Welleck et al. (2022) Sean Welleck, Ximing Lu, Peter West, Faeze Brahman, Tianxiao Shen, Daniel Khashabi, and Yejin Choi. 2022.  Generating sequences by learning to self-correct.
  * Wood et al. (1976) David Wood, Jerome S Bruner, and Gail Ross. 1976.  The role of tutoring in problem solving.  _Child Psychology & Psychiatry & Allied Disciplines_.
  * Xu et al. (2020) Benfeng Xu, Licheng Zhang, Zhendong Mao, Quan Wang, Hongtao Xie, and Yongdong Zhang. 2020.  Curriculum learning for natural language understanding.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 6095–6104, Online. Association for Computational Linguistics.
  * Yang et al. (2023) Chengrun Yang, Xuezhi Wang, Yifeng Lu, Hanxiao Liu, Quoc V. Le, Denny Zhou, and Xinyun Chen. 2023.  Large language models as optimizers.
  * Yu et al. (2023) Long Long Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T. Kwok, Zheng Li, Adrian Weller, and Weiyang Liu. 2023.  Metamath: Bootstrap your own mathematical questions for large language models.  _ArXiv_ , abs/2309.12284.
  * Yuan et al. (2023) Zheng Yuan, Hongyi Yuan, Cheng Li, Guanting Dong, Chuanqi Tan, and Chang Zhou. 2023.  Scaling relationship on learning mathematical reasoning with large language models.  _ArXiv_ , abs/2308.01825.
  * Zhou et al. (2023) Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuurmans, Claire Cui, Olivier Bousquet, Quoc Le, and Ed Chi. 2023.  Least-to-most prompting enables complex reasoning in large language models.
  * Zhou et al. (2022) Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. 2022.  Teaching algorithmic reasoning via in-context learning.