💬 NLP

Warm Up Before You Train: Unlocking General Reasoning in Resource-Constrained Settings

This paper proposes a sample-efficient, two-stage training strategy that "warms up" language models on Knights & Knaves logic puzzles to acquire general reasoning skills, significantly enhancing their performance and data efficiency when subsequently fine-tuned with Reinforcement Learning on small, target-domain datasets.

Original authors: Safal Shrestha, Minwu Kim, Aadim Nepal, Anubhav Shrestha, Keith Ross

Published 2026-02-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Safal Shrestha, Minwu Kim, Aadim Nepal, Anubhav Shrestha, Keith Ross

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Training AI is Expensive and Data-Hungry

Imagine you want to teach a student how to solve complex problems, like advanced math or writing computer code. Usually, to do this well, you need a massive library of textbooks, thousands of practice problems, and a lot of time.

In the world of Artificial Intelligence (AI), this "library" is called training data. Most current methods to make AI "smart" at reasoning require huge amounts of high-quality data. But what if you only have a tiny notebook of examples? That's the challenge this paper tackles: How do you teach an AI to reason well when you don't have enough data?

The Solution: The "Warm-Up" Strategy

The authors propose a two-step training method called "Warm Up Before You Train." Think of it like an athlete preparing for a specific sport.

Step 1: The "Toy" Warm-Up (Knights & Knaves)

Before the AI tackles real-world math or coding, the researchers put it through a "warm-up" session using a simple logic game called Knights & Knaves.

The Analogy: Imagine a logic puzzle where you have to figure out who is telling the truth (a Knight) and who is lying (a Knave) based on their statements.
Why this game? It doesn't require knowing math formulas or coding languages. It only requires pure logic. It's like a gym for the brain's reasoning muscles.
The Process: The researchers take a super-smart AI (the "teacher") and ask it to solve thousands of these logic puzzles, writing out its long, step-by-step thought process. Then, they teach a smaller, less experienced AI (the "student") to mimic these thought patterns.
The Result: Even though the student AI never saw a single math problem during this warm-up, it learned how to think. It learned habits like checking its own work, correcting mistakes, and testing hypotheses.

Key Finding: Just doing this warm-up made the AI significantly better at math, coding, and general knowledge questions, even without any specific training on those topics. It's like a runner doing general stretching and cardio; they get faster at running even before they start training for a specific marathon.

Step 2: The Specific Training (Resource-Constrained RLVR)

Now that the AI has "warmed up" its reasoning muscles, the second phase begins. This is where they teach it the specific task (like solving math problems) using a very small amount of data (as few as 100 examples).

The Analogy: Now that the athlete is warmed up, they practice for the specific race. Because they are already in good shape, they need far fewer practice runs to get ready than someone who started cold.
The Method: They use a technique called RLVR (Reinforcement Learning with Verifiable Rewards). Basically, the AI tries to solve a problem, gets a "reward" if it's right, and learns from its mistakes.
The Comparison: They compared two students:
1. Student A: Started cold and tried to learn math with only 100 examples.
2. Student B: Did the logic puzzle warm-up first, then learned math with the same 100 examples.

Key Finding: Student B (the warmed-up model) won easily. They learned faster, reached higher scores, and didn't need nearly as much data to get good results. In fact, the warmed-up model trained on just 100 math problems performed almost as well as a model trained on 7,500 math problems without a warm-up.

The Hidden Benefit: Not Forgetting Other Skills

Usually, when you train an AI heavily on one specific thing (like history), it gets really good at history but forgets how to do other things (like math). It becomes too specialized.

The Analogy: If you only practice playing the piano, you might get so used to the piano keys that you forget how to play the guitar.
The Paper's Discovery: The "Warm-Up" method acts like a universal adapter. Because the AI learned general reasoning skills first, it didn't lose its ability to do math or coding even after being trained on history or physics. It stayed flexible.

Summary of the "Magic"

General Skills First: Teach the AI how to think using simple logic puzzles (Knights & Knaves), not specific facts.
Specific Skills Second: Teach the AI the specific task (math, code) using very few examples.
The Payoff: The AI learns faster, needs less data, stays better at other tasks, and performs at a level usually reserved for models trained on massive datasets.

In short, the paper shows that if you want to build a smart AI in a world where data is scarce, don't just feed it facts; give it a logic puzzle warm-up first.

Technical Summary: Warm Up Before You Train

Problem Statement

Designing reasoning-capable Large Language Models (LLMs) typically relies on Reinforcement Learning with Verifiable Rewards (RLVR) or distillation using carefully curated Long Chain-of-Thought (CoT) data. However, these methods are heavily dependent on extensive, high-quality, domain-specific training data. This creates a significant bottleneck in resource-constrained settings where such data is scarce. The paper addresses the challenge of training models to acquire generalizable reasoning strategies that can rapidly adapt to new domains with minimal supervision, moving beyond task-specific optimization toward a meta-learning paradigm.

Methodology

The authors propose a sample-efficient, two-stage training framework designed to build robust reasoning LLMs under limited supervision:

Stage 1: Domain-Agnostic Warmup (Distillation)
- Objective: To induce generalizable reasoning behaviors (e.g., self-reflection, self-correction, hypothesis testing) without relying on domain-specific knowledge.
- Mechanism: The authors utilize the Knights & Knaves (K&K) logic puzzle domain as a "toy domain." This domain relies solely on basic boolean logic, avoiding the need for memorizing mathematical theorems or specialized libraries.
- Process: A strong reasoning teacher model (QwQ-32B) generates Long CoTs for K&K questions. Crucially, the authors do not apply rejection sampling to filter for correctness; instead, they distill the reasoning traces (including errors and corrections) into the base model via Supervised Fine-Tuning (SFT). This creates a "warmed-up" model.
- Rationale: The K&K domain is chosen because solving it requires extensive reasoning strategies that are broadly applicable, allowing the model to learn how to reason rather than what to know.
Stage 2: Target-Task Adaptation (RLVR)
- Objective: To specialize the warmed-up model for specific downstream tasks using a very limited dataset.
- Mechanism: The warmed-up model undergoes RLVR training using a small set of examples (e.g., $\le$ 100 samples) from the target domain (Math, Coding, or General Knowledge).
- Comparison: This approach is compared against training a standard base model directly with RLVR on the same small datasets.

Key Contributions

The paper makes three primary contributions:

Warmup as a Standalone Reasoning Activator: The authors demonstrate that the warmup phase alone, using only K&K reasoning traces, significantly improves performance across diverse benchmarks (Math, Coding, General Knowledge) without any domain-specific fine-tuning.
Enhanced Meta-Learning and Sample Efficiency: When both base and warmed-up models are subjected to RLVR on the same small dataset, the warmed-up model consistently outperforms the base model. It converges faster and achieves higher final accuracy, effectively acting as a meta-learner.
Preservation of Cross-Domain Generalizability: The paper shows that standard RLVR training on specific domains can lead to "catastrophic forgetting" of general reasoning abilities, causing the model to adopt narrow, domain-specific heuristics. The warmup phase mitigates this, allowing the model to maintain robust cross-domain generalization even after specialized RLVR training.

Experimental Results

Experiments were conducted primarily on the Qwen2.5 family of models (3B, 1.5B, and 14B) across three benchmarks: MATH, HumanEval+, and MMLU-Pro.

Warmup-Only Performance:
- The Qwen2.5-3B model saw improvements of +10.2% on MATH, +15.3% on HumanEval+, and +9.0% on MMLU-Pro after K&K warmup alone.
- The Qwen2.5-14B model achieved 77.4% accuracy on MATH after warmup, approaching the 80.2% achieved by prior work using full-scale RLVR with 7,500 domain-specific examples.
- Distillation from K&K (non-curated) often outperformed distillation from s1K (a curated, domain-rich dataset), particularly for smaller models where s1K sometimes caused performance degradation.
RLVR Adaptation (Sample Efficiency):
- On the MATH benchmark, a warmed-up model trained on just 100 random examples achieved 64.5% accuracy, matching the performance of a base model trained on the full 7,500-example dataset (63.2%).
- On HumanEval+, the warmed-up model required only 50 RLVR steps to achieve a +29.3% absolute gain over the baseline, whereas the base model required 100 steps to reach a lower gain of +24.3%.
Cross-Domain Generalization:
- Training a base model directly on MATH or HumanEval+ via RLVR resulted in significant performance drops on other domains (e.g., a -13.8% drop on MATH when the base model was trained on HumanEval+).
- In contrast, the warmed-up model maintained or improved its performance across domains even after specialized RLVR training.
- Exception: For knowledge-intensive tasks like History (MMLU-Pro), the base model trained directly with RLVR outperformed the warmed-up model. The authors note that for factual recall, short CoTs may be optimal, and the reasoning-heavy warmup can be a slight hindrance in this specific context, though the warmed-up model still retained better generalization to other reasoning tasks.

Significance and Claims

The paper claims that "warming up" a model with domain-agnostic reasoning traces is a practical and effective strategy for training robust reasoning LLMs in data-scarce environments.

Meta-Learning Paradigm: The work positions the warmup phase as a form of meta-learning, where the model learns abstract reasoning strategies that can be rapidly specialized.
Efficiency: The approach offers a path to high performance without the prohibitive cost of collecting and curating massive, domain-specific datasets.
Generalization: By decoupling the acquisition of reasoning skills from specific domain knowledge, the method prevents the "over-specialization" that often plagues RLVR-trained models, preserving the ability to solve problems in unseen domains.

The authors conclude that while K&K is a simple toy domain, it serves as a powerful prior for general reasoning, suggesting that future work could explore other synthetic environments to further enhance these transferable reasoning behaviors.