abstract: Heated debates continue over the best autonomous driving framework. The classic modular pipeline is widely adopted in the industry owing to its great interpretability and stability, whereas the end-to-end paradigm has demonstrated considerable simplicity and learnability along with the rise of deep learning. We introduce a new modularized end-to-end reinforcement learning framework (ModEL) for autonomous driving, which combines the merits of both previous approaches. The autonomous driving stack of ModEL decomposes into perception, planning, and control module, leveraging scene understanding, end-to-end reinforcement learning, and PID control respectively. Furthermore, we build a fully functional autonomous vehicle to deploy this framework. Through extensive simulation and real-world experiments, our framework has shown great generalizability to various complicated scenarios and outperforms the competing baselines.[^1]
author: Guan Wang$^{1\dag}$, Haoyi Niu$^{1\dag}$, Desheng Zhu$^{2}$, Jianming Hu$^{3}$, Xianyuan Zhan$^{3}\textsuperscript{\Letter}$ and Guyue Zhou$^{3}$ [^2] [^3] [^4] [^5]
bibliography: mylib.bib
title:  **ModEL: A Modularized End-to-end Reinforcement Learning Framework for Autonomous Driving**

# Introduction

The past decade has witnessed a surge of research interests in end-to-end autonomous driving systems. Despite its appealing simplicity and learnability, most attempts have shown limited performance or adaptability in real-world scenarios due to the large visual and dynamics gaps between simulation and real-world environment. These unsolvable issues remind us of the conventional modular pipeline that divides the system into modules for error-tracking and enables autonomous vehicles to behave predictably. However, this approach also leads to a considerable amount of human engineering in devising complicated rules and model fine-tuning. This naturally motivates us to consider a modularized end-to-end architecture , i.e. decomposing into a scene understanding module, an end-to-end driving planning module trained by imitation learning (IL) or reinforcement learning (RL), and a low-level controller.

Previous works on developing end-to-end autonomous driving systems are primarily using IL for driving planning or control . Learning an IL policy does not need a high-fidelity simulator, but involves arduous real driving data collection. It is also struggling to outstand human experts and has limited adaptability in diverse real-world scenarios. The huge success of RL has provided another appealing direction. RL has achieved superhuman performance in solving complex tasks in recent years, and also begins to show its potential on safety-critical tasks like autonomous driving . RL can alleviate some major headaches of IL, such as distribution shift problem , data bias , and causal confusion . Despite the strong capability of RL, developing a real-world end-to-end autonomous driving system using RL remains to be a highly challenging task. Considerable efforts need to be made to properly address the sim-to-real gap involved during the interactive training of RL policies with a simulator.

In this paper, we present a deployed **Mod**ularized **E**nd-to-end Reinforcement **L**earning Framework (**ModEL**) for autonomous driving. The architecture of ModEL is modularized into three parts: 1) **Perception Module** copes with the raw input from a monocular camera and uses advanced scene understanding models to perform drivable space and lane boundary identification; 2) **Planning Module** maps the scene understanding outputs to driving decisions via a trained RL policy (Soft Actor-Critic (SAC) , the state-of-the-art method of model-free RL); 3) **Control Module** uses mature PID controllers to regulate and control the low-level driving commands given the high-level planning decisions. The modularized framework of ModEL offers a nice solution to address all the aforementioned challenges. First, the decomposition of perception, planning, and control modules reduces the complexity of each individual task while improving the stability and interpretability of the system. Second, the RL algorithm is learned with post-perception outputs together with domain generalization techniques (domain randomization and data augmentation), which greatly reduces the impact of the sim-to-real issue, while fully exploiting the strength of RL policy learning.

In addition to the introduction of the ModEL framework, this study makes a major contribution by building a real autonomous vehicle that deploys ModEL to support real-world validation. We also design a new distributed training scheme that greatly accelerates the RL training and allows "learning to drive in less than one day". We examine our method against multiple baselines in CARLA simulator with intentionally introduced visual and dynamics gaps. In the real-world validation experiments, our system proves to be capable of generalizing to diverse complicated scenarios with varying road topology and lighting conditions, as well as the presence of obstacles. Through systematic simulation and real-world evaluation, ModEL demonstrates superior performance that exceeds all the competing baselines.

# Related Work

## Modular VS End-to-End Autonomous Driving

Currently, mainstream architectures for autonomous driving are structured in either a modular or an end-to-end way.

The conventional **modular pipeline** comprises a multitude of sub-systems, such as environment perception, localization, behavior prediction, planning, and control. In order to handle as many real-world scenarios as possible, the interoperation of these highly functional modules relies on human-engineered deterministic rules. The first in-depth discussions and analyses of the low-level software that encapsulates these rules are in the 2005 DARPA Grand Challenge and 2007 DARPA Urban Challenge, which form a basis for many later studies . The modular pipeline is widely adopted in the industry due to its remarkable interpretability. It endows the overall system with the ability to track down and locate sub-system malfunctions. Nevertheless, modular formalism bears several major drawbacks. Large amounts of human engineering are involved to fine-tune individual and cross-module configurations. The framework also suffers from severe compounding errors when decomposing the whole system into lots of smaller-scale but interpretable modules, as the perception uncertainty and modeling errors could be amplified and passed down through the pipeline.

Alternatively, **end-to-end** autonomous driving has been increasingly acknowledged. It is widely accepted that human driving is a behavior reflex task that requires little high-level reasoning and conscious attention, which harbors the same view with the end-to-end architecture. As a deep learning-based solution, the task-specific end-to-end learning ability brings a great reduction of human engineering efforts. However, many studies recognize the drawback of lack of interpretability in the development of end-to-end systems due to its black-box learning scheme. To alleviate this limitation, Chen *et al.* combine the probabilistic graphical modeling with RL to improve the interpretability for autonomous driving in simulation scenarios. Despite these efforts, weak interpretability remains a fatal disadvantage of the end-to-end approaches against the modular architecture.

Based on the above comparisons, there is a growing consensus among researchers that combining the merits of both modular and end-to-end formalism could be a promising solution. The raw perceptual inputs are processed using semantic segmentation methods, the intermediate outputs are then used for waypoints planning or control decision learning in an end-to-end manner. This combined framework alleviates the drawbacks of both approaches and offers enhanced learnability, robustness, and even transferability.

## Imitation Learning VS Reinforcement Learning

In end-to-end framework, **imitation learning** (IL) has been identified as a practical paradigm for autonomous driving. It uses a supervised learning scheme that imitates the human experts' demonstrations using algorithms like behavior cloning (BC) . IL is known to have several major disadvantages . First, since IL performs supervised learning on training dataset, when facing unseen scenarios during closed-loop testing, serious *distribution shift* problems take place and the agent has no idea what to do. Second, *data bias* deeply hurts the generalizability of IL policy, as the training process pays little attention to rare and risky scenarios, namely the long-tailed problem in self-driving. Finally, *causal confusion* is another problem that occurs in IL, as it performs pure data fitting and handles spurious correlations in data badly. Due to the above limitations, typical IL models can succeed in simple tasks like lane following but underperform in more complicated and rare traffic events. To improve the performance in harder scenarios, conditional imitation learning (CIL) introduces a latent state to fully explain the data, thus resulting in a better model expressiveness. However, due to the lack of the ability to perform long-term predictive "reasoning", it still has some safety issues during closed-loop testing.

Due to the ability to solve complex tasks as well as perform long-term optimization, **reinforcement learning** (RL) has become another popular choice for investigating end-to-end autonomous driving . RL learns how to map observations to optimized actions by maximizing the expected cumulative reward, and can learn strong policies in a simulator without real-world labels . Although RL has lower data efficiency than IL, it could be easily implemented in some high-fidelity simulators, e.g. CARLA, where agents can explore in more diverse scenarios. This interactive learning ability resolves the distribution shift, data bias, and causal confusion issues in IL. However, RL is also much harder to train in hard tasks. A modularized framework with a dedicated perception module that performs scene understanding can effectively ease the burden of RL, and also lead to more stable policies. Furthermore, due to the interactive learning nature of RL and its reliance on a simulator, considerable efforts need to be taken to properly address the sim-to-real gap issue in a deployable autonomous driving system.

# Detailed Implementation

The overall architecture of ModEL is described in Fig. , consisting of three modules: **(1) Perception Module**: we modify the advanced autonomous vehicle semantic segmentation model DeepLabv3+ with EfficientNet-B0 backbone to perform drivable space and lane boundary estimation based on the input image $x$ from the raw monocular camera. **(2) Planning Module**: using the learned representation $z$ from perception, we use state-of-the-art reinforcement learning algorithm, SAC , with specially designed actor, critic network structure to learn the optimized high-level vehicle control actions $\bm{a}$, consisting of throttle $\tau$ and steering angle $\theta$. Moreover, domain generalization methods are applied to improve generalization. **(3) Control Module**: lastly, the control module maps high-level actions $\bm{a}$ to low-level vehicle control commands according to the vehicle physical properties using PID controllers.

## Hardware Setup

Our autonomous vehicle is built upon an Agile.X HUNTER Unmanned Ground Vehicle (UGV) with an onboard computer (i7-9700 CPU, 32GB RAM, GeForce RTX 3060 GPU) and a front monocular RGB camera. The onboard computer runs all system modules based on Robot Operating System (ROS) , connected with UGV hardware by Controller Area Network (CAN). Additionally, as all modules are required to run in real-time, we utilize quantization and computational graph optimization techniques to improve model inference latency. The whole system has end-to-end latency (from sensor data collection to low-level control output) of 40ms and runs with a 100ms control interval. The setup of the physical system is detailed in the Appendix.

## Perception Module

The perception module acts as a powerful supervised state-encoder that maps camera captured RGB images from simulation $x_{sim}$ or real world $x_{real}$ RGB images to a semantically meaningful representation $z$, effectively bridging visual gap between simulation and real world. Compared to unsupervised state-encoders like image translation networks adopted by previous work , using supervised computer vision approaches yields human-readable intermediate representations, which contains cleaner and more precise information needed for the planning module. It is also less impacted by the visual gap, thus more reliable.

We choose the results of drivable area segmentation and lane boundary identification as the intermediate representation $z$, as it is an effective way to represent road situations from monocular RGB images and is widely used in autonomous driving systems. It labels every pixel as drivable (lane or area currently occupied by ego-vehicle), alternatively drivable (requires a lane-change), or non-drivable (blocked by obstacles). This scheme provides the vehicle with obstacles, lane boundaries, as well as lane change information.

We build our semantic segmentation model based on DeepLabv3+. It utilizes atrous convolution to enlarge receptive field and capture multi-scale context, along with encoder-decoder architecture for sharper fine-grained details. It has state-of-the-art performance on autonomous vehicle semantic segmentation datasets. As the ResNet101 backbone in DeepLabv3+ is way too heavy, which is not efficient enough to be deployed in our vehicle, we utilize a new backbone network produced by recent neural architecture search techniques, specifically, the EfficientNet-B0 . Our modification achieves superior accuracy over original DeepLabv3+ and ordinary segmentation models, as shown in the results of Table . The segmentation model is trained on the open real-world dataset BDD100k . Our model also has great computational efficiency, capable of running at 31 frames per second on the onboard computer.

|           Method           | mIoU/% (validation) | Inference time (ms) |
|:--------------------------:|:-------------------:|:-------------------:|
|      UNet (ResNet101)      |        87.2         |         64          |
|   DeepLabv3+ (ResNet101)   |        92.5         |         50          |
| **Ours (EfficientNet-B0)** |      **93.4**       |       **32**        |

Perception Model Comparison

## Planning Module

Our planning module is trained in the simulator CARLA via reinforcement learning but with specially designed domain generalization scheme to narrow the sim-to-real generalization gap. It chooses best high-level control actions (throttle and steering angle) according to the intermediate representation from the perception module.

### Reinforcement learning

We frame our autonomous vehicle planning task as a Markov Decision Process (MDP) problem and solve it via RL. The problem is represented as a tuple $(S, A, r, T, \gamma)$, where $S$ and $A$ denote the state and action set, $r(s,a)$ is the reward function, $T(s'|s,a)$ denotes the transition dynamics, and $\gamma\in (0,1)$ is the discount factor. In our problem, $S$, $A$ and $r$ are defined as follows:

**States $S$:** as illustrated in Figure , we use the drivable area segmentation image $z$ outputted from the perception module, concatenated with the vehicle speed $v$ as the state.

**Actions $A$:** We consider the steering angle $\theta$ and throttle $\tau$ of vehicle as actions $a = [\theta, \tau]$, where $\theta, \tau \in [-1, 1]$.

**Reward function $r$:** The design of our reward function involves four parts, including speed control $r_{speed}$, lane center keeping $r_{center}$, heading direction alignment $r_{heading}$ as well as collision and undesired lane crossing penalty $r_{penalty}$. For simplicity, we omit the state-action inputs $(s,a)$ in the reward function and describe each part as follows:

-   *Speed control:* Instruct the vehicle with current speed $v$ to drive in the desired speed range, $[v_{min}, v_{target}]$. If it drives too slow, or too fast, the reward will decay linearly. $v_{max}$ indicates maximum possible speed. $$\begin{aligned}
    r_{speed} &= \begin{cases}
    \frac{v}{v_{min}} & \quad\text{if}\quad v < v_{min} \\
    1 & \quad\text{if}\quad v_{min} \leqslant v \leqslant v_{target} \\
    1-\frac{v - v_{target}}{v_{max} - v_{target}} & \quad\text{if}\quad v > v_{target}
    \end{cases}
    \end{aligned}$$

-   *Lane center keeping:* Instruct the vehicle to drive in the center of a lane. $d$ is the current distance from vehicle to lane center and $d_{max}$ is the maximum in-lane distance. $$\begin{aligned}
    r_{center} = \mathrm{Clip} \left (1 - \frac{d}{d_{max}}, 0, 1\right )
    \end{aligned}$$

-   *Heading direction:* Instruct the vehicle to drive aligned with a lane. $\alpha$ is the heading angle difference between the vehicle and the lane, with the maximum value $\alpha_{max}$. $$\begin{aligned}
    r_{heading} &= \mathrm{Clip}\left(1 - \frac{\alpha}{\alpha_{max}}, 0, 1\right)
    \end{aligned}$$

-   *Collision and undesired lane crossing penalty:* Finally, we define the collision and undesired lane crossing penalty $r_{penalty}$ as follows, where $\mathbf{I}(\cdot)$ is the incident indicator: $$\begin{aligned}
    r_{penalty} = &25 \times \mathbf{I}(\mathrm{collision}) + 12 \times \mathbf{I}(\text{cross solid line}) \\ & + 15 \times \mathbf{I}(\text{cross double solid line})
    \end{aligned}$$

The total reward is the multiplication of $r_{speed}, r_{center},$ $r_{heading}$ and subtracted by collision and undesired lane crossing penalty. The multiplication of the first three objectives is to enforce a soft binary AND logic , expecting the agent to do well in all these aspects. The final reward is given as: $$\begin{aligned}
r = r_{speed} \cdot r_{center} \cdot r_{heading} - r_{penalty}
\end{aligned}$$

We solve this MDP problem using the Soft Actor-Critic (SAC) algorithm, which is a state-of-the-art model-free RL algorithm for continuous action control. As an actor-critic algorithm, SAC learns an actor policy $\pi_\phi (s)$ and an ensemble of critic Q-functions $Q^\pi_{\psi_i}, i=1,2$. SAC is derived from the maximum entropy reinforcement learning, which encourages stronger exploration thus leads to better performance. It aims to find an optimal policy that maximizes the following objective: $$J_\pi(\phi) = \mathbb{E}_s\mathbb{E}_{a\sim \pi_\phi}[Q_\psi^\pi(s,a)-\lambda \log \pi_\phi(s)]$$ SAC simultaneously trains the critics $Q^\pi_{\psi_1}$, $Q^\pi_{\psi_2}$ to minimize the soft Bellman residual: $$\begin{aligned}
J_Q(\psi_i) &= \mathbb{E}_{s,a}\left[Q_{\psi_i}^\pi(s,a) - (r+\gamma Q_{target}) \right]^2 \\
Q_{target} &= \min_{i=1,2} [Q_{\psi_i}^\pi(s',a') - \lambda \log \pi_\phi(s')]
\end{aligned}$$

To further improve the function approximation capability and model expressiveness, we design a light-weighted ResNet architecture for both actor and critic networks (see Fig. ). As our RL task involves understanding complex image-based state information, a more expressive network structure is critical for enhancing the RL agent's generalization ability . Previous deep network architectures used in image-based RL tasks typically place over 90% parameters on fully-connected layers. We utilize global average pooling to reduce the fully-connected parameters' size to less than 10%, placing most parameters on more trainable residual layers, thus improving learning efficiency and generalization.

### Domain Generalization

It is known that RL is prone to overfit the training environment and generalizes poorly to unseen scenarios . As the RL policy learned in a simulator needs to be deployed to unseen real-world scenarios, we focus on improving its generalizability by introducing various domain generalization methods during training.

We randomize the configurations in the simulation environment every epoch, to force the agent to adapt to environments with diverse properties, and also avoid overfitting on certain configurations. At sensor level, camera position is randomly chosen from a pre-specified range, to reflect sensor installation errors in real world. Vehicle physical properties (size, mass, etc.) are also randomly chosen from a pre-determined set to reduce the sensitivity to vehicle dynamics. Last but not the least, we apply data augmentations to agent observations, such as rotating the observed image, which contributes a lot to data efficiency and generalization .

### RL Training Acceleration

Training an RL agent to solve complex tasks like autonomous driving involves extensive amount of computation. We develop a new distributed framework, which enables highly efficient training of off-policy RL algorithms. Compared to current distributed RL frameworks that synchronize the execution of training procedures, our framework separates environment, actor and optimizer as individual nodes, exchange messages among nodes, while executing asynchronously as fast as possible. Detailed information about the distributed computation framework is discussed in the Appendix.

With our distributed acceleration techniques, we can successfully train the whole system in less than one day on a workstation with 2 GPUs. The training times of different modules are reported in Table .

| ModEL  | Perception  | Planning (RL)  |
|:------:|:-----------:|:--------------:|
| Time/h |     5.1     |      15.0      |

Wall-clock Training Time of Modules

## Low-level Control Module

On our test autonomous vehicle, we first apply a low-pass filter on high-level control actions $\theta, \tau$ from the planning module to smooth out noisy signals. The outputs are then mapped to vehicle wheel speed $\omega_1$ and steering servo angle $\omega_2$ based on physical properties of the vehicle. Finally, we utilize two PID controllers to follow low-level control targets.

# Evaluation

In this section, we present the baselines, evaluation metrics information and detailed experiment results quantitatively.

## Experiment Settings

### Baselines

We consider three baselines in this study, including modularized imitation learning (IL) , end-to-end IL and RL methods. We re-implement the IL and RL modules of these baselines in our framework for a thorough comparison.

For the training of all IL-based baselines, we follow the treatment described in. We collect 28h of driving data using a privileged modular pipeline as expert (has access to ground-truth map and obstacle information) provided by CARLA , and add noise to 20% of the expert control outputs to improve the robustness of the learned policy. For RL-based baselines, all environment and agent configurations are set the same as our proposed approach.

### Evaluation Metrics

For performance evaluation, we use the following metrics:

-   **MPI (m)**: meters per intervention, indicating the level of autonomy of the vehicle, which is widely adopted in real-world autonomous vehicle testing. Interventions are performed when collisions happen or the vehicle doesn't move in more than one minute.

-   **SR (%)**: success rate, referring to the proportion of travelled distance from the start to the first intervention with respect to the whole journey in one trial.

-   **Std\[$\theta$\] ()**: the standard deviation of the steering angle, reflecting the lateral smoothness of the trajectory.

-   **Std\[$v$\] (m/s)**: the standard deviation of the velocity, revealing the longitudinal smoothness of the trajectory.

## Comparison in Simulation

The main challenge of deploying a simulator-trained system to real world is the visual, vehicle dynamics and scenario gap. Previous simulation-based autonomous driving benchmarks, such as *NoCrash* and CARLA benchmark , do not reflect such challenges effectively. For example, the training and test scenarios are visually similar, and not suited well for comprehensive real-world deployment validation.

We propose a new *NoGap* benchmark to measure the sim-to-real gaps explicitly and provide better real-world generalization evaluations. In *NoGap* benchmark. the autonomous driving system is trained in "simulator" and then tested in "real-world" (simulator with different configurations). The evaluation settings are described as follows:

-   **Visual Gap**: We use different simulator rendering modes to create intentionally introduced visual gaps. During training, low-quality coarse rendering is used, producing game-like images. While in testing, the simulator is switched to photo-realistic rendering, producing images that look like real-world. Please refer to Fig. for examples.

-   **Dynamics Gap**: Vehicle physical properties are randomly chosen at the beginning of every testing episode, to validate generalization between different vehicle dynamics.

-   **Scenario Gap**: 7 CARLA maps are used for training, and 1 reserved for testing, to reflect scenario difference.

We use the MPI metric in simulation evaluations. Inspired by , the environment is reset to move the vehicle to a safe state after an intervention for precise intervention counting. Moreover, background vehicles controlled by modular pipeline in CARLA are generated to emulate real-world traffic and dynamic obstacles.

As shown in Table , end-to-end frameworks fail to generalize from training to testing scenarios due to large visual gaps in RGB inputs. Furthermore, training RL with RGB images is also a much harder task because of too much information contained, which has poorer performance compared with end-to-end IL. Modularized frameworks generalize well under visual gaps and different maps. They also have higher performance in training because of better state representation. Due to stronger model capability and better exploration than IL, modularized RL method achieves the best performance with high degree of generalization.

|         Framework         |   Input    | Planning |    MPI    |           |
|:-------------------------:|:----------:|:--------:|:---------:|:---------:|
|            4-5            |            |          |   Train   |   Test    |
|       End-to-end IL       |    RGB     |    IL    |   134.6   |   16.2    |
|       End-to-end RL       |    RGB     |    RL    |   24.9    |    0.4    |
|      Modularized IL       | Perception |    IL    |   306.6   |   180.9   |
| **Modularized RL (Ours)** | Perception |    RL    | **449.4** | **332.6** |

Evaluation results in simulation, metrics are averaged over 50 runs.

```latex
\begin{table*}[t]\centering
\caption{Real world evaluation ($>$ means no intervention over the whole trajectory with specified total length.)}
\begin{tabular}{c|c|c||>{\columncolor{gray!40}}c>{\columncolor{gray!30}}c>{\columncolor{gray!20}}c>{\columncolor{gray!10}}c|>{\columncolor{gray!40}}c>{\columncolor{gray!30}}c>{\columncolor{gray!20}}c>{\columncolor{gray!10}}c}
\toprule

\rowcolor{white}
 \multicolumn{3}{c||}{\multirow{2}*{\textbf{Real World Task}}} & \multicolumn{8}{c}{\multirow{2}*{\textbf{Framework}}} \\
 \multicolumn{3}{c||}{}  & \multicolumn{8}{c}{}\\%\cmidrule{4-11}

 \rowcolor{White}
Road & Obstacle & Lighting & \multicolumn{4}{c|}{\textbf{Modularized IL}  (MIL \cite{Muller2018})} & \multicolumn{4}{c}{\textbf{Modularized RL} (ModEL)} \\

Topology  & Setup  & Condition  & MPI ($m$) & SR ($\%$) & Std[$\theta$] ($^\circ$) & Std[$v$] ($m \cdot s^{-1}$) & MPI ($m$) & SR ($\%$) & Std[$\theta$] ($^\circ$) & Std[$v$] ($m \cdot s^{-1}$) \\ \midrule\midrule

\multirow{3}*{Straight} & \ding{55}& Day & 92.1 & 48.9 & \textbf{1.30} & 0.30 & $\bm{>}$\textbf{1163.5} & \textbf{100.0} & 1.72 & \textbf{0.19}\\

 & \ding{55}& Night & 187.1 & 49.0 & \textbf{1.67} & \textbf{0.18} & $\bm{>}$\textbf{1304.5} & \textbf{100.0} & 1.89 & 0.20 \\  %\cline{3-6}

 & \ding{51}& Day & 4.1 & 16.7 & \textbf{1.25} & 0.37 & \textbf{34.5} & \textbf{75.0} & 2.59 & \textbf{0.30} \\ \midrule

Turn     & \ding{55}& Day & 7.2 & 53.2 & \textbf{3.03} & 0.32 & $\bm{>}$\textbf{214.9} & \textbf{100.0} & 3.81 & \textbf{0.23}\\ \bottomrule
\end{tabular}
\label{Table4}
\end{table*}
```

## Real World Benchmark

To evaluate the autonomous driving framework on our test vehicle described in Section , we categorize the real world tasks based on different attributes: road topology (straight and turn), obstacles setup (with obstacle ahead), and lighting condition (day-time and night-time). We test our ModEL framework with the most competitive baseline standing out from simulation evaluation in Section , the modularized IL (MIL) , under the same set of tasks.

To assess the lane-keeping ability on straight and turning roads, we run the vehicle five trials per framework for each non-obstacle task, during each of which we set the vehicle to a collision-free position as close as possible to the place where the intervention takes place. The obstacle setup is implemented as the vehicle accelerates from 0m/s and tries to circumvent the obstacle 5m ahead. If the vehicle successfully bypasses the obstacle, we record it as a successful trial and reset the vehicle to the initial location for the next trial. Otherwise, it commits a failed trial and counts one more intervention. Table presents the comparisons about the real world generalizability between the MIL and our ModEL on tasks with different attribute settings.

Several observations can be drawn from these experimental results. ModEL achieves much higher MPI and SR on every task, superior to MIL in terms of autonomy and safety. ModEL is also generalizable to different lighting conditions, whereas surprisingly, MIL at night-time visibly outperforms itself at day-time. ModEL also has smaller velocity std in most of the tasks, while the more conservative MIL has more hold-backs during the real-world tests. Finally, it is observed that ModEL has better ability to bypass the obstacles. This results in larger std of steering angle of ModEL than the more conservative MIL, as MIL directly bumps into the obstacle and scrapes the curb at turns more often.

We also perform qualitative real-world evaluation on complex and unseen scenarios, such as sidewalk, crowded road and garage. ModEL can handle lane-following, turning and dynamic obstacle avoidance smoothly, revealing good generalization performance. Please refer to supplementary videos and Appendix for details.

# Conclusion

In this paper, we propose a novel deployed modularized end-to-end RL framework, ModEL, for autonomous driving. By combining modular pipeline with state-of-the-art reinforcement learning, we alleviate the challenging sim-to-real gap, enhance perception performance, and improve RL planning abilities. We also build a test autonomous vehicle to deploy our framework for real-world validation. Additionally, we develop a distributed RL acceleration framework to train the whole system in less than one day on a single workstation. We validate the performance and generalizability of our framework in simulation with intentionally introduced visual and dynamics gaps. Furthermore, our framework could generalize to diverse and unseen complex real-world scenarios. It also achieves superior performance compared with various end-to-end and imitation learning baselines.

# Appendix

## Qualitative real-world experiment scenarios

ModEL proves to be capable of avoiding both regular and unseen obstacles (setup as Fig. ) and also achieves superior generalizability on unseen scenarios, such as being exposed to head-on high beams, traveling on sidewalks, through dense crowds, and even in a garage (depicted in Fig. ). Please explore our accompanying video for more.

## Distributed Training Framework

We developed a distributed off-policy RL training framework (detailed in Fig. ), enabling highly-efficient RL training. As off-policy RL algorithms can learn from data generated by older policies, per-step data collection and learning procedure doesn't depend on each other. So we run the two procedures separately as fast as possible. Message exchanging between procedures (transitions storing and policy parameters updating) is asynchronous and delayed. For data collection, $N$ simulation environments and actor instances are created as independent nodes. An environment sends observations to its corresponding actor and receives predicted actions for stepping. For actor and critic learning, $K$ optimizer nodes are instantiated. Each optimizer has the same copy of actor $\pi_\phi(s)$ and critic $Q^\pi_{\psi_{1, 2}}(s, a)$ networks, and compute gradients using batches of transitions sampled from the replay buffer, then utilize distributed all-reduce to synchronize gradients for optimization.

## Hardware implementation details

The hardware system that deploys ModEL includes a USB front monocular RGB camera (640$\times$`<!-- -->`{=html}480 resolution), an onboard computer (i7-9700 CPU, 32GB RAM, GeForce RTX 3060 GPU), a remote-control unit and an Agile.X HUNTER Unmanned Ground Vehicle (UGV) with photoelectric encoders for wheel odometry, as shown in Figure . The whole system runs on the onboard computer as separated nodes with different functions, using Robot Operating System (ROS) to communicate with each other. The vehicle control node receives vehicle state and sends low-level control commands through Controller Area Network (CAN). The autonomous driving node runs the whole ModEL framework, including sensor data collection, perception, planning and control. The monitor node records vehicle safety state and trajectory statistics.

## Hyperparameter Setup

Table lists all hyperparameters used to train the ModEL framework, including both perception and planning module.

|            Parameter             |       Value       |
|:--------------------------------:|:-----------------:|
|          **Perception**          |                   |
|              model               |    DeepLabv3+     |
|             backbone             |  EfficientNet-B0  |
|            optimizer             |       Adam        |
|          learning rate           |     $10^{-3}$     |
|            batch size            |       $80$        |
|               loss               |     Dice loss     |
|         training epochs          |       $45$        |
|           **Planning**           |                   |
|            algorithm             | Soft Actor-Critic |
|          learning rate           |    $3*10^{-4}$    |
|            batch size            |       $256$       |
|       discount ($\gamma$)        |      $0.99$       |
|          target entropy          |       $-2$        |
|        replay buffer size        |      $10^6$       |
| target smoothing factor ($\tau$) |       0.02        |
|       update to data ratio       |        0.5        |

Hyperparameters

## Bridging the sim-to-real Gap

**System Identification** (SI) is a conventional supervised approach of transfer learning. In this paradigm, fine calibration is inevitable yet costly and dangerous as it needs fresh rollouts from the interaction with the physical world. Other than the inconvenience and unsafety, it is somewhat unrealistic to achieve the high-fidelity all the time due to the drastic variation of physical parameters, e.g. humidity, friction and wear-and-tear of mechanical components inside the vehicles.

**Domain Adaptation** (DA) is a class of unsupervised transfer learning techniques that maps the data from source domain and target domain to a consistent distribution, i.e. from simulators to the real world. It is noteworthy that DA is adapt at image-to-image translation across domains, e.g. end-to-end vision-based self-driving. The mapping function could be devised in diverse ways contextualized within different trade-off demands: First, simulated-to-real translation is data-efficient in the real world without the adaptation module. Conversely, real-to-simulated translation is computationally inefficient during the real-world evaluation, as an adaptation module is indispensable to translate real scenes from cameras into simulated-like images in real time. Besides, both domains could be translated into a canonical representation, e.g. semantic segmentation adopted in. In addition to training policies from the intermediate representation, Bewley *et al.* also reconstruct images intra-domain with direct and cyclic losses for better transferability.

**Domain Randomization** (DR) is a high-level data augmentation technique with special regard to the internal physical mechanisms from the perspective of generating images, like varying the texture, illustration and camera-view. In contrast, **RL with Augmented Data** (RAD) emphasizes the point that simple data manipulation on images, like rotation, cropping and flipping, could also achieve an astonishing transfer effect.

For simpilicity and efficiency, we implement Domain Randomization and Data Augmentation on the raw input and details are specified in .

[^1]: Please visit <http://carzero.xyz> for the supplementary video.

[^2]: Work done with equal contribution.

[^3]: $^{1}$Guan Wang and Haoyi Niu are with the Department of Computer Science and Technology, and Automation respectively, Tsinghua University, 100084 Beijing, China `{wangguan19, email@example.com`

[^4]: $^{2}$Desheng Zhu is with the School of Mechanical Electronic and Information Engineering, China University of Mining and Technology, Beijing, China `firstname.lastname@example.com`

[^5]: $^{3}$Jianming Hu, ^^Xianyuan Zhan, and Guyue Zhou are with Tsinghua University, Beijing, China `{hujm, zhanxianyuan, email@example.com`


abstract: Search in social networks such as Facebook poses different challenges than in classical web search: besides the query text, it is important to take into account the searcher's context to provide relevant results. Their social graph is an integral part of this context and is a unique aspect of Facebook search. While embedding-based retrieval (EBR) has been applied in web search engines for years, Facebook search was still mainly based on a Boolean matching model. In this paper, we discuss the techniques for applying EBR to a Facebook Search system. We introduce the unified embedding framework developed to model semantic embeddings for personalized search, and the system to serve embedding-based retrieval in a typical search system based on an inverted index. We discuss various tricks and experiences on end-to-end optimization of the whole system, including ANN parameter tuning and full-stack optimization. Finally, we present our progress on two selected advanced topics about modeling. We evaluated EBR on verticals[^1] for Facebook Search with significant metrics gains observed in online A/B experiments. We believe this paper will provide useful insights and experiences to help people on developing embedding-based retrieval systems in search engines.
author: Jui-Ting Huang; Ashish Sharma; Shuying Sun; Li Xia; David Zhang; Philip Pronin; Janani Padmanabhan; Giuseppe Ottaviano; Linjun Yang
bibliography: embedding.bib
title: Embedding-based Retrieval in Facebook Search

# Introduction

Search engines have been an important tool to help people access the huge amount of information online. Various techniques have been developed to improve search quality in the last decades, especially in web search engines including Bing and Google. Since it is difficult to accurately compute the search intent from query text and represent the semantic meaning of documents, search techniques are mostly based on various term matching methods , which performs well for the cases that keyword match can address. It still remains a challenging problem for *semantic matching* , which is to address desired results that are not exact match of the query text but can satisfy users' search intent.

In the last years, deep learning has made significant progress in speech recognition, computer vision, and natural language understanding . Among them *embedding*, which is also called *representation learning*, has been proven to be successful techniques contributing to the success . In essence, *embedding* is a way to represent a sparse vector of ids as a dense feature vector, which is also called *semantic embedding* in that it can often learn the semantics. Once the embeddings are learned, it can be used as a representation of query and documents to apply in various stages of a search engine. Due to the huge success of this technique in other domains including computer vision and recommendation system, it has been an active research topic in information retrieval community and search engine industry as the next generation search technology .

In general, a search engine comprises a recall layer targeting to retrieve a set of relevant documents in low latency and computational cost, usually called *retrieval* , and a precision layer targeting to rank the most desired documents on the top with more complex algorithms or models, usually called *ranking*. While embeddings can be applied to both layers, it usually has more opportunities to leverage embeddings in the retrieval layer, since it is at the bottom of the system which is often the bottleneck. The application of embeddings in retrieval is called *embedding-based retrieval* or *EBR* for short. Briefly, embedding-based retrieval is a technique to use embeddings to represent query and documents, and then convert the retrieval problem into a nearest neighbor (NN) search problem in the embedding space.

EBR is a challenging problem in search engines because of the huge scale of data being considered. Different from ranking layers which usually takes hundreds of documents into consideration per session, retrieval layer needs to process billions or trillions of documents in the index of a search engine. The huge scale imposes challenges on both training of embeddings and serving of embeddings. Second, different from embedding-based retrieval in computer vision tasks, search engine usually needs to incorporate both embedding-based retrieval and term matching based retrieval together to score documents in the retrieval layer.

Facebook search, as a social search engine, has unique challenges compared with traditional search engines. In Facebook search, the search intent does not only depend on query text but is also heavily influenced by the user who is issuing the query and the context where the searcher is. Because of this, embedding-based retrieval in Facebook search is not a text embedding problem, as is actively researched in the IR community . Instead it is a more complex problem that requires understanding of text, user, and the context altogether.

To deploy embedding-based retrieval in Facebook search, we developed approaches to address challenges on *modeling*, *serving*, and *full-stack optimization*. In modeling, we proposed *unified embedding*, which is a two sided model where one side is search request comprising query text, searcher, and context, and the other side is the document. To effectively train the model, we developed approaches to mine training data from search log and extract features from searcher, query, context, and documents. For fast model iteration, we adopted a recall metric on an offline evaluation set to compare models.

Building retrieval models for search engine has its unique challenges, such as how to build a representative training task for models to learn effectively and efficiently. We investigated two different directions, *hard mining* to address the challenge of representing and learning retrieval tasks effectively, as well as *ensemble embedding* to divide the model in multiple stages where each stage has different recall and precision tradeoff.

After the model is developed, we need to develop ways to effectively and efficiently serve the model in the retrieval stack. While it is straightforward to build a system combining the candidates from existing retrieval and embedding KNN, we found it is suboptimal because of several reasons: 1) it has huge performance cost from our initial experiment; 2) there is high maintenance cost because of dual index; 3) the two candidate sets might have significant overlap which makes it inefficient overall. Thereafter, we developed a hybrid retrieval framework to integrate embedding KNN and Boolean matching together to score documents for retrieval. To this purpose, we employed Faiss library for embedding vector quantization and integrated it with inverted index based retrieval to build a hybrid retrieval system. Besides addressing the above challenges, this system has two main advantages: 1) it enables the joint optimization of embedding and term matching to address Search retrieval problem; 2) it supports embedding KNN constrained by term matching, which not only helps address the system performance cost issue but also improves the precision of embedding KNN results.

Search is a multi-stage ranking system where retrieval is the first stage, followed by various stages of ranking and filtering models. To wholly optimize the system to return those new good results and suppress those new bad results in the end, we performed later-stage optimization. In particular, we incorporated embeddings into ranking layers and built a training data feedback loop to actively learn to identify those good and bad results from embedding-based retrieval. Figure  is an illustration of embedding-based retrieval system. We evaluated EBR on verticals for Facebook Search with significant metrics gains observed in online A/B experiments.

The paper is organized as follows. We start with modeling to present our solutions about loss function, model architecture, training data and feature engineering in Section  and Section . Next we discuss about details on model serving and system implementation in Section . We discuss the techniques we developed on later-stage optimization to unleash the power from embedding-based retrieval end to end in Section . Finally, we dedicate Section  to selected topics on advanced modeling techniques, followed by conclusions in Section .

# Model

We formulate the search retrieval task as a recall optimization problem. Specifically, given a search query, its target result set $T=\{t_1,t_2,...t_N\}$, and top $K$ results returned by a model, $\{d_1,d_2,...d_K\}$, we want to maximize recall by the top $K$ results, $$\label{eq:recall}
  recall@K=\frac{\sum_{i=1}^{K}d_i\in T}{N}.$$ The target results are the documents related to the given query based on certain criteria. For example, it could be results with user clicks, or relevant documents based on human rating.

We formulate the recall optimization as a ranking problem based on distances computed between a query and documents. The query and documents are encoded with a neural network model into dense vectors, on which we use cosine similarity as the distance metric. We propose to use triplet loss to approximate the recall objective to learn the neural network encoder, which is also called embedding model.

While semantic embedding is commonly formulated as text embedding problem in information retrieval, it is insufficient for Facebook search, which is a personalized search engine that considers not only text query but also searcher's information as well as the context in a search task to satisfy users' personalized information need. Taking people search as an example, while there might be thousands user profiles named "John Smith" on Facebook, the actual target person that a user searches for with query "John Smith" is likely to be their friends or acquaintances. To model this problem, we propose *unified embedding* which considers not only text but also user and context information in deriving embeddings.

## Evaluation Metrics

While our end goal is to deliver quality improvement end to end through online A/B test, it is important to develop offline metrics to quickly evaluate model quality before online experiments and isolate problems from complicated online experiment setup. We propose to run KNN search in the whole index and then use *recall@K* as defined in equation  as the model evaluation metric. In particular, we sampled 10000 search sessions to gather the query and target result set pairs for the evaluation set and reported averaged *recall@K* over 10000 sessions.

## Loss Function

For a given triplet $(q^{(i)},d^{(i)}_{+},d^{(i)}_{-})$, where $q^{(i)}$ is a query, $d^{(i)}_{+}$ and $d^{(i)}_{-}$ are the associated positive and negative documents, respectively, the triplet loss is defined as $$\label{eq:loss}
L= \sum_{i=1}^{N}\max(0, D(q^{(i)},d^{(i)}_{+}) - D(q^{(i)},d^{(i)}_{-}) + m),$$ where $D(u,v)$ is a distance metric between vector $u$ and $v$, $m$ is the margin enforced between positive and negative pairs, and $N$ is the total number of triplets selected from the training set. The intuition of this loss function is to separate the positive pair from the negative pair by a distance margin. We found that tuning margin value is important – the optimal margin value varies a lot across different training tasks, and different margin values result in 5-10% KNN recall variance.

We believe that using random samples to form negative pairs for the triplet loss can approximate the recall optimization task. The reason is as follows. If we sample $n$ negatives for each one positive in the training data, the model will be optimizing for recall at *top one position* when the candidate pool size is $n$. Assuming the actual serving candidate pool size is $N$, we are approximately optimizing recall at top $K \approx N/n$. In Section , we will verify this hypothesis and provide comparisons of different positive and negative label definitions.

## Unified Embedding Model

To learn embeddings that are optimizing the triplet loss, our model comprises three major components: a query encoder $E_Q=f(Q)$ which produces a query embedding, a document encoder $E_D=g(D)$ which produces a document embedding, and a similarity function $S(E_Q,E_D)$ which produces a score between query $Q$ and document $D$. An encoder is a neural network which transforms an input into a low-dimensional dense vector, also known as embedding. In our model, these two encoders $f(\cdot)$ and $g(\cdot)$ by default are two separate networks but have the option of sharing part of the parameters. As for similarity function, we choose cosine similarity as it is one of the commonly used in embedding learning : $$S(Q,D)=\cos(E_Q,E_D)=\frac{\langle E_Q,\ E_D\rangle}{\|E_Q\|\cdot\|E_D\|}.$$ The distance to be used in the loss function in Equation  is hence the cosine distance defined as $1-\cos(E_Q,E_D)$.

The inputs to the encoders are what distinguishes unified embedding from conventional text embedding model. Unified embedding encodes textual, social and other meaningful contextual features to represent query and document, respectively. For example, for query side we can include searcher location and their social connections, whereas for the document side we can include aggregated location and social clusters about a Facebook group by taking groups search as an example.

Most of features are categorical features of high cardinality, which could be either one-hot or multi-hot vectors. For each categorical feature, an embedding look-up layer is inserted to learn and output its dense vector representation before feeding into encoders. For multi-hot vectors, a weighted combination of multiple embeddings is applied for the final feature-level embedding. Figure  illustrates our unified embedding model architecture, and we will discuss more about feature engineering in Section .

## Training Data Mining

Defining positive and negative labels for a retrieval task in a search ranking system is a non-trivial problem. Here we compared several options based on the model recall metric. For negative, we experimented with the following two options of negatives in our initial study, while using click as positive:

-   *random samples*: for each query, we randomly sample documents from the document pool as negatives.

-   *non-click impressions*: for each query, we randomly sample those impressed but not clicked results in the same session as negatives.

The model trained using non-click impressions as negative has significantly worse model recall compared to using random negative: absolute 55% regression in recall for people embedding model. We believe it is because these negatives bias towards hard cases which might match the query in one or multiple factors, while the majority of documents in index are easy cases which do not match the query at all. Having all negatives being such hard negatives will change the representativeness of the training data to the real retrieval task, which might impose non-trivial bias to the learned embeddings.

We also experimented with different ways of mining positives and have interesting findings as follows:

-   *clicks*: it is intuitive to use clicked results as positives, since clicks indicates users' feedback of the result being a likely match to users' search intent.

-   *impressions*: the idea is that we treat retrieval as an approximation to ranker but can execute fast. Thereafter, we want to design the retrieval model to learn to return the same set of results that will be ranked high by the ranker. In this sense, all results shown or impressed to the users are equally positive for retrieval model learning.

Our experimental results showed that both definitions are equally effective; models trained using click vs impressions, given the same data volume, resulted in similar recalls. Additionally, we experimented with augmenting click-based training data with impression-based data, however we did not observe additional gain over the click-based model. It showed that adding impression data does not provide additional value, and the model does not benefit from increased training data volume either.

Our above study suggested that using click as positive and random as negative can provide a reasonable model performance. On top of it, we further explored hard mining strategies to improve the model's ability of differentiating between similar results. We will present more details in Section .

# Feature Engineering

One of the advantages of unified embedding model is that it can incorporate various features other than text to improve the model performance. We observed consistently across different verticals that unified embedding is more effective than text embedding. For example, there are +18% recall improvement when switching from text to unified embeddings for events search, and +16% recall improvement for groups search. The effectiveness of unified embeddings highly depends on the success of identifying and crafting informative features. Table  shows the incremental improvement by adding each new feature category to the group embedding model (with text features as baseline). In this section we discuss several important features that contributed to the major model improvements.

```latex
\begin{adjustbox}{center}
  \begin{tabular}{@{}l|l@{}}
  \toprule
  \textbf{Unified Embedding}  & \textbf{Abs. Recall Gain} \\ \midrule
  + location features         & + 2.20\%                     \\
    \ \ \ + social embedding features & + 1.77\%                     \\ \bottomrule
  \end{tabular}
  \end{adjustbox}
```

**Text features.** Character n-gram is a common approach to represent text for text embedding. Its advantages compared to word n-grams are two folds. First, due to its limited vocabulary size, the embedding lookup table has a smaller size and can be learned more effectively during training. Second, the subword representation is robust to out-of-vocabulary problem the we encounter for both query (e.g. spelling variations or errors) and document sides (due to large content inventory in Facebook). We compared models trained with character n-grams vs word n-grams and found the former can yield a better model. However, on top of character trigrams, including word n-gram representations additionally provides small but consistent model improvement (+1.5% recall gain). Note that since the cardinality of word n-grams is usually very high (e.g. 352M for query trigrams), hashing is needed to reduce the size of embedding lookup table. But even with the downside of hash collision, adding word n-ngrams still provides extra gain.

For a Facebook entity[^2], the main field to extract text features is name for people entities or title for non-people entities. Compared to Boolean term matching techniques, we found embeddings trained with purely text features are particularly good at addressing the following two scenarios:

-   Fuzzy text match. For example, the model is able to learn to match between query "kacis creations" and the *Kasie's creations* page while the term-based match cannot.

-   Optionalization. In the case of query "mini cooper nw", the model can learn to retrieve the expected group *Mini cooper owner/drivers club* by dropping "nw" for an optional term match.

```latex
\begin{adjustbox}{width=\columnwidth,center}

    \begin{tabular}{@{}llll@{}}

    \toprule

    \multicolumn{4}{c}{\textbf{Searcher Location}: Louisville, Kentucky}\\
    \multicolumn{4}{c}{\textbf{Query}: \textit{"equipment for sale"}}\\
    \midrule
\multicolumn{2}{c}{\textbf{Text Model}} & \multicolumn{2}{c}{\textbf{Text + Location Model}} \\
    \midrule

    1 & equipment for sale &1 &Kentucky Farm Equipment for sale                                            \\ %\midrule

    2 & EQUIPMENT FOR SALE WORLDWIDE &2& \begin{tabular}[c]{@{}l@{}}Pre-Owned Farm Equipment for Sale in KY, \\ Southern Indiana and Tennessee\end{tabular} \\ %\midrule

    3 & EQUIPMENT SALE &3 &Farm Equipment For Sale In Ky \\ %\midrule

    4 & Equipment for Sale or Wanted &4 &Central Ky Farm Equipment For Sale \\ %\midrule

    5 & Musical Equipment For Sale or Trade &5 &sed Farm Equipment for sale or trade East Ky. \\ %\midrule

    6 & Used Heavy Equipment For Sale in US &6 &Farm Equipment KY for Sale \\ %\midrule

    7 & Sale equipment & 7&kentucky hay and farm equipment for sale \\ \bottomrule

    \end{tabular}

    \end{adjustbox}
```

**Location features.** Location match is advantageous in many search scenarios such as searching for local business/groups/events. In order for embedding models to consider locations in generating the output embeddings, we added location features into both query and document side features. For query side, we extracted searcher city, region, country, and language. For document side, we added publicly available information, such as explicit group location tagged by the admin. Together with text features, the model was able to successfully learn implicit location match between query and results. Table  shows a side-by-side comparison of top similar documents returned by text embedding model versus text + location embedding model for groups search. We can see the model with location features can learn to fuse location signals into embeddings, ranking documents that has the same location as the searcher who is from Louisville, Kentucky to higher positions.

**Social embedding features.** To leverage the rich Facebook social graph to improve unified embedding model, we trained a separate embedding model to embed users and entities based on the social graph. This can help incorporate the comprehensive social graph into a unified embedding model which may not have full social graph information otherwise.

# Serving

## ANN

We deployed an inverted index based ANN (approximate near neighbor) search algorithms to our system because of the following advantages. First, it has smaller storage cost due to the quantization of embedding vectors. Second, it is easier to be integrated into the existing retrieval system which is based on inverted index. We employed Faiss library to quantize the vectors and then implemented the efficient NN search in our existing inverted table scanning system.

There are two major components for the embedding quantization, one is the *coarse quantization* which quantizes embedding vectors into coarse clusters typically through *K-means* algorithm, and the other is *product quantization* which does a fine-grained quantization to enable efficient calculation of embedding distances. There are several important parameters we need to tune:

-   *Coarse quantization.* There are different algorithms for coarse quantization. It is useful to compare between IMI and IVF algorithms. And it is important to tune number of coarse clusters *num_cluster*, which will affect both perf and recall.

-   *Product quantization.* There are multiple variants of product quantization algorithms, including vanilla PQ, OPQ, PQ with PCA transform. And number of bytes for PQ *pq_bytes* is an important parameter to tune.

-   *nprobe.* *nprobe* is the parameter to decide how many clusters will be assigned to the query embedding, which will further decide how many coarse clusters will be scanned. This parameter will affect the perf and recall.

```latex
\begin{adjustbox}{width=\columnwidth,center}

  \begin{tabular}{lll}

    \toprule

    \textbf{Quantization Scheme}    & \textbf{1-Recall@10}  & \textbf{Index Scanned}               \\

    \midrule
    PCA,PQ16 & 51.62\% & 0.48\% \\
    PCA,PQ16  & 54.11\% & 0.60\% \\
    PQ16 & 67.54\% & 0.58\%\\
    OPQ16,PQ16 & 70.51\% & 0.48\% \\
    OPQ16,PQ16 & 74.29\% & 0.60\% \\
    PQ32 & 74.27\% & 0.58\%\\
    PQ64 & 74.81\% & 0.58\%\\
    Flat (No Quantization) & 74.81\% & 0.58\%\\

    \bottomrule

  \end{tabular}

  \end{adjustbox}
```

We built an offline pipeline to efficiently tune these parameters. Besides, we needed to run online experiment to decide the final setting from the selected candidates based on offline tuning. Below we will share the tricks and learnings we got from ANN tuning.

-   *Tune recall against number of scanned documents.* Initially we compared recall of different coarse quantization algorithms for the same setting of *num_cluster* and *nprobe*. However, from more data analysis we found that the clusters are imbalanced, especially for IMI algorithm – around half of clusters only got a few samples. This would cause the number of scanned documents different for the same setting of *num_cluster* and *nprobe*. Therefore, we employed the number of scanned documents as a better metric to approximate perf impact in ANN tuning, as shown in Figure . We measured ANN accuracy using *1-recall@10*, which is averaged recall of getting top result from exact KNN search in top 10 results of ANN search.

-   *Tune ANN parameters when there is non-trivial model change.* We observed that ANN performance is related with model characteristics. For example, when we employed ensemble techniques with model trained with non-click impressions, we found that while the model showed a better recall than baseline, the recall was worse than baseline after applying quantization to both. It should always be considered to tune the ANN parameters when there is a non-trivial change of the model training task, e.g., add more hard negatives.

-   *Always try OPQ.* It is often useful to transform data prior to applying the quantization. We experimented with both PCA and OPQ to transform the data, and observed that OPQ is generally more effective, as shown in Table and Figure . One caveat of OPQ is: as it applied rotation of the embeddings, we might also need to retune *num_cluster* and *nprobe* to have similar documents scanned.

-   *Choose *pq_bytes* to be d/4.* Product quantizer compresses the vectors into $x$ byte codes. For the choice of $x$, it is related to the dimension $d$ of the embedding vector. Bigger $x$ results in higher search accuracy but at cost of increased memory and latency. From empirical results, we found that the accuracy improvement is limited after $x > d/4$.

-   *Tune nprobe, num_clusters, and pq_bytes online to understand the real perf impact.* While it is important to tune ANN algorithms and parameters offline to get a reasonable understanding of perf vs. recall trade-off, we found it is important to deploy several configs of the ANN algorithms and parameters online to get a better understanding of the perf impact from embedding-based retrieval to the real system. That is important to decide the capacity budget and reduce the scope of parameter search in offline tuning.

## System Implementation

In order to integrate embedding-based retrieval into our serving stack, we implemented first-class support for NN search in Unicorn , a retrieval engine powering most search products at Facebook. Unicorn represents each document as a bag of terms, which are arbitrary strings that express binary properties about the document, conventionally namespaced with their semantics. For example, a user *John* living in *Seattle* would have the terms `text:john` and `location:seattle`. Terms can have payloads attached to them.

A query can be any Boolean expression on the terms. For example, the following query would return all the people that have *john* and *smithe* in the name, and live in *Seattle* or *Menlo Park*:

    (and (or (term location:seattle)
             (term location:menlo_park))
         (and (term text:john)
              (term text:smithe)))

To support NN, we extended the document representation to include embeddings, each with a given string key, and added a `(nn <key> :radius <radius>)` query operator which matches all documents whose `<key>` embedding is within the specified radius of the query embedding.

At indexing time, each document embedding is quantized and turned into a term (for its coarse cluster) and a payload (for the quantized residual). At query time, the `(nn)` is internally rewritten into an `(or)` of the terms associated to the coarse clusters closest to the query embedding (*probes*), and for matching documents the term payload is retrieved to verify the radius constraint. The number of probes can be specified with an additional attribute `:nprobe`. By implementing NN support in terms of pre-existing primitives, instead of writing a separate system, we inherited all the features of the existing system, such as realtime updates, efficient query planning and execution, and support for multi-hop queries (see ).

The latter allows us to support top-K NN queries, where instead of matching by radius we select only the K documents closest to the query, and then evaluate the rest of the query. However, from our experimental study, we found that radius mode can give better trade-off of system performance and result quality. One possible reason is that radius mode enables a constrained NN search (constrained by other parts of the matching expression) but top $K$ mode provides a more relaxed operation which needs to scan the whole index to get top $K$ results. Hence, we use radius based matching in our current production.

### Hybrid Retrieval

By having the `(nn)` operator as part of our Boolean query language we can now support *hybrid* retrieval expressions, with arbitrary combinations of embeddings and terms. This can be used for model based fuzzy matching which can improve on cases like spelling variations, optionalization etc. while reusing and benefiting from other parts of retrieval expression. For example, say a mis-spelled query *john smithe* is looking for a person named *john smith* in *Seattle* or *Menlo Park*; the retrieval expression would look like the one above.

That expression will fail to retrieve the user in question since the term `text:smithe` will fail to match that document. We can add fuzzy matching to this expression through `(nn)` operator:

    (and (or (term location:seattle)
             (term location:menlo_park))
         (or (and (term text:john)
                  (term text:smithe))
             (nn model-141795009 :radius 0.24 :nprobe 16)))

where `model-141795009` is the key for the embedding. In this case the target user will be retrieved if the cosine distance between the query (*john smithe*) embedding and document (*john smith*) embedding is less than 0.24.

### Model Serving

We served the embedding model in the following way. After the two-sided embedding model was trained, we decomposed the model into a query embedding model and a document embedding model and then served the two models separately. For query embedding, we deployed the model in an online embedding inference service for real-time inference. For documents, we used Spark to do model inference in batch offline, and then published the generated embeddings together with other metadata into forward index. We did additional embedding quantization including coarse quantization and PQ to publish it into inverted index.

## Query and Index Selection

To improve efficiency and quality of EBR, we performed query and index selection. We applied the query selection technique to overcome problems like over-triggering, huge capacity cost and junkiness increase. We did not trigger EBR for certain queries as EBR would be poor at and provide no extra value for them, such as easy queries with which searchers are looking for a specific target searched and clicked before, or queries with clearly different query intents from what the embedding model was trained for. On index side, we did index selection to make searching faster. For example, we only chose monthly active users, recent events, popular pages and groups.

# Later-stage Optimization

Facebook search ranking is a complex multi-stage ranking system where each stage progressively refines the results from the preceding stage. At the very bottom of this stack is the retrieval layer, where embedding based retrieval is applied. Results from the retrieval layer are then sorted and filtered by a stack of ranking layers. The model at each stage should be optimized for the distribution of results returned by the preceding layer. However, since the current ranking stages are designed for existing retrieval scenarios, this could result in new results returned from embedding based retrieval to be ranked sub-optimally by the existing rankers. To solve this problem, we proposed two approaches:

-   *Embedding* as ranking feature. Propagating embedding similarities further down the funnel not only helps the ranker recognize new results from embedding-based retrieval, but also provides a generic semantic similarity measure for all results. We explored several options to extract features based on embeddings, including cosine similarity between the query and result embeddings, Hadamard product, and raw embeddings. From our experimental studies, cosine similarity feature consistently showed better performance than other options.

-   *Training data feedback loop*. While embedding-based retrieval can improve retrieval recall, it might have a lower precision in comparison with term matching. To address the precision issue, we built a closed feedback loop based on human rating pipeline. In particular, we logged the results after enabling embedding-based retrieval, and then sent these results to human raters to label whether they are relevant or not. We used these human rated data to re-train the relevance model so that it can be used to filter out the irrelevant results from embedding-based retrieval while keeping the relevant ones. This proved to be a useful technique to achieve high precision for the recall improvement in embedding-based retrieval.

# Advanced Topics

Embedding-based retrieval requires extensive research to continue improve the performance. We investigated two important areas for embedding modeling: hard mining and embedding ensemble, to continue advance the state of the art of embedding-based retrieval solutions.

## Hard Mining

The data space for a retrieval task has diverse data distribution in degrees of text/semantic/social matches, and it is important to design a training data set for an embedding model to learn efficiently and effectively on such space. To tackle this, hard mining is one major direction as well as an active research area for embedding learning. However, most of the research are from computer vision field and for the classification task , while search retrieval does not have concept of "classes" and therefore is a unique problem for which existing techniques do not necessarily work. In this direction, we divided our solutions into two parts: hard negative mining and hard positive mining.

### Hard negative mining (HNM)

When analyzing our embedding model for people search, we found that the top $K$ results from embeddings given a query were usually with the same name, and the model did not always rank the target results higher than others even though the social features are present. This motivated us believe that the model was not able to utilize social features properly yet, and it's very likely because the negative training data were too easy as they were random samples which are usually with different names. To make the model better at differentiating between similar results, we can use samples that are closer to the positive examples in the embedding space as hard negatives in training.

**Online hard negative mining.** As model training is based on mini-batch updates, hard negatives can be selected in every batch in a dynamic but efficient way. Each batch comprises $n$ positive pairs $\{(q^{(i)},d^{(i)}_+)\}_{i=1}^{n}$. Then for each query $q^{(i)}$, we formed a small document pool using all other positive documents $\{d^{(1)}_+,..., d^{(j)}_+,...,d^{(n)}_+|j\neq i\}$ and select the documents which received the highest similarity scores as the hardest negatives to create the training triplets. Enabling online hard negative mining was one major contributor to our modeling improvement. It consistently improved embedding model quality significantly across all verticals: +8.38% recall for people search; +7% recall for groups search, and +5.33% recall for events search. We also had observations that the optimal setting is at most two hard negatives per positive. Using more than two hard negatives will start to regress model quality.

One limitation of online HNM is that the probability of having any hard negative from random samples could be low and therefore cannot produce hard enough negatives. Next, we look at how to generate harder negatives based on the entire result pool, also known as offline Hard Negative Mining.

**Offline hard negative mining.** Offline hard negative mining has the following procedure:

1.  generate top $K$ results for each query.

2.  select hard negatives based on *hard selection strategy*.

3.  retrain embedding model using the newly generated triplets.

4.  the procedure can be iterative.

We performed extensive experiments to compare offline hard negative mining and online hard negative mining. One finding that may first seem counterintuitive is that models trained simply using hard negatives cannot outperform models trained with random negatives. A further analysis suggested that the "hard" model put more weights on non-text features but did worse in text match than the "easy" model. Therefore, we worked to adjust the sampling strategy and finally produce a model that can outperform online HNM model.

The first insight is about *hard selection strategy*. We found that using the hardest examples is not the best strategy. We compared sampling from different rank positions and found sampling between rank 101-500 achieved the best model recall. The second insight is about *retrieval task optimization*. Our hypothesis is that presence of easy negatives in training data is still necessary, as a retrieval model is to operate on an input space which comprises data with mixed levels of hardness, and the majority of which are very easy negatives. Therefore, we explored several ways of integrating random negatives together with hard negatives, including transfer learning from an easy model. From our empirical study the following two techniques showed highest effectiveness:

-   Mixed easy/hard training: blending random and hard negatives in training is advantageous. Increasing the ratio of easy to hard negatives continues to improve the model recall and saturated at easy:hard=100:1.

-   Transfer learning from "hard" model to "easy" model: while transfer learning from easy to hard model does not produce a better model, transfer learning from hard to easy achieved further model recall improvement.

Last but not least, computing exhaustively KNN for each data point in the training data is very time-consuming, and the total model training time will become unrealistic due to the limited computing resources. It is important to have an efficient top K generation for the offline hard negative mining algorithm. Approximate nearest neighbor search is a practical solution here to reduce the total computational time significantly. Furthermore, running ANN search on one random shard is sufficient to generate effective hard negatives, as we only rely on *semi-hard* negatives during training.

### Hard positive mining

Our baseline embedding model used clicks or impressions as positives, which the existing production can already return. To maximize the complementary gain by embedding based retrieval, one direction is to identify new results that have not been retrieved successfully by the production yet but positive. To this aim, we mined potential target results for failed search sessions from searchers' activity log. We found that the positive samples mined in this way is effective to help model training. Model trained using hard positives alone can achieve similar level of model recall as click training data while the data volume is only 4% if it. It can further improve model recall by combining both hard positives and impressions as training data.

## Embedding Ensemble

We learned from HNM experiments that both easy and hard examples are important for EBR model training – we need hard examples to improve model precision, but easy example are also important to represent the retrieval space. The model trained using random negatives simulates the retrieval data distribution and is optimized for recall at a very large $K$, but it has poor precision at top $K$ when $K$ is small. On the other hand, the model trained to optimize precision, e.g. models trained using non-click impressions as negatives or offline hard negatives, is good at ranking for smaller set of candidates but failed for retrieval tasks. Thereafter we propose to combine models trained with different levels of hardness by a *multi-stage approach*, in which the first stage model focuses on recall and the second stage model specializes at differentiating more similar results returned by the first stage model. We shared the same spirit as the cascaded embedding training in , which ensembled a set of models trained with different level of hardness in a cascaded manner. We explored different forms of ensemble embeddings, including weighted concatenation and cascade model and found both effective.

**Weighted Concatenation.** As different models provided different cosine similarity scores for (query, document) pair, we used weighted sum of cosine similarity as the metric to define how close this pair is. To be more specific, given a set of models $\{M_1, \cdots, M_n \}$ and their corresponding weights, $\alpha_1, \cdots, \alpha_n > 0$, for any query $Q$ and document $D$, we define the weighted ensemble similarity score $S_w(Q, D)$ between $Q$ and $D$ as $$S_w(Q,D) = \mathlarger{\sum}_{i=1}^{n} \alpha_i \cos(V_{Q,i},\ U_{D,i}),$$ where $V_{Q,i}, 1\le i\le n$, represent query vector of $Q$ by model $M_i$, and $U_{D,i}, 1\le i\le n$, represent document vector of $D$ by model $M_i$. For the serving purpose, we needed to ensemble multiple embedding vectors into a single representation, for query and document sides, respectively, that can satisfy the above metric property. We can prove that applying weighting multiplication to one side of normalized vector can satisfy the need. Specifically, we constructed the query vector and document vector in the following way: $$E_Q = \big(\alpha_1 \frac{V_{Q,1}}{\|V_{Q,1}\|}, \cdots, \alpha_n \frac{V_{Q,n}}{\|V_{Q,n}\|}\big),$$ and $$E_D = \big( \frac{U_{D,1}}{\|U_{D,1}\|}, \cdots, \frac{U_{Q,n}}{\|U_{Q,n}\|}\big).$$ It is easy to see that the cosine similarity between $E_Q$ and $E_D$ is proportional to $S_w(Q,D)$:

$$\begin{split}
  \cos(E_Q,\ E_D)
  & = \frac{S_w(Q,D)}{\sqrt{\sum_{i=1}^{n} \alpha_i ^ 2}\cdot\sqrt{n}}.
  \end{split}$$ With above, we served the ensemble embedding in the same way as described in Section . The choices of weights were empirical which could be based on performance on the evaluation data set.

We explored several model options for the second stage embedding models. The experiments showed that models trained using non-click impressions achieved the best kNN recall (4.39% recall improvement over the baseline model). However, it suffered from much more accuracy loss compared to single model when applying embedding quantization, and therefore the actual benefit diminished when serving it online. We found that to have a best recall after embedding quantization, the best strategy was to ensemble a relatively easier model with offline hard negative mining model, for which the hardness level of training negatives had been modified and tuned. This ensemble candidate had slightly lower offline model improvement but was able to achieve significant recall improvement online.

**Cascade Model.** Unlike the parallel combination as in weighted ensemble, cascade model runs the second stage model on the output of the first stage model, in a serial fashion. We compared different second-stage model options. We found that models trained using non-click impressions was not a good candidate; the overall improvement was much smaller than the weighted ensemble approach. Moreover, the gain diminished as number of results to be reranked by the second-stage model increases. However, using offline hard negative model for the second stage achieved 3.4% recall improvement over the baseline. It was a more suitable model candidate for cascade because the constructed training data for offline HNM were exactly based on the output of the first-stage model.

Additionally, we explored another cascade model composition. We observed that while unified embedding overall had greater model recall compared to text embedding, it generated new text match failures because of its shifted focus to social and location matches. To enhance the text matching capability of the model hence improve the model precision we adopted the cascade strategy: using text embedding to pre-select text-matching candidates and then using unified embedding model to re-rank result to return the final top candidates. It achieved significant online improvement compared to using unified embedding model alone.

# Conclusions

It has long term benefits to introduce semantic embeddings into search retrieval to address the semantic matching issues by leveraging the advancement on deep learning research. However, it is also a highly challenging problem due to the modeling difficulty, system implementation and cross-stack optimization complexity, especially for a large-scale personalized social search engine. In this paper, we presented our approach of unified embedding to model semantics for social search, and the implementation of embedding-based retrieval in a classical inverted index based search system.

It is only the first step to implement the unified embedding model and embedding-based retrieval system. There is still a long way to go to optimize the system end to end to make it perform well in terms of result quality and system performance. We introduced our experience in model improvement, serving algorithm tuning, and later-stage optimization. We believe this will be valuable experience to help people onboard embedding-based retrieval faster in real search engines. The successful deployment of embedding-based retrieval in production opens a door for sustainable improvement of retrieval quality by leveraging the latest semantic embedding learning techniques. We introduced our progress and learnings from the first step along this direction, especially on hard mining and embedding ensemble.

There are tremendous opportunities ahead to continuously improve the system. In the future, there are two main directions to pursue. One is to go *deep*. In terms of modeling we could apply the latest advanced models such as BERT  or build task-specific models to address particular segments of problems. We could investigate deeper in different stages including serving algorithm tuning and ranking model improvement, guided by full-stack failure analysis to identify the opportunities of improvements in different stacks. The other is to go *universal*. We could leverage the pre-trained text embedding models to develop a universal text embedding sub-model to be applied in different tasks. Furthermore, we could develop a universal query embedding model across all use cases.

[^1]: In Facebook search, verticals are based on result types, e.g., people, page, group, etc.

[^2]: Facebook entity includes people (profile), group, page, and event.