# Reconfigurable Intelligent Surface-Aided MISO Systems with Statistical CSI:
Channel Estimation, Analysis and Optimization
_(Invited Paper)_

Kangda Zhi1, Cunhua Pan1, Hong Ren2, Kezhi Wang3 and Maged Elkashlan1
1EECS, Queen Mary University of London, {k.zhi, c.pan,
<EMAIL_ADDRESS>
2National Mobile Communications Research Laboratory, Southeast University,
<EMAIL_ADDRESS>
3Department of Computer and Information Sciences, Northumbria University,
<EMAIL_ADDRESS>
###### Abstract

This paper investigates the reconfigurable reflecting surface (RIS)-aided
multiple-input-single-output (MISO) systems with imperfect channel state
information (CSI), where RIS-related channels are modeled by Rician fading.
Considering the overhead and complexity in practical systems, we employ the
low-complexity maximum ratio combining (MRC) beamforming at the base station
(BS), and configure the phase shifts of the RIS based on long-term statistical
CSI. Specifically, we first estimate the overall channel matrix based on the
linear minimum mean square error (LMMSE) estimator, and evaluate the
performance of MSE and normalized MSE (NMSE). Then, with the estimated
channel, we derive the closed-form expressions of the ergodic rate. The
derived expressions show that with Rician RIS-related channels, the rate can
maintain at a non-zero value when the transmit power is scaled down
proportionally to $1/M$ or $1/N^{2}$, where $M$ and $N$ are the number of
antennas and reflecting elements, respectively. However, if all the RIS-
related channels are fully Rayleigh, the transmit power of each user can only
be scaled down proportionally to $1/\sqrt{M}$ or $1/N$. Finally, numerical
results verify the promising benefits from the RIS to traditional MISO
systems.

###### Index Terms:

Intelligent reflecting surface (IRS), Reconfigurable Intelligent Surface
(RIS), statistical CSI, channel estimation, Rician fading.

##  I Introduction

As an emerging candidate for next-generation communication systems,
reconfigurable intelligent surfaces (RISs), also called intelligent reflecting
surfaces (IRSs), have attracted significant interest from both academia and
industry[1]. By intelligently tuning the phase shifts of the impinging wave,
an RIS can constructively strengthen the desired signal power or
deconstructively weaken the interference signal, which yields an appealing
passive beamforming gain[2, 3, 4].

The gains from RISs highly rely on the knowledge of the channel state
information (CSI). However, the accurate acquirement of the CSI in RIS-aided
systems is a challenging task, since RISs are typically passive without RF
chains. To cope with this challenge, researchers have turned their attention
to designing the phase shifts of RISs based on statistical CSI [5, 6, 7].
Thanks to the slowly changing nature of the statistical CSI, we only need to
configure the phase shifts of the RIS once over a large time-scale.
Accordingly, this greatly reduce the computational complexity, and reduce the
overhead of channel estimation and feedback.

However, existing statistical-CSI based contributions mainly characterized the
systems performance under the assumption that perfect instantaneous overall
CSI is available at the base station (BS). In practice, channel estimation
error is inevitable. Only a few contributions have taken imperfect CSI into
consideration when configuring the RIS based on statistical CSI[8, 9, 10].
Specifically, the impact of imperfect CSI on multi-RIS single-user systems was
evaluated in [8], where the channel estimation is performed for each
individual channel. Different from [8], the authors in [9, 10] investigated
the system performance when considering the cascaded channel estimation
instead of individual channels. However, the contributions[9, 10] only
considered a simple channel model where the RIS-BS link is fully line-of-sight
(LoS) while RIS-users links are fully non-LoS (NLoS). Although this model is
mathematically tractable, it cannot reveal the fundamental impact of the
scatter richness of the propagation environment on RIS-aided systems. With a
variety of richness of scatter, the quality of channel estimation and the
system capacity performance in RIS-aided systems will exhibit wide
differences.

Motivated by the above, we focus on the uplink transmission of an RIS-aided
multiple-input-single-output (MISO) system with imperfect CSI, where all the
RIS-related channels follow Rician distribution. We first propose a low-
complexity linear minimum mean square error (LMMSE) estimator to estimate the
channels as a whole in RIS-aided MISO systems. Based on the estimated channel,
we theoretically investigate the performance of the mean square error (MSE)
and normalized MSE (NMSE). Second, with low-complexity maximum ratio combining
(MRC) receiver, rigorous expressions for the uplink ergodic data rate are
derived in closed-form. Based on the expression, we design the phase shifts of
the RIS and then characterize the asymptotic behaviors of the rate, the power
scaling laws, and the impact of Rician factors. Finally, numerical results are
presented to validate the superiority of RIS-aided MISO systems over
traditional RIS-free MISO systems.

##  II System Model

Figure 1: An RIS-aided uplink MISO system.

As illustrated in Fig. 1, we focus on the uplink transmission of an RIS-aided
MISO system, where an RIS is deployed in the proximity of a user to assist the
transmission of a BS. The BS is equipped with $M$ antennas, the RIS is
comprised of $N$ passive reflecting elements, and the user are equipped with
single antenna. Without loss of generality, the user is denoted by user $k$.
The channels from the user to the BS, from the user to the RIS, and from the
RIS to the BS are denoted by $\mathbf{d}_{k}\in\mathbb{C}^{M\times 1}$,
$\mathbf{h}_{k}\in\mathbb{C}^{N\times 1}$ and
$\mathbf{H}_{2}\in\mathbb{C}^{M\times N}$, respectively.

The phase shift matrix of the RIS is expressed as
$\mathbf{\Phi}=\mathrm{diag}\left\\{e^{j\theta_{1}},...,e^{j\theta_{N}}\right\\}$,
where $\theta_{n}\in[0,2\pi)$ represents the phase shift of the $n$-th
reflecting element. Based on the above definitions, the cascaded user-RIS-BS
channel is obtained as $\mathbf{g}_{k}=\mathbf{H}_{2}{\bf\Phi}{\bf h}_{k}$,
and the overall channels is denoted by
$\mathbf{q}_{k}=\mathbf{g}_{k}+\mathbf{d}_{k}\in\mathbb{C}^{M\times 1}$. Then,
the received signal vector at the BS is given by $\mathbf{y}=\sqrt{p}{\bf
q}_{k}x_{k}+\mathbf{n}$ , where $p$ is the transmit power, $x_{k}$ represents
the transmit symbol, and $\mathbf{n}\sim\mathcal{CN}\left({\bf
0},\sigma^{2}\mathbf{I}_{M}\right)$ denotes the vector of noise.

Utilizing the channel estimate $\hat{\bf{q}}_{k}$, the BS performs MRC to
process the received signal as

$\displaystyle{r}_{k}=\hat{\bf{q}}_{k}^{H}\mathbf{y}=\sqrt{p}\hat{\bf{q}}_{k}^{H}\mathbf{q}_{k}{x}_{k}+\hat{\bf{q}}_{k}^{H}\mathbf{n}.$
(1)

Since rich scatters often exist on the ground, as in [5], we adopt the
Rayleigh fading to model the direct channel as
$\mathbf{d}_{k}=\sqrt{\gamma_{k}}\tilde{\mathbf{d}}_{k}$, where $\gamma_{k}$
denotes the path-loss. The entries of $\tilde{\mathbf{d}}_{k}$ are independent
and identically distributed (i.i.d.) complex Gaussian random variables, and we
have $\tilde{\mathbf{d}}_{k}\sim\mathcal{CN}\left({\bf
0},\mathbf{I}_{M}\right)$.

Considering that the RIS is often installed at the facade of tall buildings
and the BS also have a certain height, as in [5, 6, 7], we adopt Rician fading
to model the RIS-related channels as

$\displaystyle\begin{array}[]{l}\mathbf{h}_{k}\\!=\\!\sqrt{\frac{\alpha_{k}}{\varepsilon_{k}+1}}\left(\sqrt{\varepsilon_{k}}\,\overline{\mathbf{h}}_{k}+\tilde{\mathbf{h}}_{k}\right),\mathbf{H}_{2}=\sqrt{\frac{\beta}{\delta+1}}\left(\sqrt{\delta}\,\overline{\mathbf{H}}_{2}+\tilde{\mathbf{H}}_{2}\right),\end{array}$
(3)

where $\alpha_{k}$ and $\beta$ are path-loss coefficients, $\varepsilon_{k}$
and $\delta$ are Rician factors. $\overline{\mathbf{h}}_{k}$ and
$\overline{\mathbf{H}}_{2}$ denote LoS components, whereas
$\tilde{\mathbf{h}}_{k}$ and $\tilde{\mathbf{H}}_{2}$ denote NLoS components.
All the entries of NLoS components $\tilde{\mathbf{h}}_{k}$ and
$\tilde{\mathbf{H}}_{2}$ are i.i.d complex Gaussian random variables with zero
mean and unit variance. For LoS components, we utilize the uniform squared
planar array (USPA) model, and then $\overline{\mathbf{h}}_{k}$ and
$\overline{\mathbf{H}}_{2}$ are respectively expressed as

$\displaystyle{\overline{\bf h}_{k}}={\bf
a}_{N}\left({\varphi_{kr}^{a},\varphi_{kr}^{e}}\right),$
$\displaystyle{\overline{\bf H}_{2}}={\bf
a}_{M}\left({\phi_{r}^{a},\phi_{r}^{e}}\right){\bf
a}_{N}^{H}\left({\varphi_{t}^{a},\varphi_{t}^{e}}\right),$ (4)

where $\varphi_{kr}^{a}$ ($\varphi_{kr}^{e}$) is the azimuth (elevation)
angles of arrival (AoA) at the RIS from the user, $\varphi_{t}^{a}$
($\varphi_{t}^{e}$) is the azimuth (elevation) angles of departure (AoD) from
the RIS to the BS, and $\phi_{r}^{a}$ ($\phi_{r}^{e}$) is the azimuth
(elevation) AoA at the BS from the RIS, respectively. Besides, ${{\bf
a}_{X}}\left({\vartheta^{a},\vartheta^{e}}\right)\in\mathbb{C}^{X\times 1}$
denotes the array response vector, whose $\rm x$-th entry is

$\displaystyle\left[{{\bf
a}_{X}}\left({\vartheta^{a},\vartheta^{e}}\right)\right]_{\rm x}=$
$\displaystyle\exp\left\\{j2\pi\frac{d}{\lambda}\left(\lfloor\left({\rm
x}-1\right)/\sqrt{X}\rfloor\sin\vartheta^{e}\sin\vartheta^{a}\right.\right.$
$\displaystyle\left.\left.+\left(\left(\mathrm{x}-1\right)\bmod\sqrt{X}\right)\cos\vartheta^{e}\right)\right\\},$
(5)

where $d$ and $\lambda$ denote element spacing and wavelength, respectively.
To simply the notations, in the sequel of this paper, we will use notations
${\bf\bar{h}}_{k}$, ${\bf a}_{M}$ and ${\bf a}_{N}$ to represent ${\bf
a}_{N}\left({\varphi_{kr}^{a},\varphi_{kr}^{e}}\right)$, ${\bf
a}_{M}\left({\phi_{r}^{a},\phi_{r}^{e}}\right)$ and ${\bf
a}_{N}\left({\varphi_{t}^{a},\varphi_{t}^{e}}\right)$, respectively.

##  III Channel Estimation

Assume that the channel coherence interval consists of $\tau_{c}$ samples, and
let $\tau$ denote the number of samples used in the pilot transmission. During
the training part of each channel coherence interval, the user will send the
pilot sequence to the BS. The transmit pilot power is $\tau p$, and the
$1\times\tau$ pilot sequence is ${\bf s}_{k}$, with ${\bf
s}_{k}\mathbf{s}_{k}^{H}=1$. Therefore, the $M\times\tau$ pilot signal
received by the BS can be written as $\mathbf{Y}_{p}=\sqrt{\tau
p}\mathbf{q}_{k}\mathbf{s}_{k}+\mathbf{N}$, where $\bf N$ denotes the
$M\times\tau$ noise matrix whose entries are i.i.d. complex Gaussian random
variables with zero mean and $\sigma^{2}$ variance. By multiplying
$\mathbf{Y}_{p}$ with $\frac{{\bf s}^{H}_{k}}{\sqrt{\tau p}}$, the BS can
obtain the observation vector for the user as

$\displaystyle\mathbf{y}_{p}^{k}=\frac{1}{\sqrt{\tau
p}}\mathbf{Y}_{p}\mathbf{s}_{k}^{H}=\mathbf{q}_{k}+\frac{1}{\sqrt{\tau
p}}\mathbf{Ns}_{k}^{H}.$ (6)

To obtain closed-form channel estimates and enable our analysis, we will adopt
the sub-optimal LMMSE estimator to perform the channel estimation.111All
proofs in this paper are omitted due to the page limitation. Please refer to
[11] for details.

######  Theorem 1.

Using the observation vector ${\mathbf{y}}^{k}_{p}$, the LMMSE estimation
$\hat{\mathbf{q}}_{k}$ for channel ${\mathbf{q}}_{k}$ is given by
$\hat{\mathbf{q}}_{k}=\mathbf{A}_{k}\mathbf{y}_{p}^{k}+\mathbf{B}_{k}$, where

$\displaystyle\mathbf{A}_{k}=a_{k3}\mathbf{a}_{M}\mathbf{a}_{M}^{H}+a_{k4}\mathbf{I}_{M},$
(7)
$\displaystyle\mathbf{B}_{k}=\left(\mathbf{I}_{M}-\mathbf{A}_{k}\right)\sqrt{c_{k}\delta\varepsilon_{k}}\,\overline{\mathbf{H}}_{2}\mathbf{\Phi}\overline{\mathbf{h}}_{k},$
(8) $\displaystyle
c_{k}=\frac{\beta\alpha_{k}}{(\delta+1)\left(\varepsilon_{k}+1\right)},$ (9)
$\displaystyle a_{k1}=Nc_{k}\delta,\quad
a_{k2}=Nc_{k}\left(\varepsilon_{k}+1\right)+\gamma_{k},$ (10) $\displaystyle
a_{k3}=\frac{a_{k1}\frac{\sigma^{2}}{\tau
p}}{\left(a_{k2}+\frac{\sigma^{2}}{\tau
p}\right)\left(a_{k2}+\frac{\sigma^{2}}{\tau
p}+Ma_{k1}\right)},\;a_{k4}=\frac{a_{k2}}{a_{k2}+\frac{\sigma^{2}}{\tau p}},$
(11)

and the NMSE in the estimation of $\mathbf{q}_{k}$ is

$\displaystyle\operatorname{NMSE}_{k}=\frac{\frac{\sigma^{2}}{\tau
p}\left(Ma_{k1}a_{k2}+a_{k2}^{2}+\left(a_{k1}+a_{k2}\right)\frac{\sigma^{2}}{\tau
p}\right)}{\left(a_{k2}+\frac{\sigma^{2}}{\tau
p}\right)\left(a_{k2}+\frac{\sigma^{2}}{\tau
p}+Ma_{k1}\right)\left(a_{k1}+a_{k2}\right)}.$ (12)

As shown in Theorem 1, we only estimate the overall channel ${\bf
q}_{k}\in\mathbb{C}^{M\times 1}$, which possesses the same dimension as in
traditional MISO systems. Therefore, we only require that the length of pilot
sequences is larger than the number of users, i.e., $\tau\geq 1$. Compared
with the estimation method which needs the estimation of all $MN$ individual
channels[12], our method has very low overhead and complexity.

######  Remark 1.

We can see that the NMSE expression (12) is unrelated to the phase shifts of
the RIS. Therefore, the proposed LMMSE estimation method can work compatibly
with the statistical CSI-based RIS design, since the phase shifts of RIS don’t
need to be adjusted in each channel coherence interval.

######  Corollary 1.

The asymptotic NMSE performance is

$\displaystyle\lim\nolimits_{\frac{\sigma^{2}}{\tau p}\rightarrow
0}\mathrm{NMSE}_{k}\rightarrow
0,\quad\lim\nolimits_{N\rightarrow\infty}\mathrm{NMSE}_{k}\rightarrow 0.$ (13)

Note that the NMSE between $0$ (perfect estimation) and $1$ (using the mean
value of the variable as the estimate) measures the relative estimation error
per antenna [13]. In traditional MISO systems, a common method to reduce the
NMSE is to increase the length of pilots, i.e., increasing $\tau$. However,
Corollary 1 indicates that in RIS-aided MISO systems, increasing the number of
reflecting elements of RIS can play a similar role as increasing $\tau$.
Therefore, increasing RIS reflecting elements not only helps improve the
system capacity, but also helps decrease the NMSE.

To better understand the impact of increasing $N$ on channel estimation, we
present the following asymptotic results.

######  Corollary 2.

When $\tau\to\infty$, we have $\hat{\mathbf{q}}_{k}\to{\mathbf{q}}_{k}$. By
contrast, when $N\to\infty$, we have

$\displaystyle\hat{\mathbf{q}}_{k}\to{\mathbf{q}}_{k}+\frac{1}{\sqrt{\tau
p}}\mathbf{Ns}_{k}^{H},\quad{\rm MSE}_{k}\to\frac{\sigma^{2}}{\tau p}{\bf
I}_{M}.$ (14)

Although the NMSE will converge to zero with infinite $N$, Corollary 2 shows
that different from increasing $\tau$, the MSE under LMMSE estimator will
converge to a non-zero constant. It is interesting to find that the asymptotic
MSE in (14) is the same as the MSE performance of least-squares (LS)
estimator[13]. It is well-known that the LS estimator has a worse estimation
quality than the LMMSE estimator. Therefore, we can roughly deduce that the
MSE performance under LMMSE is an increasing function of $N$ which converges
to an upper bound of the LS estimator. This result will be validated in
Section V.

##  IV Analysis of Uplink Ergodic Rate

As in [13], we resort to the so called use-and-then-forget bound to
characterize the ergodic rate performance. We rewrite the received signal
$r_{k}$ in (1) as

$\displaystyle
r_{k}=\underbrace{\sqrt{p}\,\mathbb{E}\left\\{\hat{\mathbf{q}}_{k}^{H}\mathbf{q}_{k}\right\\}x_{k}}_{\text{Desired
signal
}}+\underbrace{\sqrt{p}\left(\hat{\mathbf{q}}_{k}^{H}\mathbf{q}_{k}-\mathbb{E}\left\\{\hat{\mathbf{q}}_{k}^{H}\mathbf{q}_{k}\right\\}\right)x_{k}}_{\text{Signal
leakage }}+\underbrace{\hat{\mathbf{q}}_{k}^{H}\mathbf{n}}_{\text{Noise }}.$
(15)

Then, we present the closed-form expression of ergodic rate in the following
theorem.

######  Theorem 2.

The ergodic rate of the user is lower bounded by

$\displaystyle\underline{R}_{k}=\frac{\tau_{c}-\tau}{\tau_{c}}\log_{2}\left\\{1+\frac{pE_{k}^{(\rm
signal)}}{pE_{k}^{(\rm leakage)}+\sigma^{2}E_{k}^{(\rm{noise})}}\right\\},$
(16)

with $E_{k}^{(\rm signal)}=\left(E_{k}^{(\rm{noise})}\right)^{2}$,

$\displaystyle E_{k}^{(\rm{noise})}=M\times$
$\displaystyle\left\\{\left|f_{k}(\mathbf{\Phi})\right|^{2}c_{k}\delta\varepsilon_{k}+Nc_{k}\delta
e_{k2}\right.\quad\qquad\qquad$
$\displaystyle+\left(Nc_{k}\left(\varepsilon_{k}+1\right)+\gamma_{k}\right)e_{k1}\Big{\\}},$
(17) $\displaystyle\begin{array}[]{l}E_{k}^{(\rm leakage)}\\\
=M\left|f_{k}(\mathbf{\Phi})\right|^{2}c_{k}^{2}\delta\varepsilon_{k}\left\\{N\left(M\delta+\varepsilon_{k}+1\right)\left(e_{k2}^{2}+1\right)\right.\\\
\qquad\qquad\qquad\qquad\quad\left.+2\left(Me_{k1}+e_{k2}\right)\left(e_{k2}+1\right)\right\\}\\\
+M\left|f_{k}(\mathbf{\Phi})\right|^{2}c_{k}\delta\varepsilon_{k}\left\\{\gamma_{k}+\left(\gamma_{k}+\frac{\sigma^{2}}{\tau
p}\right)e_{k2}^{2}\right\\}\\\ +M^{2}N^{2}c_{k}^{2}\delta^{2}e_{k2}^{2}\\\
+MN^{2}c_{k}^{2}\left\\{2\delta\left(\varepsilon_{k}+1\right)e_{k2}^{2}+\left(\varepsilon_{k}+1\right)^{2}e_{k3}\right\\}\\\
+M^{2}Nc_{k}^{2}\left\\{\left(2\varepsilon_{k}+1\right)e_{k1}^{2}+2\delta
e_{k1}e_{k2}\right\\}\\\ +MNc_{k}\Big{\\{}c_{k}\left(2\delta
e_{k2}^{2}+\left(2\varepsilon_{k}+1\right)e_{k3}\right)\\\
\qquad\qquad\left.+\left(2\gamma_{k}+\frac{\sigma^{2}}{\tau
p}\right)\left(\delta
e_{k2}^{2}+\left(\varepsilon_{k}+1\right)e_{k3}\right)\right\\}\\\
+M\gamma_{k}\left(\gamma_{k}+\frac{\sigma^{2}}{\tau
p}\right)e_{k3},\end{array}$ (28)

where

$\displaystyle
f_{k}({\bf\Phi})=\mathbf{a}_{N}^{H}{\bf\Phi}\overline{\mathbf{h}}_{k},\quad
e_{k1}=a_{k3}+a_{k4}\in\left[0,1\right],$ (29) $\displaystyle
e_{k2}=Ma_{3}+a_{4}\in\left[0,1\right],$ (30) $\displaystyle
e_{k3}=Ma_{k3}^{2}+2a_{k3}a_{k4}+a_{k4}^{2}\in\left[0,1\right].$ (31)

The closed-form expressions in Theorem 2 do not rely on fast changing NLoS
channel components, and do not involve computationally burdened integral and
matrix inverse operators. Therefore, in contrast to time-consuming Monte-Carlo
simulations, the evaluation of rate performance based on Theorem 2 can be
computed rapidly. Besides, comparing the expressions in Theorem 2 with those
in [6, Lemma 1], one can see that the impact of imperfect CSI is fully
characterized by the parameters $e_{k1},e_{k2},e_{k3}$ and
$\frac{\sigma^{2}}{\tau p}$. For the perfect CSI scenario without channel
estimation error, we can assume $\tau=\infty$, which leads to
$e_{k1}=e_{k2}=e_{k3}=1$ and $\frac{\sigma^{2}}{\tau p}=0$.

Next, we will design the phase shifts of the RIS. If $N=1$, the phase shifts
can be set arbitrarily[11]. Therefore, we only consider the case $N>1$ in the
following. From Theorem 2, we can see that the only phase shifts related
variable is $\left|f_{k}\left(\mathbf{\Phi}\right)\right|^{2}$. We can treat
$\left|f_{k}\left(\mathbf{\Phi}\right)\right|^{2}$ as a new variable $x$ where
$0\leq x\leq N^{2}$[11]. Then, we can rewrite the numerator and denominator of
SNR in Theorem 2 as $E_{k}^{(\text{signal})}=\left(s_{1}x+s_{2}\right)^{2}$
and
$E_{k}^{(\text{leakage})}+\frac{\sigma^{2}}{p}E_{k}^{(\text{noise})}=t_{1}x+t_{2}$,
respectively. The detailed expressions of $s_{1},s_{2},t_{1},t_{2}$ are
omitted here, since it can be readily obtained from Theorem 2. As a result,
the SNR in Theorem 2 can be rewritten as

$\displaystyle\mathrm{SNR}_{k}\left(x\right)=\frac{\left(s_{1}x+s_{2}\right)^{2}}{t_{1}x+t_{2}}.$
(32)

The rate can be maximized by maximizing the SNR. Thus, we calculate the first-
order derivative of $\mathrm{SNR}_{k}\left(x\right)$ with respect to $x$ as

$\displaystyle\frac{\partial\mathrm{SNR}_{k}\left(x\right)}{\partial
x}=\left(s_{1}x+s_{2}\right)\frac{s_{1}t_{1}x+2s_{1}t_{2}-s_{2}t_{1}}{\left(t_{1}x+t_{2}\right)^{2}}.$
(33)

Note that $s_{1},s_{2},t_{1},t_{2}>0$. Set the first-order derivative to zero.
Under the constraint $x\geq 0$, the only possible solution is

$\displaystyle\frac{s_{2}t_{1}-2s_{1}t_{2}}{s_{1}t_{1}}\triangleq x_{0}.$ (34)

Therefore, the optimal phase shifts design can be decided by discussing the
value of $x_{0}$, which is presented as follows.

######  Theorem 3.

The phase shifts of RIS can be designed under the following four cases:

  * •

Case 1: If $x_{0}\leq 0$, it is optimal to set
$\left|f_{k}\left(\mathbf{\Phi}\right)\right|=N$.

  * •

Case 2: If $x_{0}\geq N^{2}$, it is optimal to set
$\left|f_{k}\left(\mathbf{\Phi}\right)\right|=0$.

  * •

Case 3: If $0<x_{0}<N^{2}$, we should compare ${\rm SNR}_{k}\left(0\right)$
with ${\rm SNR}_{k}\left(N^{2}\right)$, and choose the maximum one.

  * •

Case 4: If $N\to\infty$, it is optimal to set
$\left|f_{k}\left(\mathbf{\Phi}\right)\right|=N$.

Theorem 3 indicates that it is feasible to design the phase shifts as
$\left|f_{k}\left(\mathbf{\Phi}\right)\right|=N$ if $N$ is large. Therefore,
in the rest of this section, we will set
$\left|f_{k}\left(\mathbf{\Phi}\right)\right|=N$ to facilitate our analysis.

Substituting Theorem 2 with $\left|f_{k}(\mathbf{\Phi})\right|=N$, we can see
that the power of desired signal is in the order of
$\mathcal{O}\left(M^{2}N^{4}\right)$, the power of signal leakage is in the
order of $\mathcal{O}\left(M^{2}N^{3}\right)$, and the power of noise term is
in the order of $\mathcal{O}\left(MN^{2}\right)$. Therefore, the rate is
bounded when $M\to\infty$, but can grow unbounded as $N\to\infty$. Besides, we
can reduce the power proportionally to $1/M$ or $1/N^{2}$ while keeping the
rate non-zero as $M,N\to\infty$.

######  Corollary 3.

If the transmit power is reduced to $p={E_{u}}/{N^{2}}$, as $N\to\infty$, the
rate is lower bounded by

$\displaystyle\underline{R}_{k}\rightarrow\frac{\tau_{c}-\tau}{\tau_{c}}\log_{2}\left(1+\frac{E_{u}}{\sigma^{2}}M\frac{\beta\alpha_{k}\delta\varepsilon_{k}}{(\delta+1)\left(\varepsilon_{k}+1\right)}\right).$
(35)

The SNR in (35) is equal to that in $M$-antennas systems with path-loss
$c_{k}\delta\varepsilon_{k}$. Besides, the SNR is uncorrelated with $\tau$ and
the same SNR can be realized by applying the corresponding power scaling in
perfect CSI-based systems[5, (17)]. In other words, the increase of $N$ could
compensate the detrimental effects of limited pilot lengths, and $\tau=1$ is
adequate in (35). Furthermore, it is easy to find that the rate is an
increasing function of Rician factors $\delta$ and $\varepsilon_{k}$.
Therefore, the LoS environment is beneficial for RIS-aided single-user
systems.

However, we can find that the power scaling laws in (35) no longer hold under
fully Rayleigh case, i.e., $\delta=\varepsilon_{k}=0$. In this case, we can
reduce the power proportionally to $1/\sqrt{M}$ or $1/N$ while keeping the
rate non-zero when $M,N\to\infty$.

######  Corollary 4.

Let $\delta=\varepsilon_{k}=0$. if the transmit power is reduced as
$p={E_{u}}/{N}$, as $N\to\infty$, the rate is lower bounded by

$\displaystyle\underline{R}_{k}^{(\mathrm{NL})}\rightarrow\frac{\tau_{c}-\tau}{\tau_{c}}\log_{2}\left(\frac{E_{u}M\beta\alpha_{k}}{E_{u}\beta\alpha_{k}+\frac{\sigma^{2}}{\tau}+\sigma^{2}\left(1+\frac{\sigma^{2}}{\tau
E_{u}\beta\alpha_{k}}\right)}\right).$ (36)

Unlike Corollary 3, it is observed that the SNR in Corollary 4 is an
increasing function of $\tau$. This observation again emphasizes that the LoS
environment is favorable for RIS-aided single-user systems.

##  V Numerical Results

In the simulations, we consider that the distance between the user and the RIS
is $d_{UI}=20$ m, and the distance between the RIS and the BS is $d_{IB}=700$
m. The distance between the user and the BS can be calculated from the
geometric relationship as
$\left(d_{k}^{\mathrm{UB}}\right)^{2}=\left(d_{\mathrm{IB}}-d_{\mathrm{UI}}\sin\left(\frac{\pi}{5}\right)\right)^{2}+\left(d_{\mathrm{UI}}\cos\left(\frac{\pi}{5}\right)\right)^{2}$.
With the distances, large-scale factors are set to
$\alpha_{k}=10^{-3}d_{\mathrm{UI}}^{-2}$,
$\beta=10^{-3}d_{\mathrm{IB}}^{-2.5}$ and
$\gamma_{k}=10^{-3}\left(d_{k}^{\mathrm{UB}}\right)^{-4}$. Besides, we set
$M=N=64$, $p=30$ dBm, $\sigma^{2}=-104$ dBm, $\tau_{c}=196$[14], $\tau=1$,
$\delta=1$ and $\varepsilon_{k}=10$.

Figure 2: MSE and NMSE versus the number of RIS elements.

We first investigate the estimation quality of overall channel ${\bf Q}$. As
our analysis, we can see that the MSE is an increasing function while the NMSE
is a decreasing function with respect to $N$. With a large $N$, the MSE
approaches an upper bound and the NMSE converges to zero. In addition, the
fully LoS case has the same performance with RIS-free traditional systems.
Therefore, the MSE and NMSE in the fully LoS case do not change with $N$.
Furthermore, compared with the NMSE with an increased $\tau$, we can see that
increasing $N$ can incur a similar impact.

Figure 3: Statistical CSI-based design versus instantaneous CSI-based design.

In Fig. 3, we compare our design with instantaneous CSI-based design, which
estimates the instantaneous channels in each coherence interval as [12,
Section $\rm V$]. With single user, the minimum pilot length is $1+N$ in
instantaneous CSI scheme[12]. Therefore, its rate (red line) should be
multiplied by a loss factor $1-\frac{1+N}{\tau_{c}}$. Meanwhile, we consider
an ideal case (green line) with a loss factor $1-\frac{1}{\tau_{c}}$, which is
the same loss factor of our scheme. We can see that if we ignore the
additional rate loss in channel estimation (green line), using instantaneous
CSI can achieve better performance, especially in large $N$ regime. However,
when considering the overhead in pilot transmission (red line), the
instantaneous CSI scheme first increase but then decrease, which performs
worse than our design. This is because instantaneous CSI scheme need longer
time for channel estimation, which is proportional to $N$. Hence, when $N$ is
large, only little time is left for data transmission. However, our design can
keep improving the rate performance as increasing $N$. Therefore, compared
with instantaneous CSI-based design, statistical CSI-based design can realize
better performance but with lower overhead and complexity in practical
systems.

Figure 4:  Rate versus $N$ with scaled-down transmit power $p=E_{u}/N^{2}$ or
$p=E_{u}/N$, where $E_{u}=20$ dB.

In Fig. 4, we characterizes the power scaling law with $N$. As presented in
Corollary 3, the rate will converge to a limit if we reduce the power
proportionally to $1/N^{2}$ under Rician channel, and this limit is maximized
under fully LoS RIS-related channels. However, under NLoS RIS-related
channels, reducing the power proportionally to $1/N^{2}$ will make the rate
zero, and it can be reduced proportionally to $1/N$ as shown in Corollary 4.
These observations validate that the LoS environment is preferable for RIS-
aided single user systems.

##  VI Conclusion

This paper has investigated RIS-aided MISO systems, taking the channel
estimation error into consideration. An LMMSE estimator has been proposed to
obtain the CSI for overall channels. Next, we have derived a closed-form
expression for the ergodic date rate under the Rician channel model. Based on
the derived expression, we have designed the phase shifts of the RIS based
only on statistical CSI, which significantly reduces the signaling overhead
and computational complexity of the system. We have also theoretically
analyzed the power scaling laws under different kinds of channel conditions,
which validates that the LoS environment is beneficial for RIS-aided MISO
systems.

## References

  * [1] C. Pan, H. Ren, K. Wang, J. F. Kolb, M. Elkashlan, M. Chen, M. Di Renzo, Y. Hao, J. Wang, A. L. Swindlehurst, X. You, and L. Hanzo, “Reconfigurable intelligent surfaces for 6G systems: Principles, applications, and research directions,” _IEEE Commun. Mag._ , vol. 59, no. 6, pp. 14–20, Jun. 2021\.
  * [2] Q. Wu and R. Zhang, “Intelligent reflecting surface enhanced wireless network via joint active and passive beamforming,” _IEEE Trans. Wireless Commun._ , vol. 18, no. 11, pp. 5394–5409, Nov. 2019.
  * [3] C. Pan _et al._ , “Multicell MIMO communications relying on intelligent reflecting surfaces,” _IEEE Trans. Wireless Commun._ , vol. 19, no. 8, pp. 5218–5233, Aug. 2020.
  * [4] C. Pan, H. Ren, K. Wang, M. Elkashlan, A. Nallanathan, J. Wang, and L. Hanzo, “Intelligent reflecting surface aided MIMO broadcasting for simultaneous wireless information and power transfer,” _IEEE J. Sel. Areas Commun._ , vol. 38, no. 8, pp. 1719–1734, Aug. 2020.
  * [5] Y. Han, W. Tang, S. Jin, C.-K. Wen, and X. Ma, “Large intelligent surface-assisted wireless communication exploiting statistical CSI,” _IEEE Trans. Veh. Technol._ , vol. 68, no. 8, pp. 8238–8242, Aug. 2019.
  * [6] K. Zhi, C. Pan, H. Ren, and K. Wang, “Power scaling law analysis and phase shift optimization of RIS-aided massive MIMO systems with statistical CSI,” 2020. [Online]. Available: https://arxiv.org/abs/2010.13525
  * [7] K. Zhi, C. Pan, H. Ren, and K. Wang, “Statistical CSI-based design for reconfigurable intelligent surface-aided massive MIMO systems with direct links,” _IEEE Wireless Commun. Lett._ , early access, 2021.
  * [8] Y. Gao, J. Xu, W. Xu, D. W. K. Ng, and M. S. Alouini, “Distributed IRS with statistical passive beamforming for MISO communications,” _IEEE Wireless Commun. Lett._ , vol. 10, no. 2, pp. 221–225, Feb. 2021.
  * [9] B. Al-Nahhas, Q. U. A. Nadeem, and A. Chaaban, “Intelligent reflecting surface assisted MISO downlink: Channel estimation and asymptotic analysis,” in _Proc. IEEE Global Commun. Conf. (Globecom)_ , 2020, pp. 1–6.
  * [10] A. Papazafeiropoulos, C. Pan, P. Kourtessis, S. Chatzinotas, and J. M. Senior, “Intelligent reflecting surface-assisted MU-MISO systems with imperfect hardware: Channel estimation, beamforming design,” 2021. [Online]. Available: https://arxiv.org/abs/2102.05333
  * [11] K. Zhi, C. Pan, H. Ren, K. Wang, M. Elkashlan, M. D. Renzo, R. Schober, H. V. Poor, and L. Hanzo, “Two-timescale design for reconfigurable intelligent surface-aided massive MIMO systems with imperfect CSI,” to be submitted to arXiv.
  * [12] Z. Wang, L. Liu, and S. Cui, “Channel estimation for intelligent reflecting surface assisted multiuser communications: Framework, algorithms, and analysis,” _IEEE Trans. Wireless Commun._ , vol. 19, no. 10, pp. 6607–6620, Oct. 2020.
  * [13] E. Björnson, J. Hoydis, and L. Sanguinetti, “Massive MIMO networks: Spectral, energy, and hardware efficiency,” _Found. Trends Signal Process._ , vol. 11, no. 3-4, pp. 154–655, Nov. 2017.
  * [14] H. Q. Ngo, E. G. Larsson, and T. L. Marzetta, “Energy and spectral efficiency of very large multiuser MIMO systems,” _IEEE Trans. Commun._ , vol. 61, no. 4, pp. 1436–1449, Apr. 2013.


# Empowering Character-level Text Infilling by Eliminating Sub-Tokens

Houxing Ren1,3 Mingjie Zhan211footnotemark: 1 Zhongyuan Wu2 Hongsheng Li3,4,5
1Shanghai Jiao Tong University 2SenseTime Research
3Shanghai Artificial Intelligence Laboratory 4CUHK MMLab 5CPII under InnoHK
<EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS>Corresponding author.

###### Abstract

In infilling tasks, sub-tokens, representing instances where a complete token
is segmented into two parts, often emerge at the boundaries of prefixes,
middles, and suffixes. Traditional methods focused on training models at the
token level, leading to sub-optimal performance in character-level infilling
tasks during the inference stage. Alternately, some approaches considered
character-level infilling, but they relied on predicting sub-tokens in
inference, yet this strategy diminished ability in character-level infilling
tasks due to the large perplexity of the model on sub-tokens. In this paper,
we introduce FIM-SE, which stands for Fill-In-the-Middle with both Starting
and Ending character constraints. The proposed method addresses character-
level infilling tasks by utilizing a line-level format to avoid predicting any
sub-token in inference. In addition, we incorporate two special tokens to
signify the rest of the incomplete lines, thereby enhancing generation
guidance. Extensive experiments demonstrate that our proposed approach
surpasses previous methods, offering a significant advantage. Code is
available at https://github.com/SenseLLM/FIM-SE.

Empowering Character-level Text Infilling by Eliminating Sub-Tokens

Houxing Ren1,3 Mingjie Zhan211footnotemark: 1 Zhongyuan Wu2 Hongsheng
Li3,4,5††thanks: Corresponding author. 1Shanghai Jiao Tong University
2SenseTime Research 3Shanghai Artificial Intelligence Laboratory 4CUHK MMLab
5CPII under InnoHK<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>

##  1 Introduction

The Transformer Vaswani et al. (2017) decoder-only architecture has proven
highly effective in various natural language processing (NLP) tasks. This
success has paved the way for the development of advanced causal decoder-only
models like GPT-4 OpenAI (2023), PaLM Chowdhery et al. (2023); Anil et al.
(2023), Llama Touvron et al. (2023a, b); Rozière et al. (2023), and Falcon
Penedo et al. (2023). These innovative models excel at generating coherent and
contextually relevant responses to natural language prompts, showcasing state-
of-the-art performance across various tasks, including question answering
Lewis et al. (2020b), logical reasoning Kojima et al. (2022), and code
synthesis Li et al. (2023); Rozière et al. (2023).

Table 1: Examples for random splitting with Llama tokenizer, where the red,
blue, and green text indicates the prefix, the middle, and the suffix,
respectively. These four rows represent the pieces after randomly splitting,
the sentence after exchanging suffix and middle, tokenized results, and token
IDs, respectively.

Pieces | A fine day.
---|---
Reorder | A f day. ine
Tokens | [‘A’, ‘_f’, ‘day’, ‘.’, ‘ine’, ‘_’]
IDs | [29909, 285, 3250, 29889, 457, 29871]
(a) The first splitting case.

Pieces | A fine day.
---|---
Reorder | A fi day. ne
Tokens | [‘A’, ‘_fi’, ‘day’, ‘.’, ‘ne’, ‘_’]
IDs | [29909, 5713, 3250, 29889, 484, 29871]
(b) The second splitting case.

Despite the success, the proficiency of these models is somewhat limited in
tasks involving text infilling, which aims to generate text at a specific
location within a prompt, while conditioning on both a prefix and a suffix.
The main reason is their intrinsic left-to-right autoregressive design. To
address this issue, CM3 Aghajanyan et al. (2022) introduced the causal masking
objective, placing a mask token at the intended fill location and completing
the fill at the end. In contrast, FIM Bavarian et al. (2022) proposed a fill-
in-the-middle technique, which randomly divides documents into three segments
and tags them with three special tokens. This technique then rearranges the
middle and suffix segments, to use the prefix and the suffix to predict the
middle segment in auto-regressive format. With these methods, decoder-only
based models can effectively handle various infilling tasks and achieve
excellent performance.

(a) Experiment.

(b) Probabilities.

Figure 1: The probabilities of prediction when inconsistent labels appear in
the training data.

However, employing the aforementioned methods may introduce inconsistencies
during training. This arises from the potential division of a single token
into multiple sub-tokens, as exemplified in Table 2(b). As we can see, due to
character-level random splitting, the same prefixes ([29909]) have
inconsistent objectives (285 and 5713, respectively) in different cases. The
inconsistent objectives will significantly impact the model’s perplexity,
especially on sub-tokens. To illustrate, we construct a simple experiment on a
classification task shown in Figure 1(a). The training data only contains
three samples and they have the same input but different labels. We train a
simple network on the training data and record the predicted probabilities of
the three labels at each training step. As shown in Figure 1(b), the predicted
probabilities for these three classes converge to 0.33, indicating a large
perplexity of the model on the inconsistent objectives. The large perplexity
on sub-tokens makes the probability of error increase when predicting a sub-
token. This phenomenon is notably detrimental in sensitive tasks such as code
completion, where even a minor error in any token can result in program
malfunction. As a result, previous approaches have yet to fully inspire the
potential of Transformer decoder-only models in infilling tasks.

To effectively address the issue, it is crucial to acknowledge and resolve an
inherent conflict. (1) We need to avoid the model predicting sub-tokens. In
the infilling training mode, the model’s perplexity in sub-tokens is large,
resulting in the low accuracy of predicting sub-tokens. (2) We need to output
a sub-token when the user only writes part of a token. Because it is necessary
to ensure that the output fits the context. If we directly drop several tokens
to make sure no sub-tokens exist, the model’s output may no longer align with
the removed context, rendering it unreasonable in practical use.

Based on these concerns, we propose FIM-SE, which stands for Fill-In-the-
Middle with both Starting and Ending character constraints. Our method
enhances the organizational framework of FIM Bavarian et al. (2022) to
concurrently address the two scenarios mentioned earlier. In simple terms, we
transfer the random-span infilling task to the multi-line infilling task.
Specifically, after random character level splitting, we utilize two distinct
special tokens to mark the _Last line of the Prefix (L-Prefix)_ and the _First
line of the Suffix (F-Suffix)_. The model is then tasked to generate text at
line level that starts with _L-Prefix_ and ends with _F-Suffix_. Their
inclusion in the prompt simplifies the task for the model, facilitating the
generation of text that seamlessly starts with _L-Prefix_ and ends with _F-
Suffix_. Overall, this method is designed to unlock the capabilities of
decoder-only models in infilling tasks.

The core contribution of the paper is that we design a novel training method
for the infilling task, a solution designed to effectively mitigate conflicts
mentioned above in infilling tasks. Our method can effectively eliminate any
potential inconsistencies and earnestly guarantee that the model’s output
aligns cohesively with the given context. Extensive experiments demonstrate
the effectiveness of the proposed method on infilling tasks while not
compromising code generation capabilities. Based on Code Llama 13B, our
approach not only achieves an 8.8% enhancement in the Humaneval random-span
infilling task, with substantial improvements of 11.5% and 10.7% in the
single-line and multi-line infilling tasks respectively, but also maintains
minimal impact on the model’s performance in code generation tasks.

##  2 Related Work

###  2.1 Large Language Models for Infilling

Various Large Language Models (LLMs) have been developed for general
generation tasks Touvron et al. (2023a); Xiong et al. (2023); Chowdhery et al.
(2023); Penedo et al. (2023); Jiang et al. (2023); Yang et al. (2023); Bai et
al. (2023). Most of these models adopt a left-to-right autoregressive
generation approach due to its effectiveness, as validated by research such as
GPT Radford et al. (2018, 2019); Brown et al. (2020); Ouyang et al. (2022);
OpenAI (2023). In the realm of code-related tasks, where infilling is
essential, LLMs are specifically trained for this task. For instance, InCoder
Fried et al. (2023) utilizes a causal masking objective, while SantaCoder
Allal et al. (2023), StarCoder Li et al. (2023), and Code Llama Rozière et al.
(2023) employ the fill-in-the-middle technique introduced by FIM Bavarian et
al. (2022).

###  2.2 Text Infilling Models

The infilling task plays a crucial role in numerous real-world applications,
including document editing111https://copilot.microsoft.com and code
completion222https://github.com/features/copilot. Three common Transformer
Vaswani et al. (2017) architectures are capable of executing this task,
_i.e.,_ encoder-only, encoder-decoder, decoder-only. In encoder-only
architectures, masked language modeling is employed as the pre-training task,
exemplified by models like BERT Devlin et al. (2019) and RoBERTa Liu et al.
(2019). These models are designed to infill brief spans, ranging from a single
token Devlin et al. (2019) to a word Cui et al. (2021), and even several
contiguous tokens Joshi et al. (2020). In encoder-decoder architectures, a
common approach involves masking several tokens in the encoder and then
tasking the model with decoding the complete sentence, as exemplified by MASS
Song et al. (2019). Additionally, models like BART Lewis et al. (2020a) and T5
Raffel et al. (2020) have introduced an infilling noising method. This
technique replaces multiple tokens with a single mask token, challenging the
model to decode the masked span. In decoder-only architectures, several
methods are employed for infilling tasks. The Insertion Transformer Stern et
al. (2019) instructs the model to first determine the location for the next
token, followed by the token prediction itself. Meanwhile, GLM Du et al.
(2022), CM3 Aghajanyan et al. (2022), and InCoder Fried et al. (2023) adopt a
different approach. They shift the target span to the end of the context,
employing left-to-right autoregressive modeling for training.

Most of these models are designed for token-level infilling tasks, which often
don’t align with real-world applications due to the incomplete nature of the
final token in actual prompts. FIM Bavarian et al. (2022) explored various
levels of spans, _i.e.,_ line level, token level, and character level. As
shown in their results, models trained with line-level or token-level spans
perform poorly on character-level infilling tasks. To enhance the performance
of models trained on token-level spans in character-level infilling tasks,
token healing was proposed to fix tokenization artifacts that normally arise
at the boundary between the end of a prompt and the beginning of a set of
generated tokens333https://github.com/guidance-
ai/guidance/blob/main/notebooks/tutorials/token_healing.ipynb. While it
effectively bridges the gap between the prefix and generated text, it falls
short in handling the transition between generated text and the suffix,
highlighting the need for further research in character-level infilling.

##  3 Preliminaries

In this section, we provide a straightforward introduction to the FIM Bavarian
et al. (2022) method and conduct a theoretical analysis of how inconsistent
labeling affects the model’s perplexity.

###  3.1 Fill-In-the-Middle (FIM)

FIM is designed to train models to complete the central sections of documents.
This approach involves joint training on a combined dataset of traditional
left-to-right sequences and data transformed by FIM, with an infilling rate
reaching as high as 90%. According to experimental results, FIM maintains the
autoregressive test losses of the left-to-right models without incurring
significant costs, and it only slightly impacts the performance in downstream
evaluations Allal et al. (2023).

In a particular document, FIM segments a document into three distinct parts:
the prefix, the middle, and the suffix. It introduces three levels of
segmentation: single-line, multi-line, and random-span. Because random-span is
more in line with actual usage conditions, previous studies Rozière et al.
(2023); Li et al. (2023) usually trained the model with random-span level
splitting. After splitting, it moves the middle piece to the end as

$\text{doc}\to\text{(pre, mid, suf)}\to\text{(pre, suf, mid)},$

then concatenate the three pieces using special tokens as

$\text{<PRE> pre <SUF> suf <MID> mid <EOT>}.$

This mode is termed Prefix-Suffix-Middle (PSM) mode. Additionally, FIM
introduced the Suffix-Prefix-Middle (SPM) mode, which interchanges the
positions of the prefix and suffix. A variant of the SPM mode is also
proposed, maintaining the same structure as the PSM mode. Detailed
descriptions of these modes are provided in Appendix A.

###  3.2 Impact of Inconsistent Labels

When the FIM method employs the random-span approach, a training sample can
contain up to four sub-tokens. This can potentially lead to inconsistent
labels that indicate the same input but with different labels. This issue
becomes particularly critical when the model is required to predict a sub-
token. In Section 1, we construct a simple experiment to illustrate that this
inconsistency can significantly affect the model’s perplexity. Here, we offer
a theoretical analysis to further elucidate this phenomenon.

We are considering a classification task involving $n$ classes, where each
sample’s label is associated with one of $m$ different categories across
various training instances. Let $\bm{y}\in\\{0,1\\}^{n}$ represent the actual
label, and $\hat{\bm{y}}\in\mathbb{R}^{n}$ represent the predicted
probabilities. In this context, the cross-entropy loss is computed as

$\mathcal{L}=-\sum_{i=0}^{n}\bm{y}_{i}\log\hat{\bm{y}}_{i}.$ (1)

Assuming that the first $m$ elements of $\bm{y}$ (i.e.,
$\bm{y}_{1},\dots,\bm{y}_{m}$) are set to 1, while the rest are 0, the loss
function is defined as

$\mathcal{L}(\hat{\bm{y}})=-\sum_{i=0}^{m}\log\hat{\bm{y}}_{i}.$ (2)

Then our objective can be expressed as

$\hat{\bm{y}}^{*}=\arg\max\mathcal{L}(\hat{\bm{y}}),~{}~{}~{}~{}s.t.\sum_{i=0}^{n}\hat{\bm{y}}_{i}=1.$
(3)

Here, we introduce the concept of the Lagrange Multiplier, which enables us to
formulate the target function as

$\mathcal{L}(\hat{\bm{y}},\lambda)=-\sum_{i=0}^{m}\log\hat{\bm{y}}_{i}-\lambda(\sum_{i=0}^{n}\hat{\bm{y}}_{i}-1).$
(4)

We then calculate the partial derivatives of
$\hat{\bm{y}}_{1},\dots,\hat{\bm{y}}_{m}$ respectively, which allows us to
obtain

$\frac{\partial\mathcal{L}(\hat{\bm{y}},\lambda)}{\partial\hat{\bm{y}}_{i}}=-\frac{1}{\hat{\bm{y}}_{i}}-\lambda~{}~{}~{}\text{when}~{}~{}~{}i=1,\dots,m.$
(5)

By setting these derivatives to zero, we obtain

$\hat{\bm{y}}_{1}^{*}=\dots=\hat{\bm{y}}_{m}^{*}=-\frac{1}{\lambda}.$ (6)

Since $\hat{\bm{y}}_{m+1},\dots,\hat{\bm{y}}_{n}$ are not included in the
objective function of Eq. (3), and given that the logarithm is a monotonically
increasing function, setting these values to 0 would maximize the objective
function. Consequently, the condition is formulated as
$\sum_{i=1}^{m}\hat{\bm{y}}_{i}=1$. By incorporating this condition, we derive

$\hat{\bm{y}}_{1}^{*}=\dots=\hat{\bm{y}}_{m}^{*}=\frac{1}{m}.$ (7)

We have now completed the proof, demonstrating that when a data point is
labeled differently across various samples, the model tends to assign an equal
probability of $\frac{1}{m}$ to each label. This behavior leads to a large
perplexity of the model, which further suggests its limited modeling
capability.

We consider this proof to be a microscopic explanation, which demonstrates
that the perplexity of sub-tokens will be higher. In contrast, macroscopically
speaking, we assume that the probability of the next token prediction obeys a
certain distribution. Then, these sub-tokens are outliers. The presence of
several outliers in each piece of training data will result in a low
confidence in the model, that is, a high degree of perplexity.

This phenomenon is notably detrimental in sensitive tasks. For example, the
initially predicted token in practical completion is usually a sub-token. The
higher perplexity of the first token has little impact on the overall quality
of the generated text, but in some sensitive tasks such as code completion,
even a minor error in any token can result in program malfunction.

Figure 2: An overview of the difference between FIM and the proposed FIM-SE.
Here, the green background indicates vanilla FIM and the blue background
indicates our FIM-SE.

##  4 Methodology

In this section, we introduce the proposed method. We begin by outlining the
training process with FIM-SE, followed by an explanation of the inference
procedure. Finally, we delve into more training details and highlight the
distinctions between our approach and the traditional FIM method.

###  4.1 FIM-SE Training

The core idea of FIM-SE is to ensure that the tokens predicted by the model
are complete, thereby circumventing the issue of large perplexity associated
with sub-tokens. Specifically, we shift from character-level to line-level
random splitting in training data construction and then reconstruct the prompt
to keep the ability of the model on the character-level infilling tasks.

As shown in Figure 2, our process for forming the final training sample from a
specific document involves three distinct steps. (1) Splitting: we split the
original document into three pieces at the character level, namely the prefix,
the middle, and the suffix (2) Refining: we distinguish between the last line
of the prefix and the first line of the suffix, denoting them as _L-Prefix_
and _F-Suffix_ , respectively. Correspondingly, we label the remaining lines
of the prefix and the suffix as _R-Prefix_ and _R-Suffix_. (3) Concatenating:
we concatenate all these sections in the following order along with their
special tokens as

<PRE> R-Prefix <SUF> R-Suffix <START> L-Prefix <END> F-Suffix
$\displaystyle\text{<MID>}~{}\emph{L-Prefix}~{}\emph{Middle}~{}\emph{F-Suffix}~{}\text{<EOT>}.$

When tokenizing a sample, we tokenize each section individually and then
concatenate them with the special tokens, which ensures that special tokens
will not be cut or merged.

###  4.2 FIM-SE Inference

During the inference stage, the model can be employed for left-to-right
generation in a standard manner. When working with an arbitrary location
within an existing document, we establish the preceding lines as _R-Prefix_
and the following lines as _R-Suffix_. For the specific line at the target
location, the text before this point is termed _L-Prefix_ , and the text
following it is named _F-Suffix_. Subsequently, a span is generated to be
inserted at this location by autoregressively sampling tokens from the
structured prompt

<PRE> R-Prefix <SUF> R-Suffix
$\displaystyle\text{<START>}~{}\emph{L-Prefix}~{}\text{<END>}~{}\emph{F-Suffix}\text{<MID>}.$

This process continues until the “<EOT>” (End of Text) token is produced.

After obtaining the generation, we verify if it begins with the _L-Prefix_ and
ends with the _F-Suffix_. If the generation does not adhere to these criteria,
we classify the infilling endeavor as unsuccessful. Conversely, if the
criteria are met, we eliminate the _L-Prefix_ from the beginning and the _F-
Suffix_ from the end, considering the remaining text as the completed segment.

###  4.3 Learning and Discussion

Training Details. We train our models using the StarCoder code
corpus444https://huggingface.co/datasets/bigcode/starcoderdata, a carefully
curated dataset sourced from The Stack, encompassing 92 languages. To ensure
consistency, we exclude categories such as GitHub issues, GitHub commits, and
Jupyter Notebooks, which possess distinct column structures. Additionally, we
remove flags marking repositories, files, and stars to maintain a focus on the
pure code content in the remaining files. After gathering the data, we process
it using the previously described method with a 90% FIM rate, following the
methodologies of existing studies Bavarian et al. (2022); Allal et al. (2023);
Rozière et al. (2023). It’s important to note that we exclusively employ the
PSM format depicted in Figure 2, as the SPM variant used in prior research
Bavarian et al. (2022) lacks a separator between the prefix and middle,
potentially leading to model confusion. We conduct experiments and give an
experimental analysis in Section 4.

Discussion. Compared to previous masked language modeling on encoder-only
models and encoder-decoder models, our method excels in character-level
infilling. While these traditional methods primarily concentrate on token-
level infilling, this approach often falls short in numerous industry
applications, as user text seldom forms complete tokens. Compared to vanilla
FIM Bavarian et al. (2022), our method also has the following two merits.
Firstly, our method ensures that tokens following “<MID>” are complete,
eliminating the need for sub-token predictions during inference and thereby
mitigating the effects of the large perplexity of the model on sub-tokens.
Secondly, our method transforms character-level infilling into line-level
infilling. This unification of formats enhances transfer across different
levels, significantly augmenting the efficacy of FIM training.

Model | Size | Training Methods | random-span | single-line | multi-line | Humaneval | MBPP
---|---|---|---|---|---|---|---
InCoder | 6B | Causal Masking | - | 69.0 | 38.6 | 15.0 | 19.4
FIM | 7B | FIM-SPM | 55.1 | 75.1 | 44.1 | - | -
code-davinci-002 | 175B | - | 74.2 | 91.6 | 69.9 | 44.5 | 55.4
StarCoder | 1B | FIM-PSM | 44.1∗ | 64.3∗ | 30.8∗ | 15.2 | 22.6∗
FIM-SE-PSM | 48.8 (+4.7) | 72.6 (+8.3) | 37.1 (+6.3) | 16.5 | 25.6
15B | FIM-PSM | 66.4∗ | 83.8∗ | 53.7∗ | 30.4 | 43.2∗
FIM-SE-PSM | 67.7 (+1.3) | 85.8 (+2.0) | 57.4 (+3.7) | 30.5 | 44.6
Code Llama | 7B | FIM-SPM | 39.0 | 83.3 | 50.8 | 33.5 | 41.4
FIM-PSM | 59.7 | 74.1 | 48.2
FIM-SE-PSM | 67.8 (+8.1) | 84.9 (+10.8) | 57.2 (+9.0) | 30.5 | 41.4
13B | FIM-SPM | 41.9 | 85.6 | 56.1 | 36.0 | 47.0
FIM-PSM | 63.6 | 75.9 | 51.0
FIM-SE-PSM | 72.4 (+8.8) | 87.4 (+11.5) | 61.7 (+10.7) | 37.2 | 50.2
Table 2: Pass@1 accuracy on Humaneval infilling datasets. Results evaluated on
our end are marked with “*”, while those unavailable are left blank. Note that
StarCoder was evaluated using a cleaned and smaller version of MBPP so we
conducted a re-evaluation.

##  5 Experiments

In this section, we construct experiments to demonstrate the effectiveness of
our method. Due to space limitations, we have constructed more experiments in
Appendix B.

###  5.1 Experimental Setup

Datasets. Following FIM Bavarian et al. (2022), we use code to test our
methods. Because we can use test suites to evaluate the correctness of samples
in our tasks even when evaluating long samples from open-ended generations.
Specifically, we use three levels of infilling benchmarks, namely random-span,
single-line, and multi-line. All of them are constructed from Humaneval
benchmarks Chen et al. (2021). Since other infilling benchmarks such as Return
Type Prediction and Docstring Generation focus on token-level infilling, we do
not use these benchmarks.

Implementation Details. We continually pre-train four models with our methods,
_i.e.,_ StarCoder-1B, StarCoder-15B Li et al. (2023), Code Llama 7B, and Code
Llama 13B Rozière et al. (2023). We employ AdamW Loshchilov and Hutter (2019)
optimizer with $\beta_{1}=0.9$, $\beta_{2}=0.95$, $\epsilon=10^{-8}$ and
weight decay of 0.1. Following previous study Gupta et al. (2023), we set the
peak learning rate to $3\times 10^{-5}$ and use a cosine schedule without
warm-up. We use a batch size of 4M tokens which are presented as sequences of
8K tokens each for StarCoder and 16K tokens each for Code Llama. We train each
model on 20B tokens in total. To efficiently train the computationally
intensive models, we simultaneously employ DeepSpeed (Rajbhandari et al.,
2020) and Flash Attention 2 (Dao, 2023). On 32 NVIDIA A800 80GB GPUs,
StarCoder-1B, StarCoder-15B, Code Llama 7B, and Code Llama 13B take 14 hours,
140 hours, 75 hours, and 138 hours, respectively.

###  5.2 Results

Baselines. We compare FIM-SE with previous state-of-the-art methods, including
InCoder Fried et al. (2023), FIM Bavarian et al. (2022), Codex Chen et al.
(2021), StarCoder Li et al. (2023) and Code Llama Rozière et al. (2023). For
other code models such as CodeGeeX Zheng et al. (2023) and OctoCoder
Muennighoff et al. (2023), we do not use them as baselines since they have not
undergone infilling pre-training. Since we focus on character-level infilling,
models focused on token-level infilling also do not be considered baselines.
Because these models cannot effectively handle the character-level infilling
task Bavarian et al. (2022).

Random-span. As shown in Table 2, our proposed method demonstrates notable
improvements in random-span infilling tasks across four models, specifically
achieving gains of 4.7%, 1.3%, 8.1%, and 8.8%. Notably, the enhancement in
StarCoder-15B is comparatively modest. This could be attributed to the fact
that StarCoder has undergone pre-training with four epochs on a total of 1TB
tokens, in contrast to Code Llama’s pre-training on 500B tokens, resulting in
a more refined model fit. Comparing StarCoder-15B with StarCoder-1B, the small
model trained on the same tokens has more gain, suggesting that the consistent
training approach of our method is particularly beneficial for smaller models
in achieving better fit. Comparing StarCoder-15B with Code Llama 13B, the
model with a similar size using fewer tokens achieves better results. This
indicates that the consistent training approach of our method accelerates the
fitting process in larger models.

Single-line and Multi-line. The proposed method demonstrates notable
improvements in both single-line and multi-line infilling tasks. For instance,
based on Code Llama 13B, our method surpasses FIM by 11.5% and 10.7% in
single-line and multi-line infilling tasks, respectively. This enhancement can
be attributed to two key factors. Firstly, our method integrates character-
level and line-level processing, significantly enhancing the model’s line-
level infilling capabilities. Secondly, it avoids the inclusion of any sub-
tokens after the “<MID>” token, which sharpens the model’s accuracy in
predicting the initial token. In contrast, in the standard FIM, the first
token following “<MID>” is typically a sub-token during training, while the
model is adopted to predict a complete token in the line-level infilling tasks
during the inference stage. A comprehensive case study is provided in Appendix
B.2 for further illustration.

Code Generation Task. We also report results on Humaneval Chen et al. (2021)
and MBPP Austin et al. (2021). As shown in Table 2, our method has minimal
impact on the model’s performance in the two code generation tasks (Note that
we cannot reproduce the result of Code Llama 7B on Humaneval, just 29.9% in
our environment). In summary, FIM-SE demonstrates a remarkable ability to
improve infilling tasks without compromising code generation capabilities.

###  5.3 Detail Analysis

Impact of Inconsistent Labels. As mentioned in Section 1 and Section 3.2, we
discussed how FIM leads to inconsistent labels during training at split
points. This phenomenon results in large perplexity on sub-tokens,
subsequently diminishing the model’s accuracy in generating sub-tokens. To
investigate this effect, we conducted an experiment based on the StarCoder-1B.
Specifically, we adjusted the temperature within the range of [0, 1.4] and
compared the performance of models trained using both FIM-SE and FIM in
generating 20 completions to estimate the Pass@1 rate. Figure 3 illustrates
that the performance gap between the FIM-SE and FIM generators widens as the
temperature increases, highlighting the larger perplexity associated with
models trained using FIM.

Figure 3: Performance on Humaneval random-span infilling task with different temperatures. The line denotes the difference between FIM-SE and FIM. Note that when the temperature surpasses 1.4, both models output noisy text and show very low performance. Methods | FIM-SE | w/o LF-Loss
---|---|---
random-span | 0.488 | 0.488
single-line | 0.726 | 0.716
multi-line | 0.371 | 0.369
Test loss | 0.847 | 0.834
Table 3: Effect of training loss on sub-tokens. The metric is Pass@1 accuracy.
Here, LF-Loss denotes the loss for tokens in _L-Prefix_ and _F-Suffix_.

Furthermore, we evaluate the impact of inconsistent labels on training.
Specifically, we mask the loss for tokens in _L-Prefix_ and _F-Suffix_ ,
ensuring that only complete tokens contribute to loss calculations. As shown
in table 3, computing losses for _L-Prefix_ and _F-Suffix_ led to a slightly
higher test loss without significantly affecting performance. This could be
attributed to the minimal proportion of sub-tokens, as the presence of up to
four sub-tokens per sample had a negligible impact on the final test results.
In summary, while the loss of sub-tokens in training scarcely affects
performance, the presence of sub-tokens in prediction objectives markedly
influences performance.

Methods | random-span | single-line | multi-line
---|---|---|---
FIM-SE | 0.488 | 0.726 | 0.371
SPM v1 | 0.492 | 0.703 | 0.374
SPM v2 | 0.013 | 0.085 | 0.088
SPM v3 | 0.090 | 0.717 | 0.383
Table 4: Comparison between different SPM format variants and FIM-SE. The
metric is Pass@1 accuracy.

Comparison with SPM Mode. In previous studies, the Suffix-Prefix-Middle
variant had better performance in most cases Bavarian et al. (2022); Rozière
et al. (2023). Here, we explore how to combine our method with SPM mode based
on StarCoder-1B. Specifically, we designed the following three prompt formats
for SPM mode. We train the model using these formats and the PSM mode, equally
distributed across 20 billion tokens.

  1. (1)

SPM v1: “<SUF> _R-Suffix_ <PRE> _R-Prefix_ <START> _L-Prefix_ <END> _F-Suffix_
<MID>”, which add the constraints before the middle to the vanilla SPM mode.

  2. (2)

SPM v2: “<PRE> <SUF> _R-Suffix_ <START> _L-Prefix_ <END> _F-Suffix_ <MID> _R-
Prefix_ ”, which add the constraints before the middle to the variant SPM in
FIM.

  3. (3)

SPM v3: “<PRE> <SUF> _R-Suffix_ <MID> _R-Prefix_ <START> _L-Prefix_ <END> _F-
Suffix_ ”, which add the constraints after prefix to the variant SPM in FIM.

Table 4 presents all comparison results of the three variants. As we can see,
SPM v2 and SPM v3 perform worse on random-span infilling tasks. This occurs
because there is no separator between the prefix and the middle, leading to
conflicts with the PSM mode, regardless of where the restriction is inserted.
In contrast, SPM v1 and PSM perform almost the same because there is no
conflict. To maintain consistency with the pre-trained models Li et al.
(2023); Rozière et al. (2023), we adopt the PSM mode.

(a) PCP Rate of FIM-SE. (b) Pass@1 performance of FIM-SE and FIM.

Figure 4: Statistics of length of _L-Prefix_ and _F-Suffix_.

Analysis of Post-Check during Inference. As we mentioned in Section 4.2, it’s
essential to verify if the generation begins with the _L-Prefix_ and ends with
the _F-Suffix_. Here, we perform statistical analysis on the success rate of
the model based on StarCoder-1B. We focus on the Post-Check Pass Rate (PCP
Rate), which quantifies the percentage of model outputs complying with the
post-check criteria, _i.e.,_ starting with the _L-Prefix_ and ending with the
_F-Suffix_. We then examine the correlation between the PCP Rate and the
average length of the _L-Prefix_ and _F-Suffix_. Additionally, we analyze the
Pass@1 rates for FIM-SE and FIM across varying lengths of these prefixes and
suffixes.

As shown in Figure 4, the PCP Rate increases with length, suggesting that
longer _L-Prefixes_ and _F-Suffixes_ provide more guidance for the model’s
text completion. Moreover, the Pass@1 metrics for both FIM-SE and FIM also
support this, showing enhanced performance with extended _L-Prefixes_ and _F-
Suffixes_. Across all lengths, FIM-SE consistently outperforms the standard
FIM, demonstrating the effectiveness of our approach.

##  6 Conclusion

In this paper, we showed that traditional infilling techniques struggle with
managing the boundaries of prefixes and suffixes. To address this, we
introduced a novel approach, referred to as FIM-SE. Our method transforms the
random-span mode to multi-line mode by removing the _L-Prefix_ and _F-Suffix_.
We further incorporated two special tokens to delineate the two incomplete
lines, thereby guiding the generation. Extensive experiments reveal that our
approach surpasses existing baselines with a clear edge. In future work, we
plan to explore the adaptation of our method to the variant SPM mode, which
holds the promise of even better performance.

##  7 Limitations

The primary limitation of this study is the inability of the proposed method
to accommodate the variant SPM mode, previously established as superior by
prior research. This challenge arises due to the absence of a distinct
delimiter between the prefix and middle, impeding our capacity to guide the
model on the commencement point for completion and to appropriately position
the prompt that instructs the model to start with _L-Prefix_ and end with _F-
Suffix_. In future endeavors, we plan to explore adapting our method to the
variant SPM mode, to achieve better performance. Another limitation of this
paper is the probability that our proposed method fails to complete tasks when
the generation neither starts with the _L-Prefix_ nor ends with the _F-
Suffix_. For example, the fail rates of StarCoder-1B and StarCoder-15B are
18.7% and 9.4%, respectively. This issue is a primary factor impacting model
performance. Future work will concentrate on improving the post-check pass
rate by developing more comprehensive prompts and refining constraint
decoding.

##  8 Ethics Statement

In this paper, we utilized the StarCoder dataset Li et al. (2023). This
dataset has been made publicly available for academic purposes. The creators
of the StarCoder dataset have transparently disclosed its derivation from The
Stack v1.2 Kocetkov et al. (2022). Importantly, The Stack v1.2 is compiled
from a collection of GitHub repositories, all of which operate under
permissive licenses. This ensures that the data’s utilization aligns with the
original authors’ intentions and the legal frameworks governing open-source
contributions. In conclusion, the application of the StarCoder dataset in our
study complies with the ethical guidelines for research data usage, aligning
with the broader principles of academic honesty and the responsible conduct of
research.

## Acknowledgment

This project is funded in part by National Key R&D Program of China Project
2022ZD0161100, by the Centre for Perceptual and Interactive Intelligence
(CPII) Ltd under the Innovation and Technology Commission (ITC)’s InnoHK, by
General Research Fund of Hong Kong RGC Project 14204021. Hongsheng Li is a PI
of CPII under the InnoHK.

## References

  * Aghajanyan et al. (2022) Armen Aghajanyan, Bernie Huang, Candace Ross, Vladimir Karpukhin, Hu Xu, Naman Goyal, Dmytro Okhonko, Mandar Joshi, Gargi Ghosh, Mike Lewis, and Luke Zettlemoyer. 2022.  CM3: A causal masked multimodal model of the internet.  _CoRR_ , abs/2201.07520.
  * Allal et al. (2023) Loubna Ben Allal, Raymond Li, Denis Kocetkov, Chenghao Mou, Christopher Akiki, Carlos Muñoz Ferrandis, Niklas Muennighoff, Mayank Mishra, Alex Gu, Manan Dey, Logesh Kumar Umapathi, Carolyn Jane Anderson, Yangtian Zi, Joel Lamy-Poirier, Hailey Schoelkopf, Sergey Troshin, Dmitry Abulkhanov, Manuel Romero, Michael Lappert, Francesco De Toni, Bernardo García del Río, Qian Liu, Shamik Bose, Urvashi Bhattacharyya, Terry Yue Zhuo, Ian Yu, Paulo Villegas, Marco Zocca, Sourab Mangrulkar, David Lansky, Huu Nguyen, Danish Contractor, Luis Villa, Jia Li, Dzmitry Bahdanau, Yacine Jernite, Sean Hughes, Daniel Fried, Arjun Guha, Harm de Vries, and Leandro von Werra. 2023.  Santacoder: don’t reach for the stars! _CoRR_ , abs/2301.03988.
  * Anil et al. (2023) Rohan Anil, Andrew M. Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, Eric Chu, Jonathan H. Clark, Laurent El Shafey, Yanping Huang, Kathy Meier-Hellstern, Gaurav Mishra, Erica Moreira, Mark Omernick, Kevin Robinson, Sebastian Ruder, Yi Tay, Kefan Xiao, Yuanzhong Xu, Yujing Zhang, Gustavo Hernández Ábrego, Junwhan Ahn, Jacob Austin, Paul Barham, Jan A. Botha, James Bradbury, Siddhartha Brahma, Kevin Brooks, Michele Catasta, Yong Cheng, Colin Cherry, Christopher A. Choquette-Choo, Aakanksha Chowdhery, Clément Crepy, Shachi Dave, Mostafa Dehghani, Sunipa Dev, Jacob Devlin, Mark Díaz, Nan Du, Ethan Dyer, Vladimir Feinberg, Fangxiaoyu Feng, Vlad Fienber, Markus Freitag, Xavier Garcia, Sebastian Gehrmann, Lucas Gonzalez, and et al. 2023.  Palm 2 technical report.  _CoRR_ , abs/2305.10403.
  * Austin et al. (2021) 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. 2021.  Program synthesis with large language models.  _CoRR_ , abs/2108.07732.
  * Bai et al. (2023) Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, Binyuan Hui, Luo Ji, Mei Li, Junyang Lin, Runji Lin, Dayiheng Liu, Gao Liu, Chengqiang Lu, Keming Lu, Jianxin Ma, Rui Men, Xingzhang Ren, Xuancheng Ren, Chuanqi Tan, Sinan Tan, Jianhong Tu, Peng Wang, Shijie Wang, Wei Wang, Shengguang Wu, Benfeng Xu, Jin Xu, An Yang, Hao Yang, Jian Yang, Shusheng Yang, Yang Yao, Bowen Yu, Hongyi Yuan, Zheng Yuan, Jianwei Zhang, Xingxuan Zhang, Yichang Zhang, Zhenru Zhang, Chang Zhou, Jingren Zhou, Xiaohuan Zhou, and Tianhang Zhu. 2023.  Qwen technical report.  _CoRR_ , abs/2309.16609.
  * Bavarian et al. (2022) Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry Tworek, and Mark Chen. 2022.  Efficient training of language models to fill in the middle.  _CoRR_ , abs/2207.14255.
  * Brown et al. (2020) Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020.  Language models are few-shot learners.  In _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_.
  * Chen et al. (2021) 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. 2021.  Evaluating large language models trained on code.  _CoRR_ , abs/2107.03374.
  * Chowdhery et al. (2023) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. 2023\.  Palm: Scaling language modeling with pathways.  _J. Mach. Learn. Res._ , 24:240:1–240:113.
  * Cui et al. (2021) Yiming Cui, Wanxiang Che, Ting Liu, Bing Qin, and Ziqing Yang. 2021.  Pre-training with whole word masking for chinese BERT.  _IEEE ACM Trans. Audio Speech Lang. Process._ , 29:3504–3514.
  * Dao (2023) Tri Dao. 2023.  Flashattention-2: Faster attention with better parallelism and work partitioning.  _CoRR_ , abs/2307.08691.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  BERT: pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers)_ , pages 4171–4186. Association for Computational Linguistics.
  * Du et al. (2022) Zhengxiao Du, Yujie Qian, Xiao Liu, Ming Ding, Jiezhong Qiu, Zhilin Yang, and Jie Tang. 2022.  GLM: general language model pretraining with autoregressive blank infilling.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022_ , pages 320–335. Association for Computational Linguistics.
  * Fried et al. (2023) Daniel Fried, Armen Aghajanyan, Jessy Lin, Sida Wang, Eric Wallace, Freda Shi, Ruiqi Zhong, Scott Yih, Luke Zettlemoyer, and Mike Lewis. 2023.  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.
  * Gupta et al. (2023) Kshitij Gupta, Benjamin Thérien, Adam Ibrahim, Mats L. Richter, Quentin Anthony, Eugene Belilovsky, Irina Rish, and Timothée Lesort. 2023.  Continual pre-training of large language models: How to (re)warm your model? _CoRR_ , abs/2308.04014.
  * Jiang et al. (2023) Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. 2023.  Mistral 7b.  _CoRR_ , abs/2310.06825.
  * Joshi et al. (2020) Mandar Joshi, Danqi Chen, Yinhan Liu, Daniel S. Weld, Luke Zettlemoyer, and Omer Levy. 2020.  Spanbert: Improving pre-training by representing and predicting spans.  _Trans. Assoc. Comput. Linguistics_ , 8:64–77.
  * Kocetkov et al. (2022) Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferrandis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, Dzmitry Bahdanau, Leandro von Werra, and Harm de Vries. 2022.  The stack: 3 TB of permissively licensed source code.  _CoRR_ , abs/2211.15533.
  * Kojima et al. (2022) Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. 2022.  Large language models are zero-shot reasoners.  In _Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022_.
  * Kudo and Richardson (2018) Taku Kudo and John Richardson. 2018.  Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, EMNLP 2018: System Demonstrations, Brussels, Belgium, October 31 - November 4, 2018_ , pages 66–71. Association for Computational Linguistics.
  * Lewis et al. (2020a) Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. 2020a.  BART: denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020_ , pages 7871–7880. Association for Computational Linguistics.
  * Lewis et al. (2020b) Patrick S. H. Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. 2020b.  Retrieval-augmented generation for knowledge-intensive NLP tasks.  In _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_.
  * Li et al. (2023) Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Marc Marone, Christopher Akiki, Jia Li, Jenny Chim, Qian Liu, Evgenii Zheltonozhskii, Terry Yue Zhuo, Thomas Wang, Olivier Dehaene, Mishig Davaadorj, Joel Lamy-Poirier, João Monteiro, Oleh Shliazhko, Nicolas Gontier, Nicholas Meade, Armel Zebaze, Ming-Ho Yee, Logesh Kumar Umapathi, Jian Zhu, Benjamin Lipkin, Muhtasham Oblokulov, Zhiruo Wang, Rudra Murthy V, Jason Stillerman, Siva Sankalp Patel, Dmitry Abulkhanov, Marco Zocca, Manan Dey, Zhihan Zhang, Nour Moustafa-Fahmy, Urvashi Bhattacharyya, Wenhao Yu, Swayam Singh, Sasha Luccioni, Paulo Villegas, Maxim Kunakov, Fedor Zhdanov, Manuel Romero, Tony Lee, Nadav Timor, Jennifer Ding, Claire Schlesinger, Hailey Schoelkopf, Jan Ebert, Tri Dao, Mayank Mishra, Alex Gu, Jennifer Robinson, Carolyn Jane Anderson, Brendan Dolan-Gavitt, Danish Contractor, Siva Reddy, Daniel Fried, Dzmitry Bahdanau, Yacine Jernite, Carlos Muñoz Ferrandis, Sean Hughes, Thomas Wolf, Arjun Guha, Leandro von Werra, and Harm de Vries. 2023.  Starcoder: may the source be with you! _CoRR_ , abs/2305.06161.
  * Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019.  Roberta: A robustly optimized BERT pretraining approach.  _CoRR_ , abs/1907.11692.
  * Loshchilov and Hutter (2019) Ilya Loshchilov and Frank Hutter. 2019.  Decoupled weight decay regularization.  In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_. OpenReview.net.
  * Muennighoff et al. (2023) Niklas Muennighoff, Qian Liu, Armel Zebaze, Qinkai Zheng, Binyuan Hui, Terry Yue Zhuo, Swayam Singh, Xiangru Tang, Leandro von Werra, and Shayne Longpre. 2023.  Octopack: Instruction tuning code large language models.  _CoRR_ , abs/2308.07124.
  * OpenAI (2023) OpenAI. 2023.  GPT-4 technical report.  _CoRR_ , abs/2303.08774.
  * Ouyang et al. (2022) Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F. Christiano, Jan Leike, and Ryan Lowe. 2022.  Training language models to follow instructions with human feedback.  In _NeurIPS_.
  * Penedo et al. (2023) Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. 2023.  The refinedweb dataset for falcon LLM: outperforming curated corpora with web data, and web data only.  _CoRR_ , abs/2306.01116.
  * Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. 2018.  Improving language understanding by generative pre-training.  _OpenAI_.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9.
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _J. Mach. Learn. Res._ , 21:140:1–140:67.
  * Rajbhandari et al. (2020) Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020.  Zero: memory optimizations toward training trillion parameter models.  In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2020, Virtual Event / Atlanta, Georgia, USA, November 9-19, 2020_ , page 20. IEEE/ACM.
  * Rozière et al. (2023) Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton-Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. 2023.  Code llama: Open foundation models for code.  _CoRR_ , abs/2308.12950.
  * Song et al. (2019) Kaitao Song, Xu Tan, Tao Qin, Jianfeng Lu, and Tie-Yan Liu. 2019.  MASS: masked sequence to sequence pre-training for language generation.  In _Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA_ , volume 97 of _Proceedings of Machine Learning Research_ , pages 5926–5936. PMLR.
  * Stern et al. (2019) Mitchell Stern, William Chan, Jamie Kiros, and Jakob Uszkoreit. 2019.  Insertion transformer: Flexible sequence generation via insertion operations.  In _Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA_ , volume 97 of _Proceedings of Machine Learning Research_ , pages 5976–5985. PMLR.
  * 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, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023a.  Llama: Open and efficient foundation language models.  _CoRR_ , abs/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, Dan Bikel, Lukas Blecher, Cristian Canton-Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. 2023b.  Llama 2: Open foundation and fine-tuned chat models.  _CoRR_ , abs/2307.09288.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pages 5998–6008.
  * Xiong et al. (2023) Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. 2023.  Effective long-context scaling of foundation models.  _CoRR_ , abs/2309.16039.
  * Yang et al. (2023) Aiyuan Yang, Bin Xiao, Bingning Wang, Borong Zhang, Ce Bian, Chao Yin, Chenxu Lv, Da Pan, Dian Wang, Dong Yan, Fan Yang, Fei Deng, Feng Wang, Feng Liu, Guangwei Ai, Guosheng Dong, Haizhou Zhao, Hang Xu, Haoze Sun, Hongda Zhang, Hui Liu, Jiaming Ji, Jian Xie, Juntao Dai, Kun Fang, Lei Su, Liang Song, Lifeng Liu, Liyun Ru, Luyao Ma, Mang Wang, Mickel Liu, MingAn Lin, Nuolan Nie, Peidong Guo, Ruiyang Sun, Tao Zhang, Tianpeng Li, Tianyu Li, Wei Cheng, Weipeng Chen, Xiangrong Zeng, Xiaochuan Wang, Xiaoxi Chen, Xin Men, Xin Yu, Xuehai Pan, Yanjun Shen, Yiding Wang, Yiyu Li, Youxin Jiang, Yuchen Gao, Yupeng Zhang, Zenan Zhou, and Zhiying Wu. 2023.  Baichuan 2: Open large-scale language models.  _CoRR_ , abs/2309.10305.
  * Zheng et al. (2023) Qinkai Zheng, Xiao Xia, Xu Zou, Yuxiao Dong, Shan Wang, Yufei Xue, Lei Shen, Zihan Wang, Andi Wang, Yang Li, Teng Su, Zhilin Yang, and Jie Tang. 2023.  Codegeex: A pre-trained model for code generation with multilingual benchmarking on humaneval-x.  In _Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD 2023, Long Beach, CA, USA, August 6-10, 2023_ , pages 5673–5684. ACM.

## Appendix

##  Appendix A Fill-In-the-Middle (FIM)

In Section 3.1, We briefly introduced the prefix-suffix-middle (PSM) mode of
FIM Bavarian et al. (2022). Here, we give a detailed description of the
suffix-prefix-middle (SPM) mode and a variant SPM mode.

For the vanilla SPM mode, it just swaps the prefix and the suffix.
Specifically, after splitting, it moves the suffix to the before:

$\text{doc}\to\text{(pre, mid, suf)}\to\text{(suf, pre, mid)},$

then concatenate the three pieces using special tokens as

$\text{<SUF> suf <PRE> pre <MID> mid <EOT>}.$

To maximize transfer between PSM mode and SPM mode, FIM proposed a novel
variant of SPM mode, which concatenates the prefix, the middle, and the suffix
pieces as

$\text{<PRE> <SUF> suf <MID> pre mid <EOT>}.$

The format occurs naturally as part of PSM training when the chosen prefix is
empty. In this way, the two modes have a consistent format and they can
transfer with each other in joint training and maximize the profits.

##  Appendix B Additional Experiments

###  B.1 Comparison with Token Healing

As discussed in Section 2, token healing is proposed as an ideal solution for
addressing tokenization that normally arise at the boundary between the end of
a prompt and the beginning of a set of generated tokens. Here, we evaluate our
approach against the token healing method based on StarCoder-1B. However,
since token healing struggles with the split points at the end of generated
tokens and the subsequent suffix, we integrate it with our method for a
comprehensive solution. Specifically, we construct the prompt as “<PRE> _R-
Prefix_ <SUF> _R-Suffix_ <START> _L-Prefix_ <END> _F-Suffix_ <MID> _L-Prefix_
” and focus solely on verifying if the generated text ends with _F-Suffix_.

Methods | random-span
---|---
FIM-SE | 0.488
Token Healing | 0.484
Table 5: Comparison with Token Healing.

Piece | Raw Text | Tokens
---|---|---
Prefix |  def so | def | _ so |  |  |
Output |  def sort(arr) | def | _ sort | ( | arr | )
Label |  def sort(arr) | def | _ sort | ( | arr | )
(a) A case can be solved by token healing.

Piece | Raw Text | Tokens
---|---|---
Prefix |  r.add(delim | r | . | add | (
delim |  |  |
Output |  r.add(delimter) | r | . | add | (
delim | ter | ) |
Label |  r.add(delimeter) | r | . | add | (
deli | meter | ) |
(b) A case cannot be solved by token healing.

Table 6: Case of token healing. The first case can be perfectly solved by
token healing. The second case cannot be solved by token healing. Here, ‘_’
denotes blank.

Table 5 presents the comparison results between our method and token healing.
Surprisingly, token healing performs slightly worse than our method. Detailed
analysis revealed that token healing struggles with complex scenarios, such as
splitting the last token into two sub-tokens and merging the latter sub-token
with the initial generated token.

Prefix
---
⬇ def largest_prime_factor(n: int):
"""Return␣the␣largest␣prime␣factor␣of␣n.␣Assume␣n␣>␣1␣and␣is␣not␣a␣prime.
␣␣␣␣>>>␣largest_prime_factor(13195) ␣␣␣␣29 ␣␣␣␣""" def is_prime(k): if k < 2:
return False for i in range(2, k - 1): if k % i == 0: return False return True
Suffix
⬇ for j in range(2, n + 1): if n % j == 0 and is_prime(j): largest =
max(largest, j) return largest
Target Middle
⬇ largest = 1
The top five choices for the initial generated token on StarCoder-1B (FIM),
along with their probabilities
‘\n’: 0.463; ‘\n_ _ _ ’: 0.225; ‘_ _ _ _’: 0.073; ‘<|endoftext|>’: 0.068; ‘_ _
_ _\n_ _ _ ’: 0.061;
The top five choices for the initial generated token on StarCoder-1B (FIM-SE),
along with their probabilities
‘_ _ _ ’: 0.829; ‘\n_ _ _ ’: 0.115; ‘_ _ _ _\n_ _ _ ’: 0.037; ‘_ _ _ _’:
0.008; ‘_ _ _ _ _ _ _ _\n_ _ _ ’ : 0.002;
Table 7: A Case to show perplexity of models on the initial generated token.
Here, ‘_’ denotes blank, and ‘\n’ denotes newline.

Here, we provide an in-depth analysis of the situation. Token healing backs up
the generation process by one or more tokens before the end of the prompt,
then constrains the first tokens generated to have a prefix that matches the
last token in the prompt. As illustrated in Table 6(a), if the last token in
the prompt is "so," token healing identifies a token that both matches this
last token’s prefix and possesses the highest probability, such as "sort."
Consequently, the first generated token is seamlessly integrated, allowing the
generation process to proceed smoothly.

However, due to the consistent integrated intrinsic features of the
SentencePiece Kudo and Richardson (2018) algorithm, it tends to merge the last
sub-token with the preceding one if possible. For example, as illustrated in
Table 6(b), the word "delimiter" is tokenized into "deli" and "meter." If a
prompt ends with "delim", the algorithm prefers tokenizing this as "delim"
instead of splitting it into "deli" and "m". Token healing does not intervene,
because there is no token starting with "delim". Consequently, when the last
sub-token can be combined with the previous one, token healing is unable to
rectify it effectively.

To effectively resolve this issue, it is necessary to revert several tokens
and subsequently engage in limited decoding, utilizing a Trie tree, until the
regenerated text encompasses the previously rolled-back tokens. Nonetheless,
this approach is time-intensive as each decoding step requires traversing the
Trie tree to identify all tokens corresponding to the given prefix. In
contrast, our method only requires modifying the prompt and doing some post-
processing. In addition, our method can handle both boundaries between the
prefix and middle as well as boundaries between the suffix and middle.

###  B.2 Case Study

Here, we present a case demonstrating the model’s large perplexity on the
initial generated token based on StarCoder-1B. Table 7 illustrates that,
despite being a single-line infilling scenario, the perplexity for the first
token is remarkably large, significantly influencing the overall generation.
This is primarily because the first token following “<MID>” tends to be a sub-
token during training, varying across samples due to random character-level
splitting. In contrast, our approach guarantees that no sub-token prediction
is required, leading to lower perplexity and enhanced performance.

Based on these concerns, we hypothesize that the superior performance of the
variant SPM mode over the PSM mode, particularly evident in the single-line
infilling task on Code Llama Rozière et al. (2023) (85.6% vs. 75.9%), can be
attributed to the specific processing format employed by Code Llama. In the
format “<PRE> <SUF> suf <MID> pre mid <EOT>”, Code Llama initially merges the
prefix and middle segments before tokenization. This approach ensures that,
following the “<MID>” token, there are no sub-tokens except for the initial
token. Consequently, this format also guarantees that no sub-token prediction
is required when the prefix is not empty, contributing to the enhanced
performance of the variant SPM mode. In contrast, FIM Bavarian et al. (2022)
adopts a different approach by tokenizing the prefix and the middle separately
before concatenating them. This leads to the presence of sub-tokens amidst the
tokens. Consequently, the performance gap between SPM and PSM modes in FIM is
narrower (61.6% vs. 62.2%) compared to that in Code Llama.