adu2021 Confidence Is All You Need: Few-Shot RL Fine-Tuning
Core Concept
Rather than requiring external reward models, preference annotations, or label engineering, RLSC uses the model's own confidence as training signal. The insight: majority voting implicitly performs mode sharpening—concentrating probability mass on the most likely response. By directly optimizing confidence, you improve agreement and reasoning quality without annotations.
Architecture Overview
- Self-confidence objective: Maximize probability that independent samples agree
- Token-level probabilities: Track confidence across generation sequence
- Minimal data requirements: Only 16 samples per question, 10-20 training steps
- No labels needed: Entirely self-supervised approach
- Strong gains: +13.4% AIME, +21.2% MATH, +21.7% Minerva Math
Implementation
Step 1: Implement Confidence Scoring
Compute confidence from token-level probabilities:
class ConfidenceScorer:
def __init__(self):
self.epsilon = 1e-10
def compute_token_confidence(self, logits: torch.Tensor) -> torch.Tensor:
"""Get confidence (max probability) for each token."""
# Convert logits to probabilities
probs = torch.softmax(logits, dim=-1)
# Confidence is max probability
confidence, _ = torch.max(probs, dim=-1)
return confidence
def compute_sequence_confidence(self,
logits: torch.Tensor,
input_length: int,
mask: torch.Tensor = None
) -> torch.Tensor:
"""Get confidence for entire generation sequence."""
# Get token-level confidence
token_conf = self.compute_token_confidence(logits)
# Remove input tokens
generation_conf = token_conf[input_length:]
# Mask padding if provided
if mask is not None:
generation_mask = mask[input_length:]
generation_conf = generation_conf * generation_mask
# Sequence confidence: product of token confidence
# (Use log for numerical stability)
log_conf = torch.log(generation_conf + self.epsilon)
log_conf = log_conf.sum(dim=0)
sequence_conf = torch.exp(log_conf)
return sequence_conf
def compute_agreement_confidence(self,
logits_list: list,
input_length: int
) -> torch.Tensor:
"""Confidence based on agreement between samples."""
# Get token probabilities for each sample
all_probs = []
for logits in logits_list:
probs = torch.softmax(logits[input_length:], dim=-1)
all_probs.append(probs)
# Compute probability of agreement
# P(sample_1 = sample_2) = sum_y p_1(y) * p_2(y)
avg_probs = torch.stack(all_probs).mean(dim=0)
agreement = (avg_probs ** 2).sum(dim=-1)
return agreement.mean()
Step 2: Create Self-Confidence Reward Signal
Design reward from confidence without external labels:
class ConfidenceReward:
def __init__(self, alpha: float = 0.1,
validity_checker = None):
self.alpha = alpha
self.validity_checker = validity_checker
self.scorer = ConfidenceScorer()
def compute_reward(self, completions: list,
logits_list: list,
input_length: int,
question: str = None
) -> dict:
"""Compute reward based on confidence.
R(y) = p(y|x) + α * validity_bonus(y)
"""
rewards = []
for completion, logits in zip(completions, logits_list):
# Base reward: sequence confidence
seq_conf = self.scorer.compute_sequence_confidence(
logits,
input_length
)
reward = seq_conf.item()
# Optional validity bonus (e.g., format checking)
if self.validity_checker is not None:
is_valid = self.validity_checker(completion)
reward += self.alpha * (1.0 if is_valid else -1.0)
rewards.append(reward)
return {
"rewards": torch.tensor(rewards),
"mean_confidence": torch.tensor(rewards).mean(),
"max_confidence": torch.tensor(rewards).max()
}
def compute_majority_vote_signal(self,
completions: list,
logits_list: list,
input_length: int
) -> torch.Tensor:
"""Reward based on agreement with majority."""
# Find majority answer
majority = max(set(completions),
key=completions.count)
rewards = []
for completion, logits in zip(completions, logits_list):
if completion == majority:
# Reward matching majority
seq_conf = self.scorer.compute_sequence_confidence(
logits,
input_length
)
rewards.append(seq_conf.item())
else:
rewards.append(0.0)
return torch.tensor(rewards)
Step 3: Implement RLSC Training Loop
Train with confidence rewards and PPO:
class RLSCTrainer:
def __init__(self, model,
num_samples: int = 16,
batch_size: int = 4):
self.model = model
self.num_samples = num_samples
self.batch_size = batch_size
self.optimizer = torch.optim.Adam(
model.parameters(),
lr=5e-6
)
self.reward_fn = ConfidenceReward(alpha=0.1)
def generate_completions(self, question: str,
num_samples: int = 16
) -> tuple:
"""Generate multiple completions."""
input_ids = self.model.tokenize(question)
input_length = input_ids.shape[-1]
completions = []
logits_list = []
for _ in range(num_samples):
output = self.model.generate(
input_ids,
max_new_tokens=256,
output_hidden_states=False,
return_dict_in_generate=True,
output_scores=True
)
# Detokenize
completion = self.model.decode(output.sequences[0])
completions.append(completion)
# Collect logits
logits = torch.stack(output.scores)
logits_list.append(logits)
return completions, input_ids, input_length, logits_list
def compute_policy_gradient(self,
logits_list: list,
rewards: torch.Tensor,
input_length: int
) -> torch.Tensor:
"""Compute policy gradient from rewards."""
# Normalize rewards for stability
rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-8)
total_loss = 0.0
for logits, reward in zip(logits_list, rewards):
# Compute log probabilities
log_probs = torch.log_softmax(logits, dim=-1)
# Get token-level log probs
generation_log_probs = log_probs[input_length:]
# Policy gradient: maximize log p * reward
loss = -(generation_log_probs.mean() * reward)
total_loss += loss
return total_loss / len(logits_list)
def train_step(self, questions: list) -> dict:
"""Single training step on batch of questions."""
total_loss = 0.0
total_reward = 0.0
total_samples = 0
for question in questions:
# Generate completions
completions, input_ids, input_length, logits_list = (
self.generate_completions(
question,
num_samples=self.num_samples
)
)
# Compute rewards
reward_dict = self.reward_fn.compute_reward(
completions,
logits_list,
input_length,
question
)
rewards = reward_dict["rewards"]
# Compute policy gradient
loss = self.compute_policy_gradient(
logits_list,
rewards,
input_length
)
# Update
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(
self.model.parameters(),
max_norm=1.0
)
self.optimizer.step()
total_loss += loss.item()
total_reward += rewards.mean().item()
total_samples += 1
return {
"avg_loss": total_loss / total_samples,
"avg_reward": total_reward / total_samples,
"num_questions": total_samples
}
Step 4: Evaluate Improvements
Assess reasoning gains without labels:
def evaluate_reasoning(model,
benchmark_questions: list,
reference_answers: list,
metric: str = "exact_match"
) -> dict:
"""Evaluate on reasoning benchmarks."""
correct = 0
avg_confidence = 0.0
answer_lengths = []
for question, ref_answer in zip(benchmark_questions,
reference_answers):
# Generate completion
completion = model.generate(
model.tokenize(question),
max_new_tokens=512
)
# Extract answer (assuming last line or similar)
answer = extract_answer(completion)
# Check correctness
is_correct = evaluate_answer(answer, ref_answer, metric)
if is_correct:
correct += 1
avg_confidence += extract_confidence(completion)
answer_lengths.append(len(completion.split()))
return {
"accuracy": correct / len(benchmark_questions),
"avg_confidence": avg_confidence / len(benchmark_questions),
"avg_answer_length": sum(answer_lengths) / len(
answer_lengths
),
"num_eval": len(benchmark_questions)
}
Practical Guidance
Sample Efficiency: 16 samples per question is sufficient due to confidence being a strong signal. More samples improve stability but require more compute.
Training Duration: 10-20 training steps on 16-sample batches is typical. Longer training may overfit to confidence metric.
Confidence as Prior: The confidence metric works because high-probability outputs tend to be correct. The method implicitly assumes model calibration.
Majority Voting Equivalence: Using agreement-based rewards is similar to majority voting but enables gradient-based optimization via confidence signals.
Emergent Conciseness: Models often produce shorter, more confident answers after RLSC training—an emergent behavior that improves interpretability.
When to Apply: Use RLSC for mathematical reasoning tasks where you want to improve performance without labels, or when external reward models are unavailable.
Reference
RLSC leverages a simple insight: confidence is sufficient training signal without external rewards. By maximizing F(p_θ) = E_y~p_θ[p_θ(y|x)], the method sharpens probability distributions toward high-confidence modes. Substantial gains (+13-21%) demonstrate confidence is a reliable proxy for reasoning quality in few-shot settings.