refer to “first order,” we shall mean, first order in $\kappa^{2}$, while
“second order” means order $\kappa^{4}$, etc.

In the first step we transform the molecular Hamiltonian,

${\bar{H}}_{{\rm mol},1}=U_{1}\,H_{\rm mol}\,U_{1}^{\dagger},$ (139)

where $U_{1}$ is a unitary transformation that is designed make the off-
diagonal terms of ${\bar{H}}_{{\rm mol},1}$ vanish to lowest order in
$\kappa^{2}$. All operators are expressed in the Born-Oppenheimer
representation, that is, as matrices (thus, for example, ${\bar{H}}_{{\rm
mol},1kl}$ and $U_{1,kl}$), but the subscripts are suppressed in (139). The
unitary operator $U_{1}$ is expressed in terms of an anti-Hermitian generator
$G_{1}$, $U_{1}=\exp(G_{1})$, so that, for any operator $A$, we have

$U_{1}AU_{1}^{\dagger}=A+[G_{1},A]+\frac{1}{2!}[G_{1},[G_{1},A]]+\ldots$ (140)

This series of iterated commutators turns into a power series in $\kappa^{2}$.

A single unitary transformation of the type shown in (139) is capable of
removing the off-diagonal terms only to first order, but there will remain
second-order terms. We can apply a second unitary transformation to remove
these, leaving behind third-order, off-diagonal terms. Thus to fully remove
these terms we must contemplate an infinite number of unitary transformations,
for which we write, $U=\ldots U_{3}U_{2}U_{1}$, where each $U_{n}$ has a
generator $G_{n}$. The generators $G_{n}$ turn out to be of order
$\kappa^{2n}$. In practice, the first generator $G_{1}$ is the most important,
and is responsible for most of what is described in the literature as
“nonadiabatic corrections” to the Born-Oppenheimer approximation.

It turns out that $G_{1}$ contains energy denominators of the form
$\epsilon_{k}(x)-\epsilon_{l}(x)$, where $k\in{I}$ and $l\notin{I}$. Thus when
this energy difference is of order $\kappa^{2}$ or smaller, the expansion
(140) breaks down and levels cannot be separated by adiabatic means. This
gives some quantitative meaning to the notion of “strongly coupled” levels,
which were discussed in Sec. III.3.

This sequence of unitary transformations produces Hamiltonians
${\bar{H}}_{{\rm mol},1}$, ${\bar{H}}_{{\rm mol},2}$, etc. We will write
simply ${\bar{H}}_{\rm mol}$ for ${\bar{H}}_{{\rm mol},\infty}$, so that
${\bar{H}}_{\rm mol}=UH_{\rm mol}U^{\dagger}$. We will refer to $U$ as the
“dressing transformation” and ${\bar{H}}_{\rm mol}$ as the “dressed
Hamiltonian.” The latter is block-diagonal to all orders in $\kappa^{2}$.

We can now distinguish what we will call the “original Born-Oppenheimer
representation,” what was called in Sec. V simply the “Born-Oppenheimer
representation,” from the “dressed Born-Oppenheimer representation.” There is,
of course, also the molecular representation, which was described in Sec. V.
Physical observables have different operators representing them in the
different representations. For example, the physical observables which are the
Jacobi vectors are represented by the operators ${\mathbf{X}}_{\alpha}$ in the
molecular representation, that is, the operators are multiplication by
${\mathbf{X}}_{\alpha}$. As described in Sec. V, these physical observables
are represented by the matrices ${\mathbf{X}}_{\alpha}\,\delta_{kl}$ in the
original Born-Oppenheimer representation, which we can write simply as
${\mathbf{X}}_{\alpha}$ if we remember that an identity matrix is implied. In
the dressed Born-Oppenheimer representation, however, they are represented by
the operators
$\bar{\mathbf{X}}_{\alpha}=U\,{\mathbf{X}}_{\alpha}\,U^{\dagger}$, which are
not the same as ${\mathbf{X}}_{\alpha}$. In fact, to first order in
$\kappa^{2}$, we have

$\bar{\mathbf{X}}_{\alpha}={\mathbf{X}}_{\alpha}+[G_{1},{\mathbf{X}}_{\alpha}]+\ldots.$
(141)

Similar statements can be made about the nuclear momenta, which in the
original Born-Oppenheimer representation involve the derivative couplings (see
(82)).

As for the total orbital angular momentum of the molecule, we have seen that
it is represented by ${\mathbf{L}}={\mathbf{L}}_{n}+{\mathbf{L}}_{e}$ in the
molecular representation and ${\mathbf{L}}_{n}\,\delta_{kl}$ in the original
Born-Oppenheimer representation, which we can abbreviate as simply
${\mathbf{L}}_{n}$ if we remember that it is multiplied by the identity
matrix. As for the dressed Born-Oppenheimer representation, the same physical
observable is represented by

$U\,{\mathbf{L}}_{n}\,U^{\dagger}=U\sum_{\alpha=1}^{N-1}{\mathbf{X}}_{\alpha}\times{\mathbf{P}}_{\alpha}\,U^{\dagger}=\sum_{\alpha=1}^{N-1}\bar{\mathbf{X}}_{\alpha}\times\bar{\mathbf{P}}_{\alpha}.$
(142)

But the dressing of ${\mathbf{L}}_{n}$ involves a series of commutators with
the generators $G_{n}$, such as shown in (140). The generators $G_{n}$ are
responsible for transforming the rotationally invariant Hamiltonian $H_{\rm
mol}$ to its diagonalized version, ${\bar{H}}_{\rm mol}$, which is also
rotationally invariant. The generators $G_{n}$ that do this are themselves
rotationally invariant, so they commute with angular momentum and all the
correction terms in power series like (140) vanish. Thus we have

$\sum_{\alpha=1}^{N-1}{\mathbf{X}}_{\alpha}\times{\mathbf{P}}_{\alpha}=\sum_{\alpha=1}^{N-1}\bar{\mathbf{X}}_{\alpha}\times\bar{\mathbf{P}}_{\alpha},$
(143)

to all orders of the Born-Oppenheimer expansion. The dressing does nothing to
the angular momentum ${\mathbf{L}}_{n}$, which represents physically the total
angular momentum of the molecule in both the original Born-Oppenheimer
representation and the dressed version of it.

In the case of single-surface problems the dressing transformation creates a
$1\times 1$ block ${\bar{H}}_{{\rm mol},k_{0}k_{0}}$ on the diagonal, that is
decoupled from all other levels to all orders of $\kappa$. Thus in the dressed
Born-Oppenheimer representation the solution of the Schrödinger equation is a
simple product form as seen in (14), to all orders of $\kappa^{2}$. For this
reason we suspect that there is a connection between Moyal perturbation
theory, as discussed here, and the method of “exact factorization” (Abedi,
Maitra, and Gross (2010, 2012); Cederbaum (2013); Scherrer  _et al._ (2015);
Schild, Agostini, and Gross (2016); Requist, Tandetzky, and Gross (2016);
Martinazzo and Burghardt (2022)). The possibility of such a connection is a
project for the future.

The dressed Hamiltonian ${\bar{H}}_{\rm mol}$ may be compared to $K_{\rm
mol}$, which was obtained in Sec. V.3 by throwing away off-diagonal terms.
Both are block-diagonal, but they are not the same Hamiltonians. This is
because the dressing transformation modifies the diagonal block, adding extra
terms to it. The first such term appears at second order. This term has
evidently been discovered independently several times (Moody, Shapere, and
Wilczek (1989); Weigert and Littlejohn (1993); Goldhaber (2005)), but it has
had no impact on the chemical literature. It is of order $\kappa^{4}$ and is
therefore small, but it is of the same order as terms that are routinely
discussed in connection with Born-Oppenheimer theory. We will say more about
this term in future publications.

##  VIII Discussion and Conclusions

We have presented an in-depth analysis of angular momentum in the Born-
Oppenheimer theory of polyatomic molecules, revealing exact equivalences among
its various representations. We have done this both in the electrostatic model
and when fine structure effects and electron spin are included. Several new
results are reported regarding the transformation laws under rotations of the
Hamiltonian, basis states and derivative couplings. A dressing transformation
that replaces the Born-Oppenheimer approximation reveals further exact
equivalences among representations of angular momentum.

Finally, note that we have not made any semiclassical approximations above,
and the exact equivalences described above hold rigorously. That being said,
the results do have clear implications for semiclassical calculations. In
particular, within surface hopping calculations (Fatehi  _et al._ (2011)),
there has been a long literature regarding questions of how to treat
electronic momentum and how to conserve momentum with electron translation
factors (Bates and McCarroll (1958); Schneiderman and Russek (1969); Delos
(1981); Illescas and Riera (1998)). These questions arise because the
electronic momentum is hidden in the phase conventions of the Born-Oppenheimer
representation. To that extent this article has pointed out that similar
questions can also be raised in the context of rotations and angular momentum.
This line of study will be pursued in a subsequent publication as well.

###### Acknowledgements.

RL and JR would like to thank the CHAMPS project (Chemistry and Mathematics in
Phase Space) and Bristol University for their support during the workshop,
“Exact Factorization, Geometric Phase and Bohmian Mechanics,” September 19–20,
2022, which helped advance the ideas presented in this article. JES was
supported by the National Science Foundation under Grant No. CHE-2102402.

##  Appendix A Time Reversal

Time reversal is covered in texts (Messiah (1966); Sakurai and Napolitano
(2011)) and specifically in molecular theory (Mead (1979, 1980a); Rösch
(1983); Mead (1987); Koizumi and Sugano (1995); Johnsson and Aitchison (1997);
Schön and Köppel (1998); Matsika and Yarkony (2001, 2002b, 2002a)). We prefer
an approach based on invariant subspaces.

There are at least two distinct time reversal operators relevant to this
article: the one that acts on electronic wave functions $\phi({\mathbf{r}})$,
relevant in the electrostatic model, and the one that acts on wave functions
$\phi({\mathbf{r}},m)$, relevant in the fine structure model. In the
electrostatic model time reversal acts on electronic wave functions according
to

$(T\phi)({\mathbf{r}})=\phi({\mathbf{r}})^{*},$ (144)

that is, by simple complex conjugation, so $T^{2}=+1$. In the fine structure
model the wave function $\phi({\mathbf{r}},m)$ depends on spin and time
reversal acts according to

$(T\phi)({\mathbf{r}},m)=\sum_{m^{\prime}_{1}}\ldots\sum_{m^{\prime}_{N_{e}}}\tau_{m_{1}m^{\prime}_{1}}\ldots\tau_{m_{N_{e}}m^{\prime}_{N_{e}}}\,\phi({\mathbf{r}},m^{\prime})^{*},$
(145)

where $m^{\prime}$ is a primed version of (6), where each magnetic quantum
number $m^{\prime}_{i}$, $i=1,\ldots,N_{e}$ ranges over $\pm 1/2$, and where
$\tau$ is the matrix

$\tau=e^{-i\pi\sigma_{y}/2}=\left(\begin{array}[]{cc}0&-1\\\
1&0\end{array}\right).$ (146)

The matrix $\tau$ is the spin rotation $u({\hat{\mathbf{y}}},\pi)$ in the
notation (165). We abbreviate an equation like (145) by writing

$(T\phi)({\mathbf{r}},m)=\sum_{m^{\prime}}(\tau\ldots\tau)_{mm^{\prime}}\,\phi({\mathbf{r}},m^{\prime})^{*}.$
(147)

###  A.1 Properties of $T$

Time reversal $T$ is an antiunitary operator, $T^{\dagger}T=TT^{\dagger}=1$,
that satisfies $T^{2}=+1$ in the electrostatic model or in the fine structure
model with $N_{e}={\rm even}$, and $T^{2}=-1$ in the fine structure model with
$N_{e}={\rm odd}$. These are the only properties of $T$ that are needed for
the rest of this appendix.

If ${\mathcal{H}}$ is a Hilbert space upon which $T$ acts and
${\mathcal{S}}\subset{\mathcal{H}}$ is a subspace, then we say that
${\mathcal{S}}$ is invariant under $T$ if for every
$|\psi\rangle\in{\mathcal{S}}$, $T|\psi\rangle\in{\mathcal{S}}$ (that is, $T$
maps ${\mathcal{S}}$ into itself). Important examples of invariant subspaces
include the entire Hilbert space (${\mathcal{S}}={\mathcal{H}}$) and
eigenspaces (possibly degenerate) of a Hamiltonian that commutes with $T$.
Many other examples appear in this article. As for energy eigenspaces, note
that if $H|\psi\rangle=E|\psi\rangle$, $T^{\dagger}HT=H$, and
$|\phi\rangle=T|\psi\rangle$, then $H|\phi\rangle=E|\phi\rangle$. That is, $T$
maps energy eigenstates into other energy eigenstates of the same energy. This
does not say whether $|\phi\rangle$ is linearly independent of $|\psi\rangle$.

Antiunitary operators map scalar products into their complex conjugates, that
is, if $|\psi^{\prime}\rangle=T|\psi\rangle$ and
$|\phi^{\prime}\rangle=T|\phi\rangle$, then
$\langle\phi^{\prime}|\psi^{\prime}\rangle=\langle\phi|\psi\rangle^{*}=\langle\psi|\phi\rangle$.
This in turn means that time reversal maps orthonormal frames into other
orthonormal frames, that is, if $\\{|n\rangle,n=1,\ldots,N\\}$ is an
orthonormal frame, $\langle n|m\rangle=\delta_{nm}$, and if
$|n^{\prime}\rangle=T|n\rangle$, then $\langle
n^{\prime}|m^{\prime}\rangle=\delta_{nm}$. The frame $\\{|n\rangle\\}$ need
not be complete (it need not span the whole Hilbert space). This does not say
whether the new frame $\\{|n^{\prime}\rangle\\}$ is linearly independent of
the old one $\\{|n\rangle\\}$.

In the following we will take subspaces that are invariant under $T$ and break
them down into smaller, mutually orthogonal subspaces that are also invariant
under $T$. First we note that if ${\mathcal{S}}$ is invariant under $T$, then
it is also invariant under $T^{\dagger}$, as follows from the fact that
$T^{\dagger}=\pm T$ (the same sign as in $T^{2}=\pm 1$).

Now let ${\mathcal{S}}\subset{\mathcal{H}}$ be an invariant subspace under
$T$, let ${\mathcal{A}}\subset{\mathcal{S}}$ be a subspace of ${\mathcal{S}}$
that is also invariant under $T$, and let ${\mathcal{B}}\subset{\mathcal{S}}$
be the space orthogonal to ${\mathcal{A}}$ inside ${\mathcal{S}}$, so that

${\mathcal{S}}={\mathcal{A}}\oplus{\mathcal{B}}.$ (148)

Then ${\mathcal{B}}$ is invariant under $T$. To prove this we note that a
vector $|\phi\rangle\in{\mathcal{B}}$ if and only if
$|\phi\rangle\in{\mathcal{S}}$ and $\langle\phi|\psi\rangle=0$ for all
$|\psi\rangle\in{\mathcal{A}}$. Now let
$|\phi\rangle\in{\mathcal{B}}\subset{\mathcal{S}}$ and let
$|\phi^{\prime}\rangle=T|\phi\rangle$, so that
$|\phi^{\prime}\rangle\in{\mathcal{S}}$. We wish to show that
$|\phi^{\prime}\rangle$ is orthogonal to all $|\psi\rangle\in{\mathcal{A}}$,
hence $|\phi^{\prime}\rangle\in{\mathcal{B}}$. First we note that
$\langle\phi^{\prime}|\psi\rangle=(\langle\phi|T^{\dagger})|\psi\rangle=[\langle\phi|(T^{\dagger}|\psi\rangle)]^{*}=\langle\phi|\psi^{\prime}\rangle^{*}$,
where $|\psi^{\prime}\rangle=T^{\dagger}|\psi\rangle$. But
$|\psi^{\prime}\rangle\in{\mathcal{A}}$ since ${\mathcal{A}}$ is invariant
under $T^{\dagger}$, and therefore the scalar product vanishes. Therefore
$|\phi^{\prime}\rangle\in{\mathcal{B}}$, and ${\mathcal{B}}$ is invariant
under $T$.

###  A.2 Case $T^{2}=+1$

Now we specialize to the case $T^{2}=+1$. Let ${\mathcal{S}}$ be a subspace
invariant under $T$, with $\dim{\mathcal{S}}\geq 1$. Then ${\mathcal{S}}$
possesses a 1-dimensional, invariant subspace. To prove this let
$|\psi\rangle\in{\mathcal{S}}$ be a nonzero vector and consider the two
vectors $|\psi\rangle$ and $T|\psi\rangle$. If these are linearly dependent,
then $|\psi\rangle$ spans a 1-dimensional, invariant subspace of
${\mathcal{S}}$. If they are linearly independent then
$|\psi\rangle+T|\psi\rangle$ is nonzero and spans a 1-dimensional, invariant
subspace.

Then the space inside ${\mathcal{S}}$, complementary and orthogonal to this
1-dimensional, invariant subspace, is also invariant under $T$, so, if its
dimensionality is $\geq 1$, it also possesses an invariant, 1-dimensional
subspace. Proceeding by induction, we see that if ${\mathcal{S}}$ is finite-
dimensional, then it can be decomposed into a set of mutually orthogonal,
1-dimensional subspaces, each invariant under $T$. We will assume that the
same holds when ${\mathcal{S}}$ is infinite-dimensional.

Now let $|e\rangle$ be a unit vector inside a 1-dimensional, invariant
subspace. Then $T|e\rangle=e^{i\alpha}|e\rangle$ for some phase factor
$e^{i\alpha}$, since both $|e\rangle$ and $T|e\rangle$ are bases inside the
1-dimensional subspace. Then defining
$|\phi\rangle=e^{i\alpha/2}|\psi\rangle$, we have

$T|\phi\rangle=e^{-i\alpha/2}\,T|\psi\rangle=|\phi\rangle,$ (149)

and $|\phi\rangle$ is invariant under $T$. That is, by a phase convention we
can make the basis in a 1-dimensional, invariant subspace invariant under time
reversal. This applies to each of the subspaces into which ${\mathcal{S}}$ of
the previous paragraph was decomposed, so we see that in the case $T^{2}=+1$,
a $T$-invariant subspace ${\mathcal{S}}$ always possesses a $T$-invariant
orthonormal basis.

In particular, bases can be chosen inside the eigenspaces of a Hamiltonian
that commutes with $T$ that are $T$-invariant, that is, such a Hamiltonian
always possesses a $T$-invariant energy eigenbasis.

It is easy to show that the matrix elements of a $T$-invariant operator such
as the Hamiltonian in a $T$-invariant basis are real.

Now let ${\mathcal{S}}\subset{\mathcal{H}}$ be a subspace of a Hilbert space
and let $\\{|e_{n}\rangle\\}$ and $\\{|f_{n}\rangle\\}$ be two orthonormal
bases in ${\mathcal{S}}$. Then these bases are connected by a unitary
transformation,

$|f_{n}\rangle=\sum_{m}|e_{m}\rangle\,U_{mn},$ (150)

where $U^{\dagger}U=UU^{\dagger}=1$. If however ${\mathcal{S}}$ is
$T$-invariant, as are the two bases, $T|e_{n}\rangle=|e_{n}\rangle$,
$|f_{n}\rangle=T|f_{n}\rangle$, then it is easy to show that
$U_{mn}=U^{*}_{mn}$, that is, $U$ is a real, orthogonal matrix. If
$N=\dim{\mathcal{S}}$, then $U\in O(N)$, the latter being the group that
connects choices of $T$-invariant, orthonormal bases in the case $T^{2}=+1$.
Conversely, if the basis $|e_{n}\rangle$ in (150) is $T$-invariant and if $U$
is real orthogonal, then the basis $|f_{n}\rangle$ is also $T$-invariant.

As a special case, if $N=1$, the group $O(1)$ consists of just two matrices
$(+1)$ and $(-1)$, so the choice of a $T$-invariant basis reduces to the
choice of a $\pm$ sign.

###  A.3 Case $T^{2}=-1$

Now let $T^{2}=-1$ and, as before, let ${\mathcal{S}}$ be a subspace invariant
under $T$ such that $\dim{\mathcal{S}}\geq 1$. Then ${\mathcal{S}}$ does not
possess any 1-dimensional, invariant subspaces but it does possess a
2-dimensional invariant subspace. To prove this let $|\psi\rangle\neq 0$ be a
vector in ${\mathcal{S}}$ and notice that $|\phi\rangle=T|\psi\rangle$ is also
a nonzero vector in ${\mathcal{S}}$ since ${\mathcal{S}}$ is $T$-invariant and
$T$ preserves norms. These vectors are orthogonal,

$\langle\phi|\psi\rangle=(\langle\psi|T^{\dagger})|\psi\rangle=[\langle\psi|(T^{\dagger}|\psi\rangle)]^{*}=-[\langle\psi|(T|\psi\rangle)]^{*}=-\langle\psi|\phi\rangle^{*}=-\langle\phi|\psi\rangle=0,$
(151)

where we use $T^{\dagger}=-T$, and therefore linearly independent. They span a
2-dimensional, invariant subspace of ${\mathcal{S}}$, since

$T(a|\psi\rangle+b|\phi\rangle)=a^{*}|\phi\rangle-b^{*}|\psi\rangle,$ (152)

where we use $T|\phi\rangle=T^{2}|\psi\rangle=-|\psi\rangle$.

Thus $\dim{\mathcal{S}}\geq 2$ and ${\mathcal{S}}$ possesses a 2-dimensional,
invariant subspace. If ${\mathcal{S}}$ is the eigenspace of a $T$-invariant
Hamiltonian, then this implies that all eigenvalues are at least 2-fold
degenerate (the usual statement of Kramers degeneracy). But this means that
the space inside ${\mathcal{S}}$ that is orthogonal and complementary to this
2-dimensional, invariant subspace is also invariant, so, by induction, we can
continue to split off 2-dimensional, invariant subspaces until ${\mathcal{S}}$
is exhausted (if ever). If ${\mathcal{S}}$ is finite-dimensional, this implies
that ${\mathcal{S}}$ can be decomposed into a set of mutually orthogonal,
2-dimensional, invariant subspaces; and we will assume that this can also be
done when ${\mathcal{S}}$ is infinite-dimensional. If ${\mathcal{S}}$ is
finite-dimensional, then $\dim{\mathcal{S}}=2N$ is even.

####  A.3.1 Quaternionic bases

If the vector $|\psi\rangle$ of the preceding paragraphs is a unit vector then
we may call it $|1\rangle$; and then $|\phi\rangle=T|\psi\rangle=T|1\rangle$
is also a unit vector, call it $|2\rangle$. Then the set
$\\{|1\rangle,|2\rangle\\}$ forms an orthonormal basis in the invariant
subspace that they span, such that $T|1\rangle=|2\rangle$ and
$T|2\rangle=-|1\rangle$. Doing the same for each of the 2-dimensional,
invariant subspaces into which an invariant subspace ${\mathcal{S}}$ is
decomposed, we obtain an orthonormal basis inside ${\mathcal{S}}$,
$\\{|k\mu\rangle,k=1,\ldots,N,\mu=1,2\\}$, such that

$T|k\mu\rangle=\sum_{\nu}|k\nu\rangle\,\tau_{\nu\mu},$ (153)

where $\tau$ is given by (146) and where $\dim{\mathcal{S}}=2N$. We shall call
such a basis quaternionic. (In equations like this we label the rows and
columns of $\tau$ by $1,2$; in other places by $1/2,-1/2$.)

Now let $|e;k\mu\rangle$ and $|f;l\nu\rangle$ be two quaternionic bases on a
$T$-invariant subspace ${\mathcal{S}}$. (The symbols $e$ and $f$ just
distinguish the two bases.) Since the bases are orthonormal, they must be
connected by a unitary matrix,

$|f;l\nu\rangle=\sum_{k\mu}|e;k\mu\rangle\,U_{k\mu,l\nu},$ (154)

where

$\sum_{n\sigma}U_{k\mu,n\sigma}\,U^{*}_{l\nu,n\sigma}=\delta_{kl}\,\delta_{\mu\nu}=\sum_{n\sigma}U^{*}_{n\sigma,k\mu}\,U_{n\sigma,l\nu}.$
(155)

In cases like this we shall view the matrix $U_{k\mu,l\nu}$ as a “major”
matrix that is composed of $2\times 2$ blocks that we will call “minor”
matrices. If we write simply $U_{kl}$, we shall mean the minor matrix whose
$(\mu\nu)$ component is $U_{k\mu,l\nu}$. With this understanding, (155) can be
written,

$\sum_{n}U_{kn}(U_{ln})^{\dagger}=\delta_{kl}=\sum_{n}(U_{nk})^{\dagger}\,U_{nl},$
(156)

where we use parentheses to make it clear, for example, that
$(U_{nk})^{\dagger}$ is the Hermitian conjugate of the minor matrix $U_{nk}$,
and where $\delta_{kl}$ is understood to be multiplied by the identity minor
matrix.

Now applying $T$ to both sides of (154) we obtain

$\sum_{\nu^{\prime}}|f;l\nu^{\prime}\rangle\,\tau_{\nu^{\prime}\nu}=\sum_{k\mu\mu^{\prime}}|e;k\mu^{\prime}\rangle\,\tau_{\mu^{\prime}\mu}\,U^{*}_{k\mu,l\nu}.$
(157)

We multiply this by $(\tau^{\dagger})_{\nu\sigma}$ and sum over $\nu$, to
obtain

$|f;l\sigma\rangle=\sum_{k\mu\mu^{\prime}\nu}|e;k\mu^{\prime}\rangle\,\tau_{\mu^{\prime}\mu}\,U^{*}_{k\mu,l\nu}\,(\tau^{\dagger})_{\nu\sigma}=\sum_{k\mu^{\prime}}|e;k\mu^{\prime}\rangle\,\bigl{(}\tau
U^{*}_{kl}\tau^{\dagger}\bigr{)}_{\mu^{\prime}\sigma}=\sum_{k\mu^{\prime}}|e;k\mu^{\prime}\rangle\,U_{k\mu^{\prime},l\sigma},$
(158)

where in the last step we have used (154) again. Then, since the vectors
$|e;k\mu^{\prime}\rangle$ are linearly independent, we obtain $U_{kl}=\tau
U^{*}_{kl}\tau^{\dagger}$, an equation connecting minor matrices. Multiplying
this on the left by $\tau^{\dagger}$ and on the right by $\tau$, we obtain

$\tau^{\dagger}U_{kl}\tau=U^{*}_{kl},$ (159)

showing that the minor matrices of $U$ are quaternions (see (164)).

Thus, $U$ belongs to the group $U(N,\mathbb{H})$, the set of $N\times N$
unitary matrices of quaternions. These matrices satisfy

$\sum_{n}U_{kn}\,\overline{U_{ln}}=\delta_{kl}=\sum_{n}\overline{U_{nk}}\,U_{nl},$
(160)

which is (156) written in quaternionic language. Conversely we can show that
if the basis $|e;k\mu\rangle$ is quaternionic and $U\in U(N,\mathbb{H})$, then
the basis $|f;l\nu\rangle$ defined by (154) is also quaternionic.

In the special case $N=1$, which applies to a single surface problem with an
odd number of electrons (a single Kramers doublet), the group
$U(1,\mathbb{H})$ consists of unit quaternions, those for which
$U_{11}\overline{U_{11}}=1$. As noted in Appendix B this is the group $SU(2)$.
Its role in this context was appreciated by Mead (1987).

####  A.3.2 Quaternionic matrix elements

Finally, let $A$ be a linear operator that commutes with time reversal,
$T^{\dagger}AT=A$. Then the matrix elements of $A$ in a quaternionic basis
form minor matrices that are quaternions. To prove this we consider the matrix
elements of $A$ with respect to a quaternionic basis $\\{|k\mu\rangle\\}$,

$\displaystyle A_{k\mu,l\nu}$ $\displaystyle=$ $\displaystyle\langle
k\mu|A|l\nu\rangle=\langle k\mu|(T^{\dagger}AT|l\nu\rangle)=[(\langle
k\mu|T^{\dagger})A(T|l\nu\rangle)]^{*}$ (161) $\displaystyle=$
$\displaystyle\left[\sum_{\mu^{\prime}\nu^{\prime}}(\langle
k\mu^{\prime}|\tau^{*}_{\mu^{\prime}\mu})A(|l\nu^{\prime}\rangle\tau_{\nu^{\prime}\nu})\right]^{*}=\left[\sum_{\mu^{\prime}\nu^{\prime}}(\tau^{\dagger})_{\mu\mu^{\prime}}\,A_{k\mu^{\prime},l\nu^{\prime}}\,\tau_{\nu^{\prime}\nu}\right]^{*},$

or, in terms of minor matrices, $A_{kl}=(\tau^{\dagger}A_{kl}\tau)^{*}$. Now
taking the complex conjugate of both sides and comparing to (164) we see that
$A_{kl}$ is a quaternion.

##  Appendix B Quaternions

It is well known that the quantum mechanics of systems with an odd number of
fermions is conveniently described in terms of quaternions (Dyson (1962);
Finkelstein  _et al._ (1962); Rösch (1983); Avron  _et al._ (1988); Johnsson
and Aitchison (1997); Zhang (1997); Saue and Aa Jensen (1999); De Leo and
Scolarici (2000); Sadovskii and Zhilinskii (2022)). Quaternions also play an
important role in representation theory (Simon (1996)). In this appendix we
summarize what is needed for this article. Our treatment is similar to that of
Rösch (1983).

For the purposes of this article a quaternion is a $2\times 2$ matrix of the
form

$q=a-i{\mathbf{b}}\cdot\bm{\sigma},$ (162)

where $a$ and ${\mathbf{b}}=(b_{1},b_{2},b_{3})$ are real. We denote the set
of quaternions by $\mathbb{H}$. Hamilton’s unit quaternions ${\mathbf{i}}$,
${\mathbf{j}}$, ${\mathbf{k}}$ are identified with the matrices
$-i\sigma_{i}$, $i=1,2,3$. By this definition the matrices (162) form a
representation of the algebra of quaternions (matrix multiplication and
inversion are equivalent to the same operations on quaternions, etc).

The quaternion conjugate to $q$, denoted ${\bar{q}}$, is obtained from $q$ by
the replacement ${\mathbf{b}}\to-{\mathbf{b}}$. Interpreted as a matrix, this
is the same as forming the Hermitian conjugate; therefore we will write
${\bar{q}}=q^{\dagger}$, and note that
$\overline{q_{1}q_{2}}=\overline{q_{2}}\,\overline{q_{1}}$. A quaternion $q$
is said to be real if ${\mathbf{b}}=0$, that is, ${\bar{q}}=q$. The square
magnitude of a quaternion is

$|q|^{2}={\bar{q}}q=q{\bar{q}}=a^{2}+b_{1}^{2}+b_{2}^{2}+b_{3}^{2}=\det q.$
(163)

As for complex conjugation, by $q^{*}$ we mean the complex conjugate of the
$2\times 2$ matrix (162).

A $2\times 2$ matrix $q$ is a quaternion, according to (162), if and only if

$\tau^{\dagger}q\tau=q^{*},$ (164)

where $\tau$ is given by (146). Notice that $\tau$ is the basis quaternion
${\mathbf{j}}$.

A unit quaternion $q$ is one for which $|q|^{2}=1$. The set of unit
quaternions, interpreted as matrices, is the group $SU(2)$. An arbitrary
quaternion can be written as $q=\rho u$, where $\rho\geq 0$ is real and $u$ is
an element of $SU(2)$, which is unique if $\rho>0$.

##  Appendix C Rotation Groups $SO(3)$ and $SU(2)$

We let $R\in SO(3)$ be a proper rotation, which we parameterize in axis-angle
form, $R=R({\hat{\mathbf{n}}},\theta)$, where the unit vector
${\hat{\mathbf{n}}}$ is the axis of the rotation and $\theta$ is the angle.
All proper rotations are covered if ${\hat{\mathbf{n}}}$ runs over the unit
sphere and $0\leq\theta\leq\pi$. If $0<\theta<\pi$ the axis-angle
representation is unique, but if $\theta=0$ then $R({\hat{\mathbf{n}}},0)=I$
(the identity) for all axes ${\hat{\mathbf{n}}}$, and if $\theta=\pi$ then
$R({\hat{\mathbf{n}}},\pi)=R(-{\hat{\mathbf{n}}},\pi)$. This shows that the
space of proper rotations, the group manifold $SO(3)$, is diffeomorphic to the
real projective space $\mathbb{R}P^{3}$. (Two spaces are diffeomorphic if
their points can be placed in a one-to-one correspondence in a smooth manner.
It means that the spaces are identical from a differentiable or topological
standpoint. The space $\mathbb{R}P^{3}$ is the 3-sphere $S^{3}$ with antipodal
points identified.)

The group $SU(2)$ consists of matrices that can be parameterized in axis-angle
form,

$u({\hat{\mathbf{n}}},\theta)=e^{-i\theta{\hat{\mathbf{n}}}\cdot\bm{\sigma}/2}=\cos(\theta/2)-i{\hat{\mathbf{n}}}\cdot\bm{\sigma}\sin(\theta/2),$
(165)

where $\bm{\sigma}$ is the vector of Pauli matrices. All of $SU(2)$ is covered
if the axis ${\hat{\mathbf{n}}}$ runs over the unit sphere and
$0\leq\theta\leq 2\pi$. The representation is unique if $0<\theta<2\pi$ but
when $\theta=0$, $u({\hat{\mathbf{n}}},0)=1$ for all ${\hat{\mathbf{n}}}$, and
when $\theta=2\pi$, $u({\hat{\mathbf{n}}},2\pi)=-1$ for all
${\hat{\mathbf{n}}}$. This means that the group manifold $SU(2)$ is
diffeomorphic to the 3-sphere $S^{3}$.

The projection from $SU(2)$ to $SO(3)$ is given by

$R_{ij}=\frac{1}{2}{\mathop{\textrm{tr}}}(u^{\dagger}\sigma_{i}u\sigma_{j}).$
(166)

One can show that the matrix $R$ defined by this equation belongs to $SO(3)$
if $u\in SU(2)$, so it defines a map or function $:SU(2)\to SO(3)$. We will
denote the function by $R(u)$; it is a group homomorphism,

$R(u_{1})R(u_{2})=R(u_{1}u_{2}),$ (167)

and thus $SO(3)$ forms a representation of $SU(2)$. The map (167) is two-to-
one, since $R(u)=R(-u)$. The map preserves the axis and angle, that is,
$R\bigl{(}u({\hat{\mathbf{n}}},\theta)\bigr{)}=R({\hat{\mathbf{n}}},\theta)$.

Sometimes it is desirable to invert (166), that is, given $R\in SO(3)$ we wish
to find $u$. The answer can be given by using the axis-angle parameterization;
if we write $R=R({\hat{\mathbf{n}}},\theta)$, then the two elements of $SU(2)$
that satisfy (166) are $\pm u({\hat{\mathbf{n}}},\theta)$.

##  Appendix D Some Proofs

We prove the statements made below (115), which concern the consequences of
that formula. First, suppose a basis is quaternionic at $x_{0}$,

$T|x_{0};k\mu\rangle=\sum_{\nu}|x_{0};k\nu\rangle\,\tau_{\nu\mu},$ (168)

see (153), and suppose that $|Rx_{0};k\mu\rangle$ is given by (115). Then we
have

$\displaystyle T|Rx_{0};k\mu\rangle$ $\displaystyle=$
$\displaystyle\sum_{\nu}U_{e}(u)\bigl{(}T|x_{0};k\nu\rangle\bigr{)}\bigl{(}u^{-1})^{*}_{\nu\mu}=\sum_{\nu\sigma}U_{e}(u)\,|x_{0};k\sigma\rangle\,\tau_{\sigma\nu}\,\bigl{(}u^{-1})^{*}_{\nu\mu}$
(169) $\displaystyle=$
$\displaystyle\sum_{\sigma}U_{e}(u)\,|x_{0};k\sigma\rangle\,\bigl{(}\tau
u^{-1*}\bigr{)}_{\sigma\mu},$

where in the first step we use the fact that $T$ commutes with rotations. But
since $u^{-1}\in SU(2)$ it is a quaternion and satisfies
$u^{-1*}=\tau^{\dagger}\,u^{-1}\,\tau$, see (164). Thus
$\tau\,u^{-1*}=u^{-1}\,\tau$, and (169) becomes

$\sum_{\nu\sigma}U_{e}(u)\,|x_{0};k\sigma\rangle\,\bigl{(}u^{-1})_{\sigma\nu}\,\tau_{\nu\mu}=\sum_{\nu}|Rx_{0};k\nu\rangle\,\tau_{\nu\mu}.$
(170)

Thus, the basis $|Rx_{0};k\mu\rangle$ at the rotated point $x=Rx_{0}$ is also
quaternionic.

Next, suppose a basis vector at $x_{0}$ is an energy eigenvector,

$H_{e}(x_{0})\,|x_{0};k\mu\rangle=\epsilon_{k}(x_{0})\,|x_{0};k\mu\rangle,$
(171)

where the energy does not depend on $\mu$ as indicated. Then

$\displaystyle H_{e}(Rx_{0})\,|Rx_{0};k\mu\rangle$ $\displaystyle=$
$\displaystyle
U_{e}(u)H_{e}(x_{0})U_{e}(u)^{\dagger}\sum_{\nu}U_{e}(u)\,|x_{0};k\nu\rangle\,\bigl{(}u^{-1}\bigr{)}_{\nu\mu}$
(172) $\displaystyle=$
$\displaystyle\epsilon_{k}(x_{0})\,\sum_{\nu}U_{e}(u)\,|x_{0};k\nu\rangle\,\bigl{(}u^{-1}\bigr{)}_{\nu\mu}=\epsilon_{k}(x_{0})\,|Rx_{0};k\mu\rangle,$

where in the first step we use (104) and (115). Thus, the rule (115) maps
energy eigenbases at $x_{0}$ into those at $x=Rx_{0}$, without changing the
eigenvalues.

To prove (116) we let $R_{1},R_{2}\in SO(3)$, corresponding to $u_{1},u_{2}\in
SU(2)$, and we write $x_{1}=R_{1}x_{0}$ and $x_{2}=R_{2}x_{1}$. Then we have

$\displaystyle|x_{2};k\mu\rangle$ $\displaystyle=$
$\displaystyle|R_{2}x_{1};k\mu\rangle=|R_{2}R_{1}x_{0};k\mu\rangle=\sum_{\nu}U_{e}(u_{2}u_{1})\,|x_{0};k\nu\rangle\,(u_{2}u_{1})^{-1}_{\nu\mu}$
(173) $\displaystyle=$
$\displaystyle\sum_{\nu\sigma}U_{e}(u_{2})\,U_{e}(u_{1})\,|x_{0};k\nu\rangle\,\bigl{(}u_{1}^{-1}\bigr{)}_{\nu\sigma}\,\bigl{(}u_{2}^{-1}\bigr{)}_{\sigma\mu}$
$\displaystyle=$
$\displaystyle\sum_{\sigma}U_{e}(u_{2})\,|x_{1};k\sigma\rangle\,\bigl{(}u^{-1}_{2}\bigr{)}_{\sigma\mu}.$

Now making the replacements $x_{1}\to x$, $R_{2}\to R$ and $u_{2}\to u$, we
obtain (116). The proof would not work if we had used $u$ instead of $u^{-1}$
in (116).

## References

  * Abedi, Maitra, and Gross (2010) Abedi, A., Maitra, N. T., and Gross, E. K. U., “Exact factorization of the time-dependent electron-nuclear wave function,” Phys. Rev. Lett. 105, 123002 (2010).
  * Abedi, Maitra, and Gross (2012) Abedi, A., Maitra, N. T., and Gross, E. K. U., “Correlated electron-nuclear dynamics: Exact factorization of the molecular wave function,” J. Chem. Phys. 137, 22A530 (2012).
  * Abrol and Kuppermann (2002) Abrol, R. and Kuppermann, A., “An optimal adiabatic-to-diabatic transformation of the $1{}^{2}A^{\prime}$ and $2{}^{2}A^{\prime}$ states of ${\rm H}_{3}$,” J. Chem. Phys. 116, 1035–1062 (2002).
  * Adhikari and Billing (2002) Adhikari, S. and Billing, G. D., “Non-adiabatic effects in chemical reactions: extended Born-Oppenheimer equations and its applications,” Adv. Chem. Phys. 124, 143–196 (2002).
  * Althorpe (2006) Althorpe, S. C., “General explanation of geometric phase effects in reactive systems: Unwinding the nuclear wave function using simple topology,” J. Chem. Phys. 124, 084105 (2006).
  * Althorpe (2012) Althorpe, S. C., “Symmetry analysis of geometric-phase effects in quantum dynamics,” in  _Conical Intersections, Theory, Computation and Experiment_ , edited by W. Domcke, D. R. Yarkony, and H. Köppel (World Scientific, Singapore, 2012) pp. 155–194.
  * Aquilanti and Cavalli (1986) Aquilanti, V. and Cavalli, S., “Coordinates for molecular dynamics: Orthogonal local systems,” J. Chem. Phys. 85, 1355–1361 (1986).
  * Atchity and Ruedenberg (1997) Atchity, G. J. and Ruedenberg, K., “Determination of diabatic states through enforcement of configurational uniformity,” Theor. Chem. Acc. 97, 47–58 (1997).
  * Avron  _et al._ (1988) Avron, J. E., Sadun, L., Segert, J., and Simon, B., “Topological invariants in Fermi systems with time-reversal invariance,” Phys. Rev. Lett. 61, 1329–1332 (1988).
  * Baer (1975) Baer, M., “Adiabatic and diabatic representations for atom-molecule collisions: treatment of the collinear arrangement,” Chem. Phys. Lett. 35, 112 (1975).
  * Ballhausen and Hansen (1972) Ballhausen, C. J. and Hansen, A. E., “Electronic spectra,” Annu. Rev. Phys. Chem. 23, 15 (1972).
  * Bates and McCarroll (1958) Bates, D. R. and McCarroll, R., “Electron capture in slow collisions,” Proc. Roy. Soc. Lond. A 245, 175–183 (1958).
  * Berry (1984) Berry, M. V., “Quantal phase factors accompanying adiabatic changes,” Proc. R. Soc. Lond. A 392, 45 (1984).
  * Bethe and Salpeter (1957) Bethe, H. A. and Salpeter, E. E.,  _Quantum Mechanics of One- and Two-Electron Atoms_ (Springer Verlag, Berlin, 1957).
  * Bian  _et al._ (2021) Bian, X., Wu, Y., Teh, H.-H., Zhou, Z., Chen, H.-T., and Subotnik, J. E., “Modeling nonadiabatic dynamics with degenerate electronic states, intersystem crossing, and spin separation: A key goal for chemical physics,” J. Chem. Phys. 154, 110901 (2021).
  * Bohm, Boya, and Kendrick (1991) Bohm, A., Boya, L. J., and Kendrick, B., “Derivation of the geometric phase,” Phys. Rev. A 43, 1206 (1991).
  * Bohm, Kendrick, and Loewe (1992) Bohm, A., Kendrick, B., and Loewe, M. E., “The Berry phase in molecular physics,” Int. J. Quantum Chem. 41, 53 (1992).
  * Bohm  _et al._ (1992) Bohm, A., Kendrick, B., Loewe, M. E., and Boya, L. J., “The Berry connection and Born-Oppenheimer method,” J. Math. Phys. 33, 977 (1992).
  * Born and Huang (1954) Born, M. and Huang, K.,  _Dynamical Theory of Crystal Lattices_ (Oxford University Press, New York, 1954).
  * Born and Oppenheimer (1927) Born, M. and Oppenheimer, R., “Zur Quantentheorie der Molekeln,” Ann. Phys. (Leipzig) 84, 457 (1927).
  * Cafiero and Adamowicz (2004) Cafiero, M. and Adamowicz, L., “Molecular structure in non-Born-Oppenheimer quantum mechanics,” Chem. Phys. Lett. 387, 136–141 (2004).
  * Cederbaum (2004) Cederbaum, L. S., “Born-Oppenheimer approximation and beyond,” in  _Conical Intersections: Electronic Structure, Dynamics and Spectroscopy_ , edited by W. Domcke, D. R. Yarkony, and H. Köppel (World Scientific, Singapore, 2004) p. 3.
  * Cederbaum (2013) Cederbaum, L. S., “The exact molecular wave function as a product of an electronic and a nuclear wavefunction,” J. Chem. Phys. 138, 224110 (2013).
  * Cederbaum, Schirmer, and Meyer (1989) Cederbaum, L. S., Schirmer, J., and Meyer, H. D., “Block diagonalization of Hermitian matrices,” J. Phys. A 22, 2427 (1989).
  * Child (2002) Child, M. S., “Early perspectives on geometric phase,” Adv. Chem. Phys. 124, 1 (2002).
  * Choi and Vaniček (2020) Choi, S. and Vaniček, J., “Which form of the molecular Hamiltonian is the most suitable for simulating the nonadiabatic quantum dynamics at a conical intersection,” J. Chem. Phys. 153, 211101 (2020).
  * Choi and Vaniček (2021) Choi, S. and Vaniček, J., “How important are the residual nonadiabatic couplings for an accurate simulation of nonadiabatic quantum dynamics in a quasiadiabatic representation?” J. Chem. Phys. 154, 124119 (2021).
  * De Leo and Scolarici (2000) De Leo, S. and Scolarici, G., “Right eigenvalue equation in quaternionic quantum mechanics,” J. Phys. A 33, 2971 (2000).
  * Delos (1981) Delos, J. B., “Theory of electronic transitions in slow atomic collisions,” Rev. Mod. Phys. 53, 287 (1981).
  * Delves (1960) Delves, L. M., “Tertial and general-order collisions II,” Nucl. Phys. 20, 275 (1960).
  * Domcke (2004) Domcke, W., “Generic aspects of the dynamics at conical intersections: Internal conversion, vibrational relaxation and photoisomerization,” in  _Conical Intersections: Electronic Structure, Dynamics and Spectroscopy_ , edited by W. Domcke, D. R. Yarkony, and H. Köppel (World Scientific, Singapore, 2004) p. 395.
  * Dyson (1962) Dyson, F. J., “Statistical theory of the energy levels of complex systems. I,” J. Math. Phys. 3, 140 (1962).
  * Faraji, Gómez-Carrasco, and Köppel (2012) Faraji, S., Gómez-Carrasco, S., and Köppel, H., “Multistate vibronic dynamics and multiple conical intersections,” in  _Conical Intersections, Theory, Computation and Experiment_ , edited by W. Domcke, D. R. Yarkony, and H. Köppel (World Scientific, Singapore, 2012) p. 249.
  * Fatehi  _et al._ (2011) Fatehi, S., Alguire, E., Shao, Y., and Subotnik, J. E., “Analytic derivative couplings between configuration-interaction-singles states with built-in electron-translation factors for translational invariance,” J. Chem. Phys. 135, 234105 (2011).
  * Fedorov and Levine (2019) Fedorov, D. A. and Levine, B., “A discontinuous basis enables numerically exact solution of the Schrödinger equation around conical intersections in the adiabatic representation,” J. Chem. Phys. 150, 054102 (2019).
  * Finkelstein  _et al._ (1962) Finkelstein, D., Jauch, J. M., Schiminovich, S., and Speiser, D., “Foundations of quaternionic quantum mechanics,” J. Math. Phys. 3, 207 (1962).
  * Frankel (1997) Frankel, T.,  _The Geometry of Physics: An Introduction_ (Cambridge University Press, Cambridge, UK, 1997).
  * Gamallo, González, and Petrongolo (2021) Gamallo, P., González, M., and Petrongolo, C., “Quantum dynamics of nonadiabatic Renner-Teller effects in atom-diatom collisions,” J. Phys. Chem. A 125, 6637 (2021).
  * Gatti  _et al._ (1998) Gatti, F., Iung, C., Menou, M., Justum, Y., Nauts, A., and Chapuisat, X., “Vector parameterization of the $N$-atom problem in quantum mechanics, I. Jacobi vectors,” J. Chem. Phys. 108, 8804–8820 (1998).
  * Goldhaber (2005) Goldhaber, A. S., “Newtonian adiabatics unified,” Phys. Rev. A 71, 062102 (2005).
  * Gonon  _et al._ (2017) Gonon, B., Perveaux, A., Gatti, F., Lauvergnat, D., and Lasorne, B., “On the applicability of wavefunction-free, energy-based procedure for generating first-order nonadiabatic couplings around conical intersections,” J. Chem. Phys. 147, 114114 (2017).
  * Gordon, Glezaku, and Yarkony (1998) Gordon, M. S., Glezaku, V.-A., and Yarkony, D. R., “Systematic location of intersecting seams of conical intersection in triatomic molecules: the $1{}^{2}A^{\prime}$–$2{}^{2}A^{\prime}$ conical intersection in ${\rm BH}_{2}$,” J. Chem. Phys. 108, 5657 (1998).
  * Herzberg and Longuet-Higgins (1963) Herzberg, G. and Longuet-Higgins, H. C., “Intersection of potential energy surfaces in polyatomic molecules,” Discuss. Faraday Soc. 35, 77 (1963).
  * Hess and Marian (2000) Hess, B. A. and Marian, C. M., “Relativistic effects in the calculation of electronic energy,” in  _Computational Molecular Spectroscopy_ , edited by P. Jensen and P. Bunker (John Wiley & Sons, Inc., Chichester, UK, 2000) pp. 169–220.
  * Howard and Moss (1970) Howard, B. J. and Moss, R. E., “The molecular Hamiltonian I. Non-linear molecules,” Mol. Phys. 19, 433 (1970).
  * Illescas and Riera (1998) Illescas, C. and Riera, A., “Classical outlook on the electron translation factor problem,” Phys. Rev. Lett. 80, 3029–3032 (1998).
  * Jasper  _et al._ (2006) Jasper, A. W., Nangia, S., Zhu, C., and Truhlar, D. G., “Non Born-Oppenheimer molecular dynamics,” Acc. Chem. Res. 39, 101–108 (2006).
  * Johnsson and Aitchison (1997) Johnsson, M. T. and Aitchison, I. J. R., “The $SU(2)$ instanton and the adiabatic evolution of two Kramers doublets,” J. Phys. A 30, 2085 (1997).
  * Juanes-Marcos, Althorpe, and Wrede (2005) Juanes-Marcos, J. C., Althorpe, S. C., and Wrede, E., “Theoretical study of geometric phase effects in the hydrogen-exchange reaction,” Science 309, 1227–1230 (2005).
  * Jungen (2019) Jungen, C., “The Renner-Teller effect revisited 40 years later,” J. Mol. Spec. 363, 111172 (2019).
  * Kendrick and Mead (1995) Kendrick, B. and Mead, C. A., “The geometric vector potential in molecular systems with arbitrarily many identical nuclei,” J. Chem. Phys. 102, 4160 (1995).
  * Kendrick (2004) Kendrick, B. K., “Geometric phase effects in chemical reaction dynamics,” in  _Conical Intersections: Electronic Structure, Dynamics and Spectroscopy_ , edited by W. Domcke, D. R. Yarkony, and H. Köppel (World Scientific, Singapore, 2004) p. 521.
  * Kendrick (2018) Kendrick, B. K., “Non-adiabatic quantum reactive scattering in hyperspherical coordinates,” J. Chem. Phys. 148, 044116 (2018).
  * Kendrick, Mead, and Truhlar (2002) Kendrick, B. K., Mead, C. A., and Truhlar, D. G., “Properties of nonadiabatic couplings and the generalized Born-Oppenheimer approximation,” Chem. Phys. 277, 31 (2002).
  * Koizumi and Sugano (1995) Koizumi, H. and Sugano, S., “Geometric phase in two Kramers doublets molecular systems,” J. Chem. Phys. 102, 4472 (1995).
  * Köppel (2004) Köppel, H., “Diabatic representation: Methods for the construction of diabatic electronic states,” in  _Conical Intersections: Electronic Structure, Dynamics and Spectroscopy_ , edited by W. Domcke, D. R. Yarkony, and H. Köppel (World Scientific, Singapore, 2004) p. 175.
  * Kuppermann and Abrol (2002) Kuppermann, A. and Abrol, R., “Quantum reaction dynamics for multiple electronic states,” Adv. Chem. Phys. 124, 283–322 (2002).
  * Littlejohn, Rawlinson, and Subotnik (2022) Littlejohn, R., Rawlinson, J., and Subotnik, J., “The parallel-transported (quasi)-diabatic basis,” J. Chem. Phys. 157, 184303 (2022).
  * Littlejohn and Flynn (1991) Littlejohn, R. G. and Flynn, W. G., “Geometric phases in the asymptotic theory of coupled wave equations,” Phys. Rev. A 44, 5239 (1991).
  * Longuet-Higgins (1975) Longuet-Higgins, H. C., “The intersection of potential energy surfaces in polyatomic molecules,” Proc. Roy. Soc. Lond. A 344, 147 (1975).
  * Martinazzo and Burghardt (2022) Martinazzo, R. and Burghardt, I., “Quantum theory of electronic friction,” Phys. Rev. A 105, 052215 (2022).
  * Matsika (2012) Matsika, S., “Three-state conical intersections,” in  _Conical Intersections, Theory, Computation and Experiment_ , edited by W. Domcke, D. R. Yarkony, and H. Köppel (World Scientific, Singapore, 2012) p. 83.
  * Matsika and Yarkony (2001) Matsika, S. and Yarkony, D. R., “On the effects of spin-orbit coupling on conical intersection seams in molecules with an odd number of electrons. I. Locating the seam,” J. Chem. Phys. 115, 2038 (2001).
  * Matsika and Yarkony (2002a) Matsika, S. and Yarkony, D. R., “Conical intersections and the spin-orbit interaction,” Adv. Chem. Phys. 124, 557 (2002a).
  * Matsika and Yarkony (2002b) Matsika, S. and Yarkony, D. R., “Spin-orbit coupling and conical intersections. IV. A perturbative determination of the electronic energies, derivative couplings and a rigorous diabatic representation near a conical intersection. The general case,” J. Phys. Chem. B 106, 8108 (2002b).
  * Matsunaga and Yarkony (1998) Matsunaga, N. and Yarkony, D. R., “Energies and derivative couplings in the vicinity of a conical intersection 3\. The ’most’ diabatic basis,” Mol. Phys. 93, 79–84 (1998).
  * McDonald (1988) McDonald, S. W., “Phase-space representations of wave equations with applications to the eikonal approximation for short-wavelength waves,” Phys. Rep. 158, 337–416 (1988).
  * Mead (1979) Mead, C. A., “The ‘noncrossing’ rule for electronic potential energy surfaces: The role of time-reversal invariance,” J. Chem. Phys. 70, 2276 (1979).
  * Mead (1980a) Mead, C. A., “Electronic spin-orbit interaction and the molecular Aharanov-Bohm effect,” Chem. Phys. 49, 33 (1980a).
  * Mead (1980b) Mead, C. A., “The molecular Aharanov-Bohm effect in bound states,” Chem. Phys. 49, 23 (1980b).
  * Mead (1983) Mead, C. A., “Electronic Hamiltonian, wave functions, and energies, and derivative coupling between Born-Oppenheimer states in the vicinity of a conical intersection,” J. Chem. Phys. 78, 807 (1983).
  * Mead (1987) Mead, C. A., “Molecular Kramers degeneracy and non-Abelian adiabatic phase factors,” Phys. Rev. Lett. 59, 161 (1987).
  * Mead (1988) Mead, C. A., “The Born-Oppenheimer approximation in molecular quantum mechanics,” in  _Mathematical Frontiers in Computational Chemical Physics_ , edited by D. G. Truhlar (Springer-Verlag, New York, 1988) p. 1.
  * Mead (1992) Mead, C. A., “The geometric phase in molecular systems,” Rev. Mod. Phys. 64, 51 (1992).
  * Mead and Moscowitz (1967) Mead, C. A. and Moscowitz, A., “Dipole length versus dipole velocity in the calculation of infrared intensities with Born-Oppenheimer wave functions,” Int. J. Quantum Chem. 1, 243 (1967).
  * Mead and Truhlar (1979) Mead, C. A. and Truhlar, D. G., “On the determination of Born-Oppenheimer nuclear motion wave functions including complications due to conical intersections and identical nuclei,” J. Chem. Phys. 70, 2284 (1979).
  * Messiah (1966) Messiah, A.,  _Quantum Mechanics_ (North-Holland Publishing Company, Amsterdam, 1966).
  * Moody, Shapere, and Wilczek (1989) Moody, J., Shapere, A., and Wilczek, F., “Adiabatic effective Lagrangians,” in  _Geometric phases in physics_ , edited by A. Shapere and F. Wilczek (World Scientific, Singapore, 1989) p. 160.
  * Moyal (1949) Moyal, J. E., “Quantum mechanics as a statistical theory,” Proc. Cambridge Phil. Soc. 45, 99–124 (1949).
  * Nakahara (2003) Nakahara, M.,  _Geometry, Topology and Physics_ , 2nd ed. (Taylor & Francis, New York, 2003).
  * Pacher, Cederbaum, and Köppel (1988) Pacher, T., Cederbaum, L. S., and Köppel, H., “Approximately diabatic states from block diagonalization of the electronic Hamiltonian,” J. Chem. Phys. 89, 7367 (1988).
  * Pacher, Cederbaum, and Köppel (1993) Pacher, T., Cederbaum, L. S., and Köppel, H., “Adiabatic and quasidiabatic states in a gauge theoretical framework,” Adv. Chem. Phys. 84, 293 (1993).
  * Pacher  _et al._ (1989) Pacher, T., Mead, C. A., Cederbaum, L. S., and Köppel, H., “Gauge theory and quasiadiabatic states in molecular physics,” J. Chem. Phys. 91, 7057 (1989).
  * Panati, Spohn, and Teufel (2002) Panati, G., Spohn, H., and Teufel, S., “Space-adiabatic perturbation theory in quantum dynamics,” Phys. Rev. Lett. 88, 250405 (2002).
  * Perić and Peyerimhoff (2002) Perić, M. and Peyerimhoff, S. D., “Renner-Teller effect and spin-orbit coupling in triatomic and tetraatomic molecules,” Adv. Chem. Phys. 124, 583 (2002).
  * Requist, Tandetzky, and Gross (2016) Requist, R., Tandetzky, F., and Gross, E. K. U., “Molecular geometric phase from the exact electron-nuclear factorization,” Phys. Rev. A 93, 042108 (2016).
  * Richings and Habershon (2020) Richings, G. W. and Habershon, S., “A new diabatization scheme for direct quantum dynamics: Procrustes diabatization,” J. Chem. Phys. 152, 154108 (2020).
  * Richings and Worth (2015) Richings, G. W. and Worth, G. A., “A practical diabatisation scheme for use with direct-dynamics variational multi-configuration Gaussian method,” J. Phys. Chem. A 119, 12457–12470 (2015).
  * Rösch (1983) Rösch, N., “Time-reversal symmetry, Kramers’ degeneracy and the algebraic eigenvalue problem,” Chem. Phys. 80, 1 (1983).
  * Sadovskii and Zhilinskii (2022) Sadovskii, D. A. and Zhilinskii, B. I., “Quaternionic Dirac oscillator,” J. Phys. A 55, 385204 (2022).
  * Sakurai and Napolitano (2011) Sakurai, J. J. and Napolitano, J.,  _Modern Quantum Mechanics_ , 2nd ed. (Addison-Wesley, Boston, 2011).
  * Saue and Aa Jensen (1999) Saue, T. and Aa Jensen, H. J., “Quaternion symmetry in relativistic molecular calculations: The Dirac-Hartree-Fock method,” J. Chem. Phys. 111, 6211 (1999).
  * Schaupp and Engel (2020) Schaupp, T. and Engel, V., “Born-Oppenheimer and non-Born-Oppenheimer contributions to time-dependent electron momenta,” J. Chem. Phys. 152, 204310 (2020).
  * Scherrer  _et al._ (2015) Scherrer, A., Agostini, F., Sabastiani, D., Gross, E. K. U., and Vuilleurnier, R., “Nuclear velocity perturbation theory for vibrational circular dichroism: An approach based on the exact factorization of the electron-nuclear wave function,” J. Chem. Phys. 143, 074106 (2015).
  * Schild, Agostini, and Gross (2016) Schild, A., Agostini, F., and Gross, E. K. U., “Electronic flux density beyond the Born-Oppenheimer approximation,” J. Phys. Chem. A 120, 3316–3325 (2016).
  * Schneiderman and Russek (1969) Schneiderman, S. B. and Russek, A., “Velocity-dependent orbitals in proton-on-hydrogen-atom collisions,” Phys. Rev. 181, 311–321 (1969).
  * Schön and Köppel (1998) Schön, J. and Köppel, H., “Geometric phases and quantum dynamics in spin-orbit coupled systems,” J. Chem. Phys. 108, 1503 (1998).
  * Schuurman and Yarkony (2006) Schuurman, M. S. and Yarkony, D. R., “On the characterization of three state conical intersections: A quasianalytic theory using a group homomorphism approach,” J. Chem. Phys. 124, 124109 (2006).
  * Simon (1996) Simon, B.,  _Representations of Finite and Compact Groups_ (American Mathematical Society, Providence, Rhode Island, 1996).
  * Smith (1969) Smith, F. T., “Diabatic and adiabatic representations for atomic collision problems,” Phys. Rev. 179, 111 (1969).
  * Subotnik  _et al._ (2009) Subotnik, J. E., Cave, R. J., Steele, R. P., and Shenvi, N., “The initial and final states of electron and energy transfer processes: Diabatization as motivated by system-solvent interactions,” J. Chem. Phys. 130, 234102 (2009).
  * Subotnik  _et al._ (2008) Subotnik, J. E., Yeganeh, S., Cave, R. J., and Ratner, M. A., “Constructing diabatic states from adiabatic states: Extending generalized Mulliken-Hush to multiple charge centers with Boys localization,” J. Chem. Phys. 129, 244101 (2008).
  * Teufel (2003) Teufel, S.,  _Adiabatic Perturbation Theory in Quantum Dynamics_ (Springer-Verlag, Berlin, 2003).
  * Thiel and Köppel (1999) Thiel, A. and Köppel, H., “Proposal and numerical test of a simple diabatization scheme,” J. Chem. Phys. 110, 9371 (1999).
  * Thompson and Mead (1985) Thompson, T. C. and Mead, C. A., “Adiabatic electronic energies and nonadiabatic couplings to all orders for system of three identical muclei with conical intersection,” J. Chem. Phys. 82, 2408 (1985).
  * Thompson, Truhlar, and Mead (1985) Thompson, T. C., Truhlar, D. G., and Mead, C. A., “On the form of the adiabatic and diabatic representation and the validity of the adiabatic approximation for $X_{3}$ Jahn-Teller systems,” J. Chem. Phys. 82, 2392 (1985).
  * Varshalovich, Moskalev, and Khersonskii (1988) Varshalovich, D. A., Moskalev, A. N., and Khersonskii, V. K.,  _Quantum Theory of Angular Momentum_ (World Scientific, Singapore, 1988).
  * Venghaus and Eisfeld (2016) Venghaus, F. and Eisfeld, W., “Block-diagonalization as a tool for the robust diabatization of high-dimensional potential energy surfaces,” J. Chem. Phys. 144, 114110 (2016).
  * Wang and Tucker Carrington (2000) Wang, X.-G. and Tucker Carrington, J., “A simple method for deriving kinetic energy operators,” J. Chem. Phys. 113, 7097 (2000).
  * Wang, Guan, and Yarkony (2019) Wang, Y., Guan, Y., and Yarkony, D. R., “On the impact of singularities in the two-state adiabatic to diabatic state transformation: A global treatment,” J. Phys. Chem. A 123, 9874 (2019).
  * Weigert and Littlejohn (1993) Weigert, S. and Littlejohn, R. G., “Diagonalization of multicomponent wave equations with a Born-Oppenheimer example,” Phys. Rev. A 47, 3506 (1993).
  * Wittig (2012) Wittig, C., “Geometric phase and gauge connection in polyatomic molecules,” Phys. Chem. Chem. Phys. 14, 6409 (2012).
  * Wu, Miao, and Subotnik (2020) Wu, Y., Miao, G., and Subotnik, J. E., “Chemical reaction rates for systems with spin-orbit coupling and an odd number of electrons: Does Berry’s phase lead to meaningful spin-dependent nuclear dynamics for a two state crossing?” J. Phys. Chem. A 124, 7355 (2020).
  * Wu and Subotnik (2021) Wu, Y. and Subotnik, J. E., “Electronic spin separation induced by nuclear motion near conical intersections,” Nature Communications 12, 700 (2021).
  * Yarkony (1992) Yarkony, D. R., “Spin-forbidden chemistry within the Breit-Pauli approximation,” Int. Rev. Phys. Chem. 11, 195–242 (1992).
  * Yarkony (1996) Yarkony, D. R., “Diabolical conical intersections,” Rev. Mod. Phys. 68, 985 (1996).
  * Yarkony (1997a) Yarkony, D. R., “A criterion for the confluence of two seams of conical intersection in triatomic molecules,” Theor. Chem. Acc. 98, 197 (1997a).
  * Yarkony (1997b) Yarkony, D. R., “Energies and derivative couplings in the vicinity of a conical intersection using degenerate perturbation theory and analytic gradient techniques. 1,” J. Phys. Chem. A 101, 4263 (1997b).
  * Yarkony (1999) Yarkony, D. R., “Determining the molecular Aharanov-Bohm phase angle: A rigorous approach employing a molecular properties based adiabatic to diabatic states transformation,” J. Chem. Phys. 110, 701 (1999).
  * Yarkony (2000) Yarkony, D. R., “On the adiabatic to diabatic states transformation near intersections of conical intersections,” J. Chem. Phys. 112, 2111 (2000).
  * Yarkony (2001) Yarkony, D. R., “Conical intersections: The new conventional wisdom,” J. Phys. Chem. A 105, 6277 (2001).
  * Yarkony (2004a) Yarkony, D. R., “Conical intersections: Their description and consequences,” in  _Conical Intersections: Electronic Structure, Dynamics and Spectroscopy_ , edited by W. Domcke, D. R. Yarkony, and H. Köppel (World Scientific, Singapore, 2004) p. 41.
  * Yarkony (2004b) Yarkony, D. R., “Determination of potential energy surface intersections and derviative couplings in the adiabatic representation,” in  _Conical Intersections: Electronic Structure, Dynamics and Spectroscopy_ , edited by W. Domcke, D. R. Yarkony, and H. Köppel (World Scientific, Singapore, 2004) p. 129.
  * Yarkony (2012) Yarkony, D. R., “Nonadiabatic quantum chemistry—Past, present and future,” Chem. Rev. 112, 481 (2012).
  * Zhang (1997) Zhang, F., “Quaternions and matrices of quaternions,” Linear Algebra Appl. 251, 21 (1997).
  * Zhu and Yarkony (2015) Zhu, X. and Yarkony, D. R., “On the construction of property based diabatizations: Diabolical singular points,” J. Phys. Chem. A 119, 12383 (2015).
  * Zhu and Yarkony (2016) Zhu, X. and Yarkony, D. R., “Non-adiabaticity: The importance of conical intersections,” Mol. Phys. 114, 1983 (2016).


Global.C | 81.8$\pm$0.5 | 75.9$\pm$0.4 | 83.5$\pm$0.3 | 37.3$\pm$0.8 | 23.3$\pm$2.5 | 36.7$\pm$2.7
DualT | Batch | 82.0$\pm$0.2 | 74.6$\pm$0.4 | 83.3$\pm$0.1 | 39.8$\pm$0.9 | 27.1$\pm$1.9 | 39.8$\pm$0.7
Global | 81.7$\pm$0.6 | 74.6$\pm$0.7 | 83.0$\pm$0.4 | 39.1$\pm$1.4 | 26.9$\pm$2.3 | 36.4$\pm$1.9
Global.C | 81.3$\pm$0.7 | 74.1$\pm$0.7 | 82.7$\pm$0.6 | 38.9$\pm$1.3 | 27.3$\pm$3.0 | 38.7$\pm$1.6

RENT utilizes the resampling from the mini-batch as a default setting.
However, sampling from the mini-batch could also be changed to the dataset-
level. As another ablation study on the sampling strategies of RENT, we
conduct sampling from the whole dataset level, which could be divided into
two; 1) global sampling from whole dataset, which we denote as Global, and 2)
class-wise sampling from whole dataset, which we denote as Global.C. Table 13
shows that each strategy shows robust performances over the different
experimental settings.

###  F.7 RENT vs. Sample-Selection

Apart from transition matrix based methods, sample selection methods
dynamically select a subset of training dataset during the classifier
training, by filtering out the incorrectly-labelled instances from the noisy
dataset Han et al. (2018); Yu et al. (2019); Wei et al. (2020); Bahri et al.
(2020); Cheng et al. (2020). Our method, RENT, shares the same property with
sample selection methods by recognizing the resampling of RENT as noise-
filtering procedure. Therefore, we compare the performances of RENT with
Coteaching Han et al. (2018), Jocor Wei et al. (2020), DKNN Bahri et al.
(2020) and CORES2 Cheng et al. (2020). We report the details of the baselines
as follows.

Coteaching Han et al. (2018) assumes samples with large loss are assumed to be
noisy-labelled. However, selecting samples based on the output of the trained
classifier may cause the problem of sampling only already-trained samples
again. Therefore, it utilizes two networks with same structures from different
initialization point. It selects data with small loss from other network and
train model parameter with the selected data.

Jocor Wei et al. (2020) also considers the loss value as a metric to
discriminate noisy-labelled data and they utilize two networks with same
structures, but they optimize two network at one time.

DKNN Bahri et al. (2020) considers data whose labels are different from the
K-Nearest Neighbor algorithm result as noisy-labelled data. We utilize $k=500$
as reported in the original paper.

CORES Cheng et al. (2020) select samples based on its output confidence with
regularization term making the model be more confident, since a classifier can
be uncertain for clean labels when training with noisy-labelled dataset. We
set $\beta=2$ following the original paper.

For training Coteaching and Jocor, noisy label ratio is required as input
information to decide the sample selection portion. However, since we do not
know this information, we followed the same way as DualT Yao et al. (2020) for
noisy ratio estimation, i.e. get the average of the diagonal term of $T$ as
1-noisy ratio. For $T$ estimation, we utilize the algorithm of Yao et al.
(2020).

Table 14: Comparison with Sample selection methods - Clothing1M. We did not
report results on DKNN because of too much computation.

Methods | Coteaching | Jocor | CORES2 | DKNN | RENT
---|---|---|---|---|---
Accuracy ($\%$) | 67.2$\pm$0.4 | 68.4$\pm$0.1 | 66.4$\pm$0.7 | $-$ |  69.9$\pm$0.7

Table 7 and table 14 compares test accuracies over the baselines and RENT. Our
method shows competitive performances over the baselines. It should be noted
that Coteaching and Jocor requires the training of two different classifiers,
which requires more memory space and computation than other methods.

###  F.8 When noisy dataset is class imbalanced

Solving class imbalanced dataset with noisy label can be important (Koziarski
et al., 2020; Chen et al., 2021; Huang et al., 2022). Therefore, we
experimented to show which T utilization would work well under class
imbalance. Following studies that solves class imbalanced learning Cao et al.
(2019); Zhang et al. (2023), we first make the imbalanced dataset with (the
maximum number of samples in class)/(the minimum number of samples in class)
is defined as imbalance ratio and between two classes the number of samples in
each class should increase in exponential order. Then, we flipped the label to
be noisy as we did previously.

Forward

DualT

TV

VolMinNet

Cycle

(a)

(b)

(c)

(d)

Figure 20: Performance comparison of FL, RW and RENT over class imbalance
ratio. Subcaptions represent noisy label setting and vertical captions mean
the transition matrix estimation methods. RED, GREEN and BLUE represents RENT,
RW and FL, respectively.

Figure 20 shows the performance comparison of FL, FW and RENT with several T
estimation baselines. As we can see in the figure, again, RENT shows
consistently good result over FL and RW. However, since there is no treatment
to solve class imbalance issue in all of FL, RW and RENT, the model
performances decrease with higher class imbalance ration. This direction could
be an interesting future study.

###  F.9 Time Complexity

Table 15: Iterations denotes the number of iterations per epoch.

Dataset | CIFAR10 | CIFAR100 | Clothing1M
---|---|---|---
Dataset size / Iterations | 50,000 / 391 | 50,000 / 391 | 1,000,000 / 10,000
$t_{\text{total}}$ | 15.62s | 20.44s | 2315.10s
$t_{\text{sample}}$ | 0.13s (0.83$\%$) | 0.12s (0.59$\%$) | 4.02s (0.17$\%$)

To check whether there is increment in the time complexity by RENT, Table 15
presents the wall-clock time in resampling procedure of RENT.
$t_{\text{total}}$ is the total wall-clock for a single epoch, and
$t_{\text{sample}}$ is the time only for the resampling. Given a large-scale
dataset, i.e. Clothing1M, the resampling time becomes ignorable.


#  $\mathbf{(N,K)}$-Puzzle: A Cost-Efficient Testbed for Benchmarking
Reinforcement Learning Algorithms in Generative Language Models

Yufeng Zhang  ByteDance, Inc.  Liyu Chen  ByteDance, Inc.  Boyi Liu
ByteDance, Inc.  Yingxiang Yang  ByteDance, Inc.  Qiwen Cui Work is done while
at ByteDance. University of Washington, Seattle
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>Yunzhe Tao  ByteDance, Inc.  Hongxia Yang  ByteDance, Inc.

###### Abstract

Recent advances in reinforcement learning (RL) algorithms aim to enhance the
performance of language models at scale. Yet, there is a noticeable absence of
a cost-effective and standardized testbed tailored to evaluating and comparing
these algorithms. To bridge this gap, we present a generalized version of the
24-Puzzle: the $(N,K)$-Puzzle, which challenges language models to reach a
target value $K$ with $N$ integers. We evaluate the effectiveness of
established RL algorithms such as Proximal Policy Optimization (PPO),
alongside novel approaches like Identity Policy Optimization (IPO) and Direct
Policy Optimization (DPO).

$\mathbf{(N,K)}$-Puzzle: A Cost-Efficient Testbed for Benchmarking
Reinforcement Learning Algorithms in Generative Language Models

##  1 Introduction

Reinforcement learning (RL) is an important component in the training process
of language models (LMs) (Achiam et al., 2023; Team et al., 2023; Touvron et
al., 2023a, b; Bi et al., 2024; Shao et al., 2024). Integrated into the
training process subsequent to the supervised fine-tuning (SFT) phase, RL
significantly enhances generative LMs’ performance by incorporating feedback
from sources such as human labelers (Stiennon et al., 2020; Ouyang et al.,
2022), other generative LMs (Lee et al., 2023), or automated test environments
(Le et al., 2022). However, despite its importance, there lacks a benchmark
that is targeted at isolating and testing the performance of the RL phase of
the training alone. As a matter of fact, the prevailing spectrum of benchmarks
that evaluate the LMs capabilities (Hendrycks et al., 2020, 2021; Chen et al.,
2021; Cobbe et al., 2021) contain complex and multifaceted influencing factors
ranging from the capability of the pretrained models and the SFT models to the
datasets and algorithms. These factors prevents analyzing the impact of RL in
an isolated fashion, and a cost-effective and standardized testbed for the
comparison of different RL methods applied to generative LMs is still arguably
missing.

In this paper, we introduce the $(N,K)$-Puzzle, a generalized version of the
classic 24-Puzzle, as an efficient and streamlined testbed to evaluate RL
strategies in training generative LMs. Our framework offers flexibility
through adjustable parameters $N$ and $K$, making it capable to systematically
explore the efficacy, performance, and scalability of RL strategies.
Additionally, manipulating $N$ and $K$ enables a thorough assessment of the
generalization abilities of different RL methods, which provides valuable
insights into their potential applicability and effectiveness in more complex
tasks.

##  2 Background

###  2.1 Reward Model Training

We train a reward model (RM) from a preference dataset
$\mathcal{D}_{\text{pref}}=\\{(x^{(i)},y_{w}^{(i)},y_{l}^{(i)})\\}_{i=1}^{N}$,
where $x^{(i)}$ is the prompt, $y_{w}^{(i)}$ is the chosen (win) response, and
$y_{l}^{(i)}$ is the rejected (lose) responses. We parameterize the RM as
$r_{\phi}(x,y)$ and learn the parameters $\phi$ by minimizing the negative
log-likelihood loss

$\displaystyle-\underset{(x,y_{w},y_{l})\sim\mathcal{D}_{\text{pref}}}{\mathbb{E}}\Bigl{[}\log\sigma\bigl{(}r_{\phi}(x,y_{w})-r_{\phi}(x,y_{l})\bigr{)}\Bigr{]},$

where $\sigma$ is the sigmoid function.

###  2.2 Reinforcement Learning with an RM

The RL phase trains on a dataset $\mathcal{D}=\\{x^{(i)}\\}$ where $x^{(i)}$
is a prompt. We denote the LM policy as $\pi_{\theta}$ and optimize $\theta$
through a KL regularized objective:

$\displaystyle\underset{x\sim\mathcal{D},y\sim\pi_{\theta}(\cdot|x)}{\mathbb{E}}\bigl{[}r_{\phi}(x,y)\bigr{]}-\beta\cdot\mathbb{D}_{\mathrm{KL}}(\pi_{\theta}\|\pi_{\mathrm{ref}}),$
(1)

where $\beta>0$ controls the deviance between the trained model $\pi_{\theta}$
and the reference model $\pi_{\mathrm{ref}}$. A standard approach (Ouyang et
al., 2022; Ziegler et al., 2019) to solve (1) is to initialize the LM policy
$\pi_{\theta}$ to the SFT model $\pi^{\mathrm{SFT}}$ and maximize the
objective with PPO (Schulman et al., 2017).

###  2.3 Reinforcement Learning without an RM

Recently developed RL algorithms, such as direct preference optimization (DPO)
(Rafailov et al., 2023) and identity preference optimization (IPO) (Azar et
al., 2023), optimize the policy directly from the preference dataset without
learning a reward model. Particularly, DPO combines the training of the RM and
the reinforcement learning by minimizing a unified loss

$\displaystyle\underset{(x,y_{w},y_{l})\sim\mathcal{D}_{\text{pref}}}{\mathbb{E}}\Bigl{[}\log\sigma\bigl{(}\beta\cdot
h_{\pi_{\theta},\pi_{\mathrm{ref}}}(y_{w},y_{l},x)\bigr{)}\Bigr{]},$ (2)

where

$\displaystyle
h_{\pi_{\theta},\pi_{\mathrm{ref}}}(y_{w},y_{l},x)=\log\frac{\pi_{\theta}(y_{w}|x)\pi_{\mathrm{ref}}(y_{l}|x)}{\pi_{\theta}(y_{l}|x)\pi_{\mathrm{ref}}(y_{w}|x)}.$

Alternatively, Azar et al. (2023) show that DPO suffers from weak
regularization and overfitting and generalize the DPO objective to
$\Psi$-preference optimization ($\Psi$PO). By setting $\Psi$ to identity, Azar
et al. (2023) propose IPO, which optimizes $\theta$ to minimize the following
loss,

$\displaystyle\underset{(x,y_{w},y_{l})\sim\mathcal{D}_{\text{pref}}}{\mathbb{E}}\Bigl{[}\Bigl{(}h_{\pi_{\theta},\pi_{\mathrm{ref}}}(y_{w},y_{l},x)-\frac{1}{2\beta}\Bigr{)}^{2}\Bigr{]},$
(3)

where $\beta>0$ is a regularization parameter.

In the remainder of the paper, we test and compare these aforementioned
algorithms on the $(N,K)$-Puzzle, which we now introduce.

##  3 Problem Setup: $\mathbf{(N,K)}$-Puzzle

The $(N,K)$-Puzzle generalizes the classic 24-Puzzle, where arithmetic
operations ($+$, $-$, $*$, $/$) are utilized to manipulate a set of $N$
numbers, $\\{x_{1},\ldots,x_{N}\\}$, to achieve a target value $K$. On top of
generating the correct answer, the generative LM is also expected to output a
sequence of arithmetic operations is mathematically correct. Aligning with the
classic 24-Puzzle setting, we restrict $x_{i}$ to between 1 and 13. An
illustrative example of the prompt and the response for the $(4,24)$-Puzzle
problem is presented below.

Prompt: 24;2,3,4,6:

Response: 4+6=10,10-2=8,8*3=24

This setup challenges the model’s computational abilities, its capacity for
logical reasoning, and its ability to generate coherent, step-by-step
solutions to mathematically grounded problems. In our study, we explore the
model’s performance across a range of scenarios where the number of operands
$N$ varies from 3 to 4, and the target value $K$ is chosen from the set
$\\{13,18,24,27,34\\}$. This investigation is designed to assess the model’s
adaptability and robustness in handling variations in problem complexity and
objectives.

We conduct evaluations in two distinct test set: in-distribution test set and
out-of-distribution (OOD) test set. In particular, for the in-distribution
test set, while the prompts are not presented in the training set, the pair
$(N,K)$ that defines the puzzle is presented to the model in training. For the
OOD test set, the pair $(N,K)$ is not presented to the model in training.
These two test sets allow us to comprehensively assess the model’s capability
from both the in-distribution generalization and the OOD generalization. In
particular, we assess the difficulty of each $(N,K)$-Puzzle and choose
$(3,18)$-Puzzle and $(4,27)$-Puzzle for the OOD test set, which are at the
medium difficulty among all $(N,K)$-Puzzles.

##  4 Experiments

Figure 1: Training of the model: The process of training consists of multiple
parts. Firstly, we train the model with SFT to align the model with the
$(N,K)$-Puzzle. Then, we perform RL training (PPO, DPO, IPO, and RM) based on
the SFT model.

###  4.1 Experiment Setup

Model. We utilize the GPT-2 (Radford et al., 2019) architecture. The specific
instantiation of the GPT-2 model employed in our study is defined by the
following parameters: the model comprises $n_{\text{layer}}=12$ transformer
layers, with each layer featuring $n_{\text{head}}=12$ attention heads. The
dimensionality of the embeddings and the hidden layers is set to
$n_{\text{embd}}=768$. This results in a total of approximately 124 million
trainable parameters.

Supervised fine-tuning. The initial phase of our training is the supervised
fine-tuning (SFT). We conduct two stages of SFT from scratch: format SFT and
target SFT. In format SFT, we train the model from scratch using specially
formatted data to inculcate the ability to output results in a predefined
format and to execute arithmetic operations. Note that during the format SFT
phase, the model’s outputs may not always align perfectly with the target
results. In the format SFT phase, we employ a training dataset consisting of
200,000 data points, encompassing a broad spectrum of arithmetic operations.
The model is trained for 20 epochs with a global batch size of 256. We use a
learning rate of $10^{-5}$ and a constant learning rate scheduler. Upon
completion of this phase, the model demonstrates a remarkable capability to
generate responses that outputs responses with correct arithmetic operations
and desired formats an accuracy rate of 99%.

Building on the format SFT phase, we train the model further through target
SFT. Initiated from the checkpoint from the format SFT, the objective is to
hone the model’s capabilities not only to execute arithmetic operations within
the specified format accurately but also to engage in logical reasoning
processes that ensure the final result aligns with the target value $K$.
During the target SFT phase, we employ a training dataset consisting of
300,000 data points. The model is trained for 10 epochs with a global batch
size of 128. We use a learning rate of $10^{-5}$ and a constant learning rate
scheduler. Upon completion of this phase, the model achieves an accuracy of
43.5% in the in-distribution test set and an accuracy of 8.8% in the OOD test
set.

Ground truth reward. We define the ground truth reward function $r^{*}(x,y)$
to evaluate the responses generated by the model as follows,

$\displaystyle r^{*}(x,y)=\begin{cases}1.0,\text{~{} if }y\text{ is completely
correct},\\\ 0.1,\text{~{} if }y\text{ is correct in format},\\\ 0.0,\text{~{}
otherwise}.\end{cases}$

To further illustrate the mechanism of the ground truth reward function, we
include examples of responses with different ground truth rewards in §A. This
ground truth reward function $r^{*}(x,y)$ is employed to label the preference
dataset for training RM, DPO, and IPO. Furthermore, it is utilized in
experiments involving PPO to establish a performance upper bound for PPO with
RM.

| training | test | OOD | OOD (n=20)
---|---|---|---|---
RM | $62.5$ | $60$ | $10.5$ | $12.7$
GT | $64.6$ | $61.9$ | $14.4$ | $26.9$
Table 1: Accuracy (in percentage) of Best-of-$n$ on training dataset ($n=5$),
in-distribution test set ($n=5$), and OOD test set ($n=5,20$).

###  4.2 Reward Model

The training data of our RM is generated as follows: for each prompt in the
training set of the SFT model, we collect 10 unique responses from the SFT
model, along with their ground truth rewards. If all these 10 responses have
the same reward, we sample another 10 responses to make sure that the
collected responses always provide training signals. We then train a RM using
this training set with batch size $128$ for $500$ iterations. Each batch
contains $128$ prompts, and we group all responses for each prompt into the
same batch. To test the accuracy of the RM, we perform best-of-$n$ with RM
model and the SFT model. Specifically, for each prompt in the dataset, we
sample $n$ responses from the SFT model, and pick the response with highest
reward based on the RM or ground truth reward (GT) as the final response. Our
results are shown in Table 1.

The performance gap between RM and ground truth reward is small in training
and in-distribution test set, and it is much larger in the OOD dataset,
especially when $n$ is large. We also note that best-of-$n$ greatly boosts the
performance of SFT model especially in the in-distribution test set (from
$43.4\%$ to $60\%$). Intuitively, best-of-$n$ accuracy can be treated as an
upper bound of model performance after RL training with this RM.

Figure 2: Performance of PPO with ground-truth reward and RM. While PPO with
the ground truth reward keeps boosting both the in-distribution and OOD
accuracies, PPO with RM start to see performance degradation after a short
period of training.

###  4.3 PPO

In contrast to previous studies, we set $\beta=0$ to eliminate the KL
regularization against the SFT model during PPO training. This is not a matter
of concern when training with a ground-truth reward, as it consistently
provides an accurate learning signal. However, during training with a RM, we
found KL regularization to be ineffective: a strong KL regularization impedes
the model’s learning, while insufficient KL regularization leads to a drop in
performance. In response to this, we propose a novel type of regularization.
The underlying concept is to constrain the policy to always remain within the
top-$p$ vocabulary of the reference model. We defer the full implementation
details of PPO with extra discussion to Appendix B.

Results. Results are shown in Figure 2. We see that in both in-distribution
and OOD dataset, PPO with ground truth reward consistently improves model
performance as the number of training steps increase.

PPO with RM enhances model performance in the initial 2,500 steps. However,
post this phase, performance declines below baseline as the model begins to
intentionally generate responses that were not encountered by the RM during
its training, effectively “hacking” the RM. More detailed evidence of model
“hacking” the RM is included in Appendix B.

###  4.4 DPO and IPO

The first step is to construct the preference dataset
$\mathcal{D}_{\text{pref}}=\\{(x^{(i)},y_{w}^{(i)},y_{l}^{(i)})\\}_{i=1}^{N}$.
We begin with generating responses based on a given prompt $x$ using the SFT
model with a temperature of $1$ to encourage diversity in the responses. Each
generated response is then evaluated using the ground truth reward function
$r^{*}(x,y)$ to determine its quality. If a response is assigned a reward less
than $1.0$, indicating either a format discrepancy or incorrect reasoning, it
is included in the preference dataset as a rejected response $y_{l}$. The
corresponding chosen response $y_{w}$ is sourced from the SFT dataset,
ensuring that $y_{w}$ adheres to both the correct format and reasoning
process. Through this method, we construct a preference dataset of size
$86,000$. We initialize the model from the SFT model and train it for 2 epochs
with a batch size of 128. Throughout this phase, we employ a learning rate of
$5\times 10^{-7}$ with a constant scheduler.

Regularization of DPO and IPO. In our experimental analysis of DPO and IPO
methodologies, a critical observation is regarding the regularization
practices inherent to each approach. This inspires us to conduct a detailed
investigation into the role of the regularization parameter $\beta$ within DPO
and IPO frameworks.

The results of this investigation are shown in Figure 3. We see that a larger
$\beta$ value is requisite for DPO to attain stability and mitigate
overfitting in the in-distribution test. However, even with an optimized
$\beta$, DPO’s performance in OOD test set failed to surpass that of the
baseline SFT model, showing the limitations of DPO’s OOD generalization
ability. In comparison, IPO demonstrates better performance relative to SFT
across both in-distribution and OOD test set, which is not sensitive to the
$\beta$ values explored. This resilience and performance edge of IPO can be
attributed to its inherently stronger regularization framework, which
effectively balances the model’s fidelity to the training data with its
capacity to generalize to the OOD setting.

##  5 Conclusion

Leveraging the $(N,K)$-Puzzle as a testbed revealed some valuable insights
into RL training for generative LMs. While PPO with ground truth rewards
consistently improved performance, PPO with a trained RM suffered late-stage
performance drops due to the RM being “hacked”. In contrast, although DPO and
IPO avoid RM training for simplicity of implementation, we observe limited
generalization from in-distribution to out-of-distribution prompts,
constraining their potential in developing versatile generative LMs. We see
significant potential for further exploration of the $(N,K)$-Puzzle as a
testbed, which serves as a standardized and cost-effective environment for
identifying the most effective strategies for RL training in generative LMs.

Figure 3: Test accuracy for DPO and IPO for different values of regularization
parameter $\beta$. We observe that IPO is more robust across different
$\beta$. It is worth noting that both DPO and IPO fail to enhance
generalization of the generative LMs.

##  6 Ethical Statement

In this study, we explore the application of Reinforcement Learning (RL)
within Large Language Models (LLMs) through the lens of a generalized version
of the 24-Puzzle. In our opinion, this domain appears innocuous and devoid of
direct ethical considerations due to its abstract and game-based nature.

##  7 Limitations

One of the primary limitations of the current study lies in the scale of the
LMs employed. Our research utilizes models of a relatively modest size, which,
while sufficient in the context of a generalized 24-Puzzle, may not fully
capture the complexities and capabilities inherent to larger-scale LMs. The
constrained model size can limit the depth and breadth of language
understanding and generation, potentially impacting the generalizability of
our findings to more sophisticated tasks and broader applications.

## References

  * Achiam et al. (2023) Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. 2023.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_.
  * Azar et al. (2023) Mohammad Gheshlaghi Azar, Mark Rowland, Bilal Piot, Daniel Guo, Daniele Calandriello, Michal Valko, and Rémi Munos. 2023.  A general theoretical paradigm to understand learning from human preferences.  _arXiv preprint arXiv:2310.12036_.
  * Bi et al. (2024) Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. 2024.  Deepseek llm: Scaling open-source language models with longtermism.  _arXiv preprint arXiv:2401.02954_.
  * Chen et al. (2021) Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. 2021.  Evaluating large language models trained on code.  _arXiv preprint arXiv:2107.03374_.
  * Cobbe et al. (2021) Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. 2021.  Training verifiers to solve math word problems.  _arXiv preprint arXiv:2110.14168_.
  * Hendrycks et al. (2020) Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020.  Measuring massive multitask language understanding.  _arXiv preprint arXiv:2009.03300_.
  * Hendrycks et al. (2021) Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021.  Measuring mathematical problem solving with the math dataset.  _arXiv preprint arXiv:2103.03874_.
  * Le et al. (2022) Hung Le, Yue Wang, Akhilesh Deepak Gotmare, Silvio Savarese, and Steven Chu Hong Hoi. 2022.  Coderl: Mastering code generation through pretrained models and deep reinforcement learning.  _Advances in Neural Information Processing Systems_ , 35:21314–21328.
  * Lee et al. (2023) Harrison Lee, Samrat Phatale, Hassan Mansoor, Kellie Lu, Thomas Mesnard, Colton Bishop, Victor Carbune, and Abhinav Rastogi. 2023.  Rlaif: Scaling reinforcement learning from human feedback with ai feedback.  _arXiv preprint arXiv:2309.00267_.
  * Loshchilov and Hutter (2017) Ilya Loshchilov and Frank Hutter. 2017.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_.
  * Ouyang et al. (2022) Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. 2022.  Training language models to follow instructions with human feedback.  _Advances in Neural Information Processing Systems_ , 35:27730–27744.
  * Radford et al. (2019) Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019.  Language models are unsupervised multitask learners.
  * Rafailov et al. (2023) Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D Manning, and Chelsea Finn. 2023.  Direct preference optimization: Your language model is secretly a reward model.  _arXiv preprint arXiv:2305.18290_.
  * Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_.
  * Shao et al. (2024) Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, YK Li, Y Wu, and Daya Guo. 2024.  Deepseekmath: Pushing the limits of mathematical reasoning in open language models.  _arXiv preprint arXiv:2402.03300_.
  * Stiennon et al. (2020) Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. 2020.  Learning to summarize with human feedback.  _Advances in Neural Information Processing Systems_ , 33:3008–3021.
  * Team et al. (2023) Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. 2023.  Gemini: a family of highly capable multimodal models.  _arXiv preprint arXiv:2312.11805_.
  * Touvron et al. (2023a) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. 2023a.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_.
  * Touvron et al. (2023b) Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023b.  Llama 2: Open foundation and fine-tuned chat models.  _arXiv preprint arXiv:2307.09288_.
  * Ziegler et al. (2019) Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. 2019.  Fine-tuning language models from human preferences.  _arXiv preprint arXiv:1909.08593_.

##  Appendix A Examples of Responses with Different Ground Truth Reward

Given the following prompt,

Prompt: 24;2,3,4,6:

we include examples of responses with different ground truth rewards as
follows,

  * •

Response with reward 1.0:
4+6=10,10-2=8,8*3=24

  * •

Response with reward 0.1:
4+6=10,10-2=8,8+3=11

  * •

Response with reward 0.0:
4+6=10,10-2=8,8*3=22

##  Appendix B Detail and Discussion of PPO

In all of our PPO experiments, we use policy clip range $\epsilon_{\rm
pg}=0.15$, value function clip range $\epsilon_{\rm vf}=10.0$, GAE parameter
$\lambda=1.0$, and value function parameter $0.1$. On top of the standard
formulation in (1), we employ entropy regularization on the LM policy
$\pi_{\theta}$ with a weight of $0.04$ to encourage exploration.

We train the LM policy for 2.56M episodes (prompt-response pairs). When
optimizing the clipped PPO objective, we use AdamW (Loshchilov and Hutter,
2017) with minibatches of size $512$ each. The learning rate decays at the
rate of $0.98$ every $10$ PPO epochs (or $100$ steps).

###  B.1 Training Details of PPO with Ground Truth Reward

We train the LM policy for 2.56M episodes (prompt-response pairs). We use
$5120$ episodes per PPO batch, which is optimized using AdamW (Loshchilov and
Hutter, 2017) for one epoch with $10$ minibatches of size $512$ each.

Hyperparameter | Value
---|---
learning rate | $1e-5$
learning rate decay ratio | $0.98$
discount factor $\gamma$ | $0.99$
GAE $\lambda$ | $1.0$
PPO clip ratio | $0.15$
batch size | $512\times 10$
entropy loss coefficient | $0.04$
minibatch size | $512$
Table 2: Hyperparameters of PPO with ground truth reward.

###  B.2 Training Details of PPO with RM

The hyperparameters used for training PPO with a reward model is in Table 3.
We found that KL regularization is not very effective in our task: it is
either too strong and stop the algorithm from learning, or it is too weak so
that the performance drops. To resolve this, we propose a weaker
regularization: for each response generated by the current policy, we fit it
into the SFT model and determine whether each token of the response is in SFT
model’s top-$p$ vocabulary given input before this token. If any of the token
is out of the top-$p$ vocabulary, we change the reward of current response to
the value of OOD penalty. Intuitively, such regularization ensures that the
trained policy only generates responses within the support of nucleus sampling
of the original SFT policy, but it does not force the distribution of the
trained policy to be close to the SFT policy. This ensures that the support of
the learnt policy distribution does not leave the “trust region” of the reward
model. Moreover, it does not force the policy distribution to be close to that
of reference model, which might impede learning.

Hyperparameter | Value
---|---
learning rate | $1e-6$
learning rate decay ratio | $0.98$
discount factor $\gamma$ | $0.93$
GAE $\lambda$ | $1.0$
PPO clip ratio | $0.1$
batch size | $512\times 88$
entropy loss coefficient | $0.04$
minibatch size | $512$
top-$p$ regularization | $0.95$
OOD penalty | $-40$
Table 3: Hyperparameters of PPO with RM

###  B.3 Examples of Reward Model Hacking

Figure 4: Training dynamics of PPO with reward model without KL or top-$p$
regularization.

We also include samples generated from a policy trained by PPO without KL or
top-$p$ regularization, whose performance quickly drops to zero as training
proceeds (see Figure 4). On inspecting the samples, we find that the reward
model is completely ‘hacked’: many incorrect samples have very high rewards.

episode:24;3,3,7,6:6*3=18,7-3=4,18+4=22reward model:3.284614324569702ground
truth reward:
0.0episode:13;7,7,7,2:7-7=1,7*2=14,14-1=13reward:8.917851448059082ground truth
reward:
1.0episode:24;2,2,7,7:7-2=5,7+5=12,2*12=24reward:6.0330657958984375ground
truth reward:
1.0episode:24;4,3,2,7:7-4=3,3+3=6,2*6=12reward:3.4907703399658203ground truth
reward:
0.0episode:24;3,7,4,2:3*7=21,4-2=2,21+2=23reward:4.20366096496582ground truth
reward:
0.0episode:18;5,2,4,8:8/4=2,5*2=10,2*10=20reward:1.4202626943588257ground
truth reward: 0.0

As evidence by the example rewards in cyan color, the incorrect responses can
lead to moderate-to-high rewards, which are far less separated than the ground
truth rewards. In the later stage of training, model tend to stick to some
incorrect responses with moderate-to-high rewards.