_where g_ ( _y_ ) _is given by_
            - _π_ ( _y_ )             _∇_ log _π_ ( _y_ ) _p_ ( _y ≻_ ( _π_ ) [1] _[−][β]_ ( _π_ [ref] ) _[β]_ ) _−_ _τ_ log _._
_π_ [ref] ( _y_ )
_Proof._ The gradient of Nash-MD-PG( _β_ ) is:
             -             - _π_ ( _y_ ) ��
_g_ Nash-MD-PG( _β_ ) = _−_ E _y∼π,y′∼π′_ _∇_ log _π_ ( _y_ ) _p_ ( _y ≻_ _y_ _[′]_ ) _−_ 2 [1] _[τ]_ [ log] _π_ [ref] ( _y_ )
     -      - _π_ ( _y_ )
= _−_ E _y∼π_ _∇_ log _π_ ( _y_ ) _p_ ( _y ≻_ _π_ _[′]_ ) _−_ [1]
2 _[−]_ _[τ]_ [ log] _π_ [ref] ( _y_ )
-
~~�~~ ~~��~~ ~~�~~
_g_ ( _y_ )
where we write _π_ _[′]_ = ( _π_ ) [1] _[−][β]_ ( _π_ [ref] ) _[β]_ .
Now the gradient of IPO-MD( _β_ ) is:
_g_ IPO-MD( _β_ ) = E _y,y′∼π′_
    _p_ ( _y ≻_ _y_ _[′]_ ) _∇_ _ρπ_ ( _y, y_ _[′]_ ) _−_ [1]
2 _τ_
�2 [�]
      -      = 2E _y,y_ _[′]_ _∼π_ _[′]_ _p_ ( _y ≻_ _y_ _[′]_ ) _ρπ_ ( _y, y_ _[′]_ ) _−_ [1]
2 _τ_
- _∇ρπ_ ( _y, y_ _[′]_ ) _,_
where _ρπ_ ( _y, y_ _[′]_ ) = log _ππ_ [ref] ( _y_ () _y_ ) _[−]_ [log] _ππ_ [ref] ( _y_ ( _[′]_ _y_ ) _[′]_ ) [, thus] _[ ∇][ρ][π]_ [(] _[y, y][′]_ [) =] _[ ∇]_ [log] _[ π]_ [(] _[y]_ [)] _[ −∇]_ [log] _[ π]_ [(] _[y][′]_ [)][.]
Using the anti-symmetry of the preference model, i.e., _p_ ( _y ≻_ _y_ _[′]_ ) = 1 _−_ _p_ ( _y_ _[′]_ _≻_ _y_ ), and combining terms, we have that
E _y,y_ _[′]_ _∼π_ _[′]_ [ _p_ ( _y ≻_ _y_ _[′]_ ) _∇ρπ_ ( _y, y_ _[′]_ )] = E _y∼π_ _[′]_ [ _p_ ( _y ≻_ _π_ _[′]_ ) _∇_ log _π_ ( _y_ )] _−_ E _y_ _[′]_ _∼π_ _[′]_ [ _p_ ( _π_ _[′]_ _≻_ _y_ _[′]_ ) _∇_ log _π_ ( _y_ _[′]_ )]
��
_._
     -     = 2E _y∼π′_ _∇_ log _π_ ( _y_ ) _p_ ( _y ≻_ _π_ _[′]_ ) _−_ [1]
     -     = 2E _y∼π′_ _∇_ log _π_ ( _y_ ) _p_ ( _y ≻_ _π_ _[′]_ ) _−_ [1]
Similarly:
E _y,y_ _[′]_ _∼π_ _[′]_ [ _p_ ( _y ≻_ _y_ _[′]_ ) _ρπ_ ( _y, y_ _[′]_ ) _∇ρπ_ ( _y, y_ _[′]_ )]
     -      - _π_ ( _y_ ) _π_ ( _y_ _[′]_ )
= E _y,y′∼π′_ _p_ ( _y ≻_ _y_ _[′]_ ) log
_π_ [ref] ( _y_ ) _[−]_ [log] _π_ [ref] ( _y_ _[′]_ )
-
[ _∇_ log _π_ ( _y_ ) _−∇_ log _π_ ( _y_ _[′]_ )] _,_
           - _π_ ( _y_ ) _π_ ( _y_ _[′]_ )            = E _y,y′∼π′_ _p_ ( _y ≻_ _y_ _[′]_ ) log
_π_ [ref] ( _y_ ) _[∇]_ [log] _[ π]_ [(] _[y]_ [) + (1] _[ −]_ _[p]_ [(] _[y][′][ ≻]_ _[y]_ [)) log] _π_ [ref] ( _y_ _[′]_ ) _[∇]_ [log] _[ π]_ [(] _[y][′]_ [)]
            - _π_ ( _y_ ) _π_ ( _y_ _[′]_ )
_−_ E _y,y′∼π′_ _p_ ( _y ≻_ _y_ _[′]_ ) log
_π_ [ref] ( _y_ ) _[∇]_ [log] _[ π]_ [(] _[y][′]_ [) + (1] _[ −]_ _[p]_ [(] _[y][′][ ≻]_ _[y]_ [)) log] _π_ [ref] ( _y_ _[′]_ ) _[∇]_ [log] _[ π]_ [(] _[y]_ [)]
           - _π_ ( _y_ )            = E _y,y′∼π′_ log
_π_ [ref] ( _y_ ) [(] _[∇]_ [log] _[ π]_ [(] _[y]_ [)] _[ −∇]_ [log] _[ π]_ [(] _[y][′]_ [))]
We deduce that:
[2] - _∇_ log _π_ ( _y_ ) - _p_ ( _y ≻_ _π_ _[′]_ ) _−_ [1]
_τ_ [E] _[y][∼][π][′]_ 2
_g_ IPO-MD( _β_ ) = _−_ [2]
�� - _π_ ( _y_ ) + 2E _y,y′∼π′_ log
_π_ [ref] ( _y_ ) [(] _[∇]_ [log] _[ π]_ [(] _[y]_ [)] _[ −∇]_ [log] _[ π]_ [(] _[y][′]_ [))]
[2] - _∇_ log _π_ ( _y_ ) - _p_ ( _y ≻_ _π_ _[′]_ ) _−_ [1]
_τ_ [E] _[y,y][′][∼][π][′]_ 2
= _−_ [2]
- _π_ ( _y_ )
_−_ _τ_ log
_π_ [ref] ( _y_ ) [(] _[∇]_ [log] _[ π]_ [(] _[y]_ [)] _[ −∇]_ [log] _[ π]_ [(] _[y][′]_ [))]
[2] - _∇_ log _π_ ( _y_ ) - _p_ ( _y ≻_ _π_ _[′]_ ) _−_ [1]
_τ_ [E] _[y][∼][π][′]_ 2
- - _π_ ( _y_ )
_−_ 2E _y∼π′_ log
_π_ [ref] ( _y_ )
E _y∼π′_ [ _∇_ log _π_ ( _y_ )]
~~�~~ ~~�~~  - ~~�~~
=0
= _−_ [2]
[1] _π_ ( _y_ )
2 _[−]_ _[τ]_ [ log] _π_ [ref] ( _y_ )
~~�~~ ~~�~~ - ~~�~~
= _g_ ( _y_ )


where we used that E _y∼π′_ [ _∇_ log _π_ ( _y_ )] = 1 _−_ 1 _β_ [E] _[y][∼][π][′]_ [ [] _[∇]_ [log] _[ π][′]_ [(] _[y]_ [)] = 0][, from the definition of] _[ π][′]_ [.]
We deduce that
_g_ Nash-MD-PG( _β_ ) = _−_ E _y∼π_ [ _g_ ( _y_ )]
_g_ IPO-MD( _β_ ) = _−_ _τ_ [2] [E] _[y][∼][π][′]_ [ [] _[g]_ [(] _[y]_ [)]] _[ .]_
**D. Comparison of the variance of contrastive versus non-contrastive gradient estimates**
Define the following gradient estimates based on non-contrastive vs contrastive loss functions:
         - _π_ ( _y_ ) _π_ ( _y_ _[′]_ )
_g_ ˆnon-contrastive = _−∇_ log _π_ ( _y_ ) _p_ ( _y ≻_ _y_ _[′]_ ) _−_ 2 [1] _[−]_ _[τ]_ [ log] _π_ [ref] ( _y_ ) [+] _[ τ]_ [ log] _π_ [ref] ( _y_ _[′]_ )
[1] - _p_ ( _y ≻_ _y_ _[′]_ ) _−_ [1]
2 [(] _[∇]_ [log] _[ π]_ [(] _[y]_ [)] _[ −∇]_ [log] _[ π]_ [(] _[y][′]_ [))] 2
_._
_g_ ˆcontrastive = _−_ [1]
[1] _π_ ( _y_ ) _π_ ( _y_ _[′]_ )
2 _[−]_ _[τ]_ [ log] _π_ [ref] ( _y_ ) [+] _[ τ]_ [ log] _π_ [ref] ( _y_ _[′]_ )
These two estimate resemble those of the algorithms Self-Play (as implemented by Nash-MD-PG( _β_ = 0)), which uses a
non-contrastive loss, and online-IPO (equivalent to IPO-MD( _β_ = 0)), which uses a contrastive loss, and we have
_g_ Self-Play = E _y,y′∼π_ [ˆ _g_ non-contrastive]
_g_ online-IPO = _τ_ [E] _[y,y][′][∼][π]_ [ [ˆ] _[g]_ [contrastive][]] _[ .]_
We know from the previous result that these two estimates have the same expectation. However their variance may
differ. Their respective variance depends on a non-trivial combinaison of the policy representation and the specifics of the
preference model. We now state a sufficient condition under which the contrastive gradient estimate has lower variance than
its non-contrastive counterpart.
**Proposition D.1.** _If the policy representation and the preference model are such that we have_
E _y,y′∼π_            - _∇_ log _π_ ( _y_ ) _∇_ log _π_ ( _y_ _[′]_ ) _f_ ( _y, y_ _[′]_ ) [2][�] _≥_ 0 _,_ (14)
_π_ ( _y_ ) _π_ ( _y_ _[′]_ )
_where f_ ( _y, y_ _[′]_ ) := _p_ ( _y ≻_ _y_ _[′]_ ) _−_ [1] 2 _[−]_ _[τ]_ [ log] _π_ [ref] ( _y_ ) [+] _[ τ]_ [ log] _π_ [ref] ( _y_ _[′]_ ) _[, then the variance of the contrastive gradient estimate is at]_
_least as small as that of the non-contrastive one: Var_ (ˆ _gcontrastive_ ) _≤_ _Var_ (ˆ _gnon-contrastive_ ) _._
_Proof._ Defining the random variables _X_ 1( _y, y_ _[′]_ ) := _−∇_ log _π_ ( _y_ ) _f_ ( _y, y_ _[′]_ ) and _X_ 2( _y, y_ _[′]_ ) := _∇_ log _π_ ( _y_ _[′]_ ) _f_ ( _y, y_ _[′]_ ), we have
_g_ ˆnon-contrastive = _X_ 1( _y, y_ _[′]_ )
_X_ 1( _y, y_ _[′]_ ) + _X_ 2( _y, y_ _[′]_ )
_g_ ˆcontrastive = 2 _._
Now, let us compare the variance of these estimates when _y_ and _y_ _[′]_ are independently drawn from the same distribution _π_ .
The variance of the contrastive estimate is
= [Var][(] _[X]_ [1][) +][ Var][(] _[X]_ [2][) + 2][Cov][(] _[X]_ [1] _[, X]_ [2][)]
2 [(][Var][(] _[X]_ [1][) +][ Cov][(] _[X]_ [1] _[, X]_ [2][))] _[ .]_
 - _X_ 1 + _X_ 2
Var
 - _X_ 1 + _X_ 2
Var
[) + 2][Cov][(] _[X]_ [1] _[, X]_ [2][)] = [1]
This variance is lower than Var( _X_ 1) as soon as Cov( _X_ 1 _, X_ 2) is negative (this is the principle of antithetic variates for
variance reduction).
Using the property that _f_ ( _y, y_ _[′]_ ) = _−f_ ( _y_ _[′]_ _, y_ ), and writing _c_ = E _y,y′∼π_ [ _∇_ log _π_ ( _y_ ) _f_ ( _y, y_ _[′]_ )] =
_−_ E _y,y′∼π_ [ _∇_ log _π_ ( _y_ _[′]_ ) _f_ ( _y, y_ _[′]_ )], we have:
Cov( _X_ 1 _, X_ 2) = E _y,y′∼π_ [( _−∇_ log _π_ ( _y_ ) _f_ ( _y, y_ _[′]_ ) + _c_ ) ( _∇_ log _π_ ( _y_ _[′]_ ) _f_ ( _y, y_ _[′]_ ) + _c_ )]
= E _y,y′∼π_          - _−∇_ log _π_ ( _y_ ) _∇_ log _π_ ( _y_ _[′]_ ) _f_ ( _y, y_ _[′]_ ) [2][�] _−_ _c_ [2] _._
We deduce that a sufficient condition for the contrastive estimate to have a lower variance than that of the non-contrastive
estimate is:
E _y,y′∼π_       - _∇_ log _π_ ( _y_ ) _∇_ log _π_ ( _y_ _[′]_ ) _f_ ( _y, y_ _[′]_ ) [2][�] + (E _y,y′∼π_ [ _∇_ log _π_ ( _y_ ) _f_ ( _y, y_ _[′]_ )]) [2] _≥_ 0 _._
This condition is true as soon as Equation (14) is satisfied.

Human Alignment of Large Language Models through Online Preference
  Optimisation

Daniele Calandriello
Daniel Guo
Remi Munos
Mark Rowland
Yunhao Tang
Bernardo Avila Pires
Pierre Harvey Richemond
Charline Le Lan
Michal Valko
Tianqi Liu
Rishabh Joshi
Zeyu Zheng
Bilal Piot


To build some intuition for IPO-MD, in Figure 6 we provide a plot of the trajectories obtained with a variety of values of _β_
for the game with preference probabilities as displayed below, with _τ_ = 0 _._ 1, and _π_ [ref] set to the uniform policy.


0 _._ 5 0 _._ 8 0 _._ 1
0 _._ 1 0 _._ 5 0 _._ 8
0 _._ 9 0 _._ 1 0 _._ 5

_Figure 6._ Online, offline, and MD variants of IPO.


In this section, we explore whether online DPO is related to the regularised game given Equation (9), given that Proposition 4.1 shows a similar relationship for online IPO and the regularised game.
Our strategy is to derive the gradient of the offline DPO objective (Lemma F.3), and then inspect whether different sampling
distributions _µ_ and candidate solutions _π_ satisfy the KKT conditions (Boyd & Vandenberghe, 2004) for the minimising the
objective. We can also explore online DPO, by inspecting the KKT conditions when the sampling distribution matches the
candidate solution ( _µ_ = _π_ ). For example, in Lemma F.5 we present conditions for the solution of the regularised game (see
Equation (12)) to also be a stationary point of online DPO. Given that the minimiser of the online online IPO objective is the
Nash equilibrium of the regularised game given in Equation (9) (Proposition 4.1), Lemma F.5 gives us a condition for the
online DPO problem to be equivalent (solution-wise) to the online IPO problem. However, the condition seems difficult to
satisfy: For example, there is no 2-action problem for which the condition is satisfied with the exception when preferences
are uniform ( _p_ (1 _≻_ 2) = [1] 2 [). In this sense, apart from the trivial uniform-preference case, online DPO and online IPO are]
different objectives when _|Y|_ = 2.
Regaring stationary points of online DPO, we show that, under the Bradley-Terry model assumption, the RLHF solution
(Equation (3)) is a stationary point of online DPO (Theorem F.7). We note in passing (Remark F.9) that the RLHF solution
is not the only solution of offline DPO when _µ_ ( _y_ ) = 0 for some _y ∈Y_ . In fact, there are infinitely many, with arbitrarily
small probabilities for _y_ such that _µ_ ( _y_ ) _>_ 0. So while we can say that stationary points of online DPO are also stationary
points of offline DPO, a solution for the offline DPO problem with a given _µ_ may not be a solution for corresponding the
online DPO problem, because of what happens for _y_ such that _µ_ ( _y_ ) = 0.

F. Supplementary Theoretical Study of Online DPO


We let ∆ _[◦]_ ( _Y_ ) = _[.]_ _{p ∈_ ∆( _Y_ ) : _p_ ( _y_ ) _>_ 0 _, ∀y ∈Y}_ be the interior of the simplex.
**Proposition F.1.** _For t ∈_ R _, we have_
_d_
_dt_ [log] _[ σ]_ [(] _[t]_ [) =] _[ σ]_ [(] _[−][t]_ [)] _[.]_
_Proof._ We have:
_d_ _[d]_
_dt_ [log] _[ σ]_ [(] _[t]_ [) =]
_[d]_ 1 _[d]_
_dt_ _[−]_ [log(1 +] _[ e][−][t]_ [) =] _[ −]_ 1 + _e_ _[−][t]_ _[·]_
_dt_ [(1 +] _[ e][−][t]_ [)]
_e_ _[−][t]_ 1
= 1 + _e_ _[−][t]_ [=] 1 + _e_ _[t]_ [=] _[ σ]_ [(] _[−][t]_ [)] _[.]_
**Proposition F.2.** _For t ∈_ R _, we have_
_σ_ ( _t_ ) + _σ_ ( _−t_ ) = 1
_Proof._ We have:
1 1 1 _e_ _[−][t]_
_σ_ ( _t_ ) + _σ_ ( _−t_ ) = 1 + _e_ _[−][t]_ [+] 1 + _e_ _[t]_ [=] 1 + _e_ _[t]_ [+] 1 + _e_ _[−][t]_ [= 1] _[.]_
**Lemma F.3.** _Assume that |Y| < ∞_ _and that p_ ( _y ≻_ _y_ _[′]_ ) = 1 _−_ _p_ ( _y_ _[′]_ _≻_ _y_ ) _for all y, y_ _[′]_ _∈Y. The offline DPO problem with_
_sampling distribution µ can be written as_
sup _J_ DPO( _π_ ) _,_
_π∈_ ∆( _Y_ )
_where_
    -
- _µ_ ( _y_ ) _µ_ ( _y_ _[′]_ ) _p_ ( _y ≻_ _y_ _[′]_ ) log _σ_ _τ_ log _[π]_ [(] _[y]_ [)] _[π]_ [ref] [(] _[y][′]_ [)]
_π_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
_y,y_ _[′]_
_J_ DPO( _π_ ) = _[.]_
_π_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
       + (1 _−_ _p_ ( _y ≻_ _y_ _[′]_ )) log _σ_ _−τ_ log _[π]_ [(] _[y]_ [)] _[π]_ [ref] [(] _[y][′]_ [)]
_π_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
��
_._


_Moreover, for π ∈_ ∆ _[◦]_ ( _Y_ ) _,_
_∇J_ DPO( _π_ ) _y_ = 2 _τ_ _[µ]_ [(] _[y]_ [)]
_π_ ( _y_ )
   -   _µ_ ( _y_ _[′]_ ) _p_ ( _y ≻_ _y_ _[′]_ ) _−_ _σ_ _τ_ log _[π]_ [(] _[y]_ [)] _[π]_ [ref] [(] _[y][′]_ [)]
_π_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
_y_ _[′]_
_π_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
��
_._ (15)
_Proof._ We will use the shorthands _µy,y′ ≡_ _µ_ ( _y_ ) _µ_ ( _y_ _[′]_ ), _πy ≡_ _π_ ( _y_ ) (and so forth) and _sy,y′_ = _[.]_ _τ_ log _[π]_ _π_ [(] ( _[y]_ _y_ [)] _[′]_ ) _[π]_ _π_ [ref][ref][(] ( _[y]_ _y_ _[′]_ [)] ) [.]
The offline DPO objective is obtained by taking the expectation of the DPO loss (Rafailov et al., 2023) with _Y_ [+] _, Y_ _[−]_ _∼_
( _µ, λp_ ):
_µy,y′_ ( _p_ ( _y ≻_ _y_ _[′]_ ) log _σ_ ( _sy,y′_ ) + (1 _−_ _p_ ( _y ≻_ _y_ _[′]_ )) log _σ_ ( _−sy,y′_ )) = _J_ DPO( _π_ ) _._
_y,y_ _[′]_
Before looking at _∇J_ DPO( _π_ ), we will make some simplifying observations. For _y_ _[′′]_ = _y, y_ _[′]_,
_d_
(16)
_dπy′′_ _[s][y,y][′]_ [ = 0] _[.]_
Since _µy,y′_ = _µy′,y_, _sy,y′_ = _−sy′,y_ and _p_ ( _y ≻_ _y_ _[′]_ ) = 1 _−_ _p_ ( _y_ _[′]_ _≻_ _y_ ), we have
_µy,y′_ ( _p_ ( _y ≻_ _y_ _[′]_ ) log _σ_ ( _sy,y′_ ) + (1 _−_ _p_ ( _y ≻_ _y_ _[′]_ )) log _σ_ ( _−sy,y′_ ))
(17)
= _µy′,y_ ((1 _−_ _p_ ( _y_ _[′]_ _≻_ _y_ )) log _σ_ ( _−sy′,y_ ) + _p_ ( _y_ _[′]_ _≻_ _y_ ) log _σ_ ( _sy′,y_ )) _._
Also note that
_d_ _d_
_sy,y′_ =
_dπy_ _dπy_
Therefore, for all _y ∈Y_
_τ_ log _[π][y][π]_ [ref] [(] _[y][′]_ [)]
_πy′π_ [ref] ( _y_ )
_[d]_ log _πy_ = _[τ]_
_dπy_ _π_
= _τ_ _[d]_
_._ (18)
_πy_
   _µy,y′_ _p_ ( _y ≻_ _y_ _[′]_ ) _[d]_
_dπ_
_y_ _[′]_
_d_ _J_ DPO( _π_ ) = 2 _dπy_
_[d]_ log _σ_ ( _sy,y′_ ) + (1 _−_ _p_ ( _y ≻_ _y_ _[′]_ )) _[d]_
_dπy_ _dπ_
_[d]_ log _σ_ ( _−sy,y′_ ) _,_ (Equations (16) and (17))
_dπy_
   _µy,y′_ _p_ ( _y ≻_ _y_ _[′]_ ) _σ_ ( _−sy,y′_ ) _[d]_
_dπ_
_y_ _[′]_
= 2
_[d]_ _sy,y′ −_ (1 _−_ _p_ ( _y ≻_ _y_ _[′]_ )) _σ_ ( _sy,y′_ ) _[d]_
_dπy_ _dπ_
_[d]_ _sy,y′_ (Proposition F.1)
_dπy_
= 2 - _µy,y_ _[′]_ ( _p_ ( _y ≻_ _y_ _[′]_ ) _σ_ ( _−sy,y_ _[′]_ ) _−_ (1 _−_ _p_ ( _y ≻_ _y_ _[′]_ )) _σ_ ( _sy,y_ _[′]_ )) _d_ _sy,y_ _[′]_
_dπy_
_y_ _[′]_
 - _d_
= 2 _µy,y′_ ( _p_ ( _y ≻_ _y_ _[′]_ ) _−_ _σ_ ( _sy,y′_ )) _sy,y′_ (Proposition F.2)
_dπy_
_y_ _[′]_
- _πyπy_ [ref] _[′]_
_τ_ log
_πy′πy_ [ref]
 = 2 _µy,y′_
_y_ _[′]_
_p_ ( _y ≻_ _y_ _[′]_ ) _−_ _σ_
��
_τ_
_,_ (Equation (18))
_πy_
��
= 2 _τ_ _[µ][y]_
_πy_
_µy,y′_
_y_ _[′]_
_p_ ( _y ≻_ _y_ _[′]_ ) _−_ _σ_
- _πyπy_ [ref] _[′]_
_τ_ log
_πy′πy_ [ref]
_._
_Remark_ F.4 _._ We have restricted the statement of Lemma F.3 to the interior of the simplex to stay consistent with the DPO
formulation, however Equation (15) holds for any function exp _◦f_ with _f_ : _Y →_ R. To see this, it suffices to notice that
_J_ DPO is a function of the ratios _π_ _[π]_ _j_ _[i]_ [, so for any] _[ α >]_ [ 0] _[ J]_ [DPO][(] _[π]_ [) =] _[ J]_ [DPO][(] _[α][ ·][ π]_ [)][ (where] _[ α][ ·][ π .]_ [=] _[ y][ �→]_ _[α][ ·][ π]_ [(] _[y]_ [)][).]
**Lemma F.5.** _The Nash equilibrium π_ _[∗]_ _of the regularised game given in Equation_ (9) _is a stationary point of online DPO iff:_
_p_ ( _y ≻_ _π_ _[∗]_ ) =              - _π_ _[∗]_ ( _y_ _[′]_ ) _σ_ ( _p_ ( _y ≻_ _π_ _[∗]_ ) _−_ _p_ ( _y_ _[′]_ _≻_ _π_ _[∗]_ )) _._ (19)
_y_ _[′]_


_Proof._ Let _π_ _[∗]_ be the Nash equilibrium of the regularised game given in Equation (9), the fixed point in Equation (12):
                 - 1                  _π_ _[∗]_ ( _y_ ) _∝_ _π_ [ref] ( _y_ ) exp _._ (20)
_τ_ _[p]_ [(] _[y][ ≻]_ _[π][∗]_ [)]
_π_ _[∗]_ will be a stationary point of online DPO iff _π_ = _π_ _[∗]_ is a solution of the offline DPO problem with _µ_ = _π_ _[∗]_ . Since
_π_ _[∗]_ _∈_ ∆ _[◦]_ ( _Y_ ), we can use Lemma F.3, _π_ _[∗]_ is a solution of the offline DPO problem iff
_∇J_ DPO( _π_ _[∗]_ ) = 0 _._
We have for all _y ∈Y_ :
_d_
_J_ DPO( _π_ _[∗]_ ) = 2 _τ_ _[µ]_ [(] _[y]_ [)]
_dπy_ _π_ _[∗]_ ( _y_ )
   -   _µ_ ( _y_ _[′]_ ) _p_ ( _y ≻_ _y_ _[′]_ ) _−_ _σ_ _τ_ log _[π][∗]_ [(] _[y]_ [)] _[π]_ [ref] [(] _[y][′]_ [)]
_π_ _[∗]_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
_y_ _[′]_
��
(Equation (15))
_π_ _[∗]_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
   -    _π_ _[∗]_ ( _y_ _[′]_ ) _p_ ( _y ≻_ _y_ _[′]_ ) _−_ _σ_ _τ_ log _[π][∗]_ [(] _[y]_ [)] _[π]_ [ref] [(] _[y][′]_ [)]
_π_ _[∗]_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
_y_ _[′]_
= 2 _τ_
_π_ _[∗]_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
��
( _µ_ = _π_ _[∗]_ )
= 2 _τ_ - _π_ _[∗]_ ( _y_ _[′]_ ) ( _p_ ( _y ≻_ _y_ _[′]_ ) _−_ _σ_ ( _p_ ( _y ≻_ _π_ _[∗]_ ) _−_ _p_ ( _y_ _[′]_ _≻_ _π_ _[∗]_ ))) (Equation (20))
_y_ _[′]_


 _,_
= 2 _τ_
 _p_ ( _y ≻_ _π_ _[∗]_ ) _−_
_π_ _[∗]_ ( _y_ _[′]_ ) _σ_ ( _p_ ( _y ≻_ _π_ _[∗]_ ) _−_ _p_ ( _y_ _[′]_ _≻_ _π_ _[∗]_ ))
_y_ _[′]_
and the result follows by using the fact that _τ >_ 0.
**Theorem F.6.** _No 2-action regularised game given in Equation_ (9) _has a Nash equilibrium that satisfies Equation_ (19) _,_
_except for the regularised games with p_ ( _y_ 1 _≻_ _y_ 2) = [1] 2 _[.]_
_Proof._ Let _Y_ = _{_ 1 _,_ 2 _}_ . For this two-action problem, we can write the preference matrix as
  - 12 1 _−_ _p_
_P_ = _[.]_ _p_ 12
_,_
where _Pyy_ _[′]_ = _p_ ( _y ≻_ _y_ _[′]_ ).
Let _α_ be such that _π_ _[∗]_ = ( _α,_ 1 _−_ _α_ ) _[⊤]_ . Then
      _p_ ( _y ≻_ _π_ _[∗]_ ) = _Pπ_ _[∗]_ = 1 _−_ _[α]_
[1]
2 _[−]_ _[α]_ 2
_[α]_ [1]
2 _[−]_ _[p]_ [ +] _[ αp,]_ 2
_[α]_ - _⊤_ _,_
2 [+] _[ αp]_
and
_p_ (1 _≻_ _π_ _[∗]_ ) _−_ _p_ (2 _≻_ _π_ _[∗]_ ) = [1]
2 _[−]_ _[p.]_
Then the difference between both sides of Equation (19) for _y_ = 1 is:
_p_ (1 _≻_ _π_ _[∗]_ ) _−_               - _π_ _[∗]_ ( _y_ _[′]_ ) _σ_ ( _p_ (1 _≻_ _π_ _[∗]_ ) _−_ _p_ ( _y_ _[′]_ _≻_ _π_ _[∗]_ ))
_y_ _[′]_
_[α]_ - 1
2 _[−]_ _[p]_ [ +] _[ αp][ −]_ _[ασ]_ [(0)] _[ −]_ [(1] _[ −]_ _[α]_ [)] _[σ]_ 2
= 1 _−_ _[α]_
2 _[−]_ _[p]_
         - 1          = 1 _−_ _α −_ _p_ + _αp −_ (1 _−_ _α_ ) _σ_
2 _[−]_ _[p]_
   -    - 1 ��
= (1 _−_ _α_ ) 1 _−_ _p −_ _σ_ _._
2 _[−]_ _[p]_
For _y_ = 2 we get


_p_ (2 _≻_ _π_ _[∗]_ ) _−_       - _π_ _[∗]_ ( _y_ _[′]_ ) _σ_ ( _p_ (2 _≻_ _π_ _[∗]_ ) _−_ _p_ ( _y_ _[′]_ _≻_ _π_ _[∗]_ ))
_y_ _[′]_
_[α]_ [)]
_−_ _ασ_ _p −_ [1]
2 2
= _αp_ + [(1] _[ −]_ _[α]_ [)]
_−_ (1 _−_ _α_ ) _σ_ (0)
    = _αp −_ _ασ_ _p −_ [1]
    = _αp −_ _ασ_ _p −_ [1]
     - 1     = _αp −_ _α_ + _ασ_
2 _[−]_ _[p]_
  -   - 1 ��
= _−α_ 1 _−_ _p −_ _σ_ _._
2 _[−]_ _[p]_
Now, if we let _ε_ = _[.]_ [1]
[1] [1]
2 _[−]_ _[p]_ [, we can see that for] _[ p <]_ 2
 - 1
_σ_
= _σ_ ( _ε_ ) _< σ_ (0) = [1]
2 _[−]_ _[p]_ 2
[1] [1]
2 _[<]_
2 [+] _[ ε]_ [ = 1] _[ −]_ _[p,]_
so (considering the analogous case for _p <_ [1] 2 [)]

1 - 
2 _[−]_ _[p]_
   - 1
1 _−_ _p −_ _σ_
_>_ 0 _,_ _p <_ [1]
_>_ 0 _,_ _p <_ 2 _[,]_
= 0 _,_ _p_ = 2 [1] _[,]_
_<_ 0 _,_ _p >_ [1] _[.]_

2 _[.]_
Therefore, if _p ̸_ = [1] 2 [, we cannot satisfy Equation (][19][) for] _[ y]_ [ = 1][ and] _[ y]_ [ = 2][ (note that] _[ α]_ [ = 0][ or] _[ α]_ [ = 1][ satisfy the equation for]
only one _y_ ).
**Theorem F.7.** _Assume that the preferences admit a Bradley-Terry model (Rafailov et al., 2023), that is, there exists_
_r_ : _Y →_ R _such that for all y, y_ _[′]_ _∈Y_
_p_ ( _y ≻_ _y_ _[′]_ ) = _σ_ ( _r_ ( _y_ ) _−_ _r_ ( _y_ _[′]_ )) _,_
_Then_
      - 1       _π_ _[r]_ ( _y_ ) _∝_ _π_ [ref] ( _y_ ) exp (21)
_τ_ _[r]_ [(] _[y]_ [)]
_is a critical point for offline IPO for any µ, and a stationary point of online DPO._
_Proof._ First, assume that _π_ [ref] _∈_ ∆ _[◦]_ ( _Y_ ) (we will deal with the general case at the end).
Let us consider the offline DPO problem with sampling distribution _µ_ . Assuming that the preferences admit a Bradley-Terry
model, we know from Rafailov et al. (2023) that the solution of the offline DPO problem is given by
                  - 1                   _π_ _[r]_ ( _y_ ) _∝_ _π_ [ref] ( _y_ ) exp _._ (22)
_τ_ _[r]_ [(] _[y]_ [)]
Indeed, we can verify that for any _µ_
_∇J_ DPO( _π_ _[r]_ ) = 2 _τ_ _[µ]_ [(] _[y]_ [)]
_π_ ( _y_ )
   -   _µ_ ( _y_ _[′]_ ) _p_ ( _y ≻_ _y_ _[′]_ ) _−_ _σ_ _τ_ log _[π]_ [(] _[y]_ [)] _[π]_ [ref] [(] _[y][′]_ [)]
_π_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
_y_ _[′]_
_π_ ( _y_ _[′]_ ) _π_ [ref] ( _y_ )
��
( _π_ _[r]_ _∈_ ∆ _[◦]_ ( _Y_ ), Equation (15))
= 2 _τ_ _[µ]_ [(] _[y]_ [)]
_π_ _[r]_ ( _y_ )
- _µ_ ( _y_ _[′]_ ) ( _p_ ( _y ≻_ _y_ _[′]_ ) _−_ _σ_ ( _r_ ( _y_ ) _−_ _r_ ( _y_ _[′]_ ))) (Equation (21))
_y_ _[′]_
= 2 _τ_ _[µ]_ [(] _[y]_ [)]
_π_ ( _y_ )
= 0
- _µ_ ( _y_ _[′]_ ) ( _p_ ( _y ≻_ _y_ _[′]_ ) _−_ _p_ ( _y ≻_ _y_ _[′]_ )) (Bradley-Terry model assumption)
_y_ _[′]_


It follows that _π_ _[r]_ is an offline DPO solution under any sampling distribution _µ_, and if _µ ∈_ ∆ _[◦]_ ( _Y_ ) then _π_ _[r]_ is the only
solution. In particular this means that _π_ _[r]_ is also a stationary point for online DPO (by taking _µ_ = _π_ _[r]_ ).
In the case where _π_ [ref] _∈_ ∆( _Y_ ) _−_ ∆ _[◦]_ ( _Y_ ), we can prove the result as follows. Let _Y_ _[′][ .]_ = _{y_ : _π_ [ref] ( _y_ ) = 0 _}_ . For _y ∈Y_ _[′]_, we
have _π_ _[r]_ ( _y_ ) = 0 by Equation (21). Since _π_ _[r]_ is the sampling distribution, the gradient of the offline DPO objective in _Y_ _[′]_ is
zero (and it does not matter that log _π_ _[r]_ ( _y_ ) and log _π_ [ref] ( _y_ ) are undefined since they do not appear in the objective). Now it
suffices to note that _π_ _[r]_ restricted to _Y −Y_ _[′]_ is a distribution in ∆ _[◦]_ ( _Y −Y_ _[′]_ ), so we can use the previous case ( _π_ _[r]_ _∈_ ∆ _[◦]_ ( _Y_ ))
to show the result, and it follows that _π_ _[r]_ satisfies the KKT conditions for all _y ∈Y_ .
_Remark_ F.8 _._ In the special case of rock-paper-scissors preferences (which do not admit a Bradley-Terry model) with _π_ [ref]
uniform and _τ >_ 0, we can still satisfy Equation (19), since _π_ _[∗]_ is also uniform and _p_ ( _y ≻_ _π_ _[∗]_ ) = [1] 2 [.]
_Remark_ F.9 _._ In offline DPO, the presence of sets of measure zero guarantees the solution is not unique. Let _π_ [DPO] be a
solution of offline DPO under sampling distribution _µ_ (assume it exists). _π_ [DPO] may not necessarily be _π_ _[r]_ from (21) (even
if we assume that the preferences admit a Bradley-Terry model).
To see this, assume that there exists _Y_ _[′]_ _⊂Y_ such that _Y_ _[′]_ = _∅_ and _µ_ ( _y_ ) = 0 for all _y ∈Y_ _[′]_ . For any _α ∈_ (0 _,_ 1], define
_π_ _[α]_ _∈_ ∆ _[◦]_ ( _Y_ ) by
_π_ _[α]_ ( _y_ ) = _[.]_
_α · π_ [DPO] ( _y_ ) _,_ _y /∈Y_ _[′]_ _,_
arbitrary _,_ _y ∈Y_ _[′]_ _._
By Lemma F.3, _∇J_ DPO( _π_ _[α]_ ) = 0, so _π_ _[α]_ is a minimum of the offline DPO objective. In fact, this holds for any _α >_ 0,
no matter how small, and this also means that if _µ /∈_ ∆ _[◦]_ ( _Y_ ) we can find offline DPO solutions with arbitrarily small
probabilities for sampled _y_ ( _y_ such that _µ_ ( _y_ ) _>_ 0).


**Ben Athiwaratkun** [*] **, Shiqi Wang** _[∗]_ **, Mingyue Shang** _[∗]_ **, Yuchen Tian, Zijian Wang, Sujan Kumar Gonu-**

Token Alignment via Character Matching for Subword Completion

TOKEN ALIGNMENT VIA CHARACTER MATCHING FOR SUBWORD COMPLETION

Ben Athiwaratkun
Shiqi Wang
Mingyue Shang
Yuchen Tian
Zijian Wang
Sujan Kumar Gonugondla
Sanjay Krishna Gouda
Rob Kwiatowski
Ramesh Nallapati
Bing Xiang


AWS AI Labs
{benathi,wshiqi,myshang}@amazon.com
ABSTRACT
Generative models, widely utilized in various applications, can often struggle with
prompts corresponding to partial tokens. This struggle stems from tokenization,
where partial tokens fall out of distribution during inference, leading to incorrect or
nonsensical outputs. This paper examines a technique to alleviate the tokenization
artifact on text completion in generative models, maintaining performance even
in regular non-subword cases. The method, termed _token alignment_, involves
backtracking to the last complete tokens and ensuring the model’s generation
aligns with the prompt. This approach showcases marked improvement across
many partial token scenarios, including nuanced cases like space-prefix and partial
indentation, with only a minor time increase. The technique and analysis detailed
in this paper contribute to the continuous advancement of generative models in
handling partial inputs, bearing relevance for applications like code completion
and text autocompletion.
1 INTRODUCTION
Generative models have shown remarkable efficacy in a range of applications. However, they
have been observed to falter when dealing with partially provided inputs or subwords during text
_completion_ . For instance, a generative model might struggle to predict the remaining part of the word
where a prompt ending in a subword often leads to incorrect or nonsensical outputs. This issue arises
due to the artifact of tokenization where a partial token can be out-of-distribution during inference. In
this paper, we introduce a method to address and rectify these shortcomings in generative models,
ensuring they maintain their performance on complete contexts.
Our approach involves backtracking to the last complete tokens and aligning the model’s generation
to match with the given prefix, as shown in Figure 1. For example, if the end of the prefix is “sys” we
guide the model’s generation to ensure compatibility with this prefix where all possible tokens match
with “sys”. This method not only improves the accuracy of the model’s predictions on subword data
metrics but also retains the performance on generic evaluation sets.
To illustrate our approach, consider the following programming code snippet where the end of the
prompt is right after the incomplete token “re” in Figure 2. A conventional model might generate
an incorrect or suboptimal output such as "re = []", while the correct prediction should be "return".
Our approach solves this issue by matching the characters starting from the last complete pre-token
and aligning the subsequent tokens with the prefix. The results of our approach, as shown in various
examples, demonstrate the model’s ability to reliably generate the correct output regardless of the
partial context.
Our evaluation of the token alignment approach indicates noticeable improvement in handling partial
token contexts, outperforming the traditional decoding approach without any subword handling. We
present comprehensive evaluation results on many partial token scenarios such as subword of natural
words, as well as much less obvious cases of partial tokens related to punctuations, or artifacts of
contiguous white spaces and space prefix. Our approach also does not incur much extra latency,
with an average increase of only 3-7 ms for using token alignment due to an efficient trie-based
lookup table, in addition to the number of backtracked tokens. We believe our findings provide a
*Corresponding authors (equal contribution) {benathi,wshiqi,myshang}@amazon.com
Possible Next token
Original prompt next tokens probability distribution
Figure 1: An illustration of token alignment process. We first tokenize the prompt where the last
tokens may correspond to partial tokens. We use the backtracked prompt as model input and use the
alignment prefix to filter possible tokens at byte or character level. We then mask out non-matched
tokens by zeroing their probabilities, which is later used for sampling to select the next token.
_# write a function to get three maximum_
_numbers from a list_
**def** three_max(l):
re<T> = []
**for** i **in range** ( **len** (l)):
**if** i == 0:
re.append(l[0])
(a) Without token alignment
_# write a function to get three maximum_
_numbers from a list_
**def** three_max(l):
re<T>turn **sorted** (l, reverse=True)[:3]
(b) With token alignment
Figure 2: Effects of token alignment on prompts ending with a subword. The <T> marks the end of
the prompt, after which the completion is shown. Without token alignment, the model fails to predict
“return” correctly, as the sequence of tokens “re”, “turn” is out-of-distribution. Token alignment
alleviates this constraint by backtracking to full tokens before “re”, then align subsequent generations
with the prompt.
meaningful and practical contribution to the ongoing work in improving robustness of generative
models, especially related to handling partial inputs.
2 RELATED WORK
Recent works have identified severe robustness issues for large language models. For instance, simple
perturbations like synonym substitution on one word can significantly change outputs and fool the
models (Jin et al., 2020; Zang et al., 2020). Most studies focus on general perturbations especially at
random locations of prompts for both text (Dhole et al., 2021; Nie et al., 2020; Wang et al., 2021;
Kiela et al., 2021) and code models (Wang et al., 2023; Li et al., 2022; Yang et al., 2022; Jha & Reddy,
2023). However, we see a large gap of partial versus full tokens at the end of prompts with only a few
characters missed or added (e.g., with or without a white space) due to the artifact of tokenization.
This robustness problem is barely investigated. Subword regularization tried to partially mitigate such
problems by randomly introducing sub-tokens in training and improve robustness behavior (Kudo,
2018a). However, this requires to retrain the models and also leads to inference latency increase. Our
approach, on the other hand, can work for any models with negligible latency increase.
We acknowledge concurrent work in the form of a blog post namely token healing (Lundberg,
2023), which attempts to similarly solve the partial token problems on inference time. However,
we independently develop our own token alignment algorithm in a more general design. We also
thoroughly validate and analyze the effectiveness of token alignment with extensive empirical results
on various tasks and many partial token scenarios. Other related previous work includes approximate
marginalization (Kim, 2019) for beam search. Our method however is developed for language models
which primarily use nucleus sampling (Holtzman et al., 2020) for next token prediction.
3 METHODOLOGY
Generative language models have been remarkably successful in a variety of applications, including
code completion in Integrated Development Environments (IDEs) and text autocompletion in email
or productivity tools. However, these models often encounter difficulties when dealing with partial
or incomplete inputs, which we refer to as the partial token issue. In such cases, the models tend
to generate outputs that are not compatible with the given context due to the constraint imposed by
tokenization. This limitation significantly hampers the utility and user experience of applications built
on these models. To tackle this issue, we propose a method called token alignment which makes use
of a character or byte level trie for efficient matching along with masking cache to increase efficiency
further. The underlying principle of our approach is to backtrack to the last complete token and
constrain the model’s generation in a way that aligns with the given prefix.
In our methodology, we first define the last complete token by examining the tokenization of the input
sequence. Given the tokenization of the last line as _T_ 1 _, T_ 2 _, ..., Tn_, we start generating from the _−B_ _[th]_
token, i.e., _Tn−B_ +1, where _N_ is the number of tokens we need to backtrack. The rest of the context
is the sequence _Tn−B_ +1 _, . . ., Tn_, which we refer to as the _alignment prefix_ where the subsequent
generation needs to match at a byte level (or character level if the tokenizer is not a byte-level BPE). [*]
The token alignment step ensures that the next token starts with the alignment prefix or is a prefix of
the alignment prefix itself by masking out probabilities of tokens that do not match. After each step
where the next token is selected, we deduct the generated bytes from the alignment prefix until the
alignment prefix is empty, after which we finish the token alignment process.
**Algorithm 1** TokenAlign Algorithm
1: **procedure** TOKENALIGN( _T, N_ )
2: **Define:**
3: _X_ : The input tokens [ _x_ 1 _, x_ 2 _, ..., xn_ ]
4: _B_ : The number of tokens to backtrack
5: _C_ : Tokens for model context
6: _V_ : Prebuilt character or byte-trie
7: _P_ : Alignment prefix
8: _p_ : Probability distribution of next token
9: _x_ _[′]_ : Selected next token
10: _C ←_ _X_ [: _−B_ ]
11: _P ←_ _join_ ( _X_ [ _−B_ :])
12: **while** _P_ is not empty **do**
13: _p ←_ _model_ ( _C_ )
14: _T ←_ _V_ ( _P_ ) (tokens compatible with
_P_ )
15: _pnew_ ( _T_ ) _←_ 0
16: _x_ _[′]_ _∼_ _pnew_ (sample)
17: _P ←_ _P_ [ _len_ ( _x_ _[′]_ ) :]
18: _C._ append( _x_ _[′]_ )
19: **end while**
20: **return** _C_
21: **end procedure**
The efficiency of our approach is supported by two key techniques: a character-trie for fast prefix
matching, and a mask cache. The character-trie is a tree data structure, where each node represents
a character and the paths from the root node to the leaf nodes represent complete tokens in the
vocabulary. We pre-build this trie using the model’s vocabulary before starting the generation process,
thereby enabling a fast and space-efficient representation of the vocabulary.
The mask cache is used to further accelerate the lookup process by avoiding unnecessary the trie
lookup. For common cases, such as a single space as context, we pre-build and cache a Boolean
mask. This cached mask is then used for quickly filtering tokens that share a common prefix with the
given context. These techniques significantly reduce the latency, enabling our method to be highly
efficient. The use of a trie-based lookup table and the pre-built mask cache results in a time saving
compared to a naive implementation using a dictionary lookup. In sum, our methodology presents
an effective and efficient solution to the partial token issue in generative models, enhancing their
performance in applications such as code completion and text autocompletion.
4 PARTIAL TOKEN SCENARIOS
We give an overview of many partial token scenarios where the model is susceptible to degenerate
behavior when the prompt ends with such partial token. Such partial token scenarios range from
*We use character of byte interchangeably while describing the tokenizer.
subword punctuation space-prefix contiguous spaces
Prompt
=
Truncated
text
Tokenized prompt
Prediction w.
Token alignment
if_x_==_1:
____return_True
else:
____return_False
.. __
_a_banana for_i_in_range(10):
_a_bana for_i_in_range(10)
if_x_==_1:
____return_True
.. ____
_a_banana
for_i_in_range(10):
.. ____return .. ____return_False
Prediction w/o
Token alignment _a_bana ?? for_i_in_range(10) ?? .. ____ ?? .. __ ??
Figure 3: Partial token scenarios. Note: the underscore denotes a space character.
obvious cases such as natural subwords to less obvious cases such as space prefix, as illustrated in
Figure 3.
1. **Subword** . This case refers to natural language subword. For instance, a possible subword
of “banana” is ‘banan’ or ‘bana’. We note that while it is possible for a tokenizer to use a
subword as an actual token, such as in the case of “beautifully” which can be represented as
‘beautiful’ + ‘ly’, not all possible subwords will correspond to a token, e.g., ‘beautifu’ is
unlikely to be a full token.
2. **Punctuation** . It is possible that multiple punctuations are grouped together as a token during
the tokenizer training process, which can lead to some punctuations being a partial token in
a non-obvious way. For example, in Python, it is possible that the code snippet if (x==1)
ends in a partial token ‘)’ of a full token ‘):’, which can lead to a suboptimal completion
behavior if presented to a model. In the case of punctuation, it is generally quite hard to
judge which case would correspond to a partial punctuation. We separate out this case to
emphasize the importance of token alignment as a universal method that can handle both
prompt ending in either partial token or non partial token in a prompt-agnostic way.
3. **Space prefix.** Most tokenizers use the space-prefix schema where a whitespace is often
grouped as part of a word as a prefix (Scao et al., 2022; Brown et al., 2020; Li et al., 2023;
OpenAI, 2023). This schema is primarily aimed to reduce the number of tokens required to
represent text. For example, with space-prefix, ‘I like’ can be tokenized into two tokens, ‘I’
and ‘ like’, instead of three (‘I’, ‘ ’, and ‘like’). If a prompt ‘I ’ is presented to the model,
such model can have difficulty predicting coherent text since the prompt is a sequence of
tokens (‘I’, ‘ ’), which is out-of-distribution compared to the training stage which observes
(‘I’, ‘ like’).
The space prefix’s constraint is also quite pronounced in the presence surrounding indentation
block. For instance, x= is tokenized into (‘ ’, ‘ x’, ‘=’). However, if the
prompt ends after the full block of indentation, this becomes a constraint since
(‘ ’, ‘x’) is out-of-distribution compared to (‘ ’, ‘ x’) seen during training.
We refer to this as the _space prefix with indentation_ case.
4. **Contiguous Spaces** . Modern tokenizers often intentionally group whitespaces together for
improved compression rates. For instance, Codex (Chen et al., 2021) extends the GPT-3
tokenizer (Brown et al., 2020) to include contiguous spaces of different lengths (up to 65
contiguous spaces in some version of TikToken (OpenAI, 2023), for instance). Other models
such as LLaMA (Touvron et al., 2023), StarCoder (Li et al., 2023), CodeGen (Nijkamp
et al., 2022), to name a few, also use such contiguous white spaces.
While such grouping of white spaces is compression efficient, it can lead to out-ofdistribution behavior if the prompt ends in partial token of such contiguous white spaces.
For instance, consider the following prompt ‘ if True:\n ‘ where represents a space
character. The correct syntax for Python is such that what comes after the if clause requires
another level of indentation, meaning that the model should predict extra characters. However, the model will have a hard time completing the next two since during training, the
model always observe contiguous if it were to complete the next level of indentation as
in this example.
5 EVALUATION
In this section, we empirically demonstrate that each of the partial token scenarios described in
Section 4 can unnecessarily constrain the model due to the artifact of tokenization and lead to
significant drops in evaluation scores. We perform the evaluation using both code generation tasks
where partial token completion such as in the use of code-completion tools can lead to issues if
not well handled, as well as natural language tasks such as text completion and natural language
understanding. We construct the evaluation datasets for each case described in Section 4, which
shows that token alignment method graciously handles the constraint due to all such cases, making a
universal partial token handling approach for language model inference.
5.1 DATASETS
To demonstrate the sensitivity of language models on how the prompt ends, we analyze models’
behavior by processing publicly available datasets to understand each partial token scenario such as
subword, space prefix, etc. Below, we detail our methods of processing evaluation datasets for both
code generation and natural language tasks.
5.1.1 CODE GENERATION TASKS
**Execution-Based Code Generation Benchmarks** We assess code generation abilities using the
MBXP benchmark (Athiwaratkun et al., 2022), which is a multi-lingual version of MBPP (Austin
et al., 2021), focusing on datasets in Python, Java, and JavaScript. This entails evaluating function
completion from partially given canonical solutions.
For each dataset: (1) the prompt is formed by merging the original function signature with a section
of the canonical solution, and (2) this section is selectively cut at a specific location based on different
criteria depending on the scenarios.
    - **Subword:** The cut is based on space delineation, yielding prompts ending in subwords.
    - **Punctuation:** Cuts are made within punctuation sequences to demonstrate challenges with
subtokens. For instance, ‘{};’ might be cut to produce ‘{’ or ‘{}’.
    - **Space prefix with indentation:** Here, we focus on challenges arising from prompts ending
with spaces used for indentation. For instance, in a snippet like \n return value,
a cut after the indentation spaces produces \n . Such endings can lead to suboptimal
generation behavior without token alignment due to the last space corresponds to a space
prefix of a future word such as in return.
    - **Space prefix without indentation:** Cuts are made between spaces and non-spaces
while avoiding spaces specifically used for indentation. For example, from the snippet
\n return value, a possible cut could yield \n return .
    - **Contiguous spaces:** Cuts are introduced within sequences of contiguous whitespace characters. Here, we consider, \n and \t as candidates for whitespace.
Each dataset also has a **baseline** where prompts are backtracked to end at the last full word, allowing
us to evaluate model performance without the challenge of partial tokens. For example, in the
subword case, if the prompt ends with ‘for i in rang’, we truncate it to be ‘for i in’.
In the baseline scenario, models can perform well without the use of token alignment since it does
not correspond to an ending partial token. The evaluation scores should ideally be slightly lower
due to _lower amount of information_ in the truncated prompt. The goal of the corresponding baseline
evaluation datasets is such that the prompt should not end in partial token, so that we can use them (1)
to see that token alignment does not degrade performance for prompt ending with non partial token
and (2) measure improvement on the partial token dataset due to token alignment.
Table 1: **Token alignment for subword scenario.** This table illustrates pass@1 scores (%) on
MBXP partial benchmark as well as subword version of SQuAD and Wikitext, which show clear
improvement due to token alignment.
Token Subword MBXP Baseline of Subword MBXP
Alignment
Python Java JavaScript Python Java JavaScript
with **56.58** **52.17** **49.31** **54.32** 49.54 50.16
StarCoder
w/o 30.25 25.40 30.74 53.40 **51.37** **51.12**
with **27.47** **20.82** **24.87** 23.66 **20.02** 26.36
LLaMA-7b
w/o 13.48 10.76 17.50 **25.82** 19.68 **26.68**
Token Subword SQuAD Original SQuAD
Alignment
EM ES EM ES
LLaMa-7b with **40.27** **78.81** **14.65** **52.24**
w/o 12.42 69.49 6.04 48.44
Token Subword WikiText Baseline of Subword WikiText
Alignment
Acc. ES Rouge-L Acc. ES Rouge-L
LLaMa-7b with **19.42** **47.80** **0.257** **1.45** 42.826 0.153
w/o 10.9 44.50 0.190 0.75 **44.141** **0.168**
**Metrics** We use pass@k (Kulal et al., 2019) with an unbiased estimate by Chen et al. (2021) as
the execution-based evaluation metric to measure the functional correctness of the generated code
snippet. Execution-based metrics is robust to variation in the generated code and can better reflect the
actual problem solving abilities of language models, which is highly suitable for our purposes as an
approach to test the effects of token alignment.
5.1.2 NATURAL LANGUAGE TASKS
We adapt several natural language evaluation benchmarks to reflect the impact of the subword in the
prompt.
    - Question answering: SQuAD (Rajpurkar et al., 2016; 2018) with exact match metric and
edit similarity metric.
    - Text generation: Wikitext (Merity et al., 2016) with first token accuracy and first n _[th]_ word
fuzzy-matching based metrics.
For all datasets, we use the original context and part of the ground truth as our prompt, where
the prompt ends with subword or without (control version). The control version always has one
space-delineated word fewer than the subword counterpart. In terms of partial token scenarios for
text-based evaluation, we primarily consider the subword and space prefix.
**Metrics** For question answering tasks, the ground-truth is a list of answers where each answer can
be one or more words. We use the edit similarity and exact match as the metrics. For text generation
with wikitext, we found that exact match metric is too strict for this type of open-ended generations
in open domains. Therefore, in addition to the first token accuracy, we also use fuzzy-matching based
metrics to calculate the similarity between the first n _[th]_ word of the generation and the groundtruth
with _n_ = 50 in our case. Here we use the edit similarity and Rouge-L (Lin, 2004) as the metrics.
5.2 TOKEN ALIGNMENT HANDLES MANY PARTIAL TOKEN SCENARIOS
In this section, we show that without token alignment, publicly available language models can suffer
from the partial token issue during text completion. We show that token alignment offers a simple
solution that applies to all the prompt scenarios considered. We primarily use StarCoder (Li et al.,
2023) and LLaMA 7B (Touvron et al., 2023) to perform this study.
Table 2: Token alignment for partial token of **punctuation** scenario.
Token Punc MBXP Baseline of Punc MBXP
Alignment
Python Java JavaScript Python Java JavaScript
with **59.35** **48.80** **46.03** **58.31** **51.95** **46.41**
StarCoder
w/o 45.10 31.78 24.23 58.31 49.31 46.15
with **31.75** **20.93** **26.79** **31.90** 21.44 27.05
LLaMA-7b
w/o 22.11 11.85 14.87 31.45 **24.21** **27.95**
5.2.1 SUBWORD
**Code Generation** The results with and without token alignment in Table 1 demonstrate clear
differences in the ability to handle prompts that ends with subword where the performance can drop
by _≈_ 14 _−_ 22% for pass@1 without token alignment. For all the results, we **boldface** the scenario that
obtain higher performance and highlight in red the scenario where the different is greater than 3 points
for clarity. We also show the scenarios where the prompt ends with a full word (space-delineated) as
a control experiment. In this case, we can see that token alignment yields similar scores as without
token alignment, indicating that token alignment can be used regardless of whether the prompt ends
with or without a partial token.
**Natural Language Tasks** We also show the text evaluation results in Table 1 where we observe
clear performance gain from token alignment. Observe that in the case of question answering
(SQuAD), even though we provide partial information in terms of subword, akin to providing a hint,
most models fail at answering correctly due to the constraint from the partial token artifact. This
resulting in subword SQuAD scores being lower than the baseline SQuAD even though the subword
dataset contains extra information in the prompt. However, once we use token alignment which
alleviates the tokenization constraint, the extra information helps significantly, allowing many models
to score in the range of 40% exact match versus 15% for the baseline case.
For text generation on WikiText, the next word accuracy, edit similarity, and ROUGE scores portray
obvious trends. When the prompt ends with a subword, all scores drop noticeably compared to
the baseline. Token alignment helps remove the tokenization constraint, resulting in significant
improvement for all metrics. On the control datasets which differ from the original datasets by
one fewer (sub)word, using token alignment leads to similar or slightly better results. The finding
indicates that token alignment can be used in all cases without the need to detect whether the prompt
ends with a subword or not.
5.2.2 PUNCTUATION
Table 2 demonstrates the results from the punctuation MBXP dataset, highlighting a notable improvement due to the use of token alignment. As an example, consider a prompt if x= representing a
code snippet in Python. It is unambiguous that the next character should be = since otherwise the
Python syntax would not be correct. However, this prompt corresponds to an ending partial token
= of a complete token ==; therefore, predicting an extra token ‘=’ is out-of-distribution compared
to what is seen during training, which is ‘==’. This dynamic is distinct from the natural subword
scenario where boundaries are determined by whitespace and the identification of partial tokens is
more clear. When punctuation is involved, avoiding an ending with a partial token becomes a nuanced
challenge, which token alignment effectively addresses.
5.2.3 SPACE PREFIX
The results from the space prefix MBXP datasets (for code generation) and space-SQuAD are shown
in Table 3, illustrating that a presence of a extra trailing token, while seemingly innocuous, can
noticeably degrade generation behavior without token alignment. For the code generation task, we
also have a dataset split _prefix-indent_ to ablate the case of prefix space with indentation. Table 3
illustrate the results where we observe slight improvement in the _prefix-sep_ case for code and a strong
improvement for text on SQuAD. We provide additional observations in Appendix A.2.2.
Table 3: **Space prefix** is a common design approach for many modern tokenizers for efficient
compression, but can lead to suboptimal generation if the prompt ends with a space character (due
to out-of-distribution). Token alignment resolves such problem gracefully, reflected by improved
pass@1 scores (%) on prefix-sep MBXP and prefix-indent MBXP. FS denotes few-shot prompting.
Token Prefix-sep MBXP Baseline’s
Alignment
Python Java JavaScript Python Java JavaScript
with **56.09** **50.75** **48.16** 56.73 **50.64** **48.83**
StarCoder
w/o 54.06 48.67 44.02 **57.48** 49.94 46.93
with **25.85** **21.67** **25.92** 23.82 **22.36** 25.59
LLaMA-7b
w/o 22.86 20.74 23.35 **25.00** 21.55 **26.48**
Token Prefix-indent MBXP Baseline’s
Alignment
Python Java JavaScript Python Java JavaScript
with **42.95** **50.19** **49.13** 48.82 **50.19** 46.89
StarCoder
w/o 17.10 49.31 45.33 **49.02** 50.06 **47.75**
with **45.52** **50.44** **46.89** 52.94 51.57 **46.37**
StarCoder FS
w/o 24.51 48.93 42.56 **53.14** **52.07** 46.02
with **15.45** 20.45 **28.03** 17.30 22.21 **25.95**
LLaMA-7b
w/o 12.05 **20.95** 24.39 **18.02** **22.33** 25.43
with **17.92** **23.59** **25.26** 18.74 20.83 **26.30**
LLaMA-7b FS
w/o 15.24 21.83 24.74 **19.36** **21.46** 25.78
Token Space Prefix SQuAD Baseline’s
Alignment
EM ES EM ES
LLaMA-7b with **26.82** **64.34** **9.51** **62.22**
w/o 20.45 61.34 4.38 61.88
5.2.4 CONTIGUOUS SPACES
The results on the contiguous space MBXP dataset in Table 4 show that token alignment alleviates
the partial token issue significantly, especially in Python, where the syntax is most sensitive to spaces.
We emphasize that, superficially, this case does not seem obviously linked to partial token issues (see
Figure 3). Hence, we separated it to highlight its importance and to measure the effects with and
without token alignment.
Besides this processed dataset, token alignment also impacts the _regular_ execution-based evaluation
significantly for some models that employ aggressive white space grouping in their tokenizer for high
compression. For example, in the StarCoder model, a newline character is often grouped with spaces,
such as \n, leading prompts ending with a newline character, \n, to correspond to a partial
token case. Without token alignment, we observe suboptimal behavior, as illustrated in Table 4, where
the execution scores drop from 44 _._ 0% to 27 _._ 2% without token alignment for MBXP JavaScript. For
instance, when the model encounters a JavaScript prompt ending with {\n, indicating the beginning
of a function, it might generate }\n without actually completing the function. Similar observations
occur in other programming language benchmarks.
6 ABLATION STUDY AND GENERALITY OF TOKEN ALIGNMENT
In this section, we also discuss an alternative approach to subword completion by processing training
data by subword regularization. We show that our token alignment method is complementary and can
offer a robust approach to handle subwords compatible with subword regularization. We also discuss
latency as well as ablation study on the number of backtrack tokens.
Table 4: Evaluation scores (pass@1 %) for models on the processed contiguous space MBXP and
regular MBXP datasets, considering the influence of token alignment and the partial token constraint
due to **contiguous spaces** in the tokenization.
Token Contiguous Space MBXP Baseline’s
Alignment
Python Java JavaScript Python Java JavaScript
with **43.93** 49.20 **49.41** 47.94 48.86 **50.37**
StarCoder
w/o 32.00 **50.46** 48.45 **49.28** **50.92** 49.73
with **16.56** 18.19 23.91 15.64 17.39 23.05
LLaMA-7b
w/o 14.20 **19.11** **24.01** **18.72** **18.99** **24.44**
Token MBXP with StarCoder
Alignment
Python JS Java C++ Swift TS Kotlin Go Scala
with **42.5** **44.0** **43.0** **43.3** **27.6** **42.3** **37.9** 30.8 **34.1**
w/o 7.5 27.2 42.0 39.2 8.9 7.0 16.4 30.8 22.4
6.1 COMPARISON AND COMPATIBILITY WITH SUBWORD REGULARIZATION
Subword regularization performs tokenization in a more random manner where each word can be
tokenized differently without a fixed pattern. We trained a model with subword regularization on top
of a regular language model and compare with the performance with base model before subword
regularization, with and without token alignment (details in Appendix A.2.1). We find subword
regularization alone can perform quite well on many partial token scenarios, almost matching the
scores obtained from token alignment (Table 6). Token alignment helps increase the scores further,
indicating that subword regularization itself can benefit from token alignment process. Overall, we
find that token alignment can help alleviate constraints on tenuous cases that subword regularization
does not address. We provide detailed results in Appendix A.2.1.
6.2 LATENCY AND THE NUMBER OF BACKTRACK TOKENS
The extra latency from token alignment comes mainly from two components. First, backtracking
results in more tokens needed during incremental decoding. For the usual decoding process, the
latency added is _B_ _[′]_ _· ℓd_, where _B_ _[′]_ is the number of decoding steps to finish the character or byte
matching process of token alignment, and _ℓd_ is the incremental decoding latency per token. [†] We
show the distribution of the number of steps _B_ _[′]_ is Figure 4 in Appendix, which shows that _B_ _[′]_ peaks
at the number of backtrack tokens _B_ and can be slightly lower or higher. Second, there is a minimal
cost of performing alignment _ℓT A_, which entails trie lookup and masking the next token probabilities,
which is in the order of _<_ 1 ms in most cases. After the alignment prefix (Figure 1) is fully matched
and becomes an empty string, we stop the alignment process and no longer incurs the extra latency
per token _ℓT A_ . Overall, the minimal extra latency makes it viable to use this algorithm in real time
text or code completion applications.
Throughout this experiment, we use a fixed number of tokens to backtrack _B_ = 3 and provide an
ablation study for the importance of this backtrack tokens. The detailed result in Appendix Table
7 suggests that _B_ = 1 performs slightly worse than _B_ = 3 with a difference of _−_ 3 _._ 7% on average.
Since the latency for the _B_ = 3 is not much more than that of _B_ = 1, we use _B_ = 3 as the default
setting. We note that some tokenizers use pre-token schema which offers opportunities to mark some
token as the beginning of pretoken at building time, such as in SentencePiece (Kudo & Richardson,
2018), in which case we can dynamically determine the number of backtrack tokens based on the
context. However, most tokenizers do not have this feature so we use a fixed _B_ .
†With speculative decoding, it will have a speedup due to lower amortized latency per token.
7 CONCLUSION
This paper explored the complex challenge of handling partial tokens in generative models, introducing specific scenarios such as subword, punctuation, space prefix, and contiguous spaces. Through
the application of token alignment, we demonstrated noticeable improvements across these scenarios.
In the future, as tokens become longer due to possible advancement in tokenizer compression, the
token alignment may become increasingly important as more cases may subtly fall into the categories
of partial tokens. The findings of this study emphasize the importance of understanding tokenization
artifacts and present a meaningful advancement in enhancing the robustness and practical applicability
of generative models in various real-world tasks.
8 LIMITATIONS
While token alignment works in principles on any tokenizer, it may require additional attention to
details to implement, especially with tokenizers that are not lossless which could make the token
alignment process tricky. The latency addition due to token alignment is generally small for long
generation, but can be a larger portion of overall budget if the generation is short.
REFERENCES
Ben Athiwaratkun, Sanjay Krishna Gouda, Zijian Wang, Xiaopeng Li, Yuchen Tian, Ming Tan,
Wasi Uddin Ahmad, Shiqi Wang, Qing Sun, Mingyue Shang, Sujan Kumar Gonugondla, Hantian
Ding, Varun Kumar, Nathan Fulton, Arash Farahani, Siddhartha Jain, Robert Giaquinto, Haifeng
Qian, Murali Krishna Ramanathan, Ramesh Nallapati, Baishakhi Ray, Parminder Bhatia, Sudipta
Sengupta, Dan Roth, and Bing Xiang. Multi-lingual evaluation of code generation models.
_CoRR_ [, abs/2210.14868, 2022. doi: 10.48550/arXiv.2210.14868. URL https://doi.org/10.](https://doi.org/10.48550/arXiv.2210.14868)
[48550/arXiv.2210.14868.](https://doi.org/10.48550/arXiv.2210.14868)
Jacob Austin, Augustus Odena, Maxwell I. Nye, Maarten Bosma, Henryk Michalewski, David Dohan,
Ellen Jiang, Carrie J. Cai, Michael Terry, Quoc V. Le, and Charles Sutton. Program synthesis
with large language models. _CoRR_ [, abs/2108.07732, 2021. URL https://arxiv.org/abs/](https://arxiv.org/abs/2108.07732)
[2108.07732.](https://arxiv.org/abs/2108.07732)
Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry
Tworek, and Mark Chen. Efficient training of language models to fill in the middle. _CoRR_,
abs/2207.14255, 2022. doi: 10.48550/arXiv.2207.14255. [URL https://doi.org/10.](https://doi.org/10.48550/arXiv.2207.14255)
[48550/arXiv.2207.14255.](https://doi.org/10.48550/arXiv.2207.14255)
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal,
Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel
Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler,
Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott
Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya
Sutskever, and Dario Amodei. Language models are few-shot learners. _CoRR_, abs/2005.14165,
[2020. URL https://arxiv.org/abs/2005.14165.](https://arxiv.org/abs/2005.14165)
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and
John Jumper. Accelerating large language model decoding with speculative sampling. _CoRR_,
abs/2302.01318, 2023. doi: 10.48550/arXiv.2302.01318. [URL https://doi.org/10.](https://doi.org/10.48550/arXiv.2302.01318)
[48550/arXiv.2302.01318.](https://doi.org/10.48550/arXiv.2302.01318)
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Pondé de Oliveira Pinto, Jared
Kaplan, Harrison Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri,
Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan,
Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian,
Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios
Chantzis, Elizabeth Barnes, Ariel Herbert-Voss, William Hebgen Guss, Alex Nichol, Alex Paino,
Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders,
Christopher Hesse, Andrew N. Carr, Jan Leike, Joshua Achiam, Vedant Misra, Evan Morikawa,
Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob
McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating
large language models trained on code. _CoRR_ [, abs/2107.03374, 2021. URL https://arxiv.](https://arxiv.org/abs/2107.03374)
[org/abs/2107.03374.](https://arxiv.org/abs/2107.03374)
Kaustubh D Dhole, Varun Gangal, Sebastian Gehrmann, Aadesh Gupta, Zhenhao Li, Saad Mahamood,
Abinaya Mahendiran, Simon Mille, Ashish Srivastava, Samson Tan, et al. Nl-augmenter: A
framework for task-sensitive natural language augmentation. _ArXiv preprint_, abs/2112.02721,
[2021. URL https://arxiv.org/abs/2112.02721.](https://arxiv.org/abs/2112.02721)
Daniel Fried, Armen Aghajanyan, Jessy Lin, Sida Wang, Eric Wallace, Freda Shi, Ruiqi Zhong,
Scott Yih, Luke Zettlemoyer, and Mike Lewis. Incoder: A generative model for code infilling and
synthesis. In _The Eleventh International Conference on Learning Representations, ICLR 2023,_
_Kigali, Rwanda, May 1-5, 2023_ [. OpenReview.net, 2023. URL https://openreview.net/](https://openreview.net/pdf?id=hQwb-lbM6EL)
[pdf?id=hQwb-lbM6EL.](https://openreview.net/pdf?id=hQwb-lbM6EL)
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text
degeneration. In _8th International Conference on Learning Representations, ICLR 2020, Addis_
_Ababa, Ethiopia, April 26-30, 2020_ [. OpenReview.net, 2020. URL https://openreview.](https://openreview.net/forum?id=rygGQyrFvH)
[net/forum?id=rygGQyrFvH.](https://openreview.net/forum?id=rygGQyrFvH)
Akshita Jha and Chandan K. Reddy. Codeattack: Code-based adversarial attacks for pre-trained
programming language models. In Brian Williams, Yiling Chen, and Jennifer Neville (eds.),
_Thirty-Seventh AAAI Conference on Artificial Intelligence, AAAI 2023, Thirty-Fifth Conference on_
_Innovative Applications of Artificial Intelligence, IAAI 2023, Thirteenth Symposium on Educational_
_Advances in Artificial Intelligence, EAAI 2023, Washington, DC, USA, February 7-14, 2023_, pp.
[14892–14900. AAAI Press, 2023. URL https://ojs.aaai.org/index.php/AAAI/](https://ojs.aaai.org/index.php/AAAI/article/view/26739)
[article/view/26739.](https://ojs.aaai.org/index.php/AAAI/article/view/26739)
Di Jin, Zhijing Jin, Joey Tianyi Zhou, and Peter Szolovits. Is BERT really robust? A strong baseline
for natural language attack on text classification and entailment. In _The Thirty-Fourth AAAI_
_Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications_
_of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational_
_Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020_, pp. 8018–
[8025. AAAI Press, 2020. URL https://aaai.org/ojs/index.php/AAAI/article/](https://aaai.org/ojs/index.php/AAAI/article/view/6311)
[view/6311.](https://aaai.org/ojs/index.php/AAAI/article/view/6311)
Douwe Kiela, Max Bartolo, Yixin Nie, Divyansh Kaushik, Atticus Geiger, Zhengxuan Wu, Bertie
Vidgen, Grusha Prasad, Amanpreet Singh, Pratik Ringshia, Zhiyi Ma, Tristan Thrush, Sebastian
Riedel, Zeerak Waseem, Pontus Stenetorp, Robin Jia, Mohit Bansal, Christopher Potts, and Adina
Williams. Dynabench: Rethinking benchmarking in NLP. In _North American Association for_
_Computational Linguistics (NAACL)_, pp. 4110–4124, 2021.
Gyuwan Kim. Subword language model for query auto-completion. In Kentaro Inui, Jing Jiang,
Vincent Ng, and Xiaojun Wan (eds.), _Proceedings of the 2019 Conference on Empirical Methods_
_in Natural Language Processing and the 9th International Joint Conference on Natural Language_
_Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019_, pp. 5021–5031.
[Association for Computational Linguistics, 2019. doi: 10.18653/V1/D19-1507. URL https:](https://doi.org/10.18653/v1/D19-1507)
[//doi.org/10.18653/v1/D19-1507.](https://doi.org/10.18653/v1/D19-1507)
Taku Kudo. Subword regularization: Improving neural network translation models with multiple subword candidates. In _Proceedings of the 56th Annual Meeting of the Association_
_for Computational Linguistics (Volume 1: Long Papers)_, pp. 66–75, Melbourne, Australia,
July 2018a. Association for Computational Linguistics. doi: 10.18653/v1/P18-1007. URL
[https://aclanthology.org/P18-1007.](https://aclanthology.org/P18-1007)
Taku Kudo. Subword regularization: Improving neural network translation models with multiple
subword candidates. In Iryna Gurevych and Yusuke Miyao (eds.), _Proceedings of the 56th Annual_
_Meeting of the Association for Computational Linguistics, ACL 2018, Melbourne, Australia, July_
_15-20, 2018, Volume 1: Long Papers_, pp. 66–75. Association for Computational Linguistics, 2018b.
[doi: 10.18653/v1/P18-1007. URL https://aclanthology.org/P18-1007/.](https://aclanthology.org/P18-1007/)
Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword
tokenizer and detokenizer for neural text processing. In Eduardo Blanco and Wei Lu (eds.),
_Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing,_
_EMNLP 2018: System Demonstrations, Brussels, Belgium, October 31 - November 4, 2018_, pp.
66–71. Association for Computational Linguistics, 2018. doi: 10.18653/v1/d18-2012. URL
[https://doi.org/10.18653/v1/d18-2012.](https://doi.org/10.18653/v1/d18-2012)
Sumith Kulal, Panupong Pasupat, Kartik Chandra, Mina Lee, Oded Padon, Alex Aiken, and Percy
Liang. Spoc: Search-based pseudocode to code. _CoRR_ [, abs/1906.04908, 2019. URL http:](http://arxiv.org/abs/1906.04908)
[//arxiv.org/abs/1906.04908.](http://arxiv.org/abs/1906.04908)
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative
decoding. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan
Sabato, and Jonathan Scarlett (eds.), _International Conference on Machine Learning, ICML 2023,_
_23-29 July 2023, Honolulu, Hawaii, USA_, volume 202 of _Proceedings of Machine Learning_
_Research_ [, pp. 19274–19286. PMLR, 2023. URL https://proceedings.mlr.press/](https://proceedings.mlr.press/v202/leviathan23a.html)
[v202/leviathan23a.html.](https://proceedings.mlr.press/v202/leviathan23a.html)
Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou,
Marc Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Starcoder: may the source be with
you! _arXiv preprint
Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom
Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition-level code generation
with alphacode. _Science_, 378(6624):1092–1097, 2022.
Chin-Yew Lin. ROUGE: A package for automatic evaluation of summaries. In _Text Summarization_
_Branches Out_, pp. 74–81, Barcelona, Spain, July 2004. Association for Computational Linguistics.
[URL https://aclanthology.org/W04-1013.](https://aclanthology.org/W04-1013)
Scott Lundberg. The art of prompt design: Prompt boundaries and token healing, 2023. URL [https://towardsdatascience.com/](https://towardsdatascience.com/the-art-of-prompt-design-prompt-boundaries-and-token-healing-3b2448b0be38)
[the-art-of-prompt-design-prompt-boundaries-and-token-healing-3b2448b0be38.](https://towardsdatascience.com/the-art-of-prompt-design-prompt-boundaries-and-token-healing-3b2448b0be38)
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture
models, 2016.
Yixin Nie, Adina Williams, Emily Dinan, Mohit Bansal, J. Weston, and Douwe Kiela. Adversarial
nli: A new benchmark for natural language understanding. In _Association for Computational_
_Linguistics (ACL)_, 2020.
Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou, Silvio Savarese,
and Caiming Xiong. Codegen: An open large language model for code with multi-turn program
synthesis. _arXiv preprint
[OpenAI. tiktoken: A fast BPE tokeniser for use with OpenAI’s models. https://github.com/](https://github.com/openai/tiktoken)
[openai/tiktoken, 2023.](https://github.com/openai/tiktoken)
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100, 000+ questions for
machine comprehension of text. In Jian Su, Xavier Carreras, and Kevin Duh (eds.), _Proceedings of_
_the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin,_
_Texas, USA, November 1-4, 2016_, pp. 2383–2392. The Association for Computational Linguistics,
[2016. doi: 10.18653/v1/d16-1264. URL https://doi.org/10.18653/v1/d16-1264.](https://doi.org/10.18653/v1/d16-1264)
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don’t know: Unanswerable questions
for squad. In Iryna Gurevych and Yusuke Miyao (eds.), _Proceedings of the 56th Annual Meeting_
_of the Association for Computational Linguistics, ACL 2018, Melbourne, Australia, July 15-20,_
_2018, Volume 2: Short Papers_, pp. 784–789. Association for Computational Linguistics, 2018. doi:
[10.18653/v1/P18-2124. URL https://aclanthology.org/P18-2124/.](https://aclanthology.org/P18-2124/)
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilic, Daniel Hesslow, Roman
Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, Jonathan Tow, Alexander M.
Rush, Stella Biderman, Albert Webson, Pawan Sasanka Ammanamanchi, Thomas Wang, Benoît
Sagot, Niklas Muennighoff, Albert Villanova del Moral, Olatunji Ruwase, Rachel Bawden, Stas
Bekman, Angelina McMillan-Major, Iz Beltagy, Huu Nguyen, Lucile Saulnier, Samson Tan,
Pedro Ortiz Suarez, Victor Sanh, Hugo Laurençon, Yacine Jernite, Julien Launay, Margaret
Mitchell, Colin Raffel, Aaron Gokaslan, Adi Simhi, Aitor Soroa, Alham Fikri Aji, Amit Alfassy,
Anna Rogers, Ariel Kreisberg Nitzav, Canwen Xu, Chenghao Mou, Chris Emezue, Christopher
Klamm, Colin Leong, Daniel van Strien, David Ifeoluwa Adelani, and et al. BLOOM: A 176bparameter open-access multilingual language model. _CoRR_, abs/2211.05100, 2022. doi: 10.48550/
[arXiv.2211.05100. URL https://doi.org/10.48550/arXiv.2211.05100.](https://doi.org/10.48550/arXiv.2211.05100)
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée
Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand
Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language
models. _CoRR_ [, abs/2302.13971, 2023. doi: 10.48550/arXiv.2302.13971. URL https://doi.](https://doi.org/10.48550/arXiv.2302.13971)
[org/10.48550/arXiv.2302.13971.](https://doi.org/10.48550/arXiv.2302.13971)
Boxin Wang, Chejian Xu, Shuohang Wang, Zhe Gan, Yu Cheng, Jianfeng Gao, Ahmed Hassan
Awadallah, and Bo Li. Adversarial glue: A multi-task benchmark for robustness evaluation of
language models. In _Thirty-fifth Conference on Neural Information Processing Systems Datasets_
_and Benchmarks Track (Round 2)_, 2021.
Shiqi Wang, Zheng Li, Haifeng Qian, Chenghao Yang, Zijian Wang, Mingyue Shang, Varun Kumar,
Samson Tan, Baishakhi Ray, Parminder Bhatia, Ramesh Nallapati, Murali Krishna Ramanathan,
Dan Roth, and Bing Xiang. Recode: Robustness evaluation of code generation models. In Anna
Rogers, Jordan L. Boyd-Graber, and Naoaki Okazaki (eds.), _Proceedings of the 61st Annual_
_Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2023,_
_Toronto, Canada, July 9-14, 2023_, pp. 13818–13843. Association for Computational Linguistics,
[2023. URL https://aclanthology.org/2023.acl-long.773.](https://aclanthology.org/2023.acl-long.773)
Zhou Yang, Jieke Shi, Junda He, and David Lo. Natural attack for pre-trained models of code. In
_44th IEEE/ACM 44th International Conference on Software Engineering, ICSE 2022, Pittsburgh,_
_PA, USA, May 25-27, 2022_, pp. 1482–1493. ACM, 2022. doi: 10.1145/3510003.3510146. URL
[https://doi.org/10.1145/3510003.3510146.](https://doi.org/10.1145/3510003.3510146)
Yuan Zang, Fanchao Qi, Chenghao Yang, Zhiyuan Liu, Meng Zhang, Qun Liu, and Maosong Sun.
Word-level textual adversarial attacking as combinatorial optimization. In _Proceedings of the_
_58th Annual Meeting of the Association for Computational Linguistics_, pp. 6066–6080, Online,
2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.acl-main.540. URL
[https://aclanthology.org/2020.acl-main.540.](https://aclanthology.org/2020.acl-main.540)
A APPENDIX
A.1 FAQS
1. **Q** : Is token alignment absolutely necessary?
**A** : Yes and no. In the scenario of text _completion_ such as the scenarios in IDE environment,
token alignment is highly beneficial. In the formation of question answering in chat styles
(such as in ChatGPT), the model can typically generate the completion without being
constrained by the artifact of tokenization.
2. **Q** : Is token alignment compatible with the code insertion scenario?
**A** : Yes. Since the insert scenario (Bavarian et al., 2022; Fried et al., 2023) performs text
_completion_ where the prompt during inference uses the right context follows by left context,
the insertion is simply a continuation. This makes the token alignment method compatible
with the insertion scenario.
3. **Q** : Is it compatible with speculative decoding?
**A** : Yes. For speculative decoding (Leviathan et al., 2023; Chen et al., 2023), the draft model
can propose tokens that match the prompt during token alignment, which will be either
accepted or rejected by the main model based on the token sequence.
4. **Q** : The outputs of LLMs were uncertain. Even a minor change in a prompt could lead to
variations in the output. During the use of prompts in the paper, was the specific impact of
the prompt considered?
**A** : We use the baseline datasets where we pre-backtrack so that the prompt does not contain
partial token. Therefore, the difference between the baseline and non baseline is as low as
possible, in order to control the effect of the prompt. Moreover, we average the results over
many curated datasets to understand the aggregate behavior.
5. **Q** : Given the powerful in-context learning capabilities of large language models, it would be
worth exploring whether adding relevant knowledge to the prompt could further enhance the
proposed method.
**A** : One way to use in-context learning is to make the model repeat the entire input and
generate full tokens that match the partial tokens instead. However, if the input is long, this
will amount to high extra latency. Therefore, we do not explicitly consider this approach.
It is possible that we do in-context learning in such a way that makes the model output only
a snippet shortly before the end of the prompt. Such method is likely quite complicated and
requires processing to extract out the right portion.
Our method presents a clean way to align tokens natively by guiding the token selection
with the probabilities that the model already predicts, but masking out tokens that do not
align. We believe that this is the cleanest and simplest approach.
6. **Q** : In which areas does token alignment excel? Are there scenarios or applications where
other methods might be more appropriate?
**A** : The token alignment is very suitable for _completion_ mode. For instance, if we use
language models to help write word document, or write code within the code editor. This is
in contrast to the question-answer mode where the model does not have to output text that is
a continuation of the prompt, but would perhaps start in a different line, with a new sentence
or a disjoint piece of code instead.
7. **Q** : How does subword regularization increase inference latency?
**A** : Subword regularization breaks up a piece of text into more tokens and show this to the
model during training time. The model then has higher chance of using higher number of
tokens to produce the same text, compared to without subword regularization.
Our approach of token alignment incurs additional latency only at the beginning. In contrast,
models trained via subword regularization may use additional tokens anywhere during the
incremental decoding phase.
A.2 ADDITIONAL RESULTS
A.2.1 SUBWORD REGULARIZATION
In Table 6, we outline the results where we show the performance of a model trained without
subword regularization (the usual training strategy employed in many language models) and with
additional training with subword regularization (Kudo, 2018b). We use the implementation given in
SentencePiece (Kudo & Richardson, 2018) with _α_ = 0 _._ 05. For instance, New York is segmented
into either N, e, w, _York or New, _York, where lower alpha means higher chances of the
canonical segmentation being more common. The results indicate that subword regularization without
token alignment indeed helps increase the performance in the subword MBXP from 6.89% to 18.11%,
but using token alignment with subword regularization model increases the scores further to 18.52%.
This finding highlights the generality of token alignment where it can be used with or without subword
regularization, and can even help address cases where subword regularization alone may not handle
based on the consistent improved scores of token alignment. In addition, subword regularization may
introduce additional latency cost due to finer grain tokenization of text in the training stage, which
influences how the tokens are generated during inference. We leave this direction of research on
investigating the interplay of subword regularization and tradeoff for future work.
A.2.2 ADDITIONAL OBSERVATIONS
Note that on the _prefix-indent_ split, with token alignment, the generation’s quality improved significantly but not quite matching the baseline (e.g. 42.95 for partial token + token alignment versus
Table 5: Pass@1 scores for subword regularization and token alignment with 650M model.
Token Alignment Partial Word MBXP Baseline’s
Subword-regularized model with **18.52%** 15.43%
w/o 18.11% 17.39%
Baseline model with **16.67%** 12.96%
w/o 6.89% 13.89%
Table 6: Effects of subword regularization and token alignment
Token Partial Word Baseline Word Partial Punc Baseline Punc
Model
Alignment MBPP MBPP MBPP MBPP
Subword with **18.52%** 15.43% 17.96% **19.44%**
Regularized w/o 18.11% **17.39%** **18.10%** 19.14%
Baseline with **16.67%** 12.96% **18.84%** 18.25%
600M w/o 6.89% 13.89% 11.72% **18.69%**
48.82 for baseline for Python StarCoder). We found that the subpar score compared to baseline is
due some examples corresponding to a function signature and extra spaces, which can confuse the
model that is not fully preference aligned. We include the few-shot prompting version (FS) which
includes examples of complete functions prepended at the beginning, which is to help steer the model
to complete the function and observe noticeable improvement, but still does not quite close the gap.
We hypothesize with a model that is more preference aligned will handle such cases better, where
this partial token dataset can be used as once of such benchmark to measure instruction following
abilities.
A.2.3 NUMBER OF BACKTRACK TOKENS
We show detailed results on evaluating with varying number of backtrack tokens _B_ in Table 7. While
the results for individual datasets vary, on average, the overall difference between the scores from
_B_ = 1 and _B_ = 3 is _−_ 3 _._ 69%, indicating that _B_ = 1 is less accurate since it may not backtrack
enough to a full token.
In terms of latency, the factor that affects the latency the most is the number of steps it takes to get
out of the character matching mode during token alignment. Figure 4 shows the distribution of the
number of steps in the case of _B_ = 3, where the distribution peaks around _B_ steps, with some long
tail towards higher the number of steps. We emphasize that the number of steps required to complete
the character matching need not be exactly equal to the number of backtrack tokens. This is because
given a string that corresponds to the backtracked prompt, the model can generate any tokens as long
as it matches, which could result in breaking the string down to coarser or finer granularity.
Table 7: Execution-based evaluation on various partial token versions of MBXP (pass@1). The
difference ∆ denotes the scores of B=-1 minus the scores of B=3. The total overall difference is
-3.69%.
Token Partial Word MBXP Baseline’s
Alignment Python Java JavaScript Python Java JavaScript
B=1 56.58% 52.17% 49.31% 54.32% 49.54% 50.16%
B=3 54.01% 47.37% 43.76% 54.12% 52.29% 50.59%
∆ 2.57% 4.81% 5.55% 0.21% -2.75% -0.43%
Partial Punc MBXP Baseline’s
B=1 59.35% 48.80% 46.03% 58.31% 51.95% 46.41%
B=3 59.64% 50.82% 47.44% 58.90% 50.95% 44.74%
∆ -0.30% -2.02% -1.41% -0.59% 1.01% 1.67%
Prefix-sep MBXP Baseline’s
B=1 56.09% 50.75% 48.16% 56.73% 50.64% 48.83%
B=3 58.23% 51.22% 48.16% 57.48% 50.17% 47.71%
∆ -2.14% -0.46% 0.00% -0.75% 0.46% 1.12%
Prefix-indent MBXP Baseline’s
B=1 42.95% 50.19% 49.13% 48.82% 50.19% 46.89%
B=3 44.08% 50.31% 46.71% 49.54% 53.32% 47.23%
∆ -1.13% -0.13% 2.42% -0.72% -3.14% -0.35%
Contiguous Space MBXP Baseline’s
B=1 43.93% 49.20% 49.41% 47.94% 48.86% 50.37%
B=3 45.88% 50.23% 48.99% 49.38% 49.66% 49.73%
∆ -1.95% -1.03% 0.43% -1.44% -0.80% 0.64%
|Token Alignment Step Distribution (StarCoder)|Col2|Col3|Col4|Col5|Col6|Col7|
|---|---|---|---|---|---|---|
|Token Alignment Step Distribution (StarCoder)|Token Alignment Step Distribution (StarCoder)|Token Alignment Step Distribution (StarCoder)|Token Alignment Step Distribution (StarCoder)|Token Alignment Step Distribution (StarCoder)|Token Alignment Step Distribution (StarCoder)|Token Alignment Step Distribution (StarCoder)|
|||||||MBPP<br>|
||||B<br>|asel<br>|ine~~-~~<br>|MBPP<br>|
||||S|ubword~~-~~|ubword~~-~~|MBPP|
||||||||
||||||||
||||||||
||||||||
||||||||
||||||||
Figure 4: The density estimation of the number of steps during matching process in token alignment.

gondla, Sanjay Krishna Gouda, Rob Kwiatowski, Ramesh Nallapati, Bing Xiang


Swapnaja Achintalwar, Ioana Baldini, Djallel Bouneffouf, Joan Byamugisha, Maria Chang, Pierre Dognin,
Eitan Farchi, Ndivhuwo Makondo, Aleksandra Mojsilovi´c, Manish Nagireddy, Karthikeyan Natesan Ramamurthy,
Inkit Padhi, Orna Raz, Jesus Rios, Prasanna Sattigeri, Moninder Singh, Siphiwe Thwala, Rosario A. Uceda-Sosa,
and Kush R. Varshney
_**Abstract**_ **—The alignment of large language models is usually**

Alignment Studio: Aligning Large Language Models to Particular
  Contextual Regulations

Swapnaja Achintalwar
Ioana Baldini
Djallel Bouneffouf
Joan Byamugisha
Maria Chang
Pierre Dognin
Eitan Farchi
Ndivhuwo Makondo
Aleksandra Mojsilovic
Manish Nagireddy
Karthikeyan Natesan Ramamurthy
Inkit Padhi
Orna Raz
Jesus Rios
Prasanna Sattigeri
Moninder Singh
Siphiwe Thwala
Rosario A. Uceda-Sosa
Kush R. Varshney


**common or universally understood across use cases and contexts.**

done by model providers to add or control behaviors that are


**work in concert to control the behavior of a language model.**

Studio architecture: Framers, Instructors, and Auditors that


**conduct guidelines.**
_**Index Terms**_ **—AI alignment, knowledge representation, fine-**

a company’s internal-facing enterprise chatbot to its business


INTRODUCTION
RE-TRAINED large language models (LLMs) are usually
tuned by model providers to endow them with different