| --- |
| base_model: google/gemma-2-2b-it |
| library_name: jax |
| license: gemma |
| language: |
| - en |
| pipeline_tag: text-generation |
| tags: |
| - math |
| - reasoning |
| - grpo |
| - reinforcement-learning |
| - gsm8k |
| - gemma-2 |
| - jax |
| - lora |
| - tpu |
| - google-tunix-hackathon |
| - kaggle |
| --- |
| |
| # Gemma 2 2B - GRPO Math Solver 🧮🧠 |
|
|
| This model is a fine-tuned version of `google/gemma-2-2b-it`, specifically optimized for solving mathematical word problems. It was trained on the **GSM8K** benchmark using **Group Relative Policy Optimization (GRPO)** via the JAX-native Tunix library. |
|
|
| The model was developed and submitted as part of the **Google Tunix Hackathon** hosted on Kaggle. |
|
|
| ## 🚀 Model Details |
| * **Base Architecture:** Gemma 2 (2B Parameters, Instruction Tuned) |
| * **Task:** Mathematical Reasoning & High-Precision Problem Solving |
| * **Dataset:** GSM8K (Grade School Math 8K) |
| * **Training Algorithm:** GRPO (Group Relative Policy Optimization) |
| * **Hardware:** Trained on a Google v5e-8 TPU. |
| * **Format:** JAX/Tunix checkpoints. |
|
|
| ## ⚙️ Training Hyperparameters (LoRA & GRPO) |
| To reduce hallucinations and encourage logical step-by-step thinking ("Show your work"), the model was fine-tuned using Low-Rank Adaptation (LoRA) integrated with a low-temperature GRPO setup. |
|
|
| **LoRA Config:** |
| - **Rank:** 64 |
| - **Alpha:** 64.0 |
|
|
| **GRPO & Generation Config:** |
| - **Max Prompt Length:** 256 |
| - **Total Generation Steps:** 512 |
| - **Temperature:** 0.6 *(Optimized for math: less chaotic, highly logical, yet diverse enough for GRPO)* |
| - **Top-P:** 0.95 *(Cuts off the 5% least likely tokens to prevent hallucinations)* |
| - **Top-K:** 50 |
| - **Num Generations:** 4 |
| - **Beta / Epsilon:** 0.08 / 0.2 |
|
|
| **Training Details:** |
| - **Learning Rate:** 3e-6 (AdamW, Cosine Scheduler) |
| - **Weight Decay:** 0.1 |
| - **Max Grad Norm:** 0.1 |
| - **Training Steps:** 1000 batches (Micro-batch size: 2) |
|
|
| ## 🧠 Why GRPO? |
| Instead of standard Supervised Fine-Tuning (SFT), this model utilizes GRPO, a Reinforcement Learning (RL) algorithm that enhances reasoning abilities. By generating multiple responses for a given prompt and evaluating them via a reward model, it calculates a relative advantage to update the policy without the heavy memory footprint of a separate value model. |
|
|
| ## 🏆 Hackathon Context |
| * **Competition:** [Google Tunix Hackathon - Train a model to show its work](https://www.kaggle.com/competitions/google-tunix-hackathon) |
| * **Objective:** Enhance step-by-step reasoning and mathematical capabilities of smaller LLMs using JAX-native tools. |
| * **Author:** Tomasz Słapczyński |
