README.md

TinyReasoner

A reasoning model under 1 million parameters capable of tool calling.

Architecture

• 2-layer LSTM
• Hidden size: 256
• Embedding size: 128
• Character-level tokenizer with special tokens for capabilities.

Progress

Session 1: Pretraining

• Implemented the model, tokenizer, and pretraining script.
• Pretrained on NLTK Gutenberg corpus (approx. 1M characters).
• Saved checkpoint: models/pretrained.pt.

Session 2: Supervised Fine-Tuning (SFT)

• Implemented src/capabilities.py with dictionary (NLTK WordNet) and math (SymPy) tools.
• Implemented src/sampler.py with capability call detection and result injection.
• Created src/generate_sft_data.py to generate 2000 synthetic reasoning traces.
• Created src/sft_train.py for instruction tuning using the SOAP optimizer.
• Fine-tuned the model on reasoning traces.
• Saved checkpoint: models/sft_model.pt.
• Verified that the model starts to use [DEFINE] and [SYMPY] tokens correctly.

Session 3: Reinforcement Learning (GRPO)

• Implemented GRPO training in src/grpo_train.py.
• Created multi-faceted reward functions in src/rewards.py.
• Expanded prompt generation in src/prompts.py.
• Updated src/sampler.py to support multi-rollout log-prob and mask tracking.
• Verified RL training loop and checkpointing.
• Saved initial RL checkpoint: models/rl_model.pt.

Session 4: Extended RL and Task Complexity

• Enhanced src/prompts.py with multi-step math problems and word length comparisons.
• Refined src/rewards.py with length penalties and rewards for utilizing tool results.
• Improved src/grpo_train.py for continuous training and increased iterations to 500.
• Successfully completed 500 iterations of GRPO training.
• Verified that the model maintains reasoning traces and tool use even with increased task complexity.
• Synced all artifacts to Hugging Face bucket.

Session 5: Addressing Mode Collapse and Real Data Integration

• Identified "mode collapse" where the model defaulted to [DEFINE]elephant for most prompts.
• Updated src/generate_sft_data.py to use a much larger vocabulary from NLTK.
• Created src/generate_real_sft_data.py to incorporate real tool-calling traces (from Hermes dataset) with injected capability calls.
• Retrained SFT model on combined synthetic and real data (3000 examples) for 10 epochs.
• Strengthened RL rewards in src/rewards.py:
  • Increased grounding reward to 0.25 per entity.
  • Added a -0.5 penalty for tool calls that don't match prompt entities (hallucinations).
  • Increased reward for utilizing tool results to 0.3.
• Continued GRPO training with the new reward structure.

Session 6: Aggressive RL and Grounding Incentives

• Drastically increased grounding rewards (from 0.5 to 5.0) and hallucination penalties (from -1.0 to -10.0) in src/rewards.py.
• Introduced a specific -20.0 penalty for the elephant hallucination mode to force exploration.
• Reduced KL penalty (beta from 0.01 to 0.0001) in grpo_train.py to allow the model to move away from the collapsed SFT state.
• Increased GRPO group_size to 16 for better advantage estimation.
• Implemented alternating exploration strategies in src/sampler.py (noise vs. temperature).
• Observed the model starting to explore other words like cat, banana, and jacket in RL logs.
• Added detailed logging of sample completions and unique completion counts in the training loop.

Session 7: Curriculum Learning and Dense Rewards

• Implemented a three-level curriculum strategy in src/prompts.py.
• Updated src/rewards.py for dense grounding rewards.
• Enhanced src/grpo_train.py with curriculum progression and persistence.

Session 8: Targeted Grounding and Curriculum Reset

• Identified persistent mode collapse (e.g., hallucinating elephant).
• Created src/generate_grounding_data.py to generate 200 high-quality grounding examples.
• Modified src/sft_train.py to support custom datasets and fine-tuning from existing RL checkpoints.
• Performed a "nudge" SFT phase on the grounding dataset to re-orient the model towards prompt entities.
• Reset GRPO curriculum to Level 0 to focus on simple grounding tasks.
• Ran 300 iterations of GRPO.
• Observed that while the model moved away from elephant, it partially collapsed into a new mode (jacket).
• Verified that tool use syntax and reasoning traces remain intact.

Session 9: Tightened Grounding Rewards and Continued RL

• Identified new mode collapse into jacket and guitar.
• Significantly tightened grounding rewards in src/rewards.py:
  • Increased reward for grounded entities to +10.0.
  • Added heavy penalties (-20.0) for individual hallucinated entities in tool payloads.
  • Added specific penalties (-20.0) for jacket mode collapse.
• Continued GRPO training for 300 iterations at Level 0.
• Monitored training logs to ensure the model explores beyond collapsed modes.
• Synced updated rl_model.pt and training logs to Hugging Face bucket.
• Verified that core architecture and tool-calling mechanics remain stable via integration tests.

Session 10: Grounding Refresh and Evaluation Tools

• Introduced src/compare_models.py to quantitatively track grounding rate and average reward.
• Enhanced src/generate_grounding_data.py with diverse reasoning templates and division operations.
• Generated an expanded dataset of 5000 grounding-focused examples.
• Performed a targeted SFT "refresh" on this data (models/sft_grounding_v2.pt) to stabilize tool payloads.
• Resumed GRPO training from the refreshed base, achieving a grounding rate of 0.55 (up from 0.3).
• Verified that math grounding is strong, while dictionary grounding remains a focus for future sessions.
• Confirmed that the model remains within the 1M parameter limit (~951k parameters).

Session 11: Quantitative Comparison and Continued RL

• Refactored src/compare_models.py to support command-line arguments for easier checkpoint evaluation.
• Evaluated all existing checkpoints (sft_model.pt, rl_model.pt, rl_model_grounding.pt) to establish the strongest baseline.
• Continued GRPO RL training starting from the best rl_model.pt for 200 iterations at Level 0.
• Achieved a stable grounding rate of 0.42 on Level 0 tasks.
• Verified that math tool calls are frequently grounded, although dictionary lookups still exhibit word hallucinations in the payload.
• Confirmed stability of the training loop and maintained model integrity via src/integration_test.py.

Session 13: Grounding Bonus and RL Stabilization

• Refined src/rewards.py with a Target Grounding Bonus (+10.0) for matching exact prompt entities in capability calls.
• Expanded the mode-collapse penalty list to include more hallucinated words ('moss', 'bat', 'sout', 'guitar', 'cat', 'banana', 'tomss', 'seet').
• Ran 300 iterations of GRPO starting from models/sft_grounding_v3.pt.
• Achieved a stable Level 0 grounding rate of 0.42, with math grounding being significantly stronger than dictionary grounding.
• Verified that the math reward often reaches +30.0 due to the new grounding bonus.
• Maintained model size under 1M parameters (~951k).

Next Steps

• Monitor the transition between curriculum levels to ensure the model maintains performance.
• Continue RL training with these aggressive rewards until grounding (using prompt words/numbers) becomes the dominant strategy.
• If the model still struggles with grounding, consider generating "near-miss" SFT data where only the grounding is changed.
• Monitor training diversity to ensure the model doesn't just collapse into a different fixed word (like cat).
• Fine-tune the reward balance as the model starts showing desired behaviors.

Usage

To test the model:

PYTHONPATH=. python3 src/sampler.py models/sft_model.pt