Skip to content

Repository files navigation

STaR: Self-Taught Reasoner — Reasoning Enhancement for LLMs on GSM8K

Implementation of the STaR (Self-Taught Reasoner) reasoning enhancement technique applied to Llama 3.2-3B-Instruct on the GSM8K mathematical reasoning benchmark. This project compares three training paradigms — Zero-Shot CoT, Vanilla SFT, and STaR — to quantify the impact of bootstrapped rationale generation on LLM reasoning accuracy.


Results

Method Exact Match (GSM8K Test) vs. Zero-Shot
Zero-Shot CoT (baseline) ~47%
Vanilla SFT ~64% +17%
STaR (ours) ~68% +21%

Key finding: STaR improves over Vanilla SFT by +4% by training on a self-curated, bootstrapped dataset of verified correct rationales — rather than the raw GSM8K annotations. 95% of the bootstrapped training data was successfully generated within 5 iterations of the hint-based retry loop.


What is STaR?

STaR (Zelikman et al., 2022) is a reasoning enhancement technique that teaches LLMs to reason by bootstrapping their own training data:

  1. Generate — Prompt the model to produce step-by-step rationales for each training problem
  2. Filter — Keep only rationales that lead to the correct final answer
  3. Fine-tune — Train the model on this verified, high-quality rationale dataset
  4. Repeat — Each round of fine-tuning improves the model's ability to generate correct rationales

The core insight is that STaR produces a self-curated dataset — automatically filtering out incorrect or suboptimal reasoning chains — resulting in cleaner training signal than using raw human annotations.


Three Approaches Compared

1. Zero-Shot CoT (Baseline)

  • Model: meta-llama/Llama-3.2-3B-Instruct (no fine-tuning)
  • Prompt: Chain-of-thought prompt asking for step-by-step reasoning followed by Final Answer: <number>
  • Result: ~47% exact match — pre-trained knowledge only, no math specialization
  • Script: run_zeroshot.py

2. Vanilla SFT

  • Model: Llama 3.2-3B fine-tuned on GSM8K train set (~7,500 samples)
  • Training: 1 epoch, LR=3e-5, gradient accumulation steps=16, linear warmup scheduler
  • Method: Trains on raw GSM8K (question + rationale + answer), masking question tokens from the loss
  • Result: ~64% exact match (+17% over zero-shot)
  • Script: Vanilla_SFT.py

3. STaR — Self-Taught Reasoner

  • Model: Llama 3.2-3B fine-tuned on bootstrapped STaR dataset
  • Training: Same config as Vanilla SFT (1 epoch, LR=3e-5, grad accum=16)
  • Key difference: Training data is self-generated and filtered — only rationales where the model produced the correct final answer are kept
  • Bootstrap: Zero-shot attempt first; if wrong, hint-based retry (up to 5 attempts); 95% of data generated within 5 iterations
  • Result: ~68% exact match (+21% over zero-shot, +4% over Vanilla SFT)
  • Scripts: Star_data_Generation.pyStar_Sft.py

Architecture & Pipeline

GSM8K Train Set (7,473 problems)
        │
        ▼
┌─────────────────────────────────┐
│   STaR Bootstrap Generation    │
│                                 │
│  1. Zero-shot CoT attempt       │
│     → If correct: keep          │
│  2. Hint-based retry (up to 5x) │
│     → If correct: keep          │
│     → If all fail: discard      │
└───────────────┬─────────────────┘
                │ ~7,000+ verified rationales
                ▼
┌─────────────────────────────────┐
│   STaR SFT Fine-tuning          │
│   Llama 3.2-3B-Instruct        │
│   Epochs: 1 | LR: 3e-5         │
│   Grad Accum: 16               │
│   Max Seq Len: 1024            │
└───────────────┬─────────────────┘
                │
                ▼
        Evaluation on GSM8K Test
        (~1,319 problems)
        EM: ~68%

File Structure

├── Star_data_Generation.py        # STaR bootstrap dataset generation (optimized, batched)
├── optGenerate_Star_data.py       # Optimized sharded generation for large-scale runs
├── 3_generate_star_datasetd1.py   # Alternative bootstrap script
├── Star_Sft.py                    # STaR fine-tuning on bootstrapped data
├── Vanilla_SFT.py                 # Vanilla SFT on raw GSM8K
├── run_zeroshot.py                # Zero-shot CoT evaluation
├── Final_normal_Zero_CoT.py       # Unified zero-shot evaluation with detailed analysis
├── Evaluation_star.py             # Evaluation script for STaR fine-tuned model
├── data/
│   ├── star_bootstrap/
│   │   ├── bootstrapped_train_sft.jsonl        # SFT-ready STaR dataset
│   │   └── bootstrapped_train_all_metadata.jsonl  # Full bootstrap metadata
│   └── star_bootstrap.jsonl                    # Raw bootstrap output
├── requirement.txt                # Python dependencies
└── README.md

Quick Start

1. Install Dependencies

pip install -r requirement.txt

requirement.txt:

datasets
transformers
wandb
llama-recipes
torch
huggingface_hub

2. Authenticate with Hugging Face

huggingface-cli login
# Accept Llama 3.2 license at: https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct

3. Run Zero-Shot Baseline

python run_zeroshot.py
# Output: outputs/zero_shot_cot_test.jsonl
# Expected EM: ~47%

4. Generate STaR Bootstrap Dataset

python Star_data_Generation.py \
    --batch_size 24 \
    --max_attempts 5 \
    --output_dir data/star_bootstrap
# Output: data/star_bootstrap/bootstrapped_train_sft.jsonl

For large-scale / multi-GPU sharded generation:

python optGenerate_Star_data.py \
    --num_shards 4 \
    --shard_idx 0 \
    --batch_size 8 \
    --out data/star_bootstrap_shard0.jsonl

5. Train Vanilla SFT

python Vanilla_SFT.py
# Checkpoint saved to: checkpoints/vanilla-sft-final
# Expected EM: ~64%

6. Train STaR Model

python Star_Sft.py
# Checkpoint saved to: checkpoints/Star-sft-final3
# Expected EM: ~68%

7. Evaluate

# Evaluate STaR model
python Evaluation_star.py

# Evaluate with detailed per-sample analysis
python Final_normal_Zero_CoT.py \
    --model_id checkpoints/Star-sft-final3 \
    --output_dir outputs/star_eval

Training Configuration

Parameter Value
Base model meta-llama/Llama-3.2-3B-Instruct
Epochs 1 (both Vanilla SFT and STaR)
Learning rate 3e-5
Batch size (micro) 1
Gradient accumulation 16 (effective batch = 16)
Max sequence length 1024
Max grad norm 1.0
LR scheduler Linear with warmup (3% warmup steps)
Precision bfloat16

Note: 1 epoch used deliberately — training on a small 7.5k dataset for more epochs risks overfitting on the Llama 3B model.


STaR Bootstrap Statistics

Metric Value
GSM8K train examples 7,473
Bootstrap success rate ~93–95%
Zero-shot pass rate ~47%
Hint-assisted pass rate ~46–48% of remaining
Examples generated within 5 iterations ~95%
Failed / discarded ~5%

Analysis

Zero-Shot (~47%): The base Llama 3.2-3B model is a generalist — it has not been specialized for multi-step mathematical reasoning. It fails on more than half of GSM8K problems.

Vanilla SFT (~64%): The +17% gain comes from rationale-tuning on GSM8K. The model is trained not just to predict the final number, but to reproduce step-by-step reasoning chains. This is the most significant single improvement.

STaR (~68%): The +4% gain over Vanilla SFT is the core contribution of the STaR method. The difference is not the training algorithm (both use SFT) — it is the quality of the training data. The STaR bootstrap process automatically filters out any rationale that does not lead to the correct answer, producing a cleaner, self-curated dataset. The model trained on STaR-generated data learns from verified, correct reasoning chains rather than raw human annotations that may contain errors or suboptimal reasoning steps.

Conclusion: STaR is an effective automated data curation technique. By bootstrapping its own high-quality training data, the model achieves meaningful gains over both zero-shot prompting and standard supervised fine-tuning on the same problem set.


References


Citation

@misc{kothandaraman2025star,
  author = {Kothandaraman, Mithilesh},
  title  = {STaR: Self-Taught Reasoner Implementation on GSM8K with Llama 3.2-3B},
  year   = {2025},
  url    = {https://github.com/Mithil-hub/star-reasoning-gsm8k}
}

About

STaR Self-Taught Reasoner implementation on GSM8K — Zero-Shot CoT vs Vanilla SFT vs STaR with Llama 3.2-3B

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages