Scale Haruka Kiyohara ([email protected]) University of Minnesota, CSE DSI Machine Learning Seminar Series September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 1
need to handle massive volume of items (e.g., movies, webpages, SNS posts) • We need to handle inference at web-latency (~0.02 second) • We’d like to unlock the capability of cutting-edge models (e.g., LLMs, diffusions) Chatbots/AI assistance (LLMs) September 2026 Creatives (GenAI) End-to-End Training of Two-Stage Decision Systems @ UMN 8
need to handle massive volume of items (e.g., movies, webpages, SNS posts) • We need to handle inference at web-latency (~0.02 second) • We’d like to unlock the capability of cutting-edge models (e.g., LLMs, diffusions) Chatbots/AI assistance (LLMs) A good Creatives (GenAI) tradeoff between “inference latency” and “model expressiveness” is the key! September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 9
a two-stage selection. Large-scale recommendation (RecSys) candidate retrieval (i.e., fast screening) September 2026 output generation (e.g., ranker, LLMs) End-to-End Training of Two-Stage Decision Systems @ UMN 10
a two-stage selection. Large-scale recommendation (RecSys) Retrieval-augmented generation (RAG) candidate retrieval (i.e., fast screening) September 2026 output generation (e.g., ranker, LLMs) End-to-End Training of Two-Stage Decision Systems @ UMN 11
use a two-stage decision. early-stage late-stage “Model A execute task 1: translation.” AI orchestration “Model B execute task 2: Q&A.” The “conductor” model provides instructions for a single/multiple AI model(s). September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN “Model C execute task 3: info-graphics.” 12
its early stage) end-to-end using user feedback (e.g., clicks, purchases)? Large-scale recommendation (RecSys) Optimize early-stage using click signals! September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 13
Decision Systems (~20min) Credit-assigned Policy Gradient for Early Stage Retrieval in Two-stage Ranking. [KCEKNPRDJW, ICML2026] 2. Orchestration for Large-Scale Decision Systems (~10min) An Off-Policy Learning Approach for Steering Sentence Generation towards Personalization. [KCSJ, RecSys2025] 3. Open Challenges in Two-Stage Decisions (~10min) Intrinsic and Extrinsic Diversity-aware Candidate Retrieval in Two-stage Decisions. [on-going work] Fast and Scalable Algorithms for Adaptable Recommendations. [on-going work] Policy Design for Two-sided Platforms with Participation Dynamics. [KYD, ICML2025] September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 14
Stage Retrieval in Two-stage Ranking. [KCEKNPRDJW, ICML2026] (Internship work at Meta) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 15
a two-stage selection. Large-scale recommendation (RecSys) Retrieval-augmented generation (RAG) candidate retrieval (i.e., fast screening) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 16
a two-stage selection. less research more research Large-scale recommendation (RecSys) How early stage contributes to the outcome? September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 17
a two-stage selection. less research more research Large-scale recommendation (RecSys) How early stage contributes to the outcome? How can we improve the overall decision process by appropriately “assign credits” for the early stage retrieval? September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 18
often derive the policy gradient of the objectives. Reinforcement Learning (RL) Supervised Learning (SL) gradient ascent gradient descent September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 21
often derive the policy gradient of the objectives. differentiate with 𝝅 Reinforcement Learning (RL) Supervised Learning (SL) gradient ascent gradient descent September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 22
the joint policy (ESR + LSR) for position 𝑙’s reward. ESR September 2026 LSR End-to-End Training of Two-Stage Decision Systems @ UMN This specific position 𝑙. 23
the joint policy (ESR + LSR) for position 𝑙’s reward. ESR LSR This specific position 𝑙. Why? – because we can calculate global policy gradient for the entire ranking by the linear sum of each position gradient. (when there is no item-item interaction) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 24
the joint policy (ESR + LSR) for position 𝑙’s reward. (marginal) prob. of having action 𝑎𝑙 at position 𝑙 September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 25
the joint policy (ESR + LSR) for position 𝑙’s reward. Rewriting the prob using two stage policies.. ESR September 2026 LSR End-to-End Training of Two-Stage Decision Systems @ UMN 26
the joint policy (ESR + LSR) for position 𝑙’s reward. Rewriting the prob using two stage policies.. Propagating the gradient to the ESR’s candidate selection prob. [Ma+,20]. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 27
scale when the candidate set size (𝐾) becomes large, due to → Even when 𝐴 = 10, the total number of combination (|𝐴|𝐾) becomes exponentially large.. ! (In practice, we may have 𝐴 = 1,000,000 or even more!) # of combinations • Variance issue • Action space scale up expotentially ≈ 𝑂(|𝐴|𝐾 ) • In contrast, we can only sample results with a single candidate set. 10 x 9 x 8 x .. candidate set size (𝐾) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 29
scale when the candidate set size (𝐾) becomes large, due to • Variance issue • Action space scale up expotentially ≈ 𝑂(|𝐴|𝐾 ) • In contrast, we can only sample results with a single candidate set. We want to avoid the dependence on candidate set September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN here! 30
for reducing the variance: What is this marginal probability? (marginal) prob. of having action 𝑎𝑙 (marginal) prob. of selecting action 𝑎𝑙 included in top-K (in some candidate) given the fact that 𝑎𝑙 is in (some) top-K September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 32
for reducing the variance: What is this marginal probability? item pool [A], [B], [C], [D] September 2026 V-PG (baseline) workflow (enumerated) candidate set ({𝐴𝐾}) target action (𝑎𝑙) ESR (marginal) prob. of having action 𝑎𝑙 (marginal)LSR prob. of selecting action 𝑎𝑙 [A,top-K B, C], [A] top-K included in (in some candidate) given the fact that 𝑎𝑙 is in (some) End-to-End Training of Two-Stage Decision Systems @ UMN 33
for reducing the variance: What is this marginal probability? where item pool [A], [B], [C], [D] September 2026 V-PG (baseline) workflow (enumerated) candidate set ({𝐴𝐾}) target action (𝑎𝑙) ESR (marginal) prob. of having action 𝑎𝑙 (marginal)LSR prob. of selecting action 𝑎𝑙 [A,top-K B, C], [A] top-K included in (in some V-PG;candidate) 𝝅𝐄𝐒𝐑(𝑨𝑲|𝒙)given the fact that 𝑎𝑙 is in (some) End-to-End Training of Two-Stage Decision Systems @ UMN 34
for reducing the variance: What is this marginal probability? where item pool [A], [B], [C], [D] September 2026 (enumerated) candidate set ({𝐴𝐾}) target action (𝑎𝑙) ESR (marginal) prob. of having action 𝑎𝑙 (marginal)LSR prob. of selecting action 𝑎𝑙 [A,top-K B, C], [A] top-K included in (in some candidate) given the fact that 𝑎𝑙 is in (some) CA-PG; 𝝅𝐄𝐒𝐑(𝑺𝑲(𝐀)|𝒙) [A, B, D], [A, C, D], (marginal) [B, C, D]prob. of having action 𝑎𝑙 included in top-K (in some candidate) End-to-End Training of Two-Stage Decision Systems @ UMN 35
for reducing the variance: effective size of action What is this marginal probability? where item pool [A], [B], [C], [D] September 2026 (enumerated) candidate set ({𝐴𝐾}) target action (𝑎𝑙) ESR (marginal) prob. of having action 𝑎𝑙 (marginal)LSR prob. of selecting action 𝑎𝑙 [A,top-K B, C], [A] top-K included in (in some candidate) given the fact that 𝑎𝑙 is in (some) CA-PG; 𝝅𝐄𝐒𝐑(𝑺𝑲(𝐀)|𝒙) [A, B, D], [A, C, D], (marginal) [B, C, D]prob. of having action 𝑎𝑙 included in top-K (in some candidate) End-to-End Training of Two-Stage Decision Systems @ UMN 36
difference in the gradient propargation workflow. Vanilla-PG sampled candidate set “Credit-Assigned”-PG marginal of target action choice CA-PG looks more efficient! September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 37
V-PG! (By avoiding the dependence on specific candidate set) • However, CA-PG is a biased policy gradient. (As it “marginalized” and ignored part of the gradient) (V-PG) = (CA-PG) + (𝐴𝐾 dependent gradient) z CA-PG ignores this term via marginalization September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 39
accurate, CA-PG can learn a good ESR! (as long as LSR preserves the policy-weighted reward ratio aligned with the reward ratio) ESR LSR implicit (reward-weighted) teacher distillation (i.e., uses the LSR action choice prob. as a part of reward) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 40
accurate, CA-PG can learn a good ESR! (as long as LSR preserves the policy-weighted reward ratio aligned with the reward ratio) expected reward ratio where policy-weighted expected reward ratio z = expected prob of the LSR chooses A provided candidate set under the ESR distribution September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 41
is (almost) well-specified. With the increased candidate set sizes (𝑲), CA-PG shows faster convergence and better stability! September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 42
larger size of candidate set (𝐾) in {50, 100, 200}, where 𝐴 = 1000. Observed a result similar to the synthetic setting, CA-PG-SwR converges faster when 𝑲 is large. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 43
retrieval of two-stage decisions. • To key challenge was high variance and credit-assignment issues of vanilla PG. • We proposed credit-assigned PG, which uses LSR prob as a reward signal. CA-PG improves data-efficiency under a mild condition about the LSR quality! September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 44
Sentence Generation towards Personalization. [KCSJ, RecSys2025] (PhD work at Cornell) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 45
use a two-stage decision. early-stage late-stage “Model A execute task 1: translation.” AI orchestration “Model B execute task 2: Q&A.” The “conductor” model provides instructions for a single/multiple AI model(s). September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN “Model C execute task 3: info-graphics.” 46
summary for recommendations: “WALL-E (2008)” short summary https://movies.disney.com/wall-e September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 47
summary for recommendations: “WALL-E (2008)” ・A robot called “WALL-E” and his adventure into space ・Animated films with beautiful picture and pretty charactors ・Science-fiction focuing on environment destruction ・Heart-warming drama about love and companionship ・Re-discovery of earth and humanity in dystopia ・Silent film without explicit quotes September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 48
summary for recommendations: “WALL-E (2008)” For sci-fi lovers, In the distant future, one little robot sparked a cosmic revolution. ・A robot called “WALL-E” and his adventure into space ・Animated films with beautiful picture and pretty charactors For romance lovers, ・Science-fiction focuing on environment destruction In a lonely world, a small robot discovers the power of connection. ・Heart-warming drama about love and companionship ・Re-discovery of earth and humanity in dystopia ・Silent film without explicit quotes September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN We’d like to personalize the sentence to each user. 49
where the sentence generation policy is defined as: optimize September 2026 fixed (frozen LLM) End-to-End Training of Two-Stage Decision Systems @ UMN 50
policy to maximize the total reward: , using the logged data collected by a logging policy 𝜋0. need to deal with the partial rewards and distribution shift September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 52
policy gradient (PG) to update the prompt policy. estimate from logged data 𝑛: data size apply importance sampling on the prompt policy prob [Swaminathan&Joachims,16] September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 54
policy gradient (PG) to update the prompt policy. estimate from logged data 𝑛: data size apply importance sampling on the prompt policy prob [Swaminathan&Joachims,16] Issues September 2026 • high variance due to rejection sampling on large prompt space • discarding rich information about generated sentence End-to-End Training of Two-Stage Decision Systems @ UMN 55
policy gradient using logged data as follows. (estimating (similarity-)marginalized sentence policy gradient) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 56
policy gradient using logged data as follows. (estimating (similarity-)marginalized sentence policy gradient) a kernel with bandwidth 𝝉 (1) simulating sentence generation using the current policy (similar to GRPO [Gao+,22]) (2) apply soft weighting based on similarity to logged sentence September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 57
the bias-variance tradeoff. (𝜋) (+/- 1) (+/- 5) DSO often achieves better bias-variance tradeoff than action-based IS. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 58
# of actions and reward noises are large. • DSO is much more data-efficient than the baselines. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 59
better than other OPL methods. • Especially, DSO is more robust to performance corruption. Note: “policy value” is the improvement observed over the sentences generated without prompt, which we call no-prompt baseline. Experiments results is from 25 different trials. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 60
prompt optimization. • To key point is to consider output distribution rather than prompt distribution. • We achieves better bias-variance tradeoff via GRPO-like data augmentation. We can improve data-efficiency by effectively using the decision structure! September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 61
we focused on aligning items based on relevance. Large-scale recommendation (RecSys) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 63
we focused on aligning items based on relevance. Large-scale recommendation (RecSys) However, in applications like news recommendation, item diversity matters. How can we learn to diversify candidate set from user feedback? Intrinsic and Extrinsic Diversity-aware Candidate Retrieval in Two-stage Decisions. [on-going work] September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 64
there is information gap between two-stages. Large-scale recommendation (RecSys) users’ static features September 2026 search query / user prompt End-to-End Training of Two-Stage Decision Systems @ UMN 65
there is information gap between two-stages. Large-scale recommendation (RecSys) users’ static features search query / user prompt How can we risk-hedge for the late stage decision under the gap? How can we reduce the information gap keeping inference latency? Fast and Scalable Algorithms for Adaptable Recommendations. [on-going work] September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 66
downstream AI models. updated! temporaly unavailable.. How can we learn to adapt to model changes train/test-time? How can we make the pipeline robust to runtime uncertainty? September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 68
in user participation in the decision-systems. Policy Design for Two-sided Platforms with Participation Dynamics. [KYD, ICML2025] September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 69
in user participation in the decision-systems. (While not limited to Orchestration,) AI decisions affect what we have in the future. How can we calibrate decisions for long-term success or safety? Policy Design for Two-sided Platforms with Participation Dynamics. [KYD, ICML2025] September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 70
Systems Credit-assigned Policy Gradient for Early Stage Retrieval in Two-stage Ranking. [KCEKNPRDJW, ICML2026] 2. Orchestration for Large-Scale Decision Systems An Off-Policy Learning Approach for Steering Sentence Generation towards Personalization. [KCSJ, RecSys2025] 3. Open Challenges in Two-Stage Decisions Intrinsic and Extrinsic Diversity-aware Candidate Retrieval in Two-stage Decisions. [on-going work] Fast and Scalable Algorithms for Adaptable Recommendations. [on-going work] Policy Design for Two-sided Platforms with Participation Dynamics. [KYD, ICML2025] September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 71
undergrads) Cornell-LinkedIn Strategic Partnership (2023-2025) September 2026 (2025-2026) End-to-End Training of Two-Stage Decision Systems @ UMN (2026-2027) 72
Kiyohara, Mihaela Curmei, Ariel Evnine, Shankar Kalyanaraman, Israel Nir, Ana-Roxana Pop, Nitzan Razin, Sarah Dean, Thorsten Joachims, Udi Weinsberg. Creditassigned Policy Gradient for Early Stage Retrieval in Two-stage Ranking. ICML, 2026. https://arxiv.org/pdf/2605.26385 [KCSJ, RecSys2025] Haruka Kiyohara, Daniel Yiming Cao, Yuta Saito, Thorsten Joachims. An OffPolicy Learning Approach for Steering Sentence Generation towards Personalization. RecSys, 2025. https://dl.acm.org/doi/epdf/10.1145/3705328.3748088 [KYD, ICML2025] Haruka Kiyohara, Fan Yao, Sarah Dean. Policy Design for Two-sided Platforms with Participation Dynamics. ICML, 2025. https://arxiv.org/pdf/2502.01792 Stay tuned for updates! (on-going works) Intrinsic and Extrinsic Diversity-aware Candidate Retrieval in Two-stage Decisions. Fast and Scalable Algorithms for Adaptable Recommendations. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 74
Early Stage Retrieval in Two-stage Ranking. [KCEKNPRDJW, ICML2026] (Internship work at Meta) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 75
is (almost) well-specified. When the LSR’s alignment condition is satisfied, CA-PG works well! September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 76
is (almost) well-specified. Results when using multiple sub-retriever to sample candidate set CA-PG gains benefits from combining multiple sub-retrievers (mixture of experts; MoE) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 77
each method.. 𝑲 𝑳 𝑲 𝑴 Combined with the SwR approximation, both the computational time does not increase with 𝑲 and 𝑳. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 78
credit-assigned policy gradient (CA-PG): where (marginal) prob. of having action 𝑎𝑙 (marginal) prob. of selecting action 𝑎𝑙 included in top-K (in some candidate) given the fact that 𝑎𝑙 is in (some) top-K A set of candidate sets 𝐴𝐾 that contains action 𝑎𝑙 Sum of the ESR prob of the candidates September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 79
two PGs? By ignoring from which candidate 𝑎𝑙 comes from, CA-PG reduces variance. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 81
alignment of actions? Alignment of the LSR’s action choice probability matters September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 84
alignment of actions? ・Any (oracle) epsilon-greedy and softmax-type policies satisfies the alignment condition of LSR. ・Even when the LSR policy makes mistake in the action alignment, the following mistakes are no problem. ・Any misalignment among top-1 to K (i.e., top) items. ・Any misalignment among top-K+1 to |A| (i.e., tail) items. ・Any misalignment whose probability ratio is bounded by the reward ratio: Alignment of the LSR’s action choice probability matters September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 85
alignment of actions? ・Any (oracle) epsilon-greedy and softmax-type policies satisfies the alignment condition of LSR. ・Even when the LSR policy makes mistake in the action alignment, the following mistakes are no problem. ・Any misalignment among top-1 to K (i.e., top) items. ・Any misalignment among top-K+1 to |A| (i.e., tail) items. ・Any misalignment whose probability ratio is bounded by the reward ratio: CA-PG works with a reasonably accurate (practical) LSR policy. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 86
a partial PG of the vanilla PG. • CA-PG enables the credit-assignment within the candidate set, by considering the marginal prob of action is being selected in (one of) top-K. • CA-PG intentionally ignore from which candidate action come from, greatly reducing variance by modifying the action space from 𝑂(|𝐴|𝐾 ) to 𝑂(|𝐴|). • CA-PG can learn the accurate alignment with a practical choice of LSR. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 87
overhead to compute 𝜋(𝑆𝐾(𝑎)|𝑥). The gradient computation of CA-PG requires 𝑂(𝐾𝐿), while that of V-PG is 𝑂(𝐾). September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 88
simplified alternative called TOP1-PG, Use 𝑆1 instead of 𝑆𝐾 only for the gradient computation (i.e., we actually sample 𝐾 actions using ESR). September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 𝑂(𝐿) 89
simplified alternative called TOP1-PG, Use 𝑆1 instead of 𝑆𝐾 only for the gradient computation (i.e., we actually sample 𝐾 actions using ESR). 𝑂(𝐿) with the commonly use Plackett-Luce policy. (We will test the performance in experiments.) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 90
simplified alternative called TOP1-PG, Use 𝑆1 instead of 𝑆𝐾 only for the gradient computation (i.e., we actually sample 𝐾 actions using ESR). with the commonly use Plackett-Luce policy. 𝑂(𝐿) TOP1-PG is equivalent to CA-PG when • Using sampling-with-replacement (SwR) approximation to compute probability • Using a single model for selecting top-K actions (i.e., not using mixture-of-expert) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 91
each PG, we have.. worst / best ※ SwR: Sampling-with-Replacement approximation MoE: mixture-of-experts September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 92
each PG, we have.. worst / best In experiments, we test • How does the performance of each PG change with varying # of candidate (K), # of outputs (L), optimality of LSR? • How does the computational time change with varying # of candidates (K), # of outputs (L)? September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 93
combined with other variance reduction methods. As a proof of concept, we combined CA-PG with GRPO: 1. Query 𝑚 samples per context and action, 𝑟𝑗(𝑥, 𝑎), 𝑗 ∈ [𝑚]. 2. Normalize the reward as 𝑟 ′ 𝑗 = (𝑟𝑗 − 𝑚𝑒𝑎𝑛(𝑟))/𝑠𝑡𝑑(𝑟). 3. (Add a constant value to scale rewards to be positive). September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 94
by recursively applying softmax on the remaining. SwR approximation calculates the probability as if applying softmax independently. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 95
choice probability) can be calculated as follows. We approximate this probability with small relative errors (~6%) in the next slides. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 96
follows. Replacing the expectation over all possible candidate set with the most likely candidate set. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 97
Steering Sentence Generation towards Personalization. [KCSJ, RecSys2025] (PhD work at Cornell) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 98
incur deficient support. (similar sentence support) because (action support) the similar sentence support is a relaxed condition of the action support. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 99
bandwidth 𝛕 is small. • This term comes from the within-neighbor reward shift. • • These terms comes from applying marginalization via kernels in the sentence space. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 101
when kernel bandwidth 𝛕 is large. • This term reduces variance by avoiding within-neighbor importance weights: • This term reduces variance by doing implicit data augmentation and soft-rejection sampling. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 103
IS [Swaminathan&Joachims,16] optimal policy uniform random • DR [Dudík+,11] the higher, the better • POTEC [Saito+,24] • DSO (ours) DR: hybrid of regression and IS POTEC: two-stage policy that uses the cluster of actions September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 104
in similar sentence sentence • sentence distribution prompt reward • reward distribution smooth, different sentences results in different rewards September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN sentence 105
8000} • number of candidate prompts: {10, 50, 100, 500, 1000} • reward noises: {0.0, 1.0, 2.0, 3.0} • For DSO, we use the Gaussian kernel with 𝜏 = 𝟏. 𝟎. value: default value September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 106
• logging marginal density: {w/ and w/o function approx.} (w/o is the monte-carlo estimation) • add noise 𝝈𝒔 = 𝟏. 𝟎 to the sentence embeddings to measure the distance value: default value September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 107
monte-carlo estimation. • Using a Gaussian kernel and the function approx. improves the robustness of DSO to the choice of bandwidth hyperparameter 𝜏. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 108
Because we use the MSE loss to fit the marginal density model. For example, when the true marginal density is 1e-5, estimating it as 1e-5 and 1e-4 does not change the MSE loss too much. In contrast, 1e-4 and 1e-5 make a significant difference. Using function approximation, we can avoid being too precise about small values of the marginal density. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 109
• DistilBert [Sanh+,19]-based reward simulator is trained on the data. (next page) • User and query (i.e., movie) are sampled from the dataset. • Candidate prompts are retrieved from RelatedWord.io. • Using Mistral-7B [Jiang+,23] as the frozen LLM to generate the sentence. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 110
• 𝑢: user id • 𝑞: item id (movie title) movie description • 𝑟: ratings (generated by Mistral-7B (zero-shot, w/o prompt)) Reward simulator user id embedding September 2026 inner product (・) DistilBert encoder movie description loss function: MSE in reward prediction End-to-End Training of Two-Stage Decision Systems @ UMN 111
movies with a rating of 5, while “negative” indicates those with ratings of 0-3. (Right) Showing the distribution of normalized reward, which indicate the improvement of expected reward gained by using the given prompt, compared to that of the sentence generated without prompts. The normalized value is multiplied by 10, so the difference become evident when running policy learning methods. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 113
following gradient in the sentence space. (true sentence policy gradient) gradient w.r.t. sentence distribution however, the issue is that the original sentence space is high dimensional.. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 114
following gradient in the marginalized sentence space. (true marginalized sentence policy gradient) 𝜙(𝑠): kernel-based neighbors of sentence 𝑠 gradient w.r.t. marginalized sentence distribution September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 115
following gradient in the marginalized sentence space. (true marginalized sentence policy gradient) 𝜙(𝑠): kernel-based neighbors of sentence 𝑠 gradient w.r.t. marginalized sentence distribution where September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 116
following gradient in the marginalized sentence space. (true marginalized sentence policy gradient) 𝜙(𝑠): kernel-based neighbors of sentence 𝑠 gradient w.r.t. marginalized sentence distribution where (probability of observing sentence within 𝜙(𝑠)) (expected reward within 𝜙(𝑠) under policy π) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 117
gradient using logged data as follows. (estimating marginalized sentence policy gradient) How we can actually estimate/implement the weighted score function? September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 118
re-sampling technique. ② ① ③ , which suggests that ① DSO does implicit data augmentation via resampling (𝑎, 𝑠′) from the policy 𝜋𝜃 . ② DSO uses soft rejection sampling using the kernel weight. ③ DSO corrects the logging distribution in the marginalized sentence space. See Appendix for the derivation. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 119
we need to estimate the logging marginal density defined as We can use function approximation trained on the following MSE loss. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 120
update the policy. regression-based [Konda&Tsitsiklis,99] important sampling-based [Swaminathan&Joachims,16] • impute regressed reward • correct the distribution shift to be unbiased • introduce bias when the regression is inaccurate • variance can be significantly high • (regression is often demanding due to partial reward and covariate shift) • (especially with a rich set of prompts) September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 123
as a control variate as follows, for the variance reduction purpose. control variate September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 124
leverage clusters among prompts. 𝑐: cluster regression-based greedy (estimating cluster policy gradient) IS w.r.t. clustering space control variate Not leveraging the information about generated sentences. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 125
When defining a DR-style estimator, the baseline term should be as follows. However, estimating the gradient involves another importance sampling, and we cannot reduce variance by using this formulation. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 126
When defining a DR-style estimator, the baseline term should be as follows. However, estimating the gradient involves another importance sampling, It would be interesting to explore how we can efficiently combine the regression and DSO as a potential future work! and we cannot reduce variance by using this formulation. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 127
Xinyang Yi, Ji Yang, Minmin Chen, Jiaxi Tang, Lichan Hong, Ed Chi. Off-policy Learning in Two-stage Recommender Systems. WWW, 2020. [Swaminathan&Joachims,16] Adith Swaminathan and Thorsten Joachims. Batch learning from logged bandit feedback through counterfactual risk minimization. JMLR, 2016. [Shao et al., 2022] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y.K. Li, Y. Wu, Daya Guo. DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. 2024. [Dudík et al.,11] Miroslav Dudík, John Langford, and Lihong Li. Doubly robust policy evaluation and learning. ICML, 2011. [Saito et al.,24] Yuta Saito, Jihan Yao, and Thorsten Joachims. Potec: Off-policy learning for large action spaces via two-stage policy decomposition. 2024. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 129
Wenqiang Lei, Jiawei Chen, Biao Li, Peng Jiang, Xiangnan He, Jiaxin Mao, Tat-Seng Chua. KuaiRec: A Fully-observed Dataset and Insights for Evaluating Recommender Systems. CIKM, 2022. [Harper&Konstan,15] F Maxwell Harper and Joseph A Konstan. The movielens datasets: History and context. TIIS, 2015. [Jiang wt al., 2023] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, William El Sayed. Mistral 7b. 2023. [Sanh et al., 2019] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. 2019. September 2026 End-to-End Training of Two-Stage Decision Systems @ UMN 130
harukakiyohara_
institudelelevage
yukinobaba
izumiyama_lab
koide3
tkng
mickey_0226
s_mizuki_nlp
gpeyre
matsumooon
upura
recsyslab
imai448
427marketing
nonsquared
thirion
tmiket
3n
jeffersonlam
iamctodd
bethany
marktimemedia
inesmontani
reverentgeek
honzajavorek
![Real-data experiment on KuaiRec [Gao+,22] better We test with a](https://files.speakerdeck.com/presentations/444f7dc57ae24625a8f7fdb841886883/slide_41.jpg){kind=link}
![better Results on MovieLens data [Harper&Konstan,15] • DSO often performs](https://files.speakerdeck.com/presentations/444f7dc57ae24625a8f7fdb841886883/slide_58.jpg){kind=link}
![List of presented papers (my own work) [KCEKNPRDJW, ICML2026] Haruka](https://files.speakerdeck.com/presentations/444f7dc57ae24625a8f7fdb841886883/slide_72.jpg){kind=link}
![Additional results combined with GRPO [Shao+,24] CA-PG can be easily](https://files.speakerdeck.com/presentations/444f7dc57ae24625a8f7fdb841886883/slide_92.jpg){kind=link}
![Synthetic experiments evaluation metric compared methods • Regression [Konda&Tsitsiklis,99] •](https://files.speakerdeck.com/presentations/444f7dc57ae24625a8f7fdb841886883/slide_102.jpg){kind=link}
![Full-LLM experiment • Semi-synthetic experiments on the MovieLens-10M dataset [Harper&Konstan,15].](https://files.speakerdeck.com/presentations/444f7dc57ae24625a8f7fdb841886883/slide_108.jpg){kind=link}
![A baseline approach: Doubly Robust (DR) [Dudík+,11] DR uses regression](https://files.speakerdeck.com/presentations/444f7dc57ae24625a8f7fdb841886883/slide_122.jpg){kind=link}
![A baseline approach: POTEC [Saito+,24] POTEC considers two-stage policies to](https://files.speakerdeck.com/presentations/444f7dc57ae24625a8f7fdb841886883/slide_123.jpg){kind=link}
![References (1/2) [Ma et al., 2020] Jiaqi Ma, Zhe Zhao,](https://files.speakerdeck.com/presentations/444f7dc57ae24625a8f7fdb841886883/slide_127.jpg){kind=link}
![References (2/2) [Gao et al., 2022] Chongming Gao, Shijun Li,](https://files.speakerdeck.com/presentations/444f7dc57ae24625a8f7fdb841886883/slide_128.jpg){kind=link}