ismamNur/qwen2.5-math-1.5b-jepa-pretrain-amcaime32k-openr1-step400

Qwen2.5-Math-1.5B JEPA Pretrain (stage 1) — amcaime32k + OpenR1-Math-220k traces

Stage-1 checkpoint of the two-stage JEPA→GRPO training pipeline. The model was JEPA-pretrained (LLM-JEPA cosine alignment + SIGReg, official recipe, no repulsion) on verified solved CoT traces, then serves as the init for stage-2 plain GRPO RL.

Training

• Base: Qwen/Qwen2.5-Math-1.5B-Instruct (hidden_size 1536, 28 layers)
• Data: solved traces filtered to ≤3072 tokens
  • amcaime32k solved traces: 14,781 rows
  • OpenR1-Math-220k (verified-correct, ≤3072): 17,000 rows
  • total 31,781, combined
• Loss: (1−λ)·mean(1−cos(p,z)) + λ·SIGReg(p), λ=0.1 (official LLM-JEPA default)
• Predictor: k=1 <|predictor_1|> token, learnable embedding delta (seed 1234 init)
• 400 steps of 994 planned epoch-1 (run cut early at cos≈0.90 convergence; epoch 2 skipped)
• bs 64, AdamW lr 1e-5 cosine w/ 50 warmup steps, grad clip 1.0, bf16 + flash_attention_2 + grad checkpointing
• Seed 31415

Contents

• model.safetensors — merged full fine-tuned weights
• predictor_delta.pt — learnable embedding delta for <|predictor_1|> (1×1536)
• manifest.json — full run config

Final training metrics (step 400)

• loss ≈ 0.19–0.22, cos(p,z) ≈ 0.86–0.90, SIGReg ≈ 0.67–0.73 (stable, no collapse)