Integrating Artificial Intelligence and Augmented Reality in Robotic Surgery: An Initial dVRK Study Using a Surgical Education Scenario

Abstract: Robot-assisted surgery has become progressively more and more popular due to its clinical advantages. In the meanwhile, the artificial intelligence and augmented reality in robotic surgery are developing rapidly and receive lots of attention. However, current methods have not discussed the coherent integration of AI and AR in robotic surgery. In this paper, we develop a novel system by seamlessly merging artificial intelligence module and augmented reality visualization to automatically generate the surgical guidance for robotic surgery education. Specifically, we first leverage reinforcement leaning to learn from expert demonstration and then generate 3D guidance trajectory, providing prior context information of the surgical procedure. Along with other information such as text hint, the 3D trajectory is then overlaid in the stereo view of dVRK, where the user can perceive the 3D guidance and learn the procedure. The proposed system is evaluated through a preliminary experiment on surgical education task peg-transfer, which proves its feasibility and potential as the next generation of robot-assisted surgery education solution.

\section{INTRODUCTION}

Artificial intelligence (AI) and augmented reality (AR) are two important and increasingly essential techniques to be developed for next-generation robotic surgery. So far, AI and AR have individually focused on different perspectives. In specific, AI concentrates on recognizing and planning surgical activities in a way similar to what surgeons could do, based on computation and analysis of collected sensory data such as endoscopic videos~\cite{jin2021temporal, zhang2021swnet, 9336292} and robotic kinematics~\cite{van2020multi, su2018real, long2021relational}.
Recent advances of AI have substantially enhanced a number of tasks such as situation awareness of surgical procedures~\cite{garrow2021machine,huaulme2021micro} and automation of some actions with surgical robots~\cite{van2021gesture,huynhnguyen2021toward}.
In the meanwhile, AR aims to augment the surgical environment in a way to facilitate surgeon's operation and decision-making, based on visualization and integration of additional information that is computed offline or in real-time~\cite{volonte2013augmented,qian2019review}.
Equipped with an immersive view in surgical robotic console, AR has shown effectiveness for the education of novice surgeons~\cite{sheik2019next,barsom2016systematic,mcknight2020virtual} and is envisaged to be very helpful if could be adopted intra-operatively \cite{londei2015intra}.

Unfortunately, to date, the advantages of AI and AR have not been merged in a sensible way for robotic surgery.
The intriguing combination of AI and AR emerges as a versatile topic and has been exemplified in a number of application scenarios, such as games \cite{turan2019using}, driver training \cite{gabbard2019ar,boboc2020application} and virtual patients \cite{daher2020physical,gonzalez2020neurological}. Conferring intelligence on AR not only can boost the virtual experience but also harnesses the strong power of learning-based algorithms in demanding tasks such as surgical education. However, there have been few attempts in combining AI and AR for surgical robots. Some authors \cite{tanzi2021real,pan2020real,sadeghi2021virtual} have proposed to use computer vision models to localize the anatomy regions of interest, and then superimpose the results in the view of camera. Nonetheless, these solutions only account for presenting existing clues perceptively, without shedding lights on human-like decision behaviors. Concurrently, reinforcement learning (RL) is widely recognized as an effective method for skill learning~\cite{amini2020learning,kurach2019google,chiu2021bimanual,richter2019open}, but its potential is not fully exploited in the surgical robotics domain. One interesting scenario for exploring these issues is surgical education, in which the RL-based smart agent is expected to reason about surgical task and generate constructive guidelines for the novice. Such embodied intelligence is promising to significantly increase accessibility and reduce the cost of surgical training. Realisation of the intelligent guidance, in the form of AR visualization on surgical robotics platforms, could further enhance usability and user experience, yet how to achieve this goal remains unclear.

\begin{figure}[t]
    \centering
    \includegraphics[width=1.0\hsize]{figures/intro.pdf}
    \caption{The integrated framework of AI and AR in robotic surgery on surgical education scenario of peg-transfer. The AI learns policy in a stimulated environment based on reinforcement learning and predict the guidance trajectory, which will be further visualized in the dVRK stereo viewer in the form of AR overlaid trajectory.}
    \label{TrainingFramework}
    \vspace{-0.50cm}
\end{figure}

In this paper, we aim to seamlessly integrate AI and AR, by augmenting RL-driven instrument movement trajectories with real-time AR visualizations, and implement the pipeline for a surgical education scenario using da Vinci Research Kit (dVRK), as shown in Fig.~\ref{TrainingFramework}.
For the AI-enabled analysis, reinforcement learning is employed to learn a policy from the expert demonstration and through the interaction with the environment. Then the system can reason about the future action based on the current observation.
This behavior can be embedded into an education process, where the automatically predicted actions could serve as helpful information to guide the trainee's movement step-by-step.
Subsequently, how to effectively visualize the information becomes very important, especially when incorporating with robotic platforms in real-time.
We will demonstrate the feasibility of overlaying the 3D guidance trajectory within the stereo video via the dVRK console.
By projecting the 3D location of the trajectory generated from RL policy to the stereo video frames in the dVRK console, we can vividly observe the surgical scene with overlaid trajectory for the education purpose.
We have implemented and evaluated our method on the typical surgical education task of peg-transfer.
To the best of our knowledge, this is the first work in exploring synergy of AI and AR with RL-based prediction and dVRK-based visualization for surgical education scenario. We hope to generate discussion and spark the potential of integrated benefits of AI and AR for surgical education and beyond.

\begin{figure*}[ht]
    \centering
    \includegraphics[width=0.79\hsize]{figures/Overview.pdf}
    \caption{The pipeline of our proposed robot-assisted surgery education system based on dVRK. The computing server a) can acquire the stereo video and kinematics b) from the dVRK PSM/ECM d) and input the command to control the movement of dVRK PSM. With the AR and AI algorithms, the computer can overlay the AI and AR information on the video c) and then import the augmented video to the dVRK console e) for user to view.}
    \label{overall_Framework}
    \vspace{-0.50cm}
\end{figure*}

\section{RELATED WORK}

\subsection{Artificial Intelligence in Robotic Surgery}
The application of artificial intelligence in robotic surgery has been actively studied over the last decade~\cite{d2021accelerating} since the da Vinci Surgical System being clinically introduced in 2000~\cite{kazanzides2014open}.
Several topics have been found important and studied widely such as surgical instrument segmentation~\cite{zhao2020learning}, gesture recognition~\cite{long2021relational}, workflow recognition~\cite{jin2021temporal} and surgical scene reconstruction~\cite{long2021dssr}. They could support intra-operative decision \cite{blum2010modeling} and provide valuable database for surgical training and evaluation \cite{padoy2012statistical,nara2009surgical,neumuth2017surgical}.
Although promising, these works just provide supplementary information without conceiving surgical plan, such as predicting the trajectories of the surgical instrument. Recently, the emergence of reinforcement learning opens the door to a new set of policy-based learning strategies \cite{rakhsha2020policy}. The efficacy of reinforcement learning has been revealed in surgical gesture classification \cite{liu2018deep,gao2020automatic}, surgical scene understanding \cite{torrents2020deep,chitsaz2009medical} and robot learning \cite{varier2020collaborative,chiu2021bimanual}. Through learning from expert demonstration, the RL agent could automatically generate meaningful solutions according to the task at hand. For example, in~\cite{chiu2021bimanual} and~\cite{richter2019open}, the authors propose to use Deep Deterministic Policy Gradients (DDPG) with Behavior Cloning (BC) to conduct surgical task of bimanual needles regrasping and autonomous blood suction. Both of them demonstrate encouraging results, showing that RL-based framework may potentially alleviate the demands on expert's guidance \cite{xie2020learning, loftus2020decision, nguyen2019new}.

Apart from the above, there have been some works using AI for surgical education, such as providing metrics and performance feedback based on the training records~\cite{sewell2008providing, nagyne2019robot, mirchi2020artificial} and differentiate expertise levels taking the stylistic characteristics into account \cite{ershad2018meaningful,liang2011surgical,ershad2019automatic}. But few of them consider its application to AR surgical education, where integrated AR and AI is highly demanded to guide the trainee vividly and automatically.
Given the superiority of RL introduced above, it is compelling to integrate RL into AR as an essential AI module, for example, a decision-maker \cite{wang2019lane, wierzbicki2000modern}, which encourages the AR system to generate content objectively.

\subsection{Augmented Reality for Robotic Surgery}
Augmented reality has been applied to robotic surgery in a variety of paradigms \cite{qian2019review}.
In intraoperative applications, augmentations are superimposed in real-time to offer assistance: (i) enhance depth perception \cite{wen2017augmented,penza2017envisors,cutolo2017new}, (ii) compensate tactile sensory \cite{zevallos2018surgical}, (iii) expand field of view \cite{wang2018stereoscopic}, (iv) provide more intuitive human-machine interface \cite{wen2017augmented}, and (v) annotate helpful cues \cite{chen2017real,7150398}.
Other applications utilize augmented reality for robotic surgery training \cite{tang2020review}.
The common displaying media of augmented reality for robotic surgery include the da Vinci console, the computer monitor, and head-mounted display \cite{qian2019review}. For the purpose of surgical education, head-mounted display-based augmented reality is an advantageous medium as it enables 3D display and interactions for multiple users and the environment.
Jarc et al. apply augmented reality to a clinical-like training scenario, where 3D semitransparent tools controlled by the proctor are augmented and overlaid in the trainee's console as guidance \cite{jarc2017proctors}.
User study involving seven proctor-trainee pairs is conducted, where they demonstrate the augmented tools as an effective mentoring approach. However, currently most of the AR system have not taken AI as the core component for generating and creating the context-aware information.

\section{METHOD}

\subsection{Overall Framework}
The whole system is based on the dVRK \cite{dVRK} platform, the first standard and general da Vinci surgical system that has been open-sourced and further developed for other researchers to explore~\cite{d2021accelerating}. It has been widely used for the research on surgical imaging and perception~\cite{long2021relational}, control and hardware design~\cite{lin2019reliable}, system simulation~\cite{xu2021surrol} and surgical task automation~\cite{lu2020learning}, which substantiate the high reliability and flexibility of the platform. In this work, we propose to fully leverage the advantage of dVRK to integrate the augmented reality and artificial intelligent to form a robotic surgery education system as shown in Fig.~\ref{overall_Framework}.
In total, our framework consists of several parts as described below:

\begin{enumerate}
\item AI+AR Computing Server, as shown in Fig.~\ref{overall_Framework} a): A high performance computer equipped with a high-end GPU for AR and AI algorithms deployment, which can acquire the stereo video and kinematics information from Patient Side Manipulator (PSM) and Endoscopic Camera Manipulator (ECM) on dVRK, receives input from other user interface, train and deploy RL policy or other AI framework and output the augmented video stream to monitor or other display devices.
\item Endoscope Video Acquisition and Kinematics Control: As depicted in Fig.~\ref{overall_Framework} b) and d). The dVRK contains a stereo endoscope on ECM which can capture the stereo video stream for 3D sensing and perception. In this case, we collect the video signal from the dVRK stereo endoscope with a video capture card to convert the video signal to the USB video stream, which can be retrieved by the computer. By using the Application Programming Interface (API) provided by dVRK Robot Operating System (ROS) packages, users can acquire the kinematics information (including tool tip position, velocity, rotation) from two PSMs on dVRK and input kinematics to control the movement of the PSMs. Also, the PSMs can be control directly by user using hands in the dVRK console.
\item Displaying AR video with dVRK console: As depicted in Fig.~\ref{overall_Framework} c) and e). Equipped with a stereo viewer in the dVRK console, users can view on the scene with overlaid AR and AI information from the computer in the surgical scenario. The stereo viewer can provide left and right views with disparity, where human can perceive the surgical video and overlaid information with 3D feelings. This function is later developed into TilePro$^\text{™}$~\cite{hyung2014tilepro} in the following generation of da vinci surgical robot, which showcases the potential possibility of applying on the real robotic surgery.
\end{enumerate}

\subsection{Generating AI Guidance with Reinforcement Learning}
Instead of guiding the novice surgeon by experienced expert, our work proposes to generate guidance through advanced reinforcement learning automatically. Considering the intriguing superiority of SurRoL \cite{xu2021surrol} in leaning from demonstration and trials on dVRK, we attempt to incorporate it as our core AI module, learning and generating the guidance trajectory using the RL algorithm based on specific task.

SurRoL incorporates reinforcement learning to teach an intelligent agent to take interactive actions so that the cumulative reward can be maximized. While the agent takes action $a_t\in\mathcal{A}$ and interacts with the environment at state $s_t\in\mathcal{S}$, it receives reward $r_t\in\mathcal{R}$ based on whether the agent achieves the goal $g$ at step $t$.   From step $t=0$ to $t=T$, the trajectory can be formulated as $\tau=\{s_0, r_0, a_0, s_1, r_1, a_1...,s_T, r_T, a_T\}$. Besides, the translation $p(s_t|s_{t+1})$ is derived from the underlying physics engine. For robotic surgery, we define $a_t$ as six degrees of freedom, including position movement $(d_x, d_y, d_z)$ in Cartesian space, orientation $d_{\textit{yaw}}/d_\textit{pitch}$ in a top-down/vertical space setting and $j$ for the open ($j\geq 0$) or close ($j<0$) status of the jaw. Reinforcement learning targets on finding the policy $\pi$ to generate action $a_t=\pi(s_t)$. To promote the learning process, the reward $r_t$ is set to be goal-based, where the success function $f(s_t, g, a_t)$ would determine the reward $r_t$ by checking whether $a_t$ achieves the ground-truth status, such as 3D positions and 6D poses. Such scenario is similar to the attempting-feedback procedure in the surgical education. Finally, the target policy $\pi$ is learned by maximizing the experimental expectation $\mathbb{E}_\pi[\sum_{t=0}^T\gamma^tr_t]$, where $\gamma\in[0, 1)$ denotes the discount factor to balance the agent's attention on distant future and immediate future.
In specific, we opt for a sample efficient learning algorithm called hindsight experience replay (HER)~\cite{andrychowicz2017hindsight} and combine it with Q-filtered behavior cloning. In practice, we will leverage a small amount of demonstration data generated from the scripted policies as the demonstration for imitation leaning. It may have great potential to be extensively developed and learn from large amount of surgery data and then generate the AI-powered guidance for surgical education.

\subsection{Real-time 3D Visualization by Augmented Reality}

To facilitate the visualization result and form a more intuitive education, we propose to use the AR techniques to design a 3D visualization system for surgical robots.

\subsubsection{Different Coordinate System}
As an essential part of the AR system, the coordinates and the transformation among them serve as the basis of how we locate and align the AR information to camera view of the stereo endoscope. In our surgical education scenario, we denote the world coordinate as $\{W\}$, the coordinate of multiple objects in the experiment environment as $\{O_i\}$ with $i$ indicating the index of the objects, Remote Center of Motion (RCM shown in Fig.~\ref{overall_Framework} (b), which is a fixed base frame of the dVRK PSM during the operation) as $\{R\}$, the instrument tip as $\{P\}$ and the endoscope coordinate as $\{C\}$. And we use $T_{src}^{dst}$ to denote as the transformation from source $\{src\}$ to destination coordinate $\{dst\}$. During the deployment, $T_W^{R}$ are fixed and the same, and we manually place the objects in the same location as in the SurRoL. When given the kinematics information, we can compute the transformation of $T_R^P$ using the forward kinematics theory and control the PSM to move the instrument tip to the target pose $\{P\}$ in the SurRoL and the real environment as below.
If we have the transformation of the endoscope relative to the world $T_W^C$, we can calculate and get the position of the object relative to the endoscope $T_C^O$ using the following equation:
\begin{equation}
    T_C^{O} = T_W^{O} * T_C^{W};~~~~
    T_C^{W} = {T_W^{C}}^{-1}
    \label{equ_t}
\end{equation}

\subsubsection{Human-involved Flexible Calibration}
Calibration process is to acquire the coordinate of the endoscope relative to the world $T_W^C$. It is an very important step~\cite{tabb2017solving} for AR, because it can align the AR information with the environment. Some methods have been proposed to calibrate the hand-eye coordinate~\cite{tsai1989new, park1994robot}, but they needs at least an expert on robotic vision who use the calibration board to complete the whole process. Using the QR code to automatically retrieve the coordinate is another attempt and solution~\cite{qian2020ar}, but we need to carefully take the detecting condition into consideration and it may be restricted in the surgical scene. In our work, we propose an user-friendly method to easily calibrate the endoscope coordinate, which is more suitable for surgical education scenario. Specifically, we will first define some specific points and collect the 3D position $[X_W^i,Y_W^i,Z_W^i]$ of these points from SurRoL ($i$ denotes the point index). Then we manually locate the 2D location $[u_i,v_i]$ of theses points in the image. We design an User Interface (UI) as depicted in Fig.~\ref{Figure_calibration}, where users can easily locate and record the 2D location by clicking on the image using the mouse. According to the pinhole model of the camera, we can have the following equation:
\begin{equation}
\begin{bmatrix} u_i \\ v_i \\ 1 \end{bmatrix} = K * J * T_W^{C}\begin{bmatrix} X_W^i\\Y_W^i\\Z_W^i\\1 \end{bmatrix}
\label{equ_pnp}
\end{equation}
\begin{equation*}
K=\begin{bmatrix} f_x&0&c_x \\ 0&f_y&c_y\\ 0&0&1 \end{bmatrix};~~
J=\begin{bmatrix} 1&0&0&0 \\ 0&1&0&0\\ 0&0&1&0 \end{bmatrix}
\label{equ_pnp2}
\end{equation*}
where $J$ represents the perspective projection model and $K$ is the camera intrinsic matrix which can be straightforwardly acquired through~\cite{zhang2000flexible}. To solve and acquire the solution for $T_W^{C}$, we leverage the stable method called iterative solvepnp algorithm~\cite{lepetit2009epnp}. It is based on Levenberg-Marquardt optimization which minimizes re-projection error and finds the best solution. The algorithm needs only 3 points for getting the solutions (more points contributes to better solution), which is convenient and efficient for surgical application.

\subsubsection{Augmented 3D Overlaying}
To overlay the 3D information on the image, we will first place the AR overlay information (3D) in the environment relative to the world coordinate $T_W^{\text{AR}}$, where we know the position $[X_W^\text{AR},Y_W^\text{AR},Z_W^\text{AR}]$. Then we will project the overlay information to the image using the projection equation same as the Eqn.~(\ref{equ_pnp}).
Specifically, we sometimes may need different views for AR overlaying. And we can use the proposed calibration method to calibrate the coordinate and overlay the new generated AR information relative to the new novel view.
This is very promising as in the surgical education scenario, multiple views may enrich the information and help the surgeon to learn better.

\subsection{Integrating AI and AR for Trajectory Guidance}\label{Sec:Deployment}

With the AI and AR components introduced above, we further propose to integrate these two modules to form a AI-powered trajectory guidance with AR visualization.
Firstly, the dVRK stereo endoscope will capture the stereo video from the surgical scenario and output for computing server. Then we will use the proposed calibration method to calculate and align the coordinate of the endoscope and world. We will then train the RL policy to learn on the specific tasks in the SurRoL. Given a well-learned policy and providing with the known state of the environment or the observation, we can generate the predicted trajectory from the action, which is represented in the coordinate of the world. Afterwards, we can project the trajectory points to the video frames using the projection equation. By lining up the points, we can finally generate a trajectory guidance. With the augmented video displayed in the stereo viewer of dVRK console, the user can perceive the 3D AR trajectory guidance for the purpose of robotic surgery education.

\section{Experiment}
\begin{figure}[t]
\vspace{+2mm}
    \centering
    \includegraphics[width=0.9\hsize]{figures/exp.pdf}
    \caption{Our proposed human-involved calibration process where we can flexibly locate the specific points on the peg board for calibrating.}
    \label{Figure_calibration}
\end{figure}
To study the preliminary results of the proposed method, we choose the peg-transfer task as the experimental task in our work, which is one of the Fundamentals
of Laparoscopic Surgery (FLS) tasks for surgical skills education~\cite{joseph2010chopstick}. In this experiment, we move the block from one peg to the other peg with one PSM on the dVRK. The whole process contains three procedures: 1) lifting the peg, 2) transferring the peg and 3) placing the peg, as depicted in Fig.~\ref{Figure_exp}. We train the RL policy on this task and generate the block transferring trajectory based on the trained model. Then we overlay the peg-transfer path as the AR trajectory on the stereo viewer for visualization.

We conduct the experiment on a computer using ubuntu 18.04 as the system,  which is equipped with 8 cores 3.60GHz Intel$^\text{®}$ Xeon(R) W-2123 CPU, NVIDIA$^\text{®}$ TITAN RTX GPU and 16G RAM. All the algorithms are written in Python. The image resolution from the endoscope we use is $640\times480$.

\begin{figure}[t]
    \centering
    \includegraphics[width=1.0\hsize]{figures/exp2.pdf}
    \caption{Two trajectories overlay visualization results with long range transfer (top) and the short range transfer (bottom).}
    \label{Figure_exp}
    \vspace{-0.4cm}
\end{figure}

\subsection{Reinforcement Learning on Peg-transfer}

In this work, we follow the peg-transfer setting in SurRoL\cite{xu2021surrol} where the goal tolerate distance is set as 0.5 cm in a workspace with size of 10 $\text{cm}^2$. Each episode lasts for 50 timesteps. When the episode ends, PSM reaches out of the workspace or the peg-transfer is completed, the environment will be reset and the positions of initial peg and target peg will be resampled. Two evaluation settings are proposed, (1) short range: the block moved by user will not encounter any obstacle peg during the transfer; (2) long range: the block needs to avoid 1 to 3 pegs to finally reach the target peg, as these two situations are similar to simple and complicated surgical scenario for education. We randomly generate 100 scripted demonstrations for RL to imitate and train for 100 epochs in the training stage. In the testing stage, for each experiment, we conduct the peg transfer for 1000 trials (with 500 short range and 500 long range trials) and evaluate the success rate. We select top 3 performed RL policy and average their results. The results are shown in the Table~\ref{table:RL}, where we can find that the RL policy can perform well on long range while achieving a higher performance on the short range peg-transfer with success rate 92.6\%. Overall, our RL can achieve promising result with average success rate 86.4\% on peg-transfer.

At the deployment stage, we export the complete predicted trajectory from the RL policy and input the trajectory location to control the PSM of real dVRK to conduct peg-transfer. We show a failure case from the long range setting in Fig.~\ref{Figure_fail} a), where the block is not placed to the target peg properly because of the biased trajectory generated from the RL policy. Although there exists failure trial, the above results already bring out a lot of possibilities of applying more advanced RL algorithms on the complex robotic surgery education scenario.

\begin{table}[t]
\vspace{+2mm}
    \caption{
    \textbf{The evaluation of \textit{Peg-Transfer} on SurRoL.}
    }
    \vspace{-6pt}
    \centering
    \begin{tabular}{c|cc}
    \hline
    Distance       & Trials & Success Rate (\%) \\ \hline
    Short range    & 463/500  &  92.6                \\
    Long range     & 401/500  &  80.2    \\ \hline
    All            & 864/1000  &  86.4
    \\ \hline
    \end{tabular}
    \label{table:RL}
    \vspace{-0.4cm}
\end{table}

\subsection{Visualization and Analysis on the AR results}
In the calibration process, we choose the top center points of the twelves pegs as the specific points $[u_i,v_i]$ for solving the coordinate of the endoscope. As we can see in Fig.~\ref{Figure_calibration}, the specific points are located and shown in red dot. The point locations are shown next to these points with white text. After pointing out 12 points on the peg-transfer board, we can finish calibrating the coordinate.

As shown in Fig.~\ref{Figure_exp}, we show two cases from our experiments, with the top one showing the peg-transfer with long range distance and the bottom one showing the short range distance. In the experiment, we deploy our RL policy to generate the 3D trajectory and project the trajectory to the video frames. As we can see in the figure, the white dots in the image frame represent each action step in the trajectory, and the green lines represent the trajectory. It can be observed that our method can intuitively and accurately overlay the trajectory in the video frames, where the instrument tip will follow the overlaid trajectory visually in the experiment.

To further analyse the stability and precision of the calibration, we randomly find 5 new users with engineering background to conduct the calibration process. We calculate the re-projection error (pixel) \cite{koide2019general} based on twelves specific points on the pegs. Each user conducts 10 trials and we report the mean and standard deviation of the error in Table~\ref{table:error}. From the table we can find that average error is within 1 pixel (0.76) which is precise for the visualization purpose. Besides, we only observe one failure trial where the trajectory is calculated and projected to the wrong location due to the large error from manually locating the specific points, as shown in Fig.~\ref{Figure_fail} b). The above results demonstrate the well performance of the calibration process and AR visualization for trajectory guidance on peg-transfer education scenario.

\begin{table}[tp]
\vspace{+2mm}
    \caption{
    \textbf{The evaluation of AR calibration on \textit{Peg-Transfer} scenario.}
    }
    \centering
    \vspace{-1mm}
     \resizebox{.95\columnwidth}{!}{
    \begin{tabular}{c|ccccc}
    \hline
    Re-project error (pixel) & User1 & User2 & User3 & User4 & User5 \\ \hline
    Average value       & 0.77  & 1.02  & 0.70 &  0.61  & 0.71           \\
    Standard deviation  & 0.20  & 0.14  & 0.12 &  0.10  & 0.15            \\ \hline
    \end{tabular}}
    \label{table:error}

\end{table}

\begin{figure}[t]
    \centering
    \includegraphics[width=0.8\hsize]{figures/exp_fail.pdf}
    \caption{The failure cases with (a) because of the biased
 trajectory generated from RL and (b) due to manually caused calibration error.}
    \label{Figure_fail}
    \vspace{-0.5cm}
\end{figure}

\subsection{System Latency Study}
As a systematic assessment of the whole pipeline, latency is more than important for a reliable and real-time system. In our work, we mainly focus on studying two different types of latency that may be introduced by the system. One is the video capturing and displaying latency, which is often caused by the delay of the video capture device, video signal transfer and the monitor. The other is overlaying latency, which is mainly caused by computing and projecting the points from 3D location to 2D image, as well as the time for drawing the information on the video frame.

To evaluate the video capturing and displaying delay, we place a tablet under the camera view of the endoscope which displays the timestamp on the screen. Then we capture the video from endoscope camera and display it on the screen of the computer. Afterwards, we take a photo using another camera, which will include both table screen and computer screen. The difference of two timestamps shown on two screens is the latency. The illustration can be seen in Fig.~\ref{Figure_latency} (a). Finally, we find the average latency is around 161 milliseconds with standard deviation of 17 milliseconds from 100 samples. Which need to be mentioned is that the latency may vary when given different video capture equipment or monitor. According to some studies~\cite{miller1968response, tolia2006quantifying} and our own observation among all the trials, the latency is delicate enough that the human would not be able to perceive and will not significantly affect the operating.

\begin{figure}[t]
    \centering
    \includegraphics[width=1\hsize]{figures/exp_latency_all.pdf}
    \vspace{-18pt}
    \caption{(a) Illustration for evaluating video capturing and displaying latency and (b) chart showing the latency related to the point size.}
    \label{Figure_latency}
    \vspace{-0.5cm}
\end{figure}

To further evaluate the overlaying latency, we conduct an experiment which measures the computing time of overlaying (projecting 3D trajectory points to the video frames and visualize them).  We evaluate the time cost when given different size of trajectory points (200-5000), as shown in Fig.~\ref{Figure_latency} (b), where we randomly generate the 3D locations of these points and repeat the experiment for 300 times. Finally we report the average and standard deviation of the time cost. From the figure we can observe that when given more points, the computation time may increase showing exponential growth. When around or less than 2600 points, we can achieve the real-time process ($>30$ Hz). In our peg transfer scenario, the RL will generate around 100 trajectory points which cost just about 1 ms for projecting and overlaying, it is more than enough for a real-time visualization. According to the above results, our designed system can efficiently overlay the trajectory on video frames forming a real-time AR visualization, which has great potential for more complicated real-time AR robotic surgery education scenario.

\section{CONCLUSIONS AND FUTURE WORK}
In this work, we confer the intelligence on AR by integrating AI module as a decision-maker for generating trajectory. The AI module with reinforcement learning is leveraged to learn policy-based operation and dynamically generate the trajectory based on user's task. To provide an immersive visualization and facilitate intuitive education, the trajectory plan is further projected and overlaid onto the stereo video, assisting the trainee to perceive the intelligent guidelines in 3D AR. A simple yet effective calibration with human-computer interaction is also designed to realize coordinate calibration through an user interface. We conduct the experiment with the surgical training on task peg-transfer, which demonstrates the feasibility and potential of AI-powered AR for robotic surgery education.

In our future work, we shall further exploit how to design and incorporate more AI and AR modules into the system, such as virtual tool guidance, overlaying the result of surgical workflow recognition and instrument segmentation. Furthermore, we shall integrate network communication for realizing a remote surgery education. Methods of measuring the operation of surgeon compared to the generated 'expert' trajectory and user study may also be included for evaluating the effectiveness of the proposed system.

\newpage
\bibliographystyle{IEEEtran}

Robotics (cs.RO); Artificial Intelligence (cs.AI)


Tuning Penta-Graphene Electronic Properties Through Engineered Line Defects

Abstract: Penta-graphene is a quasi-two-dimensional carbon allotrope consisting of a pentagonal lattice in which both sp2 and sp3-like carbons are present. Unlike graphene, penta-graphene exhibits a non-zero bandgap, which opens the possibility of its use in optoelectronic applications. However, as the observed bandgap is large, gap tuning strategies such as doping are required. In this work, density functional theory calculations are used to determine the effects of the different number of line defects of substitutional nitrogen or silicon atoms on the penta-graphene electronic behavior. Our results show that this doping can induce semiconductor, semimetallic, or metallic behavior depending on the doping atom and targeted hybridization (sp2 or $sp3-like carbons). In particular, we observed that nitrogen doping of sp2-like carbons atoms can produce a bandgap modulation between semimetallic and semiconductor behavior. These results show that engineering line defects can be an effective way to tune penta-graphene electronic behavior.

\section{Introduction}

Research on 2D materials has gained much attention since the discovery of graphene\cite{novoselov_Science,geim_NM}. This carbon allotrope that is composed of a hexagonal lattice was demonstrated to possess several exciting features, such as high electrical and thermal conductivity and large mechanical resistance that are aimed at developing the next generation of organic optoelectronic devices \cite{neto2009electronic,balandin2008superior}. However, when it comes to optoelectronic applications, graphene has a zero bandgap, which precludes its use, for instance, as active material in solar cells. To overcome this issue, gap opening strategies have been developed \cite{tuning_gap_graphene1, tuning_gap_graphene2, tuning_gap_graphene3}, including the cutting of graphene sheets into graphene nanoribbons \cite{son2006energy} and the use of dopants \cite{DENIS2010251,HOUMAD20151021,doi:10.1063/1.4791011}.

There are a few nature-occurring carbon allotropes, including graphite, fullerenes \cite{kroto1985c60}, and carbon nanotubes \cite{iijima1993single}. Recently, another allotrope was proposed named penta-graphene \cite{pentagraphene} (see Figure \ref{morfsp3}), which has a lattice composed of pentagons that resemble the Cairo pentagonal tiling. It is supposed to present ultrahigh ideal strength even above that of graphene and to be able to withstand temperatures of up to 1000 K \cite{1}. Penta-graphene was not sensitized yet, and it is believed that it can be at least a metastable material \cite{1,cranford20166,rajbanshi2016energetic,ewels2015predicting}. Furthermore, recent studies have suggested that chemical functionalization employing hydrogenation or fluorination could even increase its structural stability\cite{einollahzadeh2016studying,2}. Electronic structure calculations indicate that penta-graphene is a semiconductor presenting bandgap values about  2.24--4.3 eV \cite{1,rajbanshi2016energetic,einollahzadeh2016studying}. Such large bandgap suggests the need for gap tuning strategies to be developed in order to tailor this material for optoelectronic applications. In this sense, the particular topology of penta-graphene, which presents both $sp^2$ and $sp^3$-like carbon hybridizations in its lattice, raises the question of how its structural and electronic properties will behave by targeting a particular hybridization in the doping processes.

\begin{figure*}[!htb]
\centering
\includegraphics[width=1.0\linewidth]{Figure1}
\caption{(Top) Schematic representation for the $sp^3$-like carbons doping strategy. The red color represents the doping atoms (N or Si). Bond length variations for penta-graphene lattices with nitrogen and silicon doping schemes are presented in the middle and bottom panels, respectively. In the color palette on the right, the equilibrium distance is 1.575 \AA. The maximum deviations are 0.5 \AA~for nitrogen and 1.5 \AA~for silicon doping schemes.}
\label{morfsp3}
\end{figure*}

Herein, we carried out density functional theory (DFT) calculations to address the effects of the systematic substitutional doping of either $sp^2$ or $sp^3$-like carbons by nitrogen and silicon atoms in a penta-graphene lattice. Mainly, we investigate the effect of the different number of engineered line defects on their electronic and structural properties (see Figure \ref{morfsp3}). Our findings show that, in terms of morphology, nitrogen doping is responsible for increasing the stiffness of the lattice in comparison to pristine penta-graphene, whereas silicon doping results in the simultaneous stretching and compression of Si--C and C--C bonds, respectively, concerning undoped C--C bonds. The two doping schemes investigated here (Si or N) produce significantly different results in terms of electronic behavior. Silicon doping allows us to tune the bandgap when replacing $sp^3$-like carbons and produces metallic behavior when replacing $sp^2$-like ones.  Nitrogen doping replacing $sp^3$-like carbons results in a transition from semiconductor to semimetallic to a metallic character. Nitrogen doping replacing $sp^2$-like carbons produces an alternating behavior between semimetallic and semiconductor depending on the number of dopants. These results indicate that engineered line defects can be a very effective way to tune penta-graphene electronic behavior.

\section{Computational Details}
DFT calculations were carried out within the Generalized Gradient Approximation (GGA) scheme as proposed for Perdew, Burke, and Ernzerhof (GGA/PBE) \cite{perdew96} along with the DZP \cite{hohenberg64,Kohn65} basis set. Relativistic pseudopotentials parameterized within the Troullier-Martins formalism were also considered \cite{kleinman82}. These combined approximations can accurately describe the magnetic and electronic properties of materials composed of atoms with many electrons. All the calculations were performed considering spin polarization. For the bands and density of states calculations, Monkhorst pack grid of 21 x 21 x 3 was used \cite{monkhorst76}. A mesh cutoff of 200 Ry was chosen as a parameter for our calculations \cite{anglada02}. The force criteria convergence was 0.001 eV/ \AA. In order to establish a good compromise between the accuracy of our results and computational costs, the tolerance in the matrix density and total energy was set to 0.0001 and 0.00001 eV, respectively. All calculations were performed with the SIESTA software suite \cite{ordejon96,portal97}.

\section{Results and Discussion}
In the present work, we investigated the electronic and structural features of penta-graphene lattices with substitutional doping (N or Si) either at $sp^3$ or $sp^2$-like carbons forming engineered line defects. The cases of 1 up to 7 line defects were considered. In this sense, Figure \ref{morfsp3} (top panels) shows a schematic representation of these defective structures for the $sp^3$ case, which is similar to the $sp^2$ one. Figure \ref{morfsp3} presents seven scenarios, which are identified by N--XL, with N corresponding the number of line defects in the horizontal/vertical directions, and X refers to either nitrogen or silicon dopant atoms.

Significant structural differences take place in the morphology of the resulting doped structures. For the nitrogen cases, an overall decrease in bond length values was observed, as presented in Figure \ref{morfsp3} (middle panels), where the blue color represents bond length compression when compared to the original non-doped ones. As expected, these effects are more pronounced around the defect lines, but they extend to the other bonds as the number of doping atoms increases. Such a reduction in bond length values can reach 0.5 \AA, which corresponds to a roughly 30\% decrease concerning the 1.575 \AA~ equilibrium distance found in pristine penta-graphene. In contrast, for the silicon atoms, the newly formed C-Si bonds undergo a substantial expansion of up to 1.5 \AA, which would amount to 100\% increase in bond length values when compared to pristine penta-graphene C-C bonds, as shown in Figure \ref{morfsp3}-bottom. For the remaining C-C bonds, this effect is less pronounced.

\begin{figure*}[!htb]
\centering
\includegraphics[width=0.7\linewidth]{Figure2}
\caption{\label{bandasp3} Band structures for undoped and doped penta-graphene lattices, as a function of the number of line defects for the $sp^3$ case. (a) and (b) correspond to the results for nitrogen and silicon doping schemes, respectively. Here, X-NL and X-SiL denote the number (X) of dopant lines systematically inserted into the penta-graphene structure.}
\end{figure*}

It is well-known that structural modifications lead to changes in electronic properties. As such, useful information about these changes can be gained by contrasting the band structures of the doped and undoped penta-graphene sheets.  Reports in the literature for penta-graphene indicate a semiconductor material with a bandgap of about 2.4 eV \cite{pentagraphene,rajbanshi2016energetic}.  We have obtained a similar value, as shown in Figure \ref{bandasp3}. In this figure, we also present the band results for the $sp^3$ cases for nitrogen (Figure \ref{bandasp3}-a) and silicon doping (Figure \ref{bandasp3}-b) as a function of the number of line defects. As we can see from Figure \ref{bandasp3}, for the nitrogen with 1 line defect, the bandgap becomes indirect and decreases from 2.4 eV to just 1.5 eV. The Fermi level lies near the conduction band, making this doped penta-graphene structure an n-type semiconductor. As the number of line defects increases, the material no longer presents a bandgap. We can see that for 2 and 3 defect lines, the partial density of states (PDOS) near the Fermi level, presented in Figure S1a in the Electronic Supporting Information (SI), almost disappears. As the number of defect lines increases, the PDOS around the Fermi levels increases, and the doped penta-graphene becomes fully metallic.

The silicon doping of $sp^3$-like carbons, on the other hand, does not result in metallic materials. As seen in Figure \ref{bandasp3}-b and Figure S2a of the SI, doped penta-graphene lattices preserve their semiconductor nature. For the 1 line defect case, the only observed effect was the decrease in bandgap to 1.3 eV, with further doping making the bandgap indirect. Interestingly, the bandgap dependence on the number of defect lines has a convex nature, with a minimum bandgap of 0.8 eV observed for the case of 4 defect lines. After this point, the bandgap increases again, reaching 1.76 eV for the lattice with seven defect lines. The decreases in bandgap values are associated with the appearance of states in the valence band closer to the Fermi level.

\begin{figure*}[!htb]
\centering
\includegraphics[width=1.0\linewidth]{Figure3}
\caption{(Top) Schematic representation for penta-graphene lattices with $sp^2$-like carbons doping strategy. The red color represents the doping atoms (N or Si). Bond length variations for penta-graphene lattices with nitrogen and silicon doping schemes are presented in the middle and bottom panels, respectively. In the color palette on the right, the equilibrium distance is 1.575 \AA. The maximum deviations are 0.5 \AA~for nitrogen and 1.5 \AA~for silicon doping schemes.}
\label{sp2}
\end{figure*}

\begin{figure*}[!htb]
\centering
\includegraphics[width=0.7\linewidth]{Figure4}
\caption{Band structures for undoped and doped penta-graphene, as a function of the number of line defects for the $sp^2$ case. (a) and (b) correspond to the results for nitrogen and silicon doping schemes, respectively. Here, X-NL and X-SiL denote the number (X) of dopant lines systematically inserted into the penta-graphene structure.}
\label{bandassp2}
\end{figure*}

The second doping strategy considered here consists of the doping of $sp^2$-like carbons. One main difference is the possibility of having N--N and Si--Si bonds, not present for $sp^3$ case. The seven analyzed scenarios are presented in Figure \ref{sp2} (top panels). Again, nitrogen doping results in overall contraction of the bond length values up to 0.5 \AA~concerning undoped penta-graphene ones. These deviations in the bond lengths are represented in Figure \ref{sp2} (middle panels) by the blue bonds. Silicon doping, in contrast, produces different patterns. These are mostly characterized by the simultaneous expansion of Si--C bonds, by Si--Si bonds that preserve the original 1.575 \AA~bond length of pristine penta-graphene and by C-C bonds that slightly contract. The combination of these effects results in the pattern characterized by blue hexagons diagonally sliced by red Si--C bonds, as depicted in Figure \ref{sp2}-bottom, especially for the six defect lines case. However, two particular cases, 5 and 7 defect lines, break the pattern producing much more disordered configurations likely induced by symmetry breaking. The Si--C distances in the doped penta-graphene lattices can reach 3.8 \AA, considerably larger than Si--C bonds found, for instance, in disilicon carbide, which can be as large as 2.2 \AA\cite{koput2017ab}. This is suggestive that the atoms are no longer bonded and the structures undergo structural rearrangements.

In terms of electronic structure, nitrogen and silicon doping produce completely different results. Interestingly, for the $sp^2$ nitrogen doping cases, the bandgap values exhibit a bandgap modulation (alternating increasing/decreasing) behavior, as can be seen in Figure \ref{bandassp2}-a. For even values of $N$, semiconducting properties are obtained with almost direct bandgaps that decrease as $N$ grows larger from 2.4 eV for $N=0$ to 0.6 eV for $N=6$. In contrast, for odd values of $N$, the valence band maximum (VBM) touches or surpasses the Fermi levels. However, the DOS near the Fermi level is very small, increasing progressively with $N$. This behavior can be better visualized in the PDOS plots of Figure S1b in the SI. The even $N$ nitrogen-doped penta-graphene lattices display a semimetallic character, with carbon $p$ orbitals being mostly responsible for the DOS in the vicinity of the Fermi level.

Silicon doped penta-graphene structures, on the other hand, possess semiconductor properties only for $N=1$ or $N=2$, with a fast transition to bandgap closing. For larger $N$ values, the structures become fully metallic, as shown in Figure \ref{bandassp2}-b (see also Figure2-ESI in SI). It is also worth noting that the effect of the disordered morphology observed for the 5-SiL and 7-SiL cases is to make energy levels practically independent of momentum when compared to the other cases, as evidenced by the almost flat (dispersionless) bands. In this sense, silicon doping of $sp^2$-like carbons is not as effective as in cases of nitrogen doping regarding decreasing the bandgap value.

Finally, an important advantage concerning the $sp^2$ doping instead of $sp^3$ one is the effects on the structural stability of the doped structures. In Figure \ref{spec1}, we present the cohesive energy values for all doping scenarios. The cohesive energy is the difference per atom between the entire system energy and the sum of the individual energies of its constituents. Figure \ref{spec1}-top) and Figure \ref{spec1}-bottom) refer to $sp^3$ and $sp^2$ results, respectively. It can be seen from this figure that for both $sp^3$ and $sp^2$-like carbon doping, the effects on cohesive energy are similar for N and Si dopants. However, whereas $sp^3$-like carbon doping increases the cohesive energy from -8.0 eV/atom to -6.5 eV/atom, the corresponding increase in the case of $sp^2$-like carbon in less than 1.0 eV/atom.

\begin{figure}[!htb]
\centering
\includegraphics[width=0.48\textwidth]{Figure5a}
\includegraphics[width=0.48\textwidth]{Figure5b}
\caption{Cohesive energy for all cases of nitrogen and silicon doping schemes of (top) $sp^3$ carbons and (bottom) $sp^2$-like carbons.}
\label{spec1}
\end{figure}

\section{Summary and Conclusions}
In summary, we have carried out DFT calculations to assess the changes in the structural and electronic properties of penta-graphene lattices resulting from selective N and Si doping (engineered line defects) of either $sp^3$ or $sp^2$-like carbons. Regardless of the type doping, for the nitrogen cases, we observed an overall stiffening of the penta-graphene structures. On the other hand, for the silicon doping cases, we observed only the stretching of Si-C bonds and compression of the remaining C-C bonds.

From an electronic structure perspective, both the doping type and doping atom selection produce significantly different results. Silicon doping of the $sp^3$-like carbons preserves the penta-graphene semiconductor character. The smallest observed bandgap value was 0.8 eV (2.4 eV for pristine penta-graphene) for the 4 line defect case. For a larger number of line defects, the bandgaps increase again. On the other hand, for Si doping $sp^2$-like carbons, the result is a fast transition to metallic behavior.

Nitrogen doping produces more interesting results. For the $sp^3$-like carbons, the doping results in a progressive changes from semiconductor to metallic behavior. For the case of 1 line defects, the doped penta-graphene structure becomes an n-type semiconductor with a 1.5 eV bandgap. For 2 and 3 line defects, the bandgap closes, but the density of states near the Fermi level goes to zero, giving the material a semi-metallic character. From this point on, further doping leads to true metallic behavior.

For the nitrogen doping of $sp^2$-like carbons, we observed a bandgap modulation behavior (alternating increase/decrease). For even numbers of line defects, the bandgap decreases with the number of line defects. In contrast, for odd numbers of line defects, a semi-metallic behavior is observed, combining zero bandgap with near-zero density of states at the vicinity of the Fermi level. Finally, the cohesive energy values indicate that doping the $sp^2$-like carbon affects less the structural stability of the resulting doped structures than $sp^3$-like carbon ones.

These results indicate that selective doping of penta-graphene structures through engineered line defects can be an effective tool to tune their electronic behavior, being possible to create structures that vary from large bandgaps through semiconductors and even metallic or semimetallic ones. We hope the present work can stimulate further studies along these lines.

\section*{Conflicts of interest}
There are no conflicts to declare.

\section*{Acknowledgments}
The authors gratefully acknowledge the financial support from Brazilian Research Councils CNPq, CAPES, and FAP-DF. L.A.R.J. gratefully acknowledges the financial support from FAP-DF grant 00193.0000248/2019-32 and the financial support from CNPq grant 302236/2018-0.DSG thank the Center for Computational Engineering and Sciences at Unicamp for financial support through the FAPESP/CEPID Grant \#2013/08293-7.


Allocating Large Vocabulary Capacity for Cross-lingual Language Model Pre-training

Abstract: Compared to monolingual models, cross-lingual models usually require a more expressive vocabulary to represent all languages adequately. We find that many languages are under-represented in recent cross-lingual language models due to the limited vocabulary capacity. To this end, we propose an algorithm VoCap to determine the desired vocabulary capacity of each language. However, increasing the vocabulary size significantly slows down the pre-training speed. In order to address the issues, we propose k-NN-based target sampling to accelerate the expensive softmax. Our experiments show that the multilingual vocabulary learned with VoCap benefits cross-lingual language model pre-training. Moreover, k-NN-based target sampling mitigates the side-effects of increasing the vocabulary size while achieving comparable performance and faster pre-training speed. The code and the pretrained multilingual vocabularies are available atthis https URL.

\section{Introduction}
\label{sec:intro}

Pretrained cross-lingual language models~\citep{xlm, xlmr, infoxlm, mt5} have recently shown great success in improving cross-lingual transferability.
These models encode texts from different languages into universal representations with a shared multilingual vocabulary and a shared Transformer encoder~\citep{transformer}. By pre-training cross-lingual language models on the large-scale multilingual corpus, the models achieve state-of-the-art performance on various downstream tasks, e.g., cross-lingual question answering and cross-lingual sentence classification.

Although the Transformer architecture used in most pretrained monolingual and cross-lingual language models are almost identical, the vocabularies are quite different.
The vocabulary sizes in existing pretrained monolingual language models typically range from 30K to 60K subword units~\citep{bert, roberta,unilm,unilmv2}.
Meanwhile, state-of-the-art pretrained cross-lingual language models use the shared multilingual vocabulary of 250K subword units to represent more than 100 languages~\citep{xlmr, infoxlm, mt5}.
Although some subword units are shared across languages, no more than 2.5K language-specific subword units on average are allocated for each language, which is still relatively small.
Besides, the multilingual vocabulary is trained on the combined multilingual corpus with subword segmentation algorithms like BPE~\cite{bpe} and unigram language model~\cite{DBLP:conf/acl/Kudo18}.
During vocabulary construction, these algorithms tend to select more subword units shared across languages with common scripts like Latin and Cyrillic~\cite{DBLP:conf/emnlp/ChungGTR20}, but have a lower chance to select language-specific subword units.
It is hard to determine how much vocabulary capacity a particular language requires and whether the shared multilingual vocabulary has allocated enough vocabulary capacity to represent the language.

In this paper, we propose \ours{}, an algorithm to allocate large vocabulary for cross-lingual language model by separately evaluating the required vocabulary capacity of each language.
First, we use the \textit{average log probability} (ALP) to evaluate the ability of a vocabulary to represent a particular language.
We find that ALP is highly correlated to the downstream task performance, and we use it as an indicator to allocate language-specific vocabulary capacity.
In addition, the language-specific pre-training corpus size should also be considered since the pretrained model can only learn limited knowledge from low-resource languages where the pre-training data is scarce. Therefore, allocating too much vocabulary capacity for low-resource languages is inefficient.
\ours{} leverages both ALP and pre-training corpus size to evaluate the required vocabulary capacity of each language.
We finally allocate a multilingual vocabulary with 500K subword units with \ours{} and show it can significantly improve the model performance.

However, increasing the vocabulary size has two practical drawbacks: slow pre-training speed and heavy model size.
To address the pre-training speed issue, we propose $k$-NN-based target sampling, an approximate algorithm to improve the computing efficiency in the expensive softmax caused by the large vocabulary.
We pre-train the model with a small subset of the entire vocabulary constructed with $k$ nearest neighbors of the target words in current mini-batch data, evaluated with the inner product of subword embeddings.
As for the model size, we halve the embedding dimension and draw a different conclusion from~\citet{xlmr} that increasing vocabulary from 250K to 500K with a fixed capacity model can also improve the performance.

Our contributions are summarized as follows:
\begin{itemize}
\item We propose \ours{}, an algorithm to allocate appropriate vocabulary capacity for each language in the shared multilingual vocabulary of cross-lingual language models.
\item We propose $k$-NN-based target sampling, a softmax approximation algorithm to improve the computing efficiency during cross-lingual language model pre-training.
\item We evaluate our methods on the XTREME benchmark~\citep{xtreme}, including three different tasks on seven datasets.
Experiments show that \ours{} consistently outperforms previous vocabulary construction methods.
Meanwhile, our $k$-NN-based target sampling enables effective acceleration while achieving comparable performance.
\end{itemize}

\begin{figure*}[t]
\centering
\begin{minipage}[t]{0.48\textwidth}
\centering
\includegraphics[width=8cm,trim={0.5cm 0cm 1cm 1cm}]{alp.pdf}
\caption{ALP of different monolingual vocabularies with different vocabulary sizes. }
\label{fig:alp}
\end{minipage}
\quad
\begin{minipage}[t]{0.48\textwidth}
\centering
\includegraphics[width=8cm,trim={0.5cm 0cm 1cm 1cm}]{udpos.pdf}
\caption{F1 score on POS task with different vocabularies versus their ALP on the monolingual corpus.}
\label{fig:udpos}
\end{minipage}
\begin{minipage}[t]{0.48\textwidth}
\centering
\includegraphics[width=8cm,trim={0.5cm 0cm 1cm 0cm}]{panx.pdf}
\caption{F1 score on NER task with different vocabularies versus their ALP on the monolingual corpus. }
\label{fig:panx}
\end{minipage}
\quad
\begin{minipage}[t]{0.48\textwidth}
\centering
\includegraphics[width=8cm,trim={0.5cm 0cm 1cm 0cm}]{resource_alp.pdf}
\caption{
Comparison of vocabulary capacity of different-resourced languages. Shorter bars indicate larger vocabulary capacity.
}
\label{fig:res-alp}
\end{minipage}
\end{figure*}

\section{\ours{}: Language-Specific Vocabulary Capacity Allocation}
\label{sec:vocap}
We attribute the main factors that affect the performance of a particular language in a cross-lingual language model to language-specific pre-training corpus size and vocabulary capacity.
While previous work adjusts pre-training corpus size with an exponentially smoothed sampling distribution~\citep{xlm,xlmr}, few existing works have explored the effect of the language-specific vocabulary capacity in pretrained cross-lingual language models.

In this section, we first investigate the correlation between the language-specific vocabulary capacity and downstream task performance through experiments.
Then we introduce our proposed multilingual vocabulary allocation algorithm \ours{}.

\subsection{Investigating Language-Specific Vocabulary Capacity}
We start by introducing \textit{average log probability} (ALP) to quantify the language-specific vocabulary capacity in the shared multilingual vocabulary for a specific language.\footnote{For brevity and consistency, we refer to the parameterized tokenizer also as vocabulary, e.g., SentencePiece model.} Given a monolingual corpus composed of sentences $\mathcal{D}_{i}=\{s_1,...,s_{|\mathcal{D}_{i}|} \}$ from the $i$-th language and tokenized with vocabulary $V$, the average log probability is defined as follows:
\begin{align}
\textrm{ALP}(\mathcal{D}_{i}, V) = \frac{1}{|\mathcal{D}_{i}|}\sum_{j=1}^{|\mathcal{D}_{i}|}\sum_{k=1}^{|s_j|} \log p_{uni}(s_j^k)
\label{eq:alp}
\end{align}
where $s_j^k$ is the $k$-th subword of the sentence $s_j$, and $p_{uni}(\cdot)$ is the unigram distribution counted on the monolingual corpus $\mathcal{D}_{i}$.
It is difficult to count the language-specific subword units in multilingual vocabularies since the raw text contains a lot of code-switched data.
By contrast, ALP is a more convenient indicator of language-specific vocabulary capacity and it is penalized by the subword units with low-frequency.

To investigate the impact of language-specific vocabulary capacity, we first learn monolingual vocabularies in different sizes to obtain vocabularies with different ALP, i.e., language-specific vocabulary capacity.
Then we conduct pre-training with these monolingual vocabularies on their corresponding monolingual corpora.
Finally, we evaluate these monolingual models on downstream tasks and study the correlation between language-specific vocabulary capacity and downstream task performance.

\subsubsection{Setup}
To alleviate the bias from the languages' characteristics, we first select four languages with different pre-training corpus sizes from different language families, which are Hindi (hi), Persian (fa), Italian (it), Russian (ru).
We first learn thirty monolingual vocabularies for each language on the corresponding monolingual corpus, with vocabulary size ranging from 1K to 30K.
Then we pretrain monolingual language models with the corresponding monolingual vocabularies.
We evaluate these pretrained models on two downstream tasks:  NER~\cite{panx} and POS~\cite{udpos} from the XTREME benchmark since there is annotated task data for a large number of languages.
The vocabularies are learned on the reconstructed CommonCrawl corpus~\citep{infoxlm, xlmr} using SentencePiece~\citep{sentencepiece} with the unigram language model~\citep{DBLP:conf/acl/Kudo18}. The unigram distributions are also counted on the CommonCrawl corpus.
The Wikipedia corpus is used for all pre-training experiments in this paper since it is easier to run experiments due to its smaller size.
More details about the pre-training data can be found in the appendix.

\subsubsection{Observations}
\label{sec:obs}

\paragraph{Increasing vocabulary size affects ALP of different languages in varying degrees.}
In Figure~\ref{fig:alp}, we show the correlation between vocabulary size and ALP of four different languages.
We observe the ALP varies across different languages, mainly because ALP correlates with the lexicon granularity of the language, i.e., the average number of tokens per sentence.
Besides, when the vocabulary size is larger than 10,000, the gains of increasing monolingual vocabulary size in hi and fa are less than it and ru. We attribute it to that hi and fa does not have extensive compoundings.
Another observation is that for each language, every time we increase the vocabulary size by 1K, the increment in ALP is monotonically decreasing.

\paragraph{ALP correlates positively with downstream task performance.}
In Figure~\ref{fig:udpos} and Figure~\ref{fig:panx},
we illustrate downstream task performance of models pretrained with monolingual vocabularies on corresponding monolingual corpora.
We observe that ALP correlates positively with downstream task performance, making language-specific ALP a valid indicator to allocate multilingual vocabulary.
Another natural option to allocate multilingual vocabulary is directly using monolingual vocabulary size to indicate language-specific vocabulary capacity.
We compare ALP against vocabulary size and observe that ALP correlates better than vocabulary size with the downstream task performance. Besides, ALP reflects the language-specific characteristics, while vocabulary size does not. The detailed comparison is shown in the appendix.

\subsection{Allocating Multilingual Vocabulary with \ours{}
}

Based on the observations in Section~\ref{sec:obs}, we first give the implementation of our proposed vocabulary allocation algorithm \ours{}. Then we compare the multilingual vocabulary learned with \ours{} and directly learned with SentencePiece on the multilingual corpus.

\subsubsection{\ours{} Implementation}
\label{sec:vocap-imp}

\begin{algorithm}[t]
\small
\caption{\small Allocating Multilingual Vocabulary with \ours{}}
\label{alg:vocap}
\begin{algorithmic}[1]
\Require size of target multilingual vocabulary $T$; monolingual vocabularies of $N$ languages $\{V^{i}_{t_{i}}\}_{i=1}^{N}$; monolingual corpus of $N$ languages $\{D_i\}_{i=1}^{N}$
\Ensure multilingual vocabulary $V$
\For{$i \gets 1$ \textbf{to} $N$}
    \For{$j \gets 1$ \textbf{to} $50$}
        \State $a_{i,j \times 1000} \gets \textrm{ALP}(D_i, V^i_{j \times 1000})$
    \EndFor
    \State $t_{i} \gets 0$
    \State $a_{i,0} \gets -\infty$
\EndFor

\Do
    \State $j \gets 0$
    \State $\delta \gets 0$
    \For{$i \gets 1$ \textbf{to} $N$}
        \If{$\delta < a_{i, t_i+1000} - a_{i, t_i}$}
            \State $\delta \gets a_{i, t_i+1000} - a_{i, t_i}$
            \State $j \gets i$
        \EndIf
    \EndFor
    \State $t_j \gets t_j + 1000$
    \State $V \gets |\bigcup_{i=1}^{N}V^{i}_{t_i}|$
\doWhile{$|V| < T$}
\If{$|V| > T$}
\State Clip the size of $V$ to $T$
\EndIf
\end{algorithmic}
\end{algorithm}

We formulate the vocabulary construction of \ours{} as the problem of finding the optimal way to allocate language-specific vocabulary size to each language, such that the overall ALP of all languages is maximized.
In addition to language-specific vocabulary capacity measured with ALP from \eqform{eq:alp}, the language-specific pre-training corpus size also affects the downstream task performance.
Considering the two factors, the procedure of \ours{} can be formulated as follows:
\begin{align}
\underset{t_{1},...,t_{N}}{\mathrm{argmax}} \sum_{i=1}^{N}q_i^{\beta}\mathrm{ALP}(D_i, V^{i}_{t_i})
~~s.t.~~|\bigcup_{i=1}^{N}V^{i}_{t_i}| = T
\label{eq:vocap}
\end{align}
where $t_{i} \in \{x \times 1000 \mid x \leq 50, x \in N^+\}$ is the number of subword units allocated to the $i$-th language,\footnote{Since the cost of learning monolingual vocabularies with arbitrary sizes and getting the corresponding ALP is unaffordable, we learn monolingual vocabularies with vocabulary size range from 1K to 50K at intervals of 1K. } $\beta$ is a rescaling factor, $V^{i}_{t_{i}}$ is the vocabulary of the $i$-th language with $t_{i}$ subword units, $T$ is the size of the target multilingual vocabulary, and $q_{i}$ is the probability of sampling training instances from $i$-th language during pre-training~\citep{xlm, xlmr}:
\begin{align}
q_i = \frac{f_{i}^{\alpha}}{\sum_{j=1}^{N}f_{j}^{\alpha}}~~\text{with}~~f_i = \frac{n_i}{\sum_{k=1}^{N}n_k}
\label{eq:sampling}
\end{align}
where $n_i$ is the number of instances in the $i$-th language,
$\alpha$ is a rescaling factor used to alleviate the bias towards high-resource languages.
Since the increment in ALP when increasing the vocabulary size by a certain number is monotonically decreasing, \eqform{eq:vocap} can be solved with the greedy algorithm in Algorithm~\ref{alg:vocap}.

\subsubsection{Intrinsic Analysis}
We compare the multilingual vocabulary learned with \ours{} and directly learned with SentencePiece on the multilingual corpus.
The multilingual corpus to learn vocabularies in this paper is the concatenation of sentences sampled randomly from the monolingual corpora. Sentences from the $i$-th language is sampled with probability $q_i$ from \eqform{eq:sampling} and use $\alpha=0.7$. We filter languages with corpus size larger than 0.1 GB, resulting in 86 languages.

We evaluate the multilingual vocabularies with their ALP on each language's monolingual corpus, and show results of different-resourced languages in Figure~\ref{fig:res-alp}. We refer to languages with less than 1GB and more than 10GB pre-training corpus in the reconstructed CommonCrawl as low-resource and high-resource languages, respectively, otherwise mid-resource languages.
When directly learning vocabulary on the multilingual corpus using SentencePiece, the vocabulary with 500K subword units ($\joint{}_{\text{500K}}$) only has a negligible improvement compared to the vocabulary with 250K subword units ($\joint{}_{\text{250K}}$).
Meanwhile, our method ($\ours{}_{\text{500K}}$) consistently outperforms $\joint{}_{\text{500K}}$ in different-resourced languages, especially in mid and low-resource languages. The statistics of the allocated vocabulary size for each language in $\ours{}_{\text{500K}}$ are shown in the appendix.

\section{Accelerate Large-Vocabulary Language Model Pre-Training}

Although extending the multilingual vocabulary benefits cross-lingual language models, pre-training with such large vocabularies brings two practical issues: slow pre-training speed and heavy model size.
To tackle the issues, we first introduce our $k$-NN-based target sampling in Section~\ref{sec:kts}, which is a softmax approximation algorithm to improve computing efficiency.
Then we describe how we reallocate the model parameters to keep the model size fixed in Section~\ref{sec:mb}.

\subsection{$k$-NN-Based Target Sampling}
\label{sec:kts}
To reduce the expensive computation cost of the softmax function, we propose $k$-NN-based target sampling to approximate the expensive softmax.
The original masked language modeling objective minimizes the cross-entropy loss for every masked subword $w_i$ on the extensive multilingual vocabulary $V$.
The proposed $k$-NN-based target sampling instead uses a smaller vocabulary subset $V'$. The approximation of the masked language modeling loss for the masked subword $w_i$ is defined as follows:
\begin{align}
\mathcal{L}(w_i) =
-\mathrm{log}\frac{\exp(h^{\mathrm{T}}v_{w_i} + b_{w_{i}})}{\sum_{w_{j} \in V'}\exp(h^{\mathrm{T}}v_{w_j} + b_{w_j})}
\label{eq:loss}
\end{align}
where $h$ is the corresponding output vector of the penultimate network layer, i.e., the output vector of the Transformer encoder, $v_{w_{i}}$ is the embedding of the subword unit $w_{i}$, and $b_{w_{i}}$ is a bias term.
We formulate the construction of the vocabulary subset $V'$ as follows:
\begin{align}
V' &= \bigcup_{w_i \in \mathcal{W}} \mathcal{I}_{k}(w_i) \\
\mathcal{I}_{k}(w_i) &= \mathrm{top\text{-}k}(\{ v_{w_i}^{\mathrm{T}}v_{w_j} \mid w_j \in V \})
\label{eq:subvocab}
\end{align}
where $\mathcal{W}$ denotes the set of target masked subword units in the current mini-batch, and $\mathcal{I}_{k}(w_i)$ denotes the $k$ most similar subwords measured with the inner product of the subword embedding $v_{w_i}$ and $v_{w_j}$.

However, retrieving $\mathcal{I}_k(w_i)$ at every training step for every subword unit $w_i \in \mathcal{W}$ requires as much computation cost as softmax, which is unaffordable.
As an alternative, we compute $\mathcal{I}_k(w_i)$ for every subword $w_i \in V$ according to the current subword embeddings every $n$ training steps and replace the previous version of $\mathcal{I}_k(w_i)$ with the new one. We determine the value of $n$ such that $|V| \ll n \times |\mathcal{W}|$.
We illustrate the pre-training procedure with $k$-NN-based target sampling in Algorithm~\ref{alg:kts}.

\begin{algorithm}[t]
\small
\caption{\small Pre-training with $k$-NN-based target sampling}
\label{alg:kts}
\begin{algorithmic}[1]
\Require multilingual corpus $\mathcal{D}_\text{m}$; size $k$ of $k$-NN-based target sampling; multilingual vocabulary $V$; learning rate $\tau$
\Ensure model parameters ${\bm{\theta}}$

\While{not converged}
\State Sample $n$ mini-batches $\{\mathcal{X}^{(t)},  \mathcal{W}^{(t)}\}_{t=1}^{n}$ $\sim$ $\mathcal{D}_\text{m}$ \Comment{\textit{\color{gray} $\mathcal{X}^{(t)}$ is a mini-batch of monolingual text, and $\mathcal{W}^{(t)}$ is the set of masked subwords.}}
\State Update $\mathcal{I}_{k}(w_i)$ for every $w_i \in V$
\For{$t \gets 1$ \textbf{to} $n$} \Comment{\textit{\color{gray} Train the model for $n$ steps.}}
\State $V' \gets \bigcup_{w_i \in \mathcal{W}^{(t)}} \mathcal{I}_{k}(w_i)$
\State ${\bm{g}} \gets \sum_{w_i \in \mathcal{W}^{(t)}}\nabla_{\bm{\theta}} \mathcal{L}(w_i)$
\State ${\bm{\theta}} \gets {\bm{\theta}} - \tau {\bm{g}}$
\EndFor
\EndWhile
\end{algorithmic}
\end{algorithm}

From a practical point of view under the cross-lingual setting, the previous sampling-based softmax approximation methods either sample subwords from recent mini-batches or samples subwords from unigram distribution,
the task becomes simpler since a considerable part of the subword samples is from different languages.
Meanwhile, our $k$-NN-based target sampling uses subwords with similar representations like synonyms, which enforces the model focus on discriminating the ground-truth subword from a set of noise samples that are not easy to distinguish.
When using an approximate algorithm, the key point is to remain the difficult part of the original masked language modeling objective as much as possible.

\begin{table*}[t]
\centering
\small
\setlength{\tabcolsep}{1.0mm}
\begin{tabular}{lccccccccccc}
\toprule
\multirow{2}{*}{\bf Model} & \multirow{2}{*}{\bf \# Params} & \multirow{2}{*}{\bf Speed} & \multicolumn{2}{c}{\bf Pair Sentence} & \multicolumn{2}{c}{\bf Structure Prediction} & \multicolumn{3}{c}{\bf Question Answering} &  \\
& & & XNLI & PAWS-X & POS & NER & XQuAD & MLQA & TyDiQA \\ \midrule
 & & & Acc. & Acc. & F1 & F1 & F1/EM & F1/EM & F1/EM & Avg. \\ \midrule
$\text{XLM-R}_{\text{250K}}$ & 265M & 1.00x & 68.7 & 82.6 & 72.1 & 60.6 & 63.4/47.4 & 57.2/39.6 & 45.2/29.6 & 60.7  \\
$\joint{}_{\text{250K}}$ & 265M &1.00x & 69.2 & 83.3 & 72.4 & 59.7 & 63.9/47.9 & 58.9/40.7 & 45.4/29.6 & 61.1 \\
$\joint{}_{\text{500K}}$ & 448M &0.72x & 69.4 & 82.2 & 72.1 & 60.5 & 64.7/48.0 & 58.2/40.3 & 48.0/32.6 & 61.4 \\
$\ours{}_{\text{250K}}$ & 265M &1.00x & 69.3 & 82.0 & 71.4 & 60.0 & 66.2/50.3 & 60.1/42.6 & 45.6/30.6 & 61.5 \\
$\ours{}_{\text{500K}}$ & 448M &0.72x & 70.5 & 83.0 & \textbf{72.9} & \textbf{62.7} & 66.8/\textbf{50.6} & 60.9/\textbf{42.9} & 50.0/34.5 & 63.1 \\
~~+ $k$-NN & 448M &1.18x & \textbf{70.8} & 82.6 & 72.5 & 61.8 & \textbf{67.1}/49.8 & \textbf{61.4}/42.5 & \textbf{56.3/39.3} & \textbf{63.7} \\
~~+ half emb & 265M & 0.94x & 70.3 & 83.0 & 72.0 & 61.7 & 65.8/49.0 & 61.0/42.3 & 49.3/33.0 & 62.5 \\
~~+ $k$-NN \& half emb  & 265M & \textbf{1.35x} & 69.8 & \textbf{83.4} & 72.1 & 60.1 & 66.6/49.5 & 60.8/42.7 & 50.2/33.9 & 62.5 \\
\bottomrule
\end{tabular}
\caption{Evaluation results on the XTREME benchmark. ``$\text{XLM-R}_{\text{250K}}$'' denotes using the XLM-R~\citep{xlmr} vocabulary with 250K subword units. ``$k$-NN'' and ``half emb'' denote our $k$-NN-based target sampling method and using half embedding dimension, respectively. }
\label{table:benchmark-results}
\end{table*}

\subsection{Reducing the Embedding Dimension}
\label{sec:mb}

In order to keep the number of model parameters fixed while increasing the vocabulary size, we follow \cite{albert} and \cite{rembert} to reduce both the input and output embedding dimension and linearly project the embeddings to the hidden dimension of the Transformer blocks.
More precisely, we halve the embedding dimension when the vocabulary size is doubled.
This rebalancing strategy only slightly degrades the model performance but improves pre-training speed and decreases the model size.

\citet{xlmr} also studied the relation between the size of the shared multilingual vocabulary and downstream task performance with multilingual models of the fixed number of parameters.
They keep the overall number of parameters constant by adjusting the width (i.e., hidden size) of the Transformer.
Notice that we only reduce the embedding dimension while keeping the Transformer blocks untouched.

\section{Experiments}

\subsection{Setup}

\paragraph{Fine-Tuning Datasets}
To validate the effectiveness of our methods, we conduct experiments on three types of cross-lingual understanding tasks from XTREME benchmark~\cite{xtreme}, including two classification datasets: XNLI~\cite{xnli}, PAWS-X~\cite{pawsx}, three span extraction datasets: XQuAD~\cite{xquad}, MLQA~\cite{mlqa}, TyDiQA-GoldP~\cite{tydiqa}, and two sequence labeling datasets: NER~\cite{panx}, POS~\cite{udpos}.
The statistics of the datasets are shown in the appendix.

\paragraph{Implementation Details}
We adapt the Transformer architecture from the base model setting in \citet{xlmr}, i.e., 12 layers and 768 hidden dimension size.
We use masked language modeling objective to train our models for 1 million updates on eight 32GB Nvidia V100 GPUs with a batch size of 256.
We update the top-k indices for every word in the multilingual vocabulary every 1,000 training steps and use $k=50$ in $k$-NN-based target sampling.
The learning rate is scheduled with a polynomial decay with 10K warmup steps, where the peak learning rate is set as 0.0001. We adapt other hyper-parameters in pre-training from~\citet{infoxlm}.
All fine-tuning results are averaged over five random seeds.
The fine-tuning pipeline is based on the code base of~\cite{xtune}.
The fine-tuning implementation details are shown in the appendix.

\subsection{Results}
Table~\ref{table:benchmark-results} shows XTREME fine-tuning results with models pretrained using different vocabularies and acceleration strategies.
Compared to vocabulary directly learned on multilingual corpus with SentencePiece, i.e., $\text{XLM-R}_{\text{250K}}$ and $\joint{}_{\text{250K}}$, our $\ours{}_{\text{250K}}$ improves on question answering datasets but degrades on PAWS-X, POS and NER.
Then increasing the vocabulary from $\ours{}_{\text{250K}}$ to $\ours{}_{\text{500K}}$ mitigates the gap and bring improvements on six datasets except for PAWS-X, which only includes seven high-resource languages.
However, increasing the size of vocabulary directly learned with Sentencepiece from $\joint{}_{\text{250K}}$ to $\joint{}_{\text{500K}}$ does not improve the performance as our \ours{} method does, showing the importance of selecting language-specific subword units and leveraging how much vocabulary capacity each language requires.

Since increasing vocabulary size brings the issues of model size and pre-training speed, we study the proposed method to accelerate pre-training: $k$-NN-based target sampling ($k$-NN) and using half embedding dimension (half emb).
Our $k$-NN method improves pre-training speed with a 500K vocabulary so that the speed is 1.18 times that vanilla pre-training with a 250K vocabulary. Meanwhile, pre-training with our $k$-NN method does not significantly degrade the performance, it even brings improvement on XNLI, MLQA, and TyDiQA.
Then we halve the embedding dimension of the models with 500K vocabulary and results in a similar number of parameters to models with 250K vocabulary.
The overall performance degrades by 0.6-points but still consistently improves over models with 250K vocabularies while the speed is comparable.
Combining the two methods above, we achieve a 1.35-times speed-up and more than 1 point improvement with a similar model size compared to models with 250K vocabularies.

\subsection{Analysis and Discussion}
We conduct a thorough analysis to understand the impact of our proposed methods on cross-lingual language models.
To reduce the computation load, we only pre-train the cross-lingual language models for 500K steps for some of our settings.

\begin{table}[t]
\centering
\small
\begin{tabular}{lcccc}
\toprule
\textbf{Method} & \textbf{XNLI} & \textbf{POS} & \textbf{MLQA} & \textbf{Speed} \\ \midrule
$\ours{}_{\text{500K}}$ & 69.2 & \textbf{72.9} & \textbf{59.9/41.7} & 1.00x \\ \midrule
+ $k$-NN & \textbf{69.3} & 72.1 & 59.6/40.3 & \textbf{1.64x} \\
+ target sampling & 68.8 & 71.3 & 57.6/38.8 & 1.56x \\
+ NCE & 56.0 & 61.8 & 41.1/26.2 & 1.40x \\
+ NEG & 56.5 & 62.9 & 40.1/25.6 & 1.40x \\
\bottomrule
\end{tabular}
\caption{Comparison between different sampling-based softmax approximation approaches with vocabulary $\ours{}_{\text{500K}}$.  Models are pretrained for 0.5M steps.}
\label{table:sampling-results}
\end{table}

\begin{table}[t]
\centering
\small
\begin{tabular}{lcccc}
\toprule
\textbf{Method} & \textbf{XNLI} & \textbf{POS}  & \textbf{MLQA} & \textbf{Speed}\\ \midrule
$\ours{}_{\text{500K}}$ & 69.2 & 71.8 & 59.9/\textbf{41.7} & 1.00x \\ \midrule
+ $k$-NN ($k$=5) & 68.5 & 71.3 & 58.6/40.0 & \textbf{1.76x} \\
+ $k$-NN ($k$=10) & \textbf{69.3} & 71.4 & 58.9/39.6 & 1.74x \\
+ $k$-NN ($k$=25) & 69.2 & 71.7 & 59.8/40.9 & 1.69x \\
+ $k$-NN ($k$=50) & \textbf{69.3} & \textbf{72.1} & 59.6/40.3 & 1.64x \\
+ $k$-NN ($k$=100) & 69.5 & \textbf{72.1} & \textbf{60.0}/41.3 & 1.57x \\
\bottomrule
\end{tabular}
\caption{Comparison between different $k$ values in $k$-NN-based sampling method. Models are pretrained for 0.5M steps.}
\label{table:k-value}
\end{table}

\paragraph{$k$-NN-based target sampling outperforms previous sampling-based approaches.}
To verify the effectiveness of our proposed $k$-NN-based sampling method, we compare it against previous sampling-based approaches used to approximate softmax, which are target sampling~\cite{DBLP:conf/acl/JeanCMB15}, noise contrastive estimation (\citet{nce}, NCE) and negative sampling (\citet{word2vec}, NEG). The results are shown in Table~\ref{table:sampling-results}.
To make a fair comparison, since our $k$-NN-based sampling method using $k=50$ samples vocabulary subset with less than 50,000 subword units per batch on average, we here sample 50,000 negative subword units per batch for target sampling, NCE, and NEG.
Among the four methods, NCE and NEG are significantly worse than $k$-NN and target sampling. We attribute it that NCE and NEG need more training steps to converge~\citep{nce}.
Besides, the original NCE typically sample different negative samples for every target word, while we here use 50,000 negative samples for all target word in current mini-batch, which is more efficient on GPUs.

\paragraph{Effect of the value of $k$ in $k$-NN-based target sampling.}
We illustrate the downstream task performance when using different values of $k$ in our $k$-NN-based target sampling in Table~\ref{table:k-value}.
While a smaller $k$ indicates faster pre-training speed, we observe even with a small value like 5, the result does not significantly degrade compared to using the original softmax.
We attribute this to that by retrieving subword samples that are most similar to the target subword, the model can focus on the difficult part of the original masked language modeling objective.
More precisely, the model focus on discriminating the ground-truth subword from a set of noise samples that are not easy to distinguish.
Considering the overall performance, the pre-training speed, and running memory to store $k$-NN indices, we use $k=50$ in all our experiments.

\begin{table}[t]
\centering
\small
\begin{tabular}{lcccc}
\toprule
\textbf{Method} & \textbf{XNLI} & \textbf{POS} & \textbf{NER} & \textbf{MLQA}\\ \midrule
$\beta$=0 & 66.9 & \textbf{71.8} & 61.5 & 58.6/41.0 \\
$\beta$=0.3 & 69.0 & 71.7 & \textbf{61.6} & 59.2/40.1 \\
$\beta$=0.7 & 69.2 & \textbf{71.8} & 61.5 & \textbf{59.9/41.7} \\
$\beta$=1.0 & \textbf{69.5} & \textbf{71.8} & 60.9 & 58.4/40.3 \\
\bottomrule
\end{tabular}
\caption{Impact of adjusting high-resource versus low-resource vocabulary capacity trade-off with $\beta$.
$\beta=0$ indicates the vocabulary is allocated without considering pre-training corpus size. Models are pretrained for 0.5M steps.}
\label{table:beta}
\end{table}

\paragraph{Language-specific pre-training corpus should also be considered when allocating vocabulary capacity.}
The pre-training corpus size varies across different languages.
It is inefficient to allocate a large vocabulary capacity for low-resource languages with rare pre-training data since the pretrained model can only learn limited knowledge from these languages.
Here we study the value of rescaling factor $\beta$ from \eqform{eq:vocap} in multilingual vocabulary construction in Table~\ref{table:beta}.
The rescaling factor $\beta$ controls the number of selected language-specific subword units.
Increasing the value of $\beta$ improves the performance of XNLI, where most languages are high-resource languages.
However, it degrades the performance of NER, where more low-resources languages exist.
When considering overall performance, we decide to use $\beta=0.7$ in our experiments.

\begin{figure}[t]
\centering
\includegraphics[trim={2.4cm 0cm 4cm 1.4cm},scale=0.175,clip]{time_performance.pdf}
\caption{Performance on XNLI and MLQA versus the cross-lingual language models' pre-training cost.}
\label{fig:time-performance}
\end{figure}

\paragraph{The proposed acceleration strategies significantly improve the downstream task performance under the same pre-training cost.}
Increasing the vocabulary size slows the pre-training speed, even though there is almost no difference in fine-tuning speed.
We study the relationship between the downstream task performance and the pre-training cost under different model settings in Figure~\ref{fig:time-performance}.
We observe $\ours{}_{\text{500K}}$+$k$-NN achieves the best performance.
Models trained with 500K vocabulary consistently outperform 250K vocabulary on XNLI.
Besides, we observe the performance on MLQA with the model trained using 250K vocabulary degrades as the training continues while models trained using 500K vocabulary does not, indicating the sufficient vocabulary capacity is essential for question answering task.

\paragraph{\ours{} gains more improvement on mid and low-resource languages than high-resource languages.
}
In Figure~\ref{fig:res-alp} in Section~\ref{sec:vocap}, we show that the vocabulary learned with \ours{} benefits the vocabulary capacity of low-resource languages more than high-resource languages, indicating the improvements should mainly come from low-resource languages.
To verify this, we compare \ours{} against SentencePiece baseline on the performance of different-resourced languages on XNLI and NER in Figure~\ref{fig:res-performance}.
We observe that the vocabulary learned with \ours{} significantly outperforms the vocabularies directly learned with SentencePiece on mid and low-resource languages. This observation is also consistent with the ALP results in Figure~\ref{fig:res-alp}.

\begin{figure}[t]
\centering
\includegraphics[trim={1.2cm 0cm 2cm 0.7cm},scale=0.35,clip]{resource_performance.pdf}
\caption{Impact of \ours{} on the performance of different-resourced languages on XNLI and NER.}
\label{fig:res-performance}
\end{figure}

\section{Related Work}

\paragraph{Pretrained Cross-Lingual Language Models}
Recent work pre-trains Transformer models~\citep{transformer} on the large-scale multilingual corpus to obtain pretrained cross-lingual language models~\citep{xlm,xlmr,xnlg,mt6,infoxlm,xlmalign,xlme,rembert,mt5,xlmt,deltalm}.
These models are capable of encoding texts from different languages into universal representations and significantly improves cross-lingual transferability.

\paragraph{Multilingual Vocabulary Construction}
Cross-lingual language models need large vocabularies to ensure all languages are adequately represented.
Recent research work on constructing multilingual vocabulary for cross-lingual language models can be categorized into two groups. mBERT~\citep{bert}, XLM~\citep{xlm}, and XLM-R~\citep{xlmr} learn vocabularies on a combined multilingual corpus with WordPiece~\citep{wordpiece}, BPE~\citep{bpe}, and unigram language model~\citep{DBLP:conf/acl/Kudo18} from SentencePiece~\citep{sentencepiece}, respectively.
\citet{DBLP:conf/emnlp/ChungGTR20} propose to balance the trade-off between optimizing for cross-lingual subword sharing and the need for robust representation of individual languages. They first group languages into clusters and learn vocabularies individually on each cluster, then combine all cluster-vocabularies to form a single unified multilingual vocabulary.
Compared to \citet{DBLP:conf/emnlp/ChungGTR20}, our advantage is that we separately quantify the vocabulary capacity each language needs with average log probability and balance the construction procedure with pre-training corpus size.

\paragraph{Softmax Approximation}
Approximating the softmax was a core problem in training NLP tasks with a large vocabulary, e.g., neural machine translation, language modeling.
With the rise of subword representations~\citep{bpe,wordpiece,DBLP:conf/acl/Kudo18}, the vocabulary size significantly decreases, and the problem has been less studied recently.
Nevertheless, the need for training cross-lingual language models with a large multilingual vocabulary has drawn our attention again to the softmax approximation approaches.
The existing softmax approximation approaches can be grouped into softmax-based and sampling-based approaches. Softmax-based approaches includes hierarchical softmax~\citep{h-softmax}, differentiated softmax~\citep{d-softmax}, and CNN-softmax~\citep{cnn-softmax}.
However, these approaches improve the softmax efficiency by changing its architecture, which is unsuitable for either training on GPUs or multilingual settings.
Sampling-based approaches instead optimize some other easy-to-compute loss function to approximate the original softmax, including target sampling~\citep{DBLP:conf/acl/JeanCMB15}, noise contrastive estimation~\citep{nce}, negative sampling~\citep{word2vec}.
Our $k$-NN-based target sampling is also a sampling-based approach.

\section{Conclusion}

In this paper, we study pre-training cross-lingual language models with large vocabulary capacity. First, we propose \ours{} to construct large multilingual vocabulary in cross-lingual language models.
We conduct a quantitative analysis to show that average log probability is an valid indicator of vocabulary capacity for a particular language, which also correlates with downstream task performance on the language.
\ours{} uses the language-specific average log probability and pre-training corpus size to allocate appropriate vocabulary capacity for each language in the multilingual vocabulary.
Moreover, we propose $k$-NN-based target sampling to accelerate pre-training with the allocated large multilingual vocabulary by approximating the expensive softmax.
We also show that reducing the embedding dimension is an effective way to keep the improvement brought by the large vocabulary without increasing the number of model parameters.
The experiments demonstrate the effectiveness of the proposed vocabulary construction method as well as the acceleration methods.

\bibliographystyle{acl_natbib}