looped-moe-gpt2-reasoning

A GPT-2-scale looped (recurrent-depth) transformer combining sparse Mixture-of-Experts (MoE), Multi-head Latent Attention (MLA), and a per-token adaptive loop-depth router. Trained on a general-English + math-reasoning curriculum. See the project repository for the full architecture, literature survey, and training code.

Architecture

Total parameters	148,932,891
Effective depth (with looping)	5
Hidden size	1024
Attention	Multi-head Latent Attention (MLA), DeepSeek-V3-style, decoupled RoPE
Feed-forward	Sparse MoE, 8 routed experts (top-2) + 1 shared expert, auxiliary-loss-free load balancing
Looping	SharingPattern.FULL_LOOP, 3 iterations, per-token adaptive-depth routing (Mixture-of-Recursions-style)
Checkpoint step	51,727

This is not a transformers-library AutoModel-compatible checkpoint -- the architecture (block-level weight sharing, sparse MoE routing, per-token adaptive depth) has no equivalent built-in model class. Loading this checkpoint requires this project's own code.

Training data

A two-phase curriculum: general English (FineWeb-Edu) for the first 25% of training, with math-reasoning data (OpenMathInstruct-2) ramped in linearly from 25%-45% of training and held at equal weight with general text for the remainder. General text is never fully removed from the mix, to avoid the narrow-capability collapse risk of a pure-math cold start. See configs/05_curriculum_reasoning.yaml in the project repo for the exact schedule.

Evaluation

Validation loss at this checkpoint: 3.9609 (cross-entropy, on held-out text from the same corpora used for training).

Important caveat: this is a next-token-prediction loss, not a measure of correctness on reasoning tasks. A low loss on math-formatted text means the model predicts the surface structure of math solutions well (it has learned the templates/phrasing common in the training corpus) -- it does not by itself demonstrate that the model performs arithmetic or logical reasoning correctly. No answer-accuracy evaluation (e.g. GSM8K-style exact-match scoring) has been run on this checkpoint as of this upload; treat any math-reasoning claims about this model as unverified until such an evaluation is added.

How to load this checkpoint

git clone https://github.com/kameshkanna/looped-moe-gpt2.git
cd looped-moe-gpt2
pip install -e .

from pathlib import Path
from huggingface_hub import hf_hub_download
from looped_moe_gpt2.utils.config_io import load_model_config
from looped_moe_gpt2.model.gpt import LoopedMoEGPT
from looped_moe_gpt2.train.checkpoint import load_checkpoint
import torch

checkpoint_path = Path(hf_hub_download(repo_id="Kameshr/looped-moe-gpt2-reasoning", filename="checkpoint.pt"))
config_path = Path(hf_hub_download(repo_id="Kameshr/looped-moe-gpt2-reasoning", filename="config.yaml"))

model_config = load_model_config(config_path)
model = LoopedMoEGPT(model_config)
checkpoint = load_checkpoint(checkpoint_path, device=torch.device("cpu"))
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

See scripts/chat.py in the project repo for a ready-to-run interactive generation script.

License

MIT (matches the project repository's license).