NanoGEPA Guide

NanoGEPA is a 45M-parameter minimalist language model based on the JEPA architecture. Its core exploration: Does reasoning have to be performed in token space? It separates the reasoning process from text generation, performing mathematical reasoning in latent space instead of token space, aiming to verify the feasibility of latent space reasoning (not pursuing SOTA performance; it is a research prototype).

Background

Problems with Current LLMs

Modern LLMs are trained with the objective P(token_t | token_<t), learning text generation fluency rather than structured reasoning ability—when solving mathematical problems, they only mimic the appearance of thinking and easily make simple arithmetic errors.

Origin of the JEPA Architecture

Proposed by Yann LeCun, its core idea: Intelligent systems should learn abstract representations of the world and predict in latent space rather than at the pixel/token level. Traditional LLMs follow Question tokens → Answer tokens, while the JEPA style is Question latent → Answer latent → Answer tokens (reasoning is in latent space, generation is a decoding step).

Methodology

Architecture Design

Minimalist configuration:

Component	Configuration
Layers	6
Attention Heads	8
Hidden Dimension	512
Parameters	~45M
Dataset	GSM8K (~7.5k samples)

Core innovation: Custom Attention Mask

• Question→Question: Causal attention
• Answer→Answer: Causal attention (independent of Question)
• [PRED] token→Question only: Only looks at the question, not directly at the answer

Dual-Objective Training

Loss formula: L_total = L_token + λ * L_jepa

• L_token: Cross-entropy loss (stabilizes generation)
• L_jepa: Cosine similarity loss (1 − cos(pred_latent, answer_latent), aligns latent spaces)

Evidence

Training Results

Metric	Final Value
Token Loss	0.1186
JEPA Loss	0.0525
Cosine Similarity	0.9475
High cosine similarity indicates successful latent space mapping.

Ablation Experiments

• Without JEPA loss: Latent space alignment collapses; latent representations of Question and Answer have no meaningful relationship
• With JEPA loss: Representation geometry is stable; similar Questions map to adjacent regions

Performance Evaluation

Exact match accuracy on GSM8K validation set: 0.00%—authors state this is expected, as the model was trained from scratch on a small dataset and is a research prototype rather than pursuing performance.

Conclusion

Core Insights

1. Reasoning can be framed as latent representation prediction
2. JEPA loss stabilizes semantic alignment
3. Text generation ≠ reasoning
4. Standard next-token training leads to latent space geometry collapse

Comparison with Mainstream Methods

Method	Reasoning Location	Supervision Signal	Typical Scale
Standard LLM	Token space	Next-token	7B-70B+
Chain-of-Thought	Token space	Explicit reasoning steps	Same as above
NanoGEPA	Latent space	Latent representation alignment	45M

Limitations and Future Directions

Limitations

1. Scale limitations: 45M parameters + 7.5k samples
2. Single dataset: Only GSM8K
3. Generation quality: No optimization for fluency
4. No pre-training: Trained from scratch

Future Directions

1. Larger models (1B+) to validate JEPA
2. JEPA fine-tuning on pre-trained weights
3. Expansion to code/scientific reasoning
4. Exploration of latent space interpretability

Technical Implementation Highlights

• Modular design: Separation of config.py/data.py/model.py/train.py
• Complete evaluation tools: eval_alignment.py (latent alignment), evaluate_accuracy.py (exact match)
• Visualization support: Automatic generation of loss curves
• Gradio demo: Interactive latent space reasoning display
• Code style: Concise and transparent, inspired by nanoGPT