Archive untouched NeurIPS submission GRT-20M-Base checkpoint and zero-shot SLM-Bench baseline
9f377d7 verified | license: apache-2.0 | |
| tags: | |
| - grt | |
| - geometric-rotation-transformer | |
| - causal-lm | |
| - rotational | |
| datasets: | |
| - OptimalScale/ClimbMix | |
| - wikimedia/wikipedia | |
| - tiiuae/falcon-refinedweb | |
| - HuggingFaceFW/fineweb-edu | |
| - HuggingFaceTB/cosmopedia | |
| - nvidia/Nemotron-Pretraining-Dataset-sample | |
| - HuggingFaceTB/finemath | |
| language: | |
| - en | |
| library_name: transformers | |
| pipeline_tag: text-generation | |
| # GRT-20M-Base (NeurIPS Submission Pretrained Checkpoint) | |
| > **Note:** This repository preserves the original, untouched `16.7M`-parameter (`10.3M` non-embedding backbone) **GRT-20M** base checkpoint trained with first-generation `PairAdam` (`158,000` steps across an 8-corpus pretraining mix) for our NeurIPS submission, prior to `SymplecticAdam` / `HuygensSymplecticAdam` curriculum fine-tuning. For the curriculum-aligned comparison, see [`MatterWaveAI/grt-20m`](https://huggingface.co/MatterWaveAI/grt-20m). | |
| A Geometric Rotation Transformer pretrained on a mix of eight English corpora, | |
| including FineWeb-Edu, Cosmopedia, and FineMath. GRT replaces the dense linear | |
| projections of a standard transformer with cascades of learned Givens rotations | |
| on SO(D), so the residual path contains no `nn.Linear` at all. | |
| This is a research preview at small scale. It is published to demonstrate that | |
| the architecture trains stably and evaluates end-to-end through the same | |
| OpenAI-HTTP serving seam as every other model, and to give CPU-only CI a cheap, | |
| always-available upstream. | |
| ## Model Details | |
| | | | | |
| |---|---| | |
| | **Total parameters** | 16,730,869 (16.7M) | | |
| | **Non-embedding parameters** | 10,280,869 (10.3M) | | |
| | **d_model** | 256 | | |
| | **n_blocks** | 12 | | |
| | **n_heads** | 4 | | |
| | **n_qkv_rotations** | 11 | | |
| | **ffn_n_rotations** | 21 | | |
| | **Vocabulary** | 50,000 (BPE) | | |
| | **Max sequence length** | 1,024 | | |
| | **Precision** | float32 | | |
| Parameter counts are measured from the published `model.safetensors`. "20m" is | |
| the size-class label, not a parameter count. The embedding dominates: 6.45M of | |
| 16.7M parameters are the token embedding (50,000 tokens x 128 phase angles, | |
| plus a log-magnitude per token). | |
| ## Training | |
| | | | | |
| |---|---| | |
| | **Data** | ClimbMix, Wikipedia, RefinedWeb, FineWeb-Edu, Cosmopedia, Nemotron, FineMath, StackExchange | | |
| | **Steps** | 158,000 | | |
| | **Optimizer** | PairAdam | | |
| | **Peak LR** | 1e-3, cosine decay to 10% | | |
| | **Training loss** | 2.114 | | |
| ## Sample output (greedy, temperature 0) | |
| The three reference completions, generated by the published weights through | |
| `transformers` f32, and reproduced token-for-token identically by the vLLM and | |
| llama.cpp serving paths: | |
| | Prompt | Greedy continuation (20 tokens) | | |
| |---|---| | |
| | `The capital of France is` | ` the capital of the capital of the capital of the capital of the capital of the capital of the capital` | | |
| | `Once upon a time` | `, there was a great idea to help us understand the complexities of the concept of the concept of the` | | |
| | `In the beginning` | `, the first term is the same, but the second term is the same.\n\nThe problem` | | |
| ## Benchmarks (Pre-Fine-Tune Zero-Shot & Standard `PairAdam` Baseline) | |
| Evaluated via `grt-slm-bench` (`training/slm_bench.py`) across 9 standard multiple-choice benchmarks (`13,144` held-out examples) and `GSM8K` held-out perplexity: | |
| | Task (`SLM-Bench`) | Split (`n`) | Random | `GRT-20M-Base` (Zero-Shot) | Standard `PairAdam` Fine-Tune | | |
| |---|---|---:|---:|---:| | |
| | **BoolQ** (`google/boolq`) | `val` (`3,270`) | `50.00%` | **`62.17%`** | `53.43%` | | |
| | **PIQA** (`baber/piqa`) | `val` (`1,838`) | `50.00%` | **`53.97%`** | `53.37%` | | |
| | **WinoGrande** (`allenai/winogrande`) | `val` (`1,267`) | `50.00%` | `51.93%` | **`52.41%`** | | |
| | **QNLI** (`nyu-mll/glue`) | `val` (`1,000`) | `50.00%` | `50.50%` | **`55.70%`** | | |
| | **OpenBookQA** (`allenai/openbookqa`) | `test` (`500`) | `25.00%` | **`34.20%`** | `25.60%` | | |
| | **HellaSwag** (`Rowan/hellaswag`) | `val` (`1,000`) | `25.00%` | `28.70%` | **`31.80%`** | | |
| | **ARC-Easy** (`allenai/ai2_arc`) | `test` (`2,376`) | `25.00%` | `26.43%` | **`26.85%`** | | |
| | **ARC-Challenge** (`allenai/ai2_arc`) | `test` (`1,172`) | `25.00%` | `25.34%` | **`25.43%`** | | |
| | **CommonsenseQA** (`tau/commonsense_qa`) | `val` (`1,221`) | `20.00%` | `19.16%` | **`23.10%`** | | |
| | **9-Task Multiple-Choice Macro Average** | `13,144` total | `35.56%` | **`39.16%`** | `38.63%` | | |
| | Held-Out Math Perplexity | Split (`n`) | `GRT-20M-Base` (Zero-Shot) | Standard `PairAdam` Fine-Tune | | |
| |---|---:|---:|---:| | |
| | **GSM8K (`openai/gsm8k`)** | `test` (`500`) | `60.12` PPL | `12.52` PPL | | |
| Reproduce with: | |
| ```bash | |
| python -m training.slm_bench --checkpoint ./ --tasks boolq,arc_easy --device cuda | |
| ``` | |
| Zero-shot, via lm-evaluation-harness: | |
| ```bash | |
| lm_eval --model hf --model_args pretrained=MatterWaveAI/grt-20m-base \ | |
| --tasks boolq,arc_easy --batch_size 64 | |
| ``` | |
| ## Usage | |
| ```python | |
| import torch | |
| import grt # registers model_type "grt" with the transformers Auto classes | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| model = AutoModelForCausalLM.from_pretrained("MatterWaveAI/grt-20m-base", dtype=torch.float32) | |
| tokenizer = AutoTokenizer.from_pretrained("MatterWaveAI/grt-20m-base") | |
| input_ids = tokenizer("The geometry of language is", return_tensors="pt").input_ids | |
| output = model.generate(input_ids, max_new_tokens=50, do_sample=False) | |
| print(tokenizer.decode(output[0], skip_special_tokens=True)) | |
| ``` | |
| The package is pure PyTorch and runs on CPU. Optimized CUDA kernels are optional | |
| and load automatically when installed. | |
| > **`trust_remote_code=True` is not supported for this repository.** The model | |
| > code is distributed as a pip package rather than as Hub remote code, because | |
| > the implementation spans a package with absolute imports that a flat Hub | |
| > snapshot cannot resolve. Install `geometric-rotation-transformer` first; | |
| > `import grt` then registers `model_type: "grt"` so the Auto classes resolve | |
| > these weights. | |
| ### GGUF | |
| `grt-20m-f32.gguf` is the same weights in GGUF form, loadable by the MatterWave | |
| llama.cpp fork. Verify any pull against `SHA256SUMS` in this repo: | |
| ```bash | |
| sha256sum --ignore-missing -c SHA256SUMS | |
| ``` | |
| ## Architecture | |
| GRT replaces the standard Q/K/V projections and FFN weight matrices with | |
| cascades of learned 2x2 Givens rotations in a butterfly pattern. Each stage is | |
| modulated by data-dependent Clifford features: | |
| ``` | |
| feat = [S0, V1, V2, B12] # scalar energy, direction, cross-pair bivector | |
| theta = angles + tanh(modulation) # base angle + content-dependent nudge | |
| r0 = x0*cos(theta) - x1*sin(theta) | |
| r1 = x0*sin(theta) + x1*cos(theta) | |
| ``` | |
| Components: | |
| - **RotationalEmbedding** — tokens stored as phase angles on the unit hypersphere | |
| - **RotationalAttention** — Q/K/V via Givens cascades | |
| - **RotationalFFN** — feed-forward via meta-coupled rotation stages | |
| - **PearlNorm** — pair-radius normalization preserving per-pair phase | |
| - **RadialGELU** — activation gating pair magnitudes, preserving direction | |
| - **PairAdam** — AdamW with factored second moments: 1D params (angles, scales) | |
| use standard diagonal moments, 2D params (coupling matrices) use a row factor | |
| [m,1] and col factor [1,n] reconstructed as an outer product, at O(m+n) | |
| instead of O(mn). Optionally clamps per-step angular updates. | |
| Because every projection is a rotation, the transform is norm-preserving by | |
| construction rather than by normalization. | |
| ## Intended Use & Limitations | |
| Intended for research into structured/rotational architectures and as a | |
| deterministic CI upstream. Not intended for production use, deployment, or any | |
| task requiring factual output. | |
| - **No instruction tuning.** This is a base model trained on next-token | |
| prediction only. | |
| - **No chat template.** There is no `chat_template` in this repo. | |
| - **No tool calling, no structured/schema-constrained output, no vision.** | |
| - English only. The mix includes web crawl (RefinedWeb, ClimbMix), encyclopedic | |
| text, synthetic textbooks (Cosmopedia), math (FineMath) and Q&A | |
| (StackExchange), so it inherits the coverage gaps and biases of all of them. | |
| - Generated text is frequently incoherent past a clause or two. Do not expect | |
| instruction-following, factual reliability, or coherent long-form output. | |
| - No safety tuning, RLHF, or content filtering of any kind. | |
| - Context is capped at 1,024 tokens. | |
| ## Citation | |
| ```bibtex | |
| @article{grt2026, | |
| title={Geometric Rotation Transformer: Structured SO(n) Rotations for Efficient Sequence Modeling}, | |
| author={MatterWave AI}, | |
| year={2026}, | |
| url={https://github.com/matterwaveai/grt} | |
| } | |
| ``` | |
| ## License | |
| Apache 2.0 | |