Improving the process for large-scale recommender systems: an interview with Haruka Kiyohara

In their paper Credit-assigned Policy Gradient for Early Stage Retrieval in Two-stage Ranking, Haruka Kiyohara, Mihaela Curmei, Ariel Evnine, Shankar Kalyanaraman, Israel Nir, Ana-Roxana Pop, Nitzan Razin, Sarah Dean, Thorsten Joachims, and Udi Weinsberg study how to improve early-stage retrieval in two-stage decision processes. We caught up with first author Haruka to find out more. This is a project that she worked on during an internship with the Central Applied Science (CAS) team at Meta.

What is the topic of the research in your paper and why is it an interesting area for study?

This paper studies how to improve early-stage retrieval in two-stage decision processes. In large-scale recommender systems (e.g., e-commerce, social network feeds), platforms often need to handle massive volumes of items at web latency. To meet this severe requirement, systems often employ two-stage decision processes: an early-stage ranker (ESR), which filters up to billions of items cheaply into a smaller candidate set, and a more expensive late-stage ranker (LSR) that generates the final set of recommendations from this candidate set, as illustrated in Figure 1.

While there are many existing papers discussing how to train the late-stage policy effectively, the studies on efficient end-to-end training of the early-stage policy have been underinvestigated. In particular, the training of the early-stage policy traditionally used supervised learning methods, where a model aims to predict how users will interact with items, e.g. clicks and purchases, usually through regression models. However, optimizing prediction accuracy is not always best for maximizing such user interaction rewards. For example, when there is a positional factor that affects the reward observation (called “position bias”), aligning items based on the simple reward regression results often fails in identifying the best candidate sets. In contrast, the end-to-end training means that we use the user interaction signals to recursively improve the candidate set to directly optimize the early-stage decision. This end-to-end training is effective, because the policy is able to eventually learn the optimal candidate selection to maximize the reward signal (e.g., clicks) by evaluating the candidate set itself.

However, the end-to-end training process also entails the following challenges. The first challenge is that we need some “exploration” to identify the best selection. This is problematic because exploring the best candidate can be more complicated than simply regressing the reward obtained from a static policy. The second challenge is that the reward for the candidate set chosen by the early-stage policy is not directly observed. Instead, we only observe the reward for the final ranking after being processed by the late-stage policy. This means that the valuation of the candidate set is affected by the late-stage decisions, even though we aim to optimize the early-stage model.

With these challenges in mind, we discuss a principled and efficient approach for training the early-stage (ESR) policy using the reinforcement learning (RL)-style policy gradient, which optimizes the policy via updating the policy parameters to the gradient direction toward improving the reward.

Figure 1: Illustration of the two-stage ranking problem in a large-scale recommender system.

What is the problem that your work looks to address?

The naive approach for calculating the policy gradient of the early-stage (ESR) policy is to derive the exact gradient with respect to the probability of choosing the candidate set sampled by the ESR policy via the chain rule. However, this approach suffers from excessive variance, especially when we increase the candidate-set size (K) and the number of items (|A|). This is because the action (candidate set) space of the ESR policy, which is the combination of K different items, can be exponentially large (i.e., |A| to the power of K). As we select only one candidate set for each query in the user interaction, this large action space makes it difficult to accurately estimate the gradient.

Now, let us give a simple example of illustrating why the gradient estimation becomes difficult using the above naive derivation of the gradient. Consider the case where there are three items, [“A”, “B”, “C”], selected in the candidate set, and the items “A” and “B” are presented in the final ranking selected by the late-stage policy. Also suppose that we get a positive reward for “A” and a negative reward for “B” as their item-wise reward. Intuitively, we expect that the probability of selecting “A” in the candidate set should be increased, while the probability of selecting “B” should be decreased by the early-stage policy. However, the naive policy gradient, which we call Vanilla Policy Gradient (V-PG), increases the joint probability of selecting all members of (“A”, “B”, “C”) when “A” receives a high reward. This credit-assignment issue has been overlooked, and this inefficient propagation of the gradient becomes the key cause of the variance. Therefore, our work aims to resolve this variance issue by presenting a more efficient policy gradient to train the ESR policy.

What is the main contribution of your paper?

The main contribution of our work is a novel policy gradient method that mitigates the above variance issue. We call the proposed method “credit-assigned policy gradient” (CA-PG).

The key point of CA-PG is to give “credit” only to the items that resulted in a high reward, as illustrated in Figure 2. CA-PG achieves this by computing the gradient with respect to the probability that the target item is chosen in any candidate set, i.e. marginalizing over all candidate sets that contain it. That is, instead of calculating the probability of selecting a specific candidate set as done by the naive version (V-PG), CA-PG calculates the gradient only w.r.t. the probability that the target item is selected in any candidate set, without distinguishing individual candidate sets that contain the target item. By doing so, we can avoid the dependence of the policy gradient on each specific candidate set, reducing the action space from combination of K items to only the target item (i.e., from “|A| to the power of K” to “|A|”).

However, due to the above marginalization operation, CA-PG is not an unbiased estimation of the true ESR’s policy gradient. Thus, we further examine in the theoretical analysis when CA-PG can learn the accurate alignment of the items. The analysis reveals that CA-PG uses the late stage (LSR) model choice as a part of reward signal, and CA-PG learns the accurate alignment of items when LSR successfully partitions top-K items over the other items. Complex LSR models, such as those consisting with large language models (LLMs), are often able to satisfy this requirement. Therefore, the proposed CA-PG can often learn accurate alignment, significantly reducing variance compared to the vanilla policy gradient.

Figure 2: Illustration of the reason the proposed method is called “credit-assigned” policy gradient – Vanilla policy gradient (V-PG) increases the joint probability of the candidate set, regardless of the contribution of each item in the candidate set to the final reward. In contrast, CA-PG increases the (marginal) probability of only the item that resulted in a high reward, efficiently assigning the “credit” to the target item.

How did you go about testing your method?

We empirically verify the theoretical results on both synthetic and real-recommender datasets. Both dataset has simulated/ground-truth valuation per each (user, item) pair, indicating how an item is effective for the given user. We calculate “policy value” – the average reward that users get from the joint policy (ESR + LSR), to evaluate the quality of the early-stage policy at each training step. A smaller training step uses smaller cumulative data for training, and the policy used for the data collection is repeatedly updated via the corresponding policy gradient method during the training phase.

We compare the proposed CA-PG against the vanilla policy gradient (V-PG). To enable fair comparison of the policy gradient (i.e., loss function), we use the same model architecture for both methods. For the computation of the policy gradient, we use the default CA-PG and V-PG, and its simplified approximation called “sampling-with-replacement” (SwR) (We refer readers to Section 3.2 and Appendix B.2 of the original paper for the details about the SwR approximation). Therefore, the compared methods are V-PG, V-PG-SwR, CA-PG, and CA-PG-SwR.

Figure 3 compares the training stability and performance of the policies learned by each policy gradient method across varying candidate set size K (i.e., how many items are chosen by the early stage (ESR) policy) in the synthetic experiment. As we see in the figure, the SwR variants stabilize the training process for both V-PG and CA-PG. However, even if we use the SwR approximation, V-PG-SwR slows down the convergence speed as we increase the size of candidates (K), while CA-PG-SwR shows rapid convergence of the policy in the large candidate set size regime (e.g., K=20). This trend is confirmed by the following observation: CA-PGs perform much better than V-PG-SwR when we compare the performance with a small number of online interactions (@50K), and achieve 95% of the final policy value of V-PG-SwR about 3x faster than V-PG-SwR when K=20. These differences are significant in a practical situation, as a poor performance of the policy in the initial training phase can harm both user satisfaction and business metrics, and the fast convergence can help early adaptation to user’s reward distribution shifts (e.g., user’s preference changes). A similar trend is observed on the real-data experiment on the KuaiRec dataset, and additional results on the synthetic experiments also confirms that CA-PG(-SwR) learns well when the late-stage (LSR) policy is reasonably accurate, as theoretical analysis suggested.

Figure 3: Comparing the training stability and convergence performance (policy value) of each PG with varying # of candidates (K) in the synthetic experiment. Each line shows the learning curve for a single random seed (i.e., 10 random seed results can overlap). The vertical lines indicate training termination; any line appearing before 500K gradient steps signifies an interruption due to gradient overflow. Two horizontal lines show the 100% (top) and 95% (bottom) lines of the V-PG-SwR policy values, to compare the convergence speed of each method.

Are you planning to extend this work further?

Although this may be done as a project independent of the internship, these topics are very interesting and open, and I would definitely continue work on the follow-up research questions.

The first future direction is to explore the off-policy learning (OPL) setting, where we aim to optimize the (two-stage) policy using logged data. This would help practitioners to learn a new policy without carefully auditing the online training phase.

The second direction is to investigate how to incorporate reward interaction within ranking. Our paper primarily focuses on the setting where each reward is affected only by the corresponding item. However, reward interactions such as diversity within the final ranking may arise in applications like news recommendation. Exploring how to assign credit to each item under item-item interactions can also be impactful.

Finally, our framework conceptually fits large language model (LLM) applications, such as those using retrieval-augmented generation (RAG) framework. However, we did not conduct experiments on LLM/RAG applications, and I am excited to see how a related idea can be integrated in these other cutting-edge applications as well.

About Haruka

Haruka Kiyohara is a fourth-year Computer Science Ph.D. candidate at Cornell University. Her research interest lies in evaluating and optimizing decision-making systems using causal inference and machine learning, particularly learning from logged data and optimizing for long-term social goods in large-scale recommender systems. Her work has been published at machine learning and data mining conferences, including ICML, NeurIPS, ICLR, KDD, WSDM, and RecSys. Prior to Cornell, she received a B.E. in Industrial Engineering and Economics from Tokyo Institute of Technology. She is supported by the Funai Overseas Scholarship, Quad Fellowship, and LinkedIn PhD Award for her graduate study at Cornell.

Useful links

• ICML paper page
• arXiv preprint
• Slides
• GitHub

Lucy Smith is Senior Managing Editor for AIhub.

Lucy Smith is Senior Managing Editor for AIhub.