Add an intermediate RL stage between pretraining and post-training using dynamic token budgeting, curriculum sampling, and dual training. Trigger: reduce reasoning steps while maintaining or improving performance in post-training.
Standardmäßig ist der Prompt ausgewählt, der zuerst die Quelle prüft. Sie können zu einem direkten Befehl wechseln oder eine lokale Kopie herunterladen.
Quelldateien prüfen
Lesen Sie SKILL.md und alle von SkillsMP angezeigten Begleitdateien, bevor Sie sich für eine Installation entscheiden.
Mit Codex oder Claude installieren Kopieren Sie diesen Prompt, fügen Sie ihn in Codex, Claude oder einen anderen Assistant ein und lassen Sie die Skill-Seite prüfen und installieren.
Ein direkter Befehl überspringt den Prüf-Prompt. Prüfen Sie die Quelle, bevor Sie ihn ausführen.
Add an intermediate RL stage between pretraining and post-training using dynamic token budgeting, curriculum sampling, and dual training. Trigger: reduce reasoning steps while maintaining or improving performance in post-training.
Reinforcement Mid-Training (RMT): Three-Stage LLM Development
Core Concept
Standard LLM development has two stages: pretraining (next-token prediction) and post-training (instruction-tuning + RL). RMT inserts a critical middle stage that applies reinforcement learning to discourage unnecessary reasoning steps while focusing training on high-value tokens. This achieves up to 64.91% performance improvement while using only 21% of reasoning tokens.
The key insight: Not all tokens contribute equally to performance. Mid-training learns to allocate compute efficiently before the model develops bad reasoning habits during post-training.
Architecture Overview
Dynamic Token Budgeting: Learns when to stop generating reasoning and produce answers
Curriculum-Based Sampling: Progressive training from simple to complex tokens
Dual Training Strategy: Combines RL with next-token prediction for token importance weighting
Create a mechanism to reward efficient reasoning and penalize excessive token generation.
classTokenBudgetingReward:
def__init__(self, max_tokens=4096, target_ratio=0.5):
self.max_tokens = max_tokens
self.target_ratio = target_ratio # Ideally use 50% of budgetdefcompute_reward(self, thinking_tokens, final_answer_tokens, is_correct):
"""
Reward correctness, but penalize verbose thinking.
Args:
thinking_tokens: Number of reasoning tokens generated
final_answer_tokens: Number of tokens in final answer
is_correct: Boolean whether answer is correct
Returns:
Scalar reward signal
"""# Base correctness reward
correctness_reward = 1.0if is_correct else -1.0# Efficiency bonus: reward using <50% of budget
efficiency_ratio = thinking_tokens / self.max_tokens
efficiency_bonus = 0.5if efficiency_ratio < self.target_ratio else -0.2# Penalize excessive reasoningif thinking_tokens > self.max_tokens:
excessive_penalty = -0.5else:
excessive_penalty = 0# Combined reward
total_reward = (
correctness_reward +
0.3 * efficiency_bonus +
excessive_penalty
)
return total_reward
defrecord_trajectory(self, problem, thinking, answer, is_correct):
"""Log a trajectory with token counts and correctness."""
thinking_tokens = len(thinking.split())
answer_tokens = len(answer.split())
reward = self.compute_reward(
thinking_tokens,
answer_tokens,
is_correct
)
return {
"problem": problem,
"thinking": thinking,
"answer": answer,
"thinking_tokens": thinking_tokens,
"answer_tokens": answer_tokens,
"is_correct": is_correct,
"reward": reward
}
3. Implement Curriculum-Based Sampling
Start with easy problems and progressively increase difficulty. This helps the model learn efficient reasoning on simpler tasks before tackling complex ones.
classCurriculumSampler:
def__init__(self, dataset, config):
self.dataset = dataset
self.config = config
self.current_stage = 0self.stage_step = 0defsample_batch(self, batch_size):
"""
Sample a batch weighted by curriculum.
"""# Determine current difficulty stage based on training progress
stage_config = self.config.difficulty_stages[self.current_stage]
# Sample from current stage with high probability# Sample from adjacent stages with lower probability
difficulties = ["easy", "medium", "hard"]
ifself.current_stage == 0:
difficulty_weights = [0.8, 0.2, 0.0]
elifself.current_stage == 1:
difficulty_weights = [0.2, 0.6, 0.2]
else:
difficulty_weights = [0.0, 0.3, 0.7]
# Weighted sampling
batch = []
for _ inrange(batch_size):
difficulty = np.random.choice(
difficulties,
p=difficulty_weights
)
example = self.dataset.sample_by_difficulty(difficulty)
batch.append(example)
# Progress through stagesself.stage_step += 1ifself.stage_step > stage_config.get("duration_steps", 30000):
self.current_stage = min(self.current_stage + 1, 2)
self.stage_step = 0return batch
defget_curriculum_weight(self):
"""Return curriculum weight for loss computation."""returnself.config.difficulty_stages[self.current_stage]["curriculum_weight"]
4. Implement Dual Training (RL + Supervised)
Combine RL losses (based on correctness) with supervised losses (predicting important tokens).
defcompute_mid_training_loss(
model,
batch,
token_budgeter,
curriculum_weight,
rl_weight=0.7,
supervised_weight=0.3):
"""
Compute combined RL and supervised loss for mid-training.
Args:
model: LLM to train
batch: List of training examples
token_budgeter: Reward computer
curriculum_weight: Importance weight from curriculum
Returns:
Scalar loss value
"""
total_loss = 0
batch_size = len(batch)
for example in batch:
# Generate thinking and answer
output = model.generate(
example["problem"],
max_thinking_tokens=4096,
return_intermediate=True
)
thinking = output["thinking"]
answer = output["answer"]
# Evaluate correctness
is_correct = evaluate_answer(answer, example["ground_truth"])
# Compute reward (RL component)
reward = token_budgeter.compute_reward(
len(thinking.split()),
len(answer.split()),
is_correct
)
log_prob = model.compute_log_prob(thinking + answer)
# RL loss: policy gradient
rl_loss = -reward * log_prob
# Supervised loss: predict important tokens# Identify "important" tokens (problem-solving steps, key insights)
important_tokens = identify_important_tokens(
thinking,
example["solution_tokens"]
)
supervised_loss = compute_token_importance_loss(
model,
thinking,
important_tokens
)
# Combine losses with curriculum weight
combined_loss = (
rl_weight * rl_loss +
supervised_weight * supervised_loss
) * curriculum_weight
total_loss += combined_loss
return total_loss / batch_size
