$\displaystyle\text{Pr}\left[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c_{n,1},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+2})=c_{n,2},X_{m+1}=X_{m+2}\right]$
$\displaystyle\quad=\mathbbm{1}\\{l_{1}=l_{2}=:l,c_{n,1}=c_{n,2}=c\\}\binom{m}{l}$
$\displaystyle\quad\quad\times\text{Pr}\left[X_{1}=\dots,X_{l}=X_{m+1}=X_{m+2},\\{X_{l+1},\dots,X_{m}\\}\cap\\{X_{m+1}\\}=\emptyset,\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c\right]$
$\displaystyle\quad=\mathbbm{1}\\{l_{1}=l_{2}=:l,c_{n,1}=c_{n,2}=c\\}\binom{m}{l}$
$\displaystyle\quad\quad\times\text{Pr}\left[X_{1}=\dots,X_{l}=X_{m+1}=X_{m+2},\\{X_{l+1},\dots,X_{m}\\}\cap\\{X_{m+1}\\}=\emptyset,\sum_{i=l+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c-l\right]$

which is determined by the distribution of $(X_{1},\dots,X_{m+2})$. In view of
Equation 3.5 of Sangalli (2006)

$\text{Pr}[X_{1}\in C_{1},\dots,X_{m+2}\in
C_{m+2}]=\sum_{k=1}^{m+2}\frac{\theta^{k}}{(\theta)_{(m+2)}}\sum_{(\pi_{1},\dots,\pi_{k})\in\Pi(m+2,k)}\prod_{i=1}^{k}(|\pi_{i}|-1)!\
\nu(\cap_{r\in\pi_{i}}C_{r})\ .$

We set
$D(m,l):=\\{X_{1}=\dots,X_{l}=X_{m+1}=X_{m+2},\\{X_{l+1},\dots,X_{m}\\}\cap\\{X_{m+1}\\}=\emptyset\\}$,
and we define $\mu_{(\pi_{1},\dots,\pi_{k})}$ as the probability measure on
$(\mathcal{V}^{m+2},\mathscr{V}^{n+2})$ generated by the following identity

$\nu_{\pi_{1},\dots,\pi_{k}}(C_{1}\times\dots\times
C_{m+2}):=\prod_{i=1}^{k}\nu(\cap_{r\in\pi_{i}}C_{r})\ ,$

It is clear that such measures attach to $D(m,l)$ a probability value that is
either 0 or 1. In particular, $\nu_{\pi_{1},\dots,\pi_{k}}(D(m,l))=1$ if and
only if one of the $\pi$’s (e.g. $\pi_{k}$, being these partitions given up to
the order) is exactly equal to the set $\\{1,\dots,l,m+1,m+2\\}$. Accordingly,
we write

$\displaystyle\text{Pr}\left[D(m,l),\sum_{i=l+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c-l\right]$
$\displaystyle\quad=\sum_{k=2}^{m-l+1}\frac{\theta^{k}}{(\theta)_{(m+2)}}\sum_{(\pi_{1},\dots,\pi_{k-1})\in\Pi(m-l,k-1)}(l+1)!\prod_{i=1}^{k-1}(|\pi_{i}|-1)!\nu_{\pi_{1},\dots,\pi_{k}}\left(\sum_{i=l+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c-l\right)$
$\displaystyle\quad=\frac{\theta(\theta)_{(m-l)}}{(\theta)_{(m+2)}}(l+1)!\sum_{r=1}^{m-l}\frac{\theta^{r}}{(\theta)_{(m-l)}}\sum_{(\pi_{1},\dots,\pi_{r})\in\Pi(m-l,r)}\prod_{i=1}^{r}(|\pi_{i}|-1)!\times$
$\displaystyle\quad\quad\times\left\\{\sum_{j=1}^{J}\nu(\\{j\\})\nu_{\pi_{1},\dots,\pi_{r}}\left(\sum_{i=l+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(j)=c-l\right)\right\\}$
$\displaystyle\quad=\frac{\theta(\theta)_{(m-l)}}{J(\theta)_{(m+2)}}(l+1)!\sum_{r=1}^{m-l}\frac{\theta^{r}}{(\theta)_{(m-l)}}\sum_{(\pi_{1},\dots,\pi_{r})\in\Pi(m-l,r)}\prod_{i=1}^{r}(|\pi_{i}|-1)!\times$
$\displaystyle\quad\quad\times\left\\{\sum_{j=1}^{J}\nu_{\pi_{1},\dots,\pi_{r}}\left(\sum_{i=l+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(j)=c-l\right)\right\\}\
.$

Hence,

$\displaystyle\text{Pr}\left[f_{X_{n+1}}=l_{1},f_{X_{m+2}}=l_{2},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c_{n,1},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+2})=c_{n,2},X_{m+1}=X_{m+2}\right]$
(42)
$\displaystyle\quad=\mathbbm{1}\\{l_{1}=l_{2}=:l,c_{n,1}=c_{n,2}=c\\}\frac{m!}{(c-l)!(m-c)!}\frac{\theta(l+1)}{\Gamma(\theta+m+2)}\times$
$\displaystyle\quad\quad\times\frac{\Gamma(\theta)}{\Gamma(\theta/J)\Gamma(\theta(1-1/J))}\Gamma(\theta/J+c-l)\Gamma(\theta(1-1/J)+n-c).$

Now, we consider the second term on the right-hand side of the probability
(41). In particular, we write

$\displaystyle\text{Pr}\left[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c_{n,1},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+2})=c_{n,2},X_{m+1}\neq
X_{m+2}\right]$
$\displaystyle\quad=\binom{m}{l_{1},l_{2}}\text{Pr}\Bigg{[}X_{1}=\dots,X_{l_{1}}=X_{m+1},X_{l_{1}+1}=\dots,X_{l_{1}+l_{2}}=X_{m+2},X_{m+1}\neq
X_{m+2},$
$\displaystyle\quad\quad\quad\quad\quad\quad\\{X_{l_{1}+l_{2}+1},\dots,X_{m}\\}\cap\\{X_{m+1},X_{m+2}\\}=\emptyset,\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c_{n,1},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+2})=c_{n,2}\Bigg{]}$
$\displaystyle\quad=\binom{m}{l_{1},l_{2}}\text{Pr}\Bigg{[}X_{1}=\dots,X_{l_{1}}=X_{m+1},X_{l_{1}+1}=\dots,X_{l_{1}+l_{2}}=X_{m+2},X_{m+1}\neq
X_{m+2},$
$\displaystyle\quad\quad\quad\quad\quad\quad\\{X_{l_{1}+l_{2}+1},\dots,X_{m}\\}\cap\\{X_{m+1},X_{m+2}\\}=\emptyset,$
$\displaystyle\quad\quad\quad\quad\quad\quad
l_{2}\mathbbm{1}_{h_{n}(X_{l_{1}+1})}(X_{m+1})+\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c_{n,1}-l_{1},$
$\displaystyle\quad\quad\quad\quad\quad\quad
l_{1}\mathbbm{1}_{h_{n}(X_{1})}(X_{m+2})+\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+2})=c_{n,2}-l_{2}\Bigg{]}\
.$

Setting

$\displaystyle E(n,l_{1},l_{2})$
$\displaystyle:=\Bigg{\\{}X_{1}=\dots,X_{l_{1}}=X_{m+1},X_{l_{1}+1}=\dots,X_{l_{1}+l_{2}}=X_{m+2},X_{m+1}\neq
X_{m+2},$
$\displaystyle\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\\{X_{l_{1}+l_{2}+1},\dots,X_{m}\\}\cap\\{X_{m+1},X_{m+2}\\}=\emptyset\Bigg{\\}},$

we have that $\nu_{\pi_{1},\dots,\pi_{k}}(E(n,l_{1},l_{2}))=1$ if and only if
two of the $\pi$’s (e.g. $\pi_{k-1}$ and $\pi_{k}$, being these partitions
given up to the order) are exactly equal to the sets $\\{1,\dots,l_{1},m+1\\}$
and $\\{l_{1}+1,\dots,l_{1}+l_{2},m+2\\}$, respectively. Therefore, from
above, we write the following probability

$\displaystyle\text{Pr}\Bigg{[}E(n,l_{1},l_{2}),l_{2}\mathbbm{1}_{h_{n}(X_{l_{1}+1})}(X_{m+1})+\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c_{n,1}-l_{1},$
$\displaystyle\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad
l_{1}\mathbbm{1}_{h_{n}(X_{1})}(X_{m+2})+\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+2})=c_{n,2}-l_{2}\Bigg{]}$
$\displaystyle\quad=\sum_{k=3}^{n-l_{1}-l_{2}+2}\frac{\theta^{k}}{(\theta)_{(m+2)}}\sum_{(\pi_{1},\dots,\pi_{k-2})\in\Pi(m-l_{1}-l_{2},k-2)}l_{1}!l_{2}!\prod_{i=1}^{k-2}(|\pi_{i}|-1)!\times$
$\displaystyle\quad\quad\times\nu_{\pi_{1},\dots,\pi_{k}}\Bigg{(}l_{2}\mathbbm{1}_{h_{n}(X_{l_{1}+1})}(X_{m+1})+\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c_{n,1}-l_{1},$
$\displaystyle\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad
l_{1}\mathbbm{1}_{h_{n}(X_{1})}(X_{m+2})+\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+2})=c_{n,2}-l_{2}\Bigg{)}$
$\displaystyle\quad=\frac{\theta^{2}(\theta)_{(m-l_{1}-l_{2})}}{(\theta)_{(m+2)}}l_{1}!l_{2}!\sum_{r=1}^{m-l_{1}-l_{2}}\frac{\theta^{r}}{(\theta)_{(m-l_{1}-l_{2})}}\sum_{(\pi_{1},\dots,\pi_{r})\in\Pi(m-l_{1}-l_{2},r)}\prod_{i=1}^{r}(|\pi_{i}|-1)!$
$\displaystyle\quad\quad\times\Bigg{[}\sum_{(j_{1},j_{2})\in[J]^{2}}\nu(\\{j_{1}\\})\nu(\\{j_{2}\\})$
$\displaystyle\quad\quad\quad\times\nu_{\pi_{1},\dots,\pi_{r}}\Bigg{(}\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(j_{1})=c_{n,1}-l_{1}-l_{2}\mathds{1}\\{j_{1}=j_{2}\\},$
$\displaystyle\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(j_{2})=c_{n,2}-l_{2}-l_{1}\mathds{1}\\{j_{1}=j_{2}\\}\Bigg{)}\Bigg{]}\
.$

We observe that the expression within the brackets in the last term can be
split into the sum of two terms, according on whether $j_{1}=j_{2}$ or not, in
the sum over $[J]^{2}$. Therefore, we write

$\displaystyle\sum_{r=1}^{m-l_{1}-l_{2}}\frac{\theta^{r}}{(\theta)_{(m-l_{1}-l_{2})}}\sum_{(\pi_{1},\dots,\pi_{r})\in\Pi(m-l_{1}-l_{2},r)}\prod_{i=1}^{r}(|\pi_{i}|-1)!$
$\displaystyle\quad\quad\times\Bigg{[}\sum_{j_{1}=j_{2}\in[J]}\nu(\\{j_{1}\\})\nu(\\{j_{2}\\})\nu_{\pi_{1},\dots,\pi_{r}}\Bigg{(}\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(j_{1})=c_{n,1}-l_{1}-l_{2},$
$\displaystyle\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(j_{2})=c_{n,2}-l_{2}-l_{1}\Bigg{)}\Bigg{]}$
$\displaystyle\quad=\frac{1}{J}\mathds{1}\\{c_{n,1}=c_{n,2}=:c\\}\binom{m-l_{1}-l_{2}}{c-l_{1}-l_{2}}\frac{\Gamma(\theta)}{\Gamma(\theta/J)\Gamma(\theta(1-1/J))}$
$\displaystyle\quad\quad\times\frac{\Gamma(\theta/J+c-l_{1}-l_{2})\Gamma(\theta(1-1/J)+n-c)}{\Gamma(\theta+m-l_{1}-l_{2})}\
.$

On the other hand, assuming $J\geq 3$

$\displaystyle\sum_{r=1}^{m-l_{1}-l_{2}}\frac{\theta^{r}}{(\theta)_{(m-l_{1}-l_{2})}}\sum_{(\pi_{1},\dots,\pi_{r})\in\Pi(m-l_{1}-l_{2},r)}\prod_{i=1}^{r}(|\pi_{i}|-1)!$
$\displaystyle\quad\quad\times\Bigg{[}\sum_{\begin{subarray}{c}(j_{1},j_{2})\in[J]^{2}\\\
j_{1}\neq
j_{2}\end{subarray}}\nu(\\{j_{1}\\})\nu(\\{j_{2}\\})\nu_{\pi_{1},\dots,\pi_{r}}\Bigg{(}\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(j_{1})=c_{n,1}-l_{1},$
$\displaystyle\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\sum_{i=l_{1}+l_{2}+1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(j_{2})=c_{n,2}-l_{2}\Bigg{)}\Bigg{]}$
$\displaystyle\quad=\frac{J-1}{J}\binom{m-l_{1}-l_{2}}{c_{n,1}-l_{1},c_{n,2}-l_{2}}\frac{\Gamma(\theta)}{[\Gamma(\theta/J)]^{2}\Gamma(\theta(1-2/J))}$
$\displaystyle\quad\quad\times\frac{\Gamma(\theta/J+c_{n,1}-l_{1})\Gamma(\theta/J+c_{n,2}-l_{2})\Gamma(\theta(1-2/J)+m-c_{n,1}-c_{n,2})}{\Gamma(\theta+m-l_{1}-l_{2})}.$

Then,

$\displaystyle\text{Pr}\left[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c_{n,1},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+2})=c_{n,2},X_{m+1}\neq
X_{m+2}\right]$ (43)
$\displaystyle\quad=\binom{m}{l_{1},l_{2}}\frac{\theta^{2}(\theta)_{(m-l_{1}-l_{2})}}{(\theta)_{(m+2)}}l_{1}!l_{2}!$
$\displaystyle\quad\quad\times\Bigg{[}\frac{1}{J}\mathbbm{1}\\{c_{n,1}=c_{n,2}=c\\}\binom{m-l_{1}-l_{2}}{c-l_{1}-l_{2}}\frac{\Gamma(\theta)}{\Gamma(\theta/J)\Gamma(\theta(1-1/J))}\frac{\Gamma(\theta/J+c-l_{1}-l_{2})\Gamma(\theta(1-1/J)+m-\nu)}{\Gamma(\theta+m-l_{1}-l_{2})}$
$\displaystyle\quad+\frac{J-1}{J}\binom{m-l_{1}-l_{2}}{c_{n,1}-l_{1},c_{n,2}-l_{2}}\frac{\Gamma(\theta)}{[\Gamma(\theta/J)]^{2}\Gamma(\theta(1-2/J))}$
$\displaystyle\quad\quad\times\frac{\Gamma(\theta/J+c_{n,1}-l_{1})\Gamma(\theta/J+c_{n,2}-l_{2})\Gamma(\theta(1-2/J)+m-c_{n,1}-c_{n,2})}{\Gamma(\theta+m-l_{1}-l_{2})}\Bigg{]}.$

Then, by combining the probability (42) and the probability (43) we write the
following expression

$\displaystyle\text{Pr}\left[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+1})=c_{n,1},\sum_{i=1}^{m}\mathbbm{1}_{h_{n}(X_{i})}(X_{m+2})=c_{n,2}\right]$
$\displaystyle=\frac{m!}{\Gamma(\theta+m+2)}\Big{\\{}\mathbbm{1}\\{l_{1}=l_{2}=:l,c_{n,1}=c_{n,2}=c\\}\frac{\theta(l+1)}{(c-l)!(m-c)!}\beta_{1}(\theta,J)\times$
$\displaystyle\qquad\times\Gamma(\theta/J+c-l)\Gamma(\theta(1-1/J)+m-c)$
$\displaystyle+\mathbbm{1}\\{c_{n,1}=c_{n,2}=c\\}\frac{\theta^{2}}{J(c-l_{1}-l_{2})!(m-c)!}\beta_{1}(\theta,J)\Gamma(\theta/J+c-l_{1}-l_{2})\Gamma(\theta(1-1/J)+m-c)$
$\displaystyle+\left(\frac{J-1}{J}\right)\frac{\theta^{2}}{(c_{n,1}-l_{1})!(c_{n,2}-l_{2})!(m-c_{n,1}-c_{n,2})!}\beta_{2}(\theta,J)\times$
$\displaystyle\qquad\times\Gamma(\theta/J+c_{n,1}-l_{1})\Gamma(\theta/J+c_{n,2}-l_{2})\Gamma(\theta(1-2/J)+m-c_{n,1}-c_{n,2})\Big{\\}}.$
(44)

The proof is completed by combing the posterior distribution (39) with
probabilities (40) and (44).

##  Appendix G Proof of Equation (30)

Because of the independence assumption of $\mathcal{H}$, and by an application
of Bayes theorem, we write

$\displaystyle\text{Pr}[f_{X_{m+1}}=l_{1},\,f_{X_{m+1}}=l_{2}\,|\,\\{(C_{n,h_{n}(X_{m+1})},C_{n,h_{n}(X_{m+2})})\\}_{n\in[N]}=\\{(c_{n,1},c_{n,2})\\}_{n\in[N]}]$
$\displaystyle\quad=\frac{1}{\text{Pr}[\\{(C_{n,h_{n}(X_{m+1})},C_{n,h_{n}(X_{m+2})})\\}_{n\in[N]}=\\{(c_{n,1},c_{n,2})\\}]}\text{Pr}[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2}]$
$\displaystyle\quad\quad\times\prod_{n=1}^{N}\text{Pr}[(C_{n,h_{n}(X_{m+1})},C_{n,h_{n}(X_{m+2})})=(c_{n,1},c_{n,2})\,|\,f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2}]$
$\displaystyle\quad=\frac{1}{\text{Pr}[\\{(C_{n,h_{n}(X_{m+1})},C_{n,h_{n}(X_{m+2})})\\}_{n\in[N]}=\\{(c_{n,1},c_{n,2})\\}]}\text{Pr}[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2}]$
$\displaystyle\quad\quad\times\prod_{n=1}^{N}\frac{\text{Pr}[(C_{n,h_{n}(X_{m+1})},C_{n,h_{n}(X_{m+2})})=(c_{n,1},c_{n,2}),\,f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2}]}{\text{Pr}[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2}]}$
$\displaystyle\quad=\frac{1}{\text{Pr}[\\{(C_{n,h_{n}(X_{m+1})},C_{n,h_{n}(X_{m+2})})\\}_{n\in[N]}=\\{(c_{n,1},c_{n,2})\\}]}(\text{Pr}[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2}])^{1-N}$
$\displaystyle\quad\quad\times\prod_{n=1}^{N}\text{Pr}[(C_{n,h_{n}(X_{m+1})},C_{n,h_{n}(X_{m+2})})=(c_{n,1},c_{n,2})]$
$\displaystyle\quad\quad\quad\times\text{Pr}[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2}\,|\,(C_{n,h_{n}(X_{m+1})},C_{n,h_{n}(X_{m+2})})=(c_{n,1},c_{n,2})]$
$\displaystyle\quad=(\text{Pr}[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2}])^{1-N}$
$\displaystyle\quad\quad\times\prod_{n=1}^{N}\text{Pr}[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2}\,|\,(C_{n,h_{n}(X_{m+1})},C_{n,h_{n}(X_{m+2})})=(c_{n,1},c_{n,2})]$
$\displaystyle\quad\propto\prod_{n=1}^{N}\text{Pr}[f_{X_{m+1}}=l_{1},f_{X_{m+2}}=l_{2}\,|\,(C_{n,h_{n}(X_{m+1})},C_{n,h_{n}(X_{m+2})})=(c_{n,1},c_{n,2})],$

where the $n$-th term of the last expression is precisely the probability in
Theorem 3. The proof is completed.

##  Appendix H Additional experiments

We present additional experiments on the application of the CMS-PYP on
synthetic and real data. First, we recall the synthetic and real data to which
the CMS-PYP is applied. As regards synthetic data, we consider datasets of
$m=500000$ tokens from a Zipf’s distributions with parameter
$\nu=1.3,\,1.6,\,1.9,\,2,2,\,2.5$. As regards real data, we consider: i) the
20 Newsgroups dataset, which consists of $m=2765300$ tokens with $K_{m}=53975$
distinct tokens; ii) the Enron dataset, which consists of $m=6412175$ tokens
with $K_{m}=28102$ distinct tokens. Tables 6, 7, 8 and 9 report the MAE (mean
absolute error) between true frequencies and their corresponding estimates
via: i) the CMS-PYP estimate $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$; ii) the CMS
estimate $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$; iii) the CMS-DP
estimate$\hat{f}_{v}^{\text{\tiny{(DP)}}}$, the CMM estimate
$\hat{f}_{v}^{\text{\tiny{(CMM)}}}$.

| $\mathcal{Z}_{1.3}$ | $\mathcal{Z}_{1.6}$ | $\mathcal{Z}_{1.9}$ | $\mathcal{Z}_{2.2}$ | $\mathcal{Z}_{2.5}$
---|---|---|---|---|---
Bins $v$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$
(0,1] | 1,061.3 | 161.72 | 0.86 | 629.40 | 62.19 | 3.10 | 308.11 | 81.10 | 92.35 | 51.65 | 1.04 | 3.40 | 32.65 | 1.02 | 30.69
(1,2] | 1,197.9 | 169.74 | 1.71 | 514.31 | 102.42 | 1.85 | 154.20 | 2.00 | 29.62 | 289.50 | 2.04 | 88.25 | 48.15 | 2.01 | 45.35
(2,4] | 1,108.3 | 116.37 | 3.22 | 474.82 | 52.10 | 1.52 | 2,419.51 | 2,215.85 | 1,155.12 | 134.05 | 3.40 | 16.47 | 54.34 | 10.50 | 50.85
(4,8] | 1,275.9 | 378.04 | 6.14 | 786.73 | 214.46 | 7.34 | 460.13 | 258.90 | 115.39 | 118.40 | 6.44 | 8.15 | 69.85 | 6.03 | 66.67
(8,16] | 1,236.1 | 230.32 | 11.68 | 719.84 | 232.24 | 8.73 | 380.05 | 139.50 | 92.64 | 413.13 | 129.03 | 129.23 | 80.80 | 13.10 | 77.75
(16,32] | 1,256.8 | 221.98 | 24.13 | 831.70 | 79.73 | 14.66 | 288.59 | 23.90 | 59.95 | 503.60 | 364.30 | 160.75 | 9.86 | 9.86 | 7.55
(32,64] | 1,312.8 | 235.87 | 42.80 | 783.90 | 184.99 | 34.63 | 415.58 | 54.82 | 155.87 | 217.81 | 82.92 | 94.49 | 10.22 | 10.22 | 6.79
(64,128] | 1,721.7 | 766.29 | 94.68 | 950.31 | 304.36 | 78.05 | 1,875.50 | 1,762.20 | 1,557.88 | 64.01 | 64.01 | 81.62 | 13.75 | 13.75 | 11.17
(128,256] | 1,107.7 | 334.57 | 178.55 | 1,727.19 | 1,488.38 | 347.06 | 202.09 | 163.61 | 99.75 | 46.80 | 46.80 | 134.87 | 17.51 | 17.51 | 14.37

Table 6: Synthetic data: MAE for $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$,
$\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ and $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$,
case $J=320,N=2$.

| $\mathcal{Z}_{1.3}$ | $\mathcal{Z}_{1.6}$ | $\mathcal{Z}_{1.9}$ | $\mathcal{Z}_{2.2}$ | $\mathcal{Z}_{2.5}$
---|---|---|---|---|---
Bins $v$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$
(0,1] | 2,212.1 | 590.48 | 0.60 | 1,262.0 | 146.11 | 1.00 | 424.80 | 130.90 | 0.96 | 154.70 | 47.10 | 29.34 | 56.70 | 1.01 | 15.07
(1,2] | 2,339.8 | 359.57 | 1.00 | 1,332.7 | 63.21 | 2.00 | 552.00 | 65.00 | 1.95 | 182.70 | 2.01 | 21.77 | 48.20 | 2.03 | 5.49
(2,4] | 2,270.9 | 69.42 | 1.70 | 1,277.8 | 301.89 | 3.60 | 487.30 | 163.55 | 3.36 | 184.70 | 97.15 | 16.62 | 57.80 | 14.35 | 6.80
(4,8] | 2,234.6 | 339.95 | 3.90 | 1,375.7 | 579.94 | 6.00 | 545.20 | 243.08 | 6.16 | 252.50 | 62.70 | 40.26 | 51.10 | 8.30 | 4.66
(8,16] | 2,213.3 | 313.37 | 7.70 | 1,165.7 | 152.53 | 11.90 | 493.20 | 196.20 | 11.66 | 247.30 | 29.70 | 102.23 | 24.10 | 14.11 | 6.04
(16,32] | 2,283.0 | 23.30 | 12.60 | 1,217.2 | 22.94 | 20.60 | 535.50 | 154.30 | 22.95 | 295.90 | 190.92 | 55.98 | 25.00 | 23.20 | 15.09
(32,64] | 2,305.7 | 133.09 | 34.10 | 1,284.6 | 209.13 | 44.80 | 637.80 | 150.05 | 43.65 | 120.60 | 71.86 | 32.92 | 31.70 | 31.70 | 24.33
(64,128] | 2,244.5 | 102.43 | 73.20 | 1,120.2 | 118.42 | 95.60 | 425.10 | 198.60 | 96.05 | 180.30 | 113.75 | 66.87 | 29.20 | 29.20 | 24.68
(128,256] | 2,237.4 | 294.43 | 130.60 | 1,141.3 | 573.12 | 175.10 | 525.90 | 267.15 | 186.95 | 129.70 | 129.70 | 93.83 | 32.10 | 32.10 | 38.13

Table 7: Synthetic data: MAE for $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$,
$\hat{f}_{v}^{\text{\tiny{(CMM)}}}$ and $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$,
case $J=160,N=4$.

| 20 Newsgroups | Enron
---|---|---
Bins $v$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(DP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(DP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$
(0,1] | 46.4 | 46.39 | 1.16 | 12.2 | 12.20 | 0.99
(1,2] | 16.6 | 16.60 | 1.96 | 13.8 | 13.80 | 1.99
(2,4] | 38.4 | 38.40 | 2.93 | 61.5 | 61.49 | 3.76
(4,8] | 59.4 | 59.39 | 6.10 | 88.4 | 88.39 | 7.50
(8,16] | 54.3 | 54.29 | 11.65 | 23.4 | 23.40 | 11.97
(16,32] | 17.8 | 17.80 | 21.14 | 55.1 | 55.09 | 20.78
(32,64] | 40.8 | 40.79 | 45.85 | 128.5 | 128.48 | 43.78
(64,128] | 26.0 | 25.99 | 88.47 | 131.1 | 131.08 | 81.38
(128,256] | 13.6 | 13.59 | 170.27 | 50.7 | 50.68 | 171.99

Table 8: Real data ($J=12000$ and $N=2$): MAE for
$\hat{f}_{v}^{\text{\tiny{(PYP)}}}$, $\hat{f}_{v}^{\text{\tiny{(DP)}}}$ and
$\hat{f}_{v}^{\text{\tiny{(CMS)}}}$.

| 20 Newsgroups | Enron
---|---|---
Bins $v$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(DP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(CMS)}}}$ | $\hat{f}_{v}^{\text{\tiny{(DP)}}}$ | $\hat{f}_{v}^{\text{\tiny{(PYP)}}}$
(0,1] | 53.4 | 53.39 | 1.00 | 71.0 | 70.98 | 1.00
(1,2] | 30.5 | 30.49 | 2.00 | 47.4 | 47.38 | 2.00
(2,4] | 32.5 | 32.49 | 3.30 | 52.5 | 52.49 | 3.90
(4,8] | 38.7 | 38.69 | 6.70 | 53.1 | 53.08 | 6.80
(8,16] | 25.3 | 25.29 | 12.60 | 57.0 | 56.98 | 11.90
(16,32] | 25.0 | 24.99 | 21.90 | 90.0 | 89.98 | 20.60
(32,64] | 39.7 | 39.69 | 43.70 | 108.4 | 108.37 | 47.80
(64,128] | 22.1 | 22.09 | 92.19 | 55.7 | 55.67 | 88.10
(128,256] | 25.8 | 25.79 | 207.58 | 80.8 | 80.76 | 179.30

Table 9: Real data ($J=8000$ and $N=4$): MAE for
$\hat{f}_{v}^{\text{\tiny{(PYP)}}}$, $\hat{f}_{v}^{\text{\tiny{(DP)}}}$ and
$\hat{f}_{v}^{\text{\tiny{(CMS)}}}$.


# Now It Sounds Like You: Learning Personalized Vocabulary On Device

Sid Wang
<EMAIL_ADDRESS>
&Ashish Shenoy
<EMAIL_ADDRESS>
Pierce Chuang
<EMAIL_ADDRESS>
&John Nguyen
<EMAIL_ADDRESS>

###### Abstract

In recent years, Federated Learning (FL) has shown significant advancements in
its ability to perform various natural language processing (NLP) tasks. This
work focuses on applying personalized FL for on-device language modeling. Due
to limitations of memory and latency, these models cannot support the
complexity of sub-word tokenization or beam search decoding, resulting in the
decision to deploy a closed-vocabulary language model. However, closed-
vocabulary models are unable to handle out-of-vocabulary (OOV) words belonging
to specific users. To address this issue, We propose a novel technique called
"OOV expansion" that improves OOV coverage and increases model accuracy while
minimizing the impact on memory and latency. This method introduces a
personalized "OOV adapter" that effectively transfers knowledge from a central
model and learns word embedding for personalized vocabulary. OOV expansion
significantly outperforms standard FL personalization methods on a set of
common FL benchmarks.

##  1 Introduction

Figure 1: The quantile plot of word frequency for top 10k words from 2
datasets that demonstrates the long-tail phenomenon of OOV.

Federated learning (FL) is a distributed machine learning paradigm that
enables model training on decentralized datasets without exchanging or sharing
sensitive data McMahan et al. (2016). Personalized FL considers models unique
to each client under heterogeneous data settings Hanzely and Richtárik (2020).
During personalization, the server sends a global model to the clients for
local fine-tuning then the client uses that model for inference. Our work
studies a class of on-device language models for next-word prediction task
trained with personalized FL.

The resource-constrained edge devices Chen et al. (2019a); Qiu et al. (2022);
Mathur et al. (2021) limits the usage of subword-level tokenizers. On the one
hand, subword tokenizers with large vocabulary require a large memory
footprint, making deployment infeasible. On the other hand, a smaller
vocabulary leads to longer tokenized sequences and increases generation
latency. In order to satisfy these constraints, a _closed vocabulary_ Qin and
Rudnicky (2013), a word-level vocabulary with a white-space tokenizer that
treats all unknown words as a special token, must be used.

Consequently, the model cannot handle out-of-vocabulary (OOV) words Chen et
al. (2019b), making it more challenging to understand the communication style
of an individual user. The heavy tail of OOV (shown in Figure 1) motivates the
need for personalized vocabulary. In this work, we propose _OOV Expansion_ to
personalize on-device vocabulary, tailoring the model toward users’ unique
wording habits and spelling patterns. _OOV Expansion_ uses an adapter – an MLP
with residual connections inserted to different submodules to help the model
adapt to new knowledge domain Houlsby et al. (2019) – to compute the embedding
output of the OOV words.

Our contributions We highlight the main contributions:

$\bullet$ We propose _OOV Expansion_ , a novel personalized FL method, to
extract useful features for user-specific vocabulary and address the long tail
OOV issue.

$\bullet$ We perform comprehensive experiments, showing superior results over
that of the baselines, across a set of standard FL datasets.

$\bullet$ We demonstrate that our technique significantly outperforms previous
methods with up to $5.6\%$ relative improvements on next-word prediction
accuracy and reducing the averaged unknown-word-rate by at least $97\%$.

##  2 Method

This section describes the model architecture and outlines the training
techniques for the on-device next-word prediction task.

Figure 2: The mechanism flowchart of the OOV Expansion stage; here $\|$
denotes vector concatenation.

###  2.1 Character-aware language model

The limitations on the memory and latency of on-device modeling confines the
model architecture and rules out some common choices such as the transformers
Vaswani et al. (2017). Instead, we use the character-aware LSTM based language
model, a commonly adopted architecture for on-device applications Kim et al.
(2015); Hard et al. (2018).

The model consists of 3 major components: (1) a character level CNN that
computes word embeddings, which will be denoted as CharCNN from now on, (2) an
LSTM encoder that provides the input representations, (3) an MLP decoder that
gives the scores over vocabulary. See the non-yellow stack in Figure 2.

###### Remark 2.1

The term “score” is an important and recurring notion throughout the section.
The scores of multiple sources of vocabulary may be joined and passed through
a LogSoftmax transform to produce logits over the vocabulary union, which are
eventually used to compute the cross entropy loss for next-word prediction.

The model hyperparameters are defined as follows: the CharCNN contains a utf-8
embedding of shape $\mathbb{R}^{256\times E}$ and a CNN with $D$ channels and
kernel size $K$. The LSTM encoder contains $M$ layers with both input and
output dimensions equal to $D$. The decoder is a linear layer that has input
dimension $D$ and output dimension $|\mathcal{V}|$, with $\mathcal{V}$ being
the closed vocabulary.

###  2.2 Baseline 1: OOV-as-UNK

This is the traditional FL personalization (i.e. involving only the non-yellow
stack in Figure 2). To be precise, the model is first pretrained on a general
dataset using the standard “next-word prediction” task, and then sent to
client devices for federated learning. Lastly, perform personalization without
any additional treatments for OOVs rather than replacing them with the special
token [UNK].

###  2.3 Baseline 2: OOV-oracle

The second baseline assumes to have the knowledge of all users’ OOV on server.
However, because of the tight memory budgets in practice we cannot afford full
vocabulary training, we instead expand the OOV-as-UNK vocabulary by the $N$
most frequent OOV words. Upon expansion, the rest stages (including
pretraining, FL, and personalization) proceed in exactly the same way as in
Subsection 2.2.

###  2.4 Our method: OOV Expansion

We start with the unpersonalized OOV-as-UNK (ref. Subsection 2.2) and perform
personalization as demonstrated in Figure 2. To begin, given an input sentence
the model forwards in exactly the same fashion as in subsection 2.1 to get the
vocabulary scores (ref. Remark 2.1). Next, retrieve the client’s top $n$ OOV
words and compute their representations from CharCNN, and further pass them
through an adapter (i.e. a residual MLP that is randomly initialized at the
beginning of personalization on each device) with hidden dimensions
$\vec{\textbf{H}}=(H_{1},H_{2},\cdots,H_{L})$ where $L$ is the number of
layers in the adapter. This intermediate result could be interpreted as
“adapted OOV word embeddings”. Then take the inner product between the adapter
outputs and the LSTM encoder outputs to get OOV scores (ref. Remark 2.1).
Finally, concatenating the OOV scores with the vocabulary scores, we obtain
the full scores over the client personalized vocabulary for next word
prediction.

To understand the functionality of adapter, note that the model never learns
to compute OOV representations during pretraining or federated learning stage
because unknown words are ignored and replaced by a special token. During
personalization, the CharCNN module mutlitasks on (1) providing inputs for
LSTM encoder and (2) computing OOV embeddings. Morally speakng, OOV adapter
alleviates the interference between the two tasks as well as adapting the
prior knowledge of CharCNN to the new “OOV task”.

Lastly, our design guarantees that no sensitive OOV information ever leaves
the client device and hence mitigates the privacy leakage risk while
preserving the model’s ability to learn OOV features.

Figure 3: $\text{EMR}_{3}$ of OOV Expansion and two baselines on 2 datasets.

##  3 Experiments

###  3.1 Model setup

The model hyperparameters in Subsection 2.1 are chosen as $E=100$, $D=200$,
$K=4$, $M=2$.

To initialize OOV-as-UNK (ref. Subsection 2.2) we extract the 5k most frequent
words from a centralized FL dataset (see subsection 3.2) to form a closed
vocabulary, and pretrain the model on wikitext-103 Merity et al. (2016) using
next-word prediction task for 100 epochs with learning rate as $10^{-4}$,
batch size equals to 32, across 8 GPU devices, and evaluated by the exact
match metric $\text{EMR}_{3}$ (see Subsection 3.3 for the precise definition).
The OOV-oracle baseline (ref. Subsection 2.3) follows the same configuration
as OOV-as-UNK, except for having a 10k-sized vocabulary. In other words, OOV-
oracle expands the vocabulary of OOV-as-UNK by $N=5000$ most frequent OOVs.
The OOV-as-UNK (resp. OOV-Oracle) model has 1.76M (resp. 2.76M) parameters in
total, and took up to 36 hours on 8 Nvidia V100 16GB GPUs to complete
pretraining.

During OOV expansion (see Subsection 2.4), we personalize up to 1000 most
frequent out-of-vocabulary words from each decentralized dataset.

Figure 4: $\text{EMR}_{3}$ of two baselines before and after personalizations
on each dataset.

###  3.2 Datasets

We train, validate, and test our approach on 2 publicly available benchmark
datasets for federated learning: pushift.io’s Reddit Caldas et al. (2018),
Stack Overflow Authors. (2019). Pushshift.io Reddit is a previously existing
dataset extracted and obtained by a third party that contains preprocessed
comments posted on the social network Reddit and hosted by pushift.io. Stack
Overflow consists of questions and answers from Stack Overflow. We experiment
the next-word-prediction task with federated learning, using the natural non-
IID partitionings of all datasets.

Each dataset admits a training/validation/test split on client IDs. The
training clients are used for federated learning, validation clients for
early-stopping/hyperparameter search, and test clients for metric report and
personalization.

Each test client’s dataset is further divided into training, validation, and
test segments with ratio $8:1:1$, where we (1) personalize the global model
locally on the training segment, (2) use the validation segment for early-
stopping and hyperparameter search, and (3) report model performance on the
test segments.

###  3.3 Evaluation metric

The model quality is measured by how often the actual next word appears among
the top $K$ word predictions, denoted by $\text{EMR}_{K}$ (Exact Match Rate).
More precisely, given a dataset $\mathcal{D}$ of sentences,

$\text{EMR}_{K}(\mathcal{D})=\frac{\sum_{S\in\mathcal{D}}\sum_{w\in
S}\mathbbm{1}_{\\{w\in\text{Top}_{K}\text{pred}(w)\\}}}{\sum_{S\in\mathcal{D}}|S|}.$
(1)

This is a natural metric for real-time language model in production, where $K$
next-word suggestions are surfaced when a user is typing. In our experiments
we set $K=3$.

The final metric is given below:

$\displaystyle\begin{split}\text{EMR}_{3}&=\frac{\sum_{u\in\mathcal{U}}\sum_{S\in\mathcal{D}_{u}}\sum_{w\in
S}\mathbbm{1}_{\\{w\in\text{Top}_{3}\text{pred}(w)\\}}}{\sum_{u\in\mathcal{U}}\sum_{S\in\mathcal{D}_{u}}|S|}\\\
&=\frac{\sum_{u\in\mathcal{U}}\sum_{S\in\mathcal{D}_{u}}|S|\cdot\text{EMR}_{3}(\mathcal{D}_{u})}{\sum_{u\in\mathcal{U}}\sum_{S\in\mathcal{D}_{u}}|S|}\end{split}$

where $\mathcal{U}$ is the set of all test clients and $\mathcal{D}_{u}$
stands for the test segment (see Subsection 3.2) of client $u$’s dataset.
Equivalently, the model quality is measured by $\text{EMR}_{3}$ on the
centralized test segments of all test clients.

Figure 5: $\text{EMR}_{3}$ and $\text{KEMR}_{3}$ of OOV expansion on 2
datasets, with and without adapter

###  3.4 Experimental details

The training recipes reported below are the same for all datasets unless
stated otherwise. At each FL training round, 96 clients are randomly selected
to participate in local training (evenly distributed to 8 GPU devices), with
local epochs set to be 1 and training batch size as 8. Each global training
epoch consists of $\\#\text{users}/96$ training rounds, where the total number
of global training epochs are chosen to be 6 and 3 for Pushshift.io Reddit and
Stackoverflow respectively.

We use SGD without momentum as client optimizer, and FedAdam for server side
optimizer with weight decay of $10^{-5}$, $\beta_{1}$ as 0.9, $\beta_{2}$ as
0.999, and $\epsilon$ as $10^{-8}$. We do hyperparameter search on both client
and server learning rates, with ranges $[10^{-6},0.05]$ and $[10^{-5},1]$
respectively. In Table 2 we specify the best choices for both baselines on
each dataset.

Every hyperparameter search contains 64 experiments (with 8 running in
parallel at a time), where each single experiment takes 8 Nvidia V100 16GB GPU
up to 10 hours and 12 hours for Pushshift.io Reddit and Stackoverflow
respectively.

During personalizations, we perform hyperparameter search at each client among
the following grid (if applied):

$\displaystyle\begin{split}\text{lr}&\in\\{10^{-4},10^{-3},10^{-2},0.1,1,10\\},\\\
\sigma&\in\\{0,10^{-5},10^{-4},10^{-3},10^{-2},0.1,1\\},\\\
\vec{\textbf{H}}&\in\\{(960,),(128,256,128),(256,512,256)\\}.\end{split}$

Here lr denotes the learning rate for local training; $\sigma$ stands for the
standard deviation of Gaussian distribution used to initialize the adapter
parameters (except for LayerNorm); $\vec{\textbf{H}}$ represents the hidden
dimensions of adapter. Note the grid only include $\sigma$ and
$\vec{\textbf{H}}$ during OOV expansion. For each set of hyperparameters, we
run 10 local epochs that early stops when the validation $\text{EMR}_{3}$ did
not improve.

| OOV-as-UNK | OOV-Oracle | OOV Expansion
---|---|---|---
OOV rate on Pushshift.io Reddit | 8.8% | 5.5% | 0.2%
OOV rate on Stackoverflow | 4.8% | 3.0% | 0.001%
# Model parameters | 1.76M | 2.76M | 2.12M
Table 1: Comparison of OOV rates and number of model parameters among 3
methods

###  3.5 Results

Improves overall performance As shown in Figure 3, our method achieves $2.1\%$
and $1.0\%$ (resp. $5.6\%$ and $2.5\%$) absolute (resp. relative) gains of
$\text{EMR}_{3}$ compared to the standard approach (i.e. OOV-as-UNK) on
Pushshift.io Reddit and Stackoverflow respectively. It shows competitive
performance even when compared with the stronger baseline OOV-Oracle, seeing
$3.7\%$ and $1.7\%$ relative $\text{EMR}_{3}$ gain on Pushshift.io Reddit and
Stackoverflow respectively. The fact that OOV expansion yields better
improvements on Pushshift.io Reddit relative to Stackoverflow partially
attributes to the higher heterogeneity of Pushshift.io Reddit data, which is
reflected by the larger OOV rates (see Table 1) and the worse long tail
behavior (ref. Figure 1).

Increases word-coverage rates As Table 1 shows, the OOV expansion approach
reduces the OOV rate of OOV-as-UNK (resp. OOV-oracle) by more than $97.7\%$
and $99.9\%$ (resp. $96.4\%$ and $99.9\%$) on Pushshift.io Reddit and
Stackoverflow respectively.

More parameter efficient As Table 1 indicates, OOV expansion uses $24\%$ less
parameters than the OOV-Oracle during personalization. In addition, our
approach does not require any extra training parameters during pretraining or
FL. Consequently, it is $36\%$ more parameter-efficient in FL compared to OOV-
Oracle.

Personalization is essential Taking Pushshift.io Reddit for instance, the
standard personalization can already improve the global model accuracy by
$6.5\%$ relatively on $\text{EMR}_{3}$ (see Figure 4). With OOV expansion we
can further boost this quality gain up to $12.5\%$.

Adapter is necessary As can be seen from the first plot in Figure 5, adapter
yields $15\%$ relative $\text{EMR}_{3}$ gain over trivial adapter (i.e.
identity block) on Pushshift.io Reddit. For Stackoverflow dataset, inserting
adapter achieves 2.6 times accuracy of the trivial adapter approach. In the
second plot we present quality comparison under a new metric denoted by
$\text{KEMR}_{K}$ which represents top $K$ known word exact match rates, and
(similar to formula (1)) it is defined by

| OOV-as-UNK | OOV-oracle
---|---|---
client lr | server lr | client lr | server lr
Pushshift.io Reddit | 0.840 | 0.003 | 0.258 | 0.004
Stackoverflow | 0.168 | 0.005 | 0.129 | 0.008
Table 2: Best learning rates for two baselines on 2 datasets

$\displaystyle\begin{split}\text{KEMR}_{K}(\mathcal{D})=\frac{\sum_{S\in\mathcal{D}}\sum_{w\in
S\cap\mathcal{V}}\mathbbm{1}_{\\{w\in\text{Top}_{K}\text{pred}(w)\\}}}{\sum_{S\in\mathcal{D}}|S\cap\mathcal{V}|}.\end{split}$

where $\mathcal{V}$ is the closed vocabulary.

From Figure 5, we see that adapter not only helps achieves better OOV-
understanding (i.e. higher EMR), the fact that it consistently improves
accuracy on known word implies the important role of adapter in suppressing
the issue of forgetting the knowledge from pretraining and federated learning
stage.

##  4 Related work

Federated learning McMahan et al. (2016) has achieved significant impact in
the field of natural language processing (NLP) in recent years Chen et al.
(2019b, c); Wu et al. (2020); Lin et al. (2021); Hilmkil et al. (2021); Liu et
al. (2021); Ro et al. (2022). Unlike non-federated language models where
subword-based tokenizations such as WordPiece Schuster and Nakajima (2012),
Byte Pair Encoding (BPE) Sennrich et al. (2015) or SentencePiece Kudo and
Richardson (2018) have become major choices, word-level tokenizer with a
closed vocabulary is the default choice for edge device settings due to its
low capacity and real-time nature. This gives rise to the commonly known OOV
issue Chen et al. (2019a, c), which has been investigated by several preceding
work in federated learning: Chen et al. (2019b) trains a separate character-
level generative model to sample new words from, but by design requires an
extra training stage; Singhal et al. (2021) introduces a new but more complex
FL algorithm named FedRecon based on partial personalization; Bagdasaryan et
al. (2022) proposes to iteratively update a subword-level tokenizer using the
token sequences sampled from an FL-trained language model. Whereas Khandelwal
et al. (2019) and Bekal et al. (2021) use memorization to dynamically update
the vocabulary without retraining. However, these approaches either focuses on
a different problem, or is not suitable for our settings due to their
requirements of high training budgets, system complexity, and latency/memory
costs. Instead, our approach is based on FL personalization, a technique that
has driven a large body of work to address the presumed presence of
heterogeneity Wang et al. (2019); Hanzely and Richtárik (2020); Yu et al.
(2020); Fallah et al. (2020); Dinh et al. (2020); Li et al. (2020); Singhal et
al. (2021); Pillutla et al. (2022); Kulkarni et al. (2020). Another key
ingredient of our work is adapter Houlsby et al. (2019); Stickland and Murray
(2019), which has been widely studied in a large volume of work in the non-
federated setting Mahabadi et al. (2021); Hu et al. (2021); Li and Liang
(2021); He et al. (2021); Pfeiffer et al. (2020).

## Conclusion

This paper proposes a personalized federated learning method that enables out-
of-vocabulary words understanding for a class of on-device language models
with closed vocabulary. By evaluating on two public benchmarks, we show that
our method significantly outperforms the commonly used personalization
approach in terms of next-word-prediction accuracy and drastically reduce the
unknown-word rate on average while preserving user privacy and being parameter
efficient.

## Limitations

One limitation of the results herein is the lack of study using subword-based
language models. This was due to the tight memory, compute and latency budget
of our experiments which made sub-word tokenization or beam search decoding
infeasible. Another aspect that could be further explored is the trade off
between accuracy and privacy by adding differential privacy to FL training.
Last but not least, our method falls shy at cold-start problem (i.e. when a
user’s on-device historical data is insufficient or unavailable) because we
assume to know all possible OOVs on-device.

## Ethics Statement

The proposed method in this work has the potential to improve the performance
of language models for individual users by personalizing the vocabulary to
their specific needs. This can have a positive impact in various applications
including next word prediction, sentence completion and automatic speech
recognition. Additionally, when used with federated learning, our approach can
also help to protect the privacy of users by keeping personal data on
individual devices and only sharing model updates with a centralized server.
However, as is typical of any generative model, we also recognize that there
are potential downsides to our approach. For example, if not implemented
correctly, our approach could perpetuate existing biases in the data,
specifically vocabulary and text usage patterns, and create unfair or biased
models. Therefore, it is important for practitioners to be aware of these
issues and take steps to mitigate them when implementing our method.

## References

  * Authors. (2019) The TensorFlow Federated Authors. 2019.  Tensorflow federated stack overflow dataset.
  * Bagdasaryan et al. (2022) Eugene Bagdasaryan, Congzheng Song, Rogier C. van Dalen, Matthew Stephen Seigel, and ’Aine Cahill. 2022.  Training a tokenizer for free with private federated learning.  _ArXiv_ , abs/2203.09943.
  * Bekal et al. (2021) Dhanush Bekal, Ashish Shenoy, Monica Sunkara, Sravan Bodapati, and Katrin Kirchhoff. 2021.  Remember the context! asr slot error correction through memorization.  In _2021 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU)_ , pages 236–243.
  * Caldas et al. (2018) Sebastian Caldas, Peter Wu, Tian Li, Jakub Konecný, H. B. McMahan, Virginia Smith, and Ameet S. Talwalkar. 2018.  Leaf: A benchmark for federated settings.  _ArXiv_ , abs/1812.01097.
  * Chen et al. (2019a) Mia Xu Chen, Benjamin Lee, Gagan Bansal, Yuan Cao, Shuyuan Zhang, Justin Lu, Jackie Tsay, Yinan Wang, Andrew M. Dai, Z. Chen, Timothy Sohn, and Yonghui Wu. 2019a.  Gmail smart compose: Real-time assisted writing.  _Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_.
  * Chen et al. (2019b) Mingqing Chen, Rajiv Mathews, Tom Y. Ouyang, and Françoise Beaufays. 2019b.  Federated learning of out-of-vocabulary words.  _ArXiv_ , abs/1903.10635.
  * Chen et al. (2019c) Mingqing Chen, Ananda Theertha Suresh, Rajiv Mathews, Adeline Wong, Cyril Allauzen, Françoise Beaufays, and Michael Riley. 2019c.  Federated learning of n-gram language models.  In _Proceedings of the 23rd Conference on Computational Natural Language Learning (CoNLL)_ , pages 121–130, Hong Kong, China. Association for Computational Linguistics.
  * Dinh et al. (2020) Canh T. Dinh, Nguyen H. Tran, and Tuan Dung Nguyen. 2020.  Personalized federated learning with moreau envelopes.  _ArXiv_ , abs/2006.08848.
  * Fallah et al. (2020) Alireza Fallah, Aryan Mokhtari, and Asuman E. Ozdaglar. 2020.  Personalized federated learning: A meta-learning approach.  _ArXiv_ , abs/2002.07948.
  * Hanzely and Richtárik (2020) Filip Hanzely and Peter Richtárik. 2020.  Federated learning of a mixture of global and local models.  _ArXiv_ , abs/2002.05516.
  * Hard et al. (2018) Andrew Hard, Kanishka Rao, Rajiv Mathews, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. 2018.  Federated learning for mobile keyboard prediction.  _ArXiv_ , abs/1811.03604.
  * He et al. (2021) Junxian He, Chunting Zhou, Xuezhe Ma, Taylor Berg-Kirkpatrick, and Graham Neubig. 2021.  Towards a unified view of parameter-efficient transfer learning.  _ArXiv_ , abs/2110.04366.
  * Hilmkil et al. (2021) Agrin Hilmkil, Sebastian Callh, Matteo Barbieri, Leon René Sütfeld, Edvin Listo Zec, and Olof Mogren. 2021.  Scaling federated learning for fine-tuning of large language models.  In _International Conference on Applications of Natural Language to Data Bases_.
  * Houlsby et al. (2019) Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin de Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. 2019.  Parameter-efficient transfer learning for nlp.  In _International Conference on Machine Learning_.
  * Hu et al. (2021) Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, and Weizhu Chen. 2021.  Lora: Low-rank adaptation of large language models.  _ArXiv_ , abs/2106.09685.
  * Khandelwal et al. (2019) Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2019\.  Generalization through memorization: Nearest neighbor language models.  _CoRR_ , abs/1911.00172.
  * Kim et al. (2015) Yoon Kim, Yacine Jernite, David A. Sontag, and Alexander M. Rush. 2015.  Character-aware neural language models.  In _AAAI Conference on Artificial Intelligence_.
  * 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 _Conference on Empirical Methods in Natural Language Processing_.
  * Kulkarni et al. (2020) Viraj Kulkarni, Milind V. Kulkarni, and Aniruddha Pant. 2020.  Survey of personalization techniques for federated learning.  _2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainability (WorldS4)_ , pages 794–797.
  * Li et al. (2020) Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. 2020.  Ditto: Fair and robust federated learning through personalization.  In _International Conference on Machine Learning_.
  * Li and Liang (2021) Xiang Lisa Li and Percy Liang. 2021.  Prefix-tuning: Optimizing continuous prompts for generation.  _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , abs/2101.00190.
  * Lin et al. (2021) Bill Yuchen Lin, Chaoyang He, ZiHang Zeng, Hulin Wang, Yufen Huang, Mahdi Soltanolkotabi, Xiang Ren, and Salman Avestimehr. 2021.  Fednlp: Benchmarking federated learning methods for natural language processing tasks.  In _NAACL-HLT_.
  * Liu et al. (2021) Ming Liu, Stella Ho, Mengqi Wang, Longxiang Gao, Yuan Jin, and Heng Zhang. 2021\.  Federated learning meets natural language processing: A survey.  _ArXiv_ , abs/2107.12603.
  * Mahabadi et al. (2021) Rabeeh Karimi Mahabadi, James Henderson, and Sebastian Ruder. 2021.  Compacter: Efficient low-rank hypercomplex adapter layers.  In _NeurIPS_.
  * Mathur et al. (2021) Akhil Mathur, Daniel J. Beutel, Pedro Porto Buarque de Gusmão, Javier Fernández-Marqués, Taner Topal, Xinchi Qiu, Titouan Parcollet, Yan Gao, and Nicholas D. Lane. 2021.  On-device federated learning with flower.  _ArXiv_ , abs/2104.03042.
  * McMahan et al. (2016) H. B. McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. 2016.  Communication-efficient learning of deep networks from decentralized data.  In _International Conference on Artificial Intelligence and Statistics_.
  * Merity et al. (2016) Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016.  Pointer sentinel mixture models.  _ArXiv_ , abs/1609.07843.
  * Pfeiffer et al. (2020) Jonas Pfeiffer, Aishwarya Kamath, Andreas Rücklé, Kyunghyun Cho, and Iryna Gurevych. 2020.  Adapterfusion: Non-destructive task composition for transfer learning.  _ArXiv_ , abs/2005.00247.
  * Pillutla et al. (2022) Krishna Pillutla, Kshitiz Malik, Abdelrahman Mohamed, Michael G. Rabbat, Maziar Sanjabi, and Lin Xiao. 2022.  Federated learning with partial model personalization.  _ArXiv_ , abs/2204.03809.
  * Qin and Rudnicky (2013) Longlu Qin and Alexander I. Rudnicky. 2013.  Finding recurrent out-of-vocabulary words.  In _INTERSPEECH_.
  * Qiu et al. (2022) Xinchi Qiu, Javier Fernández-Marqués, Pedro Gusmão, Yan Gao, Titouan Parcollet, and Nicholas Donald Lane. 2022.  Zerofl: Efficient on-device training for federated learning with local sparsity.  _ArXiv_ , abs/2208.02507.
  * Ro et al. (2022) Jae Hun Ro, Theresa Breiner, Lara McConnaughey, Mingqing Chen, Ananda Theertha Suresh, Shankar Kumar, and Rajiv Mathews. 2022.  Scaling language model size in cross-device federated learning.  _ArXiv_ , abs/2204.09715.
  * Schuster and Nakajima (2012) Mike Schuster and Kaisuke Nakajima. 2012.  Japanese and korean voice search.  _2012 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 5149–5152.
  * Sennrich et al. (2015) Rico Sennrich, Barry Haddow, and Alexandra Birch. 2015.  Neural machine translation of rare words with subword units.  _ArXiv_ , abs/1508.07909.
  * Singhal et al. (2021) K. Singhal, Hakim Sidahmed, Zachary Garrett, Shanshan Wu, Keith Rush, and Sushant Prakash. 2021.  Federated reconstruction: Partially local federated learning.  In _Neural Information Processing Systems_.
  * Stickland and Murray (2019) Asa Cooper Stickland and Iain Murray. 2019.  Bert and pals: Projected attention layers for efficient adaptation in multi-task learning.  In _International Conference on Machine Learning_.
  * Vaswani et al. (2017) Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  _ArXiv_ , abs/1706.03762.
  * Wang et al. (2019) Kangkang Wang, Rajiv Mathews, Chloé Kiddon, Hubert Eichner, Franccoise Beaufays, and Daniel Ramage. 2019.  Federated evaluation of on-device personalization.  _ArXiv_ , abs/1910.10252.
  * Wu et al. (2020) Xing Wu, Zhaowang Liang, and Jianjia Wang. 2020.  Fedmed: A federated learning framework for language modeling.  _Sensors (Basel, Switzerland)_ , 20.
  * Yu et al. (2020) Tao Yu, Eugene Bagdasaryan, and Vitaly Shmatikov. 2020.  Salvaging federated learning by local adaptation.  _ArXiv_ , abs/2002.04758.