A device-agnostic, CPU-to-GPU-scalable JAX/Flax codebase for training a Pre-LN Decoder-Only Transformer to perform mathematical operations, including arithmetic, trigonometry, logarithms, exponentials, roots, and nested compositions.
The training pipeline features an on-the-fly, adaptive curriculum data generator that dynamically updates sampling weights based on category-level training and validation losses, actively suppresses overfitted tasks, and prevents category starvation.
📖 Deep Dive Documentation: For exhaustive architectural details, mathematical sequence representations, component maps, and design decisions, consult the Codebase Memory Book.
- Unified Decoder Sequence Model: Treats formula evaluation as sequence completion (
BOS <expr> SEP <result> EOS PAD...) with target-only cross-entropy loss masking. - Digit-by-Digit Tokenization: Splits numbers into individual digit tokens to maintain a fixed 30-token vocabulary and force base-10 value representation learning.
- Adaptive Curriculum Engine: Real-time Exponential Moving Average (EMA) tracking of per-task losses, Softmax sampling probability adjustment, overfit detection, and floor-probability starvation protection.
- Safe Symbolic Tree Generator: Built-in SymPy integration for safe ground-truth calculation with domain bounds preventing zero division, complex outputs, overflow, and non-real numbers.
- Flexible Data Pipelines: Live background-threaded batch streaming or zero-memory disk streaming from sharded
.jsonlbulk datasets. - Mixed Precision & Scalability: Seamless switching between CPU development (
float32, learned position embeddings) and high-throughput GPU training (bfloat16, sinusoidal position embeddings).
Input Sequence: [BOS] "s" "i" "n" "(" "0" "." "5" ")" [SEP] "0" "." "5" [EOS] [PAD] ...
Loss Mask: 0 0 0 0 0 0 0 0 0 0 1 1 1 1 0 ...
^
└─ Gradients backpropagated ONLY here
graph TD
A[ExpressionGenerator] -->|Tree Gen & SymPy Eval| B[ExpressionSampler]
B -->|Batch Streaming / Prefetch| C[CurriculumTracker]
C -->|Softmax Sampling Weights| B
B -->|Batches: input_ids, loss_mask| D[TransformerDecoder]
D -->|JIT Loss & Gradients| E[Optax AdamW Optimizer]
E -->|Update Parameters| D
E -->|Checkpointing| F[Orbax CheckpointManager]
D -->|Greedy Auto-Regressive Decoding| G[eval.py Evaluation]
G -->|Category Val Losses| C
├── MEMORY_BOOK.md # Complete codebase reference & architectural memory
├── README.md # Project overview & quickstart guide
├── requirements.txt # System dependencies
├── configs/
│ ├── cpu_dev.yaml # Fast local CPU smoke-test configuration
│ └── gpu_train.yaml # Scaled bfloat16 GPU production training config
├── scripts/
│ └── generate_dataset.py # Sharded offline JSONL dataset generator
├── src/
│ ├── train.py # Main CLI training entry point & JIT step functions
│ ├── eval.py # Greedy auto-regressive decoder & accuracy evaluators
│ ├── tokenizer/
│ │ ├── tokenizer.py # Tokenizer class (digit-by-digit, fixed 30-token vocab)
│ │ └── vocab.json # Vocabulary mapping dictionary
│ ├── data/
│ │ ├── generator.py # ExpressionGenerator (tree generation & SymPy safety)
│ │ ├── sampler.py # ExpressionSampler (streaming, prefetching, offline sharding)
│ │ └── curriculum.py # CurriculumTracker (loss EMA, overfit mitigation, softmax weights)
│ └── model/
│ └── transformer.py # TransformerDecoder (Flax Pre-LN, Learned vs. Sinusoidal)
└── tests/ # Pytest unit test suite
├── test_tokenizer.py
├── test_data.py
├── test_curriculum.py
├── test_model.py
└── test_train.py
Install dependencies for CPU-only mode:
pip install -r requirements.txt(For GPU acceleration, install JAX with CUDA support: pip install "jax[cuda12]")
Run an end-to-end 120-step smoke test verifying data generation, forward/backward passes, adaptive curriculum updates, validation evaluation, and Orbax checkpointing:
PYTHONPATH=. python3 src/train.py --config configs/cpu_dev.yaml --device_profile cpu_devScale up seamlessly to high-throughput GPU training with bfloat16 mixed precision and sinusoidal position embeddings:
PYTHONPATH=. python3 src/train.py --config configs/gpu_train.yaml --device_profile gpu_trainThe CurriculumTracker manages task difficulty dynamically across (operator, depth) buckets (e.g. sin_d1, +_d2).
- Loss Tracking: Computes Exponential Moving Averages (EMA) of training and validation losses per category.
-
Overfitting Mitigation: If
$L_{train} < \text{threshold}$ and$L_{val} > L_{train} \times \text{ratio}$ , effective loss is scaled down by$\text{overfit_decay}$ ($0.1$ ) to prevent over-sampling memorized categories. -
Softmax Weight Allocation: Effective losses are mapped to sampling probabilities:
$$P_c \propto \exp\left(\frac{L^{eff}_c}{T}\right)$$ -
Starvation Protection: Enforces a
floor_probbaseline so solved tasks are periodically re-sampled.
To pre-compute static sharded JSONL datasets on disk:
PYTHONPATH=. python3 scripts/generate_dataset.py \
--config configs/cpu_dev.yaml \
--output_dir ./offline_dataset \
--num_examples 100000 \
--shard_size 10000The resulting dataset can be streamed with zero memory footprint using ExpressionSampler.stream_offline_dataset.
Run the full unit test suite covering tokenization, data tree generation, curriculum weight shifts, model forward passes, NaN safeguards, and batch overfitting:
python3 -m pytest testsFor in-depth explanations of every module, hyperparameter matrix, API conventions, and developer guidelines, see MEMORY_BOOK.md.