Ada-RS: Adaptive Rejection Sampling for Selective Thinking

Yirou Ge; Yixi Li; Alec Chiu; Shivani Shekhar; Zijie Pan; Avinash Thangali; Yun-Shiuan Chuang; Chaitanya Kulkarni

cs.AIFeb 23, 2026

Ada-RS: Adaptive Rejection Sampling for Selective Thinking

Authors: Yirou Ge, Yixi Li, Alec Chiu, Shivani Shekhar, Zijie Pan, Avinash Thangali, Yun-Shiuan Chuang, Chaitanya Kulkarni,

Organizations: PayPal AI

Abstract

Large language models (LLMs) are increasingly being deployed in cost and latency-sensitive settings. While chain-of-thought improves reasoning, it can waste tokens on simple requests. We study selective thinking for tool-using LLMs and introduce Adaptive Rejection Sampling (Ada-RS), an algorithm-agnostic sample filtering framework for learning selective and efficient reasoning. For each given context, Ada-RS scores multiple sampled completions with an adaptive length-penalized reward then applies stochastic rejection sampling to retain only high-reward candidates (or preference pairs) for downstream optimization. We demonstrate how Ada-RS plugs into both preference pair (e.g. DPO) or grouped policy optimization strategies (e.g. DAPO). Using Qwen3-8B with LoRA on a synthetic tool call-oriented e-commerce benchmark, Ada-RS improves the accuracy-efficiency frontier over standard algorithms by reducing average output tokens by up to 80% and reducing thinking rate by up to 95% while maintaining or improving tool call accuracy. These results highlight that training-signal selection is a powerful lever for efficient reasoning in latency-sensitive deployments.

Figures & tables

Figure 1: Comparison between thinking always reasoning and selective reasoning in a tool-calling LLM agent. The agent performs explicit reasoning even for a simple user query, resulting in unnecessary inference cost and latency (left). The agent selectively skips reasoning and directly calls the appropriate tool (right).

Figure 2: Overview of the proposed Ads-RS training framework. Off-policy Ada-RS-DPO training pipeline (left). Given an input context

x

, a teacher policy model generates multiple rollouts. Ada-RS performs pair-wise rejection sampling based on adaptive thinking reward signals for selective thinking to construct high-quality preference pairs, which are then used to optimize the student policy model via DPO training. On-policy Ada-RS-DAPO training pipeline (right). The current policy model generates multiple rollouts for the input context. Ada-RS applies group-wise rejection sampling based on adaptive thinking rewards to select informative candidate subsets. The resulting candidate group is used to update the policy through an on-policy DAPO training.

Figure 3: Tool call accuracy versus Thinking Rate across methods. The most favorable target region (high accuracy, low thinking rate) is highlighted. Points are colored by algorithm: DPO (green), DAPO (blue), SFT (orange), and no-fine-tuning/base model (red).

Figure 4: Tool call accuracy and average Output token across methods relative to the Qwen3-8B base model with thinking enabled. Numbers above the bars show percentage decrease in average amount of output tokens relative to the base model. Bars are colored by algorithm: DPO (green), DAPO (blue), and SFT (orange).

Method	Accuracy (%)	Avg Output Tokens	Thinking Rate (%)
DPO ALP + RS	63.82	451.64	100.00
Ada-RS-DPO	89.24	87.81	6.10

Table 1: Ablation study results for the effect of using NLL loss with DPO loss. For both experiments,

\beta_{rs}=0.1

and

\alpha=0.01

. RS denotes the use of the rejection sampling procedure used in Ada-RS for preference pairs.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Method	$\boldsymbol{\beta_{rs}}$	$\boldsymbol{\alpha}$	Acc.	Avg Output Tokens	Avg Reason. Tokens	Thinking Rate	Training Samples	RS Accept. Rate	Train Time Decrease
Ada-RS-DPO	1	0.01	84.10	236.52	126.54	50.60	96k	79.00	0.79x
Ada-RS-DPO	0.5	0.01	86.26	233.06	114.25	50.48	86k	71.00	0.71x
Ada-RS-DPO	0.1	0.01	89.24	87.81	0.94	6.10	53k	43.00	0.43x
Ada-RS-DPO	0.1	0.001	87.69	88.00	3.77	28.64	53k	43.00	0.43x
Ada-RS-DPO	0.01	0.01	89.68	84.81	0.83	16.97	28k	23.00	0.23x
Ada-RS-DPO	$\approx$ 0	0.01	90.04	87.21	0.87	13.03	12k	10.00	0.1x

Appendix

Table 2: Hyperparameter sweep for

\boldsymbol{\beta_{rs}}

and

\boldsymbol{\alpha}

. Sample numbers are reported to the nearest thousand. RS Acceptance Rate denotes the empirical acceptance rate from rejection sampling. Training time decrease denotes the fold change relative to using all produced samples (120k samples or preference pairs).

Efficient Language Model Reasoning

Efficient language model reasoning studies how to improve problem-solving accuracy while reducing inference, memory, latency, and training costs.

See all 370 papers

Learning When to Sample: Confidence-Aware Selective Sampling for Efficient Chain-of-Thought Reasoning
Mar 9, 2026Juming Xiong, Kevin Guo, Congning Ni +7CoT Reasoning Efficient Language Model Reasoning
1Vanderbilt University · 2Vanderbilt University Medical Center · 3Intuit AI Research
BOSS: Token-Budgeted Online Sampling and Selection for LLM Reasoning
Oct 19, 2026Yupeng Chang, Yahan Li, Chenlu Guo +3Efficient Language Model Reasoning Inference-Time Compute Allocation
School of Artificial Intelligence, Jilin University · Xingchen AGI Lab, China Telecom Artificial Intelligence (Beijing) Co., Ltd., China · Bilibili Inc., China +3
Adaptive Latent Agentic Reasoning
Jun 1, 2026Dongwon Jung, Peng Shi, Yi Zhang +2Agentic Reasoning LLM Agents
University of California, Davis · University of Waterloo · Greenshoe, Inc.

Mar 9, 2026J/K move · Enter open · S save

Learning When to Sample: Confidence-Aware Selective Sampling for Efficient Chain-of-Thought Reasoning

Juming Xiong, Kevin Guo, Congning Ni, Weixin Liu+6

1Vanderbilt University · 2Vanderbilt University Medical Center · 3Intuit AI Research

CoT Reasoning Efficient Language Model Reasoning LLM Inference Acceleration Adaptive Sampling LLM Uncertainty Estimation LLM Reasoning+1

Open paper AskPDF arXiv

Large language models (LLMs) can achieve strong reasoning performance through chain-of-thought (CoT) reasoning, yet they often generate unnecessarily long reasoning paths that incur high inference cost. Self-consistency-based approaches push accuracy higher still, but they require sampling and aggregating multiple reasoning trajectories, leading to substantial computational overhead. In this paper, we introduce a confidence-aware selective sampling framework that, at inference time, analyzes a single reasoning trajectory to adaptively determine whether to rely on that trajectory alone or trigger multi-path sampling. The framework uses trajectory-level numeric features and sentence-level linguistic features extracted from reasoning states to guide selective multi-path reasoning. We train it on MedQA and evaluate it in-domain on MedQA and under calibration-only transfer on MathQA, MedMCQA, and MMLU, without further fine-tuning. Experimental results show that the proposed framework maintains comparable performance to full and efficient multi-path reasoning baselines, with accuracy changes of

-0.41 \pm 0.58

and

-0.31 \pm 0.58

percentage points, respectively, while reducing token usage by

71.7 \pm 5.0%

and

36.6 \pm 9.1%

. These findings demonstrate that reasoning trajectories contain rich signals for uncertainty estimation, enabling a simple, transferable mechanism to balance accuracy and efficiency in LLM reasoning.

Abstract

Efficient Language Model Reasoning

Explore similar work

Learning When to Sample: Confidence-Aware Selective Sampling for Efficient Chain-of-Thought Reasoning

BOSS: Token-Budgeted Online Sampling and Selection for LLM Reasoning

Adaptive Latent Agentic Reasoning