Release Sol Milkshake 3M Base

Private initial release of the validated 900M recovery checkpoint with standalone MLX runtime, tokenizer, provenance, and raw evaluations.

LICENSE

@@ -0,0 +1,12 @@
+Creative Commons Attribution 4.0 International
+
+Copyright (c) 2026 Sol Labs
+
+Licensed under the Creative Commons Attribution 4.0 International License.
+You may obtain a copy of the License at:
+
+https://creativecommons.org/licenses/by/4.0/legalcode
+
+Unless required by applicable law or agreed to in writing, material distributed
+under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+CONDITIONS OF ANY KIND, either express or implied.

README.md

@@ -0,0 +1,181 @@
+---
+license: cc-by-4.0
+language:
+  - en
+pipeline_tag: text-generation
+tags:
+  - causal-lm
+  - decoder-only
+  - small-language-model
+  - recurrent-depth
+  - mlx
+  - ngpt
+  - research
+---
+
+![Sol Milkshake](./sol-milkshake-banner.svg)
+
+# Sol Milkshake 3M Base
+
+Sol Milkshake is an exact **2,990,000-parameter** recurrent, hyperspherical decoder-only language model trained from scratch and continued through **2,526,565,888 total token exposures**.
+
+It combines recurrent block reuse, grouped-query attention, parameter-free XSA value subtraction, value residuals, tensorized 2–5-gram memory, completed-chunk rolling memory, adaptive recurrent routing, and nGPT-style hyperspherical optimization.
+
+**This is a base model, not an instruction-tuned assistant.**
+
+## Model summary
+
+| Property | Value |
+| --- | --- |
+| Parameters | **2,990,000** |
+| Total token exposures | **2,526,565,888** |
+| Architecture | Decoder-only recurrent causal language model |
+| Physical / effective blocks | 5 / 11 |
+| Residual width | 192 |
+| Vocabulary | 2,048-token byte-level BPE |
+| Maximum context | 2,048 tokens |
+| Weight format | MLX NPZ |
+| License | CC BY 4.0 |
+
+## Quick start
+
+Install the runtime dependencies:
+
+```bash
+pip install "mlx>=0.29" "tokenizers>=0.22" huggingface_hub
+```
+
+Load the model and generate text on Apple Silicon:
+
+```python
+from huggingface_hub import snapshot_download
+import sys
+
+model_dir = snapshot_download("solintellegence/Sol-Milkshake-3M-Base")
+sys.path.insert(0, model_dir)
+
+from modeling_sol_milkshake import load_model, generate
+
+model, tokenizer = load_model(model_dir)
+
+text = generate(
+    model,
+    tokenizer,
+    "The future of efficient language models is",
+    max_new_tokens=64,
+)
+
+print(text)
+```
+
+The repository includes the complete standalone MLX implementation and does not require the original training repository.
+
+## Results
+
+| Benchmark | Examples | Score |
+| --- | ---: | ---: |
+| HellaSwag | 10,042 | 25.02% |
+| ARC-Easy | 2,376 | 29.76% |
+| ARC-Challenge | 1,172 | 22.70% |
+| PIQA | 1,838 | 52.50% |
+| ArithMark-3 | 1,000 | 30.60% |
+| Intelligence Index | — | **3.158** |
+
+The four language-model tasks were evaluated zero-shot on their complete evaluation splits with `lm-eval` 0.4.12 using normalized accuracy.
+
+ArithMark-3 used its independent-tokenization normalized-accuracy protocol.
+
+Raw evaluation outputs are included under `evals/`.
+
+These measurements are not independently verified, and no leaderboard position is claimed. The release checkpoint was selected from recovery checkpoints using this evaluation suite, so the results should not be treated as a pristine held-out estimate.
+
+## Architecture
+
+| Component | Configuration |
+| --- | --- |
+| Stored parameters | **2,990,000** |
+| Physical / effective blocks | 5 / 11 |
+| Recurrent layout | 1 prelude + 3 middle blocks × 3 passes + 1 coda |
+| Residual width | 192 |
+| Attention | 6 query heads, 2 key/value heads, head dimension 32 |
+| Attention modification | XSA value subtraction after causal attention |
+| Attention normalization | Unit-normalized Q/K with RoPE |
+| Recurrent routing | Full first pass, then 75% and 50% token capacity |
+| MLP | Gated SiLU MLP, width 512 |
+| Vocabulary | 2,048-token byte-level BPE |
+| Learned memory | Rank-26 tensorized 2–5-gram memory |
+| Rolling memory | 32-token chunks, 32 slots, width 64 |
+| Embeddings | Tied input/output table |
+| Maximum context | 2,048 tokens |
+
+## Training data
+
+Initial pretraining used a staged English curriculum built from:
+
+- FineWeb-Edu
+- FinePDFs-Edu
+- English UltraFineWeb multi-domain and question-answer subsets
+- Cosmopedia
+- FineMath
+
+Recovery training used a **70% / 15% / 15%** mixture of the original frozen curriculum, Cosmopedia-v2 English, and FinePhrase.
+
+The tokenizer and prepared training streams were frozen before training. Existing quality filtering, deduplication, and decontamination rules were retained.
+
+## Intended use
+
+Sol Milkshake is intended for research on highly parameter-efficient language modeling, including:
+
+- extremely small language models
+- recurrent-depth architectures
+- nGPT-style hyperspherical optimization
+- compact learned and rolling memory
+- grouped-query attention at small scale
+- Apple Silicon and MLX inference research
+
+Because this is a base model, users should expect raw next-token completion behavior rather than instruction-following or assistant-style responses.
+
+## Limitations
+
+At approximately 3M parameters, Sol Milkshake has severe capability limitations.
+
+Outputs may be:
+
+- inconsistent
+- factually incorrect
+- repetitive
+- incoherent
+- unsafe
+
+The model is not instruction tuned and should not be treated as a conversational assistant.
+
+Do not use it for high-stakes decisions.
+
+Users are responsible for evaluating outputs and determining whether the model is suitable for their application.
+
+## Repository files
+
+| File | Purpose |
+| --- | --- |
+| `model.npz` | MLX model weights |
+| `modeling_sol_milkshake.py` | Standalone model, loader, and generator |
+| `sol_config.py` | Frozen architecture configuration |
+| `config.json` | Model metadata |
+| `tokenizer.json` | Frozen tokenizer |
+| `tokenizer_config.json` | Tokenizer metadata |
+| `training_state.json` | Training and provenance summary |
+| `evals/` | Raw benchmark results |
+
+## Reproducibility and evaluation notes
+
+The repository includes raw evaluation outputs under `evals/`.
+
+Reported benchmark values should be interpreted as reference measurements for this release. They have not been independently verified, and checkpoint selection used the same suite.
+
+No leaderboard position is claimed.
+
+## License
+
+Sol Milkshake 3M Base is released under the **CC BY 4.0** license.
+
+Dataset licenses and terms remain with their respective owners.

config.json

@@ -0,0 +1,16 @@
+{
+  "architectures": ["SolMilkshake"],
+  "model_type": "sol_milkshake",
+  "model_name": "Sol Milkshake 3M Base",
+  "parameter_count": 2990000,
+  "vocab_size": 2048,
+  "hidden_size": 192,
+  "max_position_embeddings": 2048,
+  "num_hidden_layers": 5,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "head_dim": 32,
+  "intermediate_size": 512,
+  "tie_word_embeddings": true,
+  "framework": "mlx"
+}

evals/open_slm_lm_eval_0.4.12.json

@@ -0,0 +1,334 @@
+{
+  "results": {
+    "hellaswag": {
+      "name": "hellaswag",
+      "alias": "hellaswag",
+      "sample_len": 10042,
+      "acc,none": 0.2577175861382195,
+      "acc_stderr,none": 0.004364838000336011,
+      "acc_norm,none": 0.2502489543915555,
+      "acc_norm_stderr,none": 0.004322710911026276
+    },
+    "arc_easy": {
+      "name": "arc_easy",
+      "alias": "arc_easy",
+      "sample_len": 2376,
+      "acc,none": 0.2857744107744108,
+      "acc_stderr,none": 0.009270380606981266,
+      "acc_norm,none": 0.2975589225589226,
+      "acc_norm_stderr,none": 0.009381226721815406
+    },
+    "arc_challenge": {
+      "name": "arc_challenge",
+      "alias": "arc_challenge",
+      "sample_len": 1172,
+      "acc,none": 0.17320819112627986,
+      "acc_stderr,none": 0.011058694183280343,
+      "acc_norm,none": 0.22696245733788395,
+      "acc_norm_stderr,none": 0.012240491536132965
+    },
+    "piqa": {
+      "name": "piqa",
+      "alias": "piqa",
+      "sample_len": 1838,
+      "acc,none": 0.5315560391730142,
+      "acc_stderr,none": 0.011642567710878958,
+      "acc_norm,none": 0.5250272034820457,
+      "acc_norm_stderr,none": 0.011651200967623562
+    }
+  },
+  "group_subtasks": {},
+  "configs": {
+    "arc_challenge": {
+      "task": "arc_challenge",
+      "dataset_path": "allenai/ai2_arc",
+      "dataset_name": "ARC-Challenge",
+      "training_split": "train",
+      "validation_split": "validation",
+      "test_split": "test",
+      "doc_to_text": "Question: {{question}}\nAnswer:",
+      "doc_to_target": "{{choices.label.index(answerKey)}}",
+      "unsafe_code": false,
+      "doc_to_choice": "{{choices.text}}",
+      "description": "",
+      "target_delimiter": " ",
+      "fewshot_delimiter": "\n\n",
+      "fewshot_config": {
+        "sampler": "default",
+        "split": null,
+        "process_docs": null,
+        "fewshot_indices": null,
+        "samples": null,
+        "doc_to_text": "Question: {{question}}\nAnswer:",
+        "doc_to_choice": "{{choices.text}}",
+        "doc_to_target": "{{choices.label.index(answerKey)}}",
+        "gen_prefix": null,
+        "fewshot_delimiter": "\n\n",
+        "target_delimiter": " "
+      },
+      "num_fewshot": 0,
+      "metric_list": [
+        {
+          "metric": "acc",
+          "aggregation": "mean",
+          "higher_is_better": true
+        },
+        {
+          "metric": "acc_norm",
+          "aggregation": "mean",
+          "higher_is_better": true
+        }
+      ],
+      "output_type": "multiple_choice",
+      "repeats": 1,
+      "should_decontaminate": true,
+      "doc_to_decontamination_query": "Question: {{question}}\nAnswer:",
+      "metadata": {
+        "version": 1.0,
+        "config_source": "/Users/jonah/Documents/ChatGPT/Training/.venv/lib/python3.14/site-packages/lm_eval/tasks/arc/arc_challenge.yaml"
+      }
+    },
+    "arc_easy": {
+      "task": "arc_easy",
+      "dataset_path": "allenai/ai2_arc",
+      "dataset_name": "ARC-Easy",
+      "training_split": "train",
+      "validation_split": "validation",
+      "test_split": "test",
+      "doc_to_text": "Question: {{question}}\nAnswer:",
+      "doc_to_target": "{{choices.label.index(answerKey)}}",
+      "unsafe_code": false,
+      "doc_to_choice": "{{choices.text}}",
+      "description": "",
+      "target_delimiter": " ",
+      "fewshot_delimiter": "\n\n",
+      "fewshot_config": {
+        "sampler": "default",
+        "split": null,
+        "process_docs": null,
+        "fewshot_indices": null,
+        "samples": null,
+        "doc_to_text": "Question: {{question}}\nAnswer:",
+        "doc_to_choice": "{{choices.text}}",
+        "doc_to_target": "{{choices.label.index(answerKey)}}",
+        "gen_prefix": null,
+        "fewshot_delimiter": "\n\n",
+        "target_delimiter": " "
+      },
+      "num_fewshot": 0,
+      "metric_list": [
+        {
+          "metric": "acc",
+          "aggregation": "mean",
+          "higher_is_better": true
+        },
+        {
+          "metric": "acc_norm",
+          "aggregation": "mean",
+          "higher_is_better": true
+        }
+      ],
+      "output_type": "multiple_choice",
+      "repeats": 1,
+      "should_decontaminate": true,
+      "doc_to_decontamination_query": "Question: {{question}}\nAnswer:",
+      "metadata": {
+        "version": 1.0,
+        "config_source": "/Users/jonah/Documents/ChatGPT/Training/.venv/lib/python3.14/site-packages/lm_eval/tasks/arc/arc_easy.yaml"
+      }
+    },
+    "hellaswag": {
+      "task": "hellaswag",
+      "dataset_path": "Rowan/hellaswag",
+      "training_split": "train",
+      "validation_split": "validation",
+      "process_docs": "def process_docs(dataset: datasets.Dataset) -> datasets.Dataset:\n    def _process_doc(doc):\n        ctx = doc[\"ctx_a\"] + \" \" + doc[\"ctx_b\"].capitalize()\n        out_doc = {\n            \"query\": preprocess(doc[\"activity_label\"] + \": \" + ctx),\n            \"choices\": [preprocess(ending) for ending in doc[\"endings\"]],\n            \"gold\": int(doc[\"label\"]),\n        }\n        return out_doc\n\n    return dataset.map(_process_doc)\n",
+      "doc_to_text": "{{query}}",
+      "doc_to_target": "{{label}}",
+      "unsafe_code": false,
+      "doc_to_choice": "choices",
+      "description": "",
+      "target_delimiter": " ",
+      "fewshot_delimiter": "\n\n",
+      "fewshot_config": {
+        "sampler": "default",
+        "split": null,
+        "process_docs": "<function process_docs at 0x136493740>",
+        "fewshot_indices": null,
+        "samples": null,
+        "doc_to_text": "{{query}}",
+        "doc_to_choice": "choices",
+        "doc_to_target": "{{label}}",
+        "gen_prefix": null,
+        "fewshot_delimiter": "\n\n",
+        "target_delimiter": " "
+      },
+      "num_fewshot": 0,
+      "metric_list": [
+        {
+          "metric": "acc",
+          "aggregation": "mean",
+          "higher_is_better": true
+        },
+        {
+          "metric": "acc_norm",
+          "aggregation": "mean",
+          "higher_is_better": true
+        }
+      ],
+      "output_type": "multiple_choice",
+      "repeats": 1,
+      "should_decontaminate": false,
+      "metadata": {
+        "version": 1.0,
+        "config_source": "/Users/jonah/Documents/ChatGPT/Training/.venv/lib/python3.14/site-packages/lm_eval/tasks/hellaswag/hellaswag.yaml"
+      }
+    },
+    "piqa": {
+      "task": "piqa",
+      "dataset_path": "baber/piqa",
+      "training_split": "train",
+      "validation_split": "validation",
+      "doc_to_text": "Question: {{goal}}\nAnswer:",
+      "doc_to_target": "label",
+      "unsafe_code": false,
+      "doc_to_choice": "{{[sol1, sol2]}}",
+      "description": "",
+      "target_delimiter": " ",
+      "fewshot_delimiter": "\n\n",
+      "fewshot_config": {
+        "sampler": "default",
+        "split": null,
+        "process_docs": null,
+        "fewshot_indices": null,
+        "samples": null,
+        "doc_to_text": "Question: {{goal}}\nAnswer:",
+        "doc_to_choice": "{{[sol1, sol2]}}",
+        "doc_to_target": "label",
+        "gen_prefix": null,
+        "fewshot_delimiter": "\n\n",
+        "target_delimiter": " "
+      },
+      "num_fewshot": 0,
+      "metric_list": [
+        {
+          "metric": "acc",
+          "aggregation": "mean",
+          "higher_is_better": true
+        },
+        {
+          "metric": "acc_norm",
+          "aggregation": "mean",
+          "higher_is_better": true
+        }
+      ],
+      "output_type": "multiple_choice",
+      "repeats": 1,
+      "should_decontaminate": true,
+      "doc_to_decontamination_query": "goal",
+      "metadata": {
+        "version": 1.0,
+        "config_source": "/Users/jonah/Documents/ChatGPT/Training/.venv/lib/python3.14/site-packages/lm_eval/tasks/piqa/piqa.yaml"
+      }
+    }
+  },
+  "versions": {
+    "arc_challenge": 1.0,
+    "arc_easy": 1.0,
+    "hellaswag": 1.0,
+    "piqa": 1.0
+  },
+  "n-shot": {
+    "arc_challenge": 0,
+    "arc_easy": 0,
+    "hellaswag": 0,
+    "piqa": 0
+  },
+  "higher_is_better": {
+    "arc_challenge": {
+      "acc": true,
+      "acc_norm": true
+    },
+    "arc_easy": {
+      "acc": true,
+      "acc_norm": true
+    },
+    "hellaswag": {
+      "acc": true,
+      "acc_norm": true
+    },
+    "piqa": {
+      "acc": true,
+      "acc_norm": true
+    }
+  },
+  "n-samples": {
+    "hellaswag": {
+      "original": 10042,
+      "effective": 10042
+    },
+    "arc_easy": {
+      "original": 2376,
+      "effective": 2376
+    },
+    "arc_challenge": {
+      "original": 1172,
+      "effective": 1172
+    },
+    "piqa": {
+      "original": 1838,
+      "effective": 1838
+    }
+  },
+  "config": {
+    "model": "MLXCheckpointLM",
+    "model_args": null,
+    "batch_size": 2,
+    "batch_sizes": [],
+    "device": null,
+    "use_cache": null,
+    "limit": null,
+    "bootstrap_iters": 100000,
+    "gen_kwargs": null,
+    "random_seed": 0,
+    "numpy_seed": 1234,
+    "torch_seed": 1234,
+    "fewshot_seed": 1234
+  },
+  "git_hash": "6bdcd14",
+  "date": 1789940047.395892,
+  "pretty_env_info": "PyTorch version: 2.13.0\nIs debug build: False\nCUDA used to build PyTorch: None\nROCM used to build PyTorch: N/A\n\nOS: macOS 27.0 (arm64)\nGCC version: Could not collect\nClang version: 21.0.0 (clang-2100.1.1.101)\nCMake version: version 4.3.3\nLibc version: N/A\n\nPython version: 3.14.7 (main, Aug  5 2026, 10:29:49) [Clang 21.0.0 (clang-2100.1.1.101)] (64-bit runtime)\nPython platform: macOS-27.0-arm64-arm-64bit-Mach-O\nIs CUDA available: False\nCUDA runtime version: No CUDA\nCUDA_MODULE_LOADING set to: N/A\nGPU models and configuration: No CUDA\nNvidia driver version: No CUDA\ncuDNN version: No CUDA\nIs XPU available: False\nHIP runtime version: N/A\nMIOpen runtime version: N/A\nIs XNNPACK available: False\nCaching allocator config: N/A\n\nCPU:\nApple M5 Max\n\nVersions of relevant libraries:\n[pip3] numpy==2.5.2\n[pip3] torch==2.13.0\n[conda] Could not collect",
+  "transformers_version": "5.16.1",
+  "lm_eval_version": "0.4.12",
+  "upper_git_hash": null,
+  "tokenizer_pad_token": [
+    null,
+    "None"
+  ],
+  "tokenizer_eos_token": [
+    "<|eos|>",
+    "1"
+  ],
+  "tokenizer_bos_token": [
+    null,
+    "None"
+  ],
+  "eot_token_id": 1,
+  "max_length": 2048,
+  "xsa_evaluation": {
+    "checkpoint": "/Users/jonah/Documents/ChatGPT/Training/runs/sol_milkshake_recovery_500m_v1/checkpoint_000900005888",
+    "checkpoint_tokens_seen": null,
+    "lm_eval_version": "0.4.12",
+    "tasks": [
+      "hellaswag",
+      "arc_easy",
+      "arc_challenge",
+      "piqa"
+    ],
+    "num_fewshot": 0,
+    "metric": "acc_norm",
+    "max_length": 2048,
+    "batch_size": 2,
+    "elapsed_seconds": 174.03599381446838
+  }
+}

model.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8bd1efee28989bd49c85c70d72824c91c3872b871a34f48ccad21903b8ef36a
+size 11980144

modeling_sol_milkshake.py

@@ -0,0 +1,169 @@
+"""Exact 2,990,000-parameter Sol Milkshake model implemented in MLX."""
+from __future__ import annotations
+import math
+from collections import OrderedDict
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten
+from sol_config import DEPLOYED_PARAMS, SolConfig
+
+def unit_norm(x):
+    xf=x.astype(mx.float32)
+    return (xf*mx.rsqrt(mx.maximum(mx.sum(xf*xf,axis=-1,keepdims=True),mx.array(1e-12)))).astype(x.dtype)
+
+def _count(module): return sum(int(x.size) for _,x in tree_flatten(module.trainable_parameters()))
+
+def _rope(q,k,theta):
+    n,d=q.shape[-2:]; inv=mx.exp(mx.arange(0,d,2,dtype=mx.float32)*(-math.log(theta)/d))
+    ang=mx.arange(n,dtype=mx.float32)[:,None]*inv[None,:]; cs,sn=mx.cos(ang)[None,None],mx.sin(ang)[None,None]
+    def rot(x):
+        xf=x.astype(mx.float32); a,b=xf[...,::2],xf[...,1::2]
+        return mx.stack((a*cs-b*sn,a*sn+b*cs),axis=-1).reshape(x.shape).astype(x.dtype)
+    return rot(q),rot(k)
+
+class TiedTokens(nn.Module):
+    def __init__(self,c):
+        super().__init__(); self.weight=unit_norm(mx.random.normal((c.vocab_size,c.d_model)))
+    def encode(self,ids): return self.weight[ids]
+    def decode(self,x,scale): return x.astype(mx.float32)@unit_norm(self.weight).astype(mx.float32).T*mx.exp(scale)
+
+class TNGram(nn.Module):
+    """Causal hashed CP retrieval for orders 2--5 (exactly 573,460 parameters)."""
+    def __init__(self,c):
+        super().__init__(); r=c.tn_rank; self.token_factor=mx.random.normal((c.vocab_size,r))*.02
+        self.hash_tables=mx.random.normal((4,4950,r))*.02; self.order_factors=mx.ones((4,r))
+        self.rank_to_hidden=mx.random.normal((r,c.d_model))*.02; self.output_bias=mx.zeros((c.d_model,))
+        self.gates=mx.full((4,),-2.0); self.order_mix=mx.eye(4,r); self.context_controls=mx.zeros((16,))
+    def __call__(self,ids):
+        b,t=ids.shape; out=mx.broadcast_to(self.output_bias,(b,t,self.output_bias.shape[0]))
+        for oi,order in enumerate(range(2,6)):
+            h=mx.zeros_like(ids); prod=mx.ones((b,t,self.token_factor.shape[1]),dtype=mx.float32); valid=mx.arange(t)>=order-1
+            for lag in range(1,order):
+                shifted=mx.zeros((b,t),dtype=ids.dtype) if lag>=t else mx.concatenate((mx.zeros((b,lag),dtype=ids.dtype),ids[:,:-lag]),axis=1)
+                h=(h*(131+oi*6)+shifted)%4950; prod=prod*self.token_factor[shifted].astype(mx.float32)
+            rank=prod*self.hash_tables[oi,h]*self.order_factors[oi]
+            rank=rank+self.order_mix[oi][None,None,:]*mx.mean(mx.tanh(self.context_controls[oi*4:(oi+1)*4]))
+            out=out+(rank@self.rank_to_hidden)*mx.sigmoid(self.gates[oi])*valid[None,:,None]
+        return out
+
+class ShakeMemory(nn.Module):
+    """Completed-chunk rolling memory (exactly 51,520 parameters)."""
+    def __init__(self,c):
+        super().__init__(); w=c.memory_width; self.chunk_size=c.chunk_size; self.slots=c.memory_slots
+        self.summary_router=nn.Linear(c.d_model,1,bias=False); self.compress=nn.Linear(c.d_model,w,bias=False)
+        self.query=nn.Linear(c.d_model,w,bias=False); self.key=nn.Linear(w,w,bias=False)
+        self.value=nn.Linear(w,w,bias=False); self.output=nn.Linear(w,c.d_model,bias=False)
+        self.slot_positions=mx.random.normal((c.memory_slots,c.d_model))*.01
+        self.decay=mx.zeros((w,)); self.update_gate=mx.zeros((w,))
+    def __call__(self,x):
+        b,t,d=x.shape; chunks=[]; reads=[]
+        for start in range(0,t,self.chunk_size):
+            end=min(t,start+self.chunk_size); length=end-start
+            if chunks:
+                mem=mx.stack(chunks[-self.slots:],axis=1); count=mem.shape[1]
+                mem=mem+self.compress(self.slot_positions[:count])[None]*mx.sigmoid(self.decay)[None,None]
+                a=mx.softmax((self.query(x[:,start:end]).astype(mx.float32)@self.key(mem).astype(mx.float32).transpose(0,2,1))/math.sqrt(mem.shape[-1]),axis=-1)
+                reads.append(self.output((a@self.value(mem).astype(mx.float32)).astype(x.dtype)))
+            else: reads.append(mx.zeros((b,length,d),dtype=x.dtype))
+            if length==self.chunk_size:
+                weights=mx.softmax(self.summary_router(x[:,start:end]).astype(mx.float32).squeeze(-1),axis=-1)
+                summary=mx.sum(weights[:,:,None]*x[:,start:end].astype(mx.float32),axis=1)
+                chunks.append(self.compress(summary.astype(x.dtype))*mx.sigmoid(self.update_gate))
+        return mx.concatenate(reads,axis=1)
+
+class Attention(nn.Module):
+    def __init__(self,c):
+        super().__init__(); d=c.d_model; self.q=nn.Linear(d,d,bias=False); self.k=nn.Linear(d,c.n_kv_heads*c.head_dim,bias=False)
+        self.v=nn.Linear(d,c.n_kv_heads*c.head_dim,bias=False); self.o=nn.Linear(d,d,bias=False)
+        self.hq=c.n_q_heads; self.hk=c.n_kv_heads; self.hd=c.head_dim; self.theta=c.rope_theta
+        self.value_current=mx.zeros((1,)); self.value_prior=mx.zeros((1,))
+    def __call__(self,x,prior=None):
+        b,t,_=x.shape; q=unit_norm(self.q(x).reshape(b,t,self.hq,self.hd)).transpose(0,2,1,3)
+        k=unit_norm(self.k(x).reshape(b,t,self.hk,self.hd)).transpose(0,2,1,3); v=self.v(x).reshape(b,t,self.hk,self.hd).transpose(0,2,1,3)
+        if prior is not None: v=mx.sigmoid(self.value_current)*v+mx.sigmoid(self.value_prior)*prior
+        q,k=_rope(q,k,self.theta); k=mx.repeat(k,self.hq//self.hk,axis=1); vr=mx.repeat(v,self.hq//self.hk,axis=1)
+        scores=(q.astype(mx.float32)@k.astype(mx.float32).transpose(0,1,3,2))/math.sqrt(self.hd)
+        pos=mx.arange(t); mask=mx.where(pos[:,None]>=pos[None,:],0.0,-1e9)
+        attended=(mx.softmax(scores+mask[None,None],axis=-1)@vr.astype(mx.float32)).astype(x.dtype)
+        vh=unit_norm(vr); attended=attended-mx.sum(attended.astype(mx.float32)*vh.astype(mx.float32),axis=-1,keepdims=True).astype(x.dtype)*vh
+        return self.o(attended.transpose(0,2,1,3).reshape(b,t,-1)),v
+
+class Block(nn.Module):
+    def __init__(self,c):
+        super().__init__(); self.attention=Attention(c); self.gate=nn.Linear(c.d_model,c.ffn_hidden,bias=False)
+        self.up=nn.Linear(c.d_model,c.ffn_hidden,bias=False); self.down=nn.Linear(c.ffn_hidden,c.d_model,bias=False)
+        self.attn_step=mx.zeros((c.d_model,)); self.ffn_step=mx.zeros((c.d_model,))
+    def __call__(self,x,prior=None):
+        a,v=self.attention(unit_norm(x),prior); x=unit_norm(x+mx.sigmoid(self.attn_step)*(unit_norm(a)-x))
+        f=self.down(nn.silu(self.gate(x))*self.up(x)); return unit_norm(x+mx.sigmoid(self.ffn_step)*(unit_norm(f)-x)),v
+
+class SolMilkshake(nn.Module):
+    def __init__(self,c=None):
+        super().__init__(); self.config=c or SolConfig(); self.tokens=TiedTokens(self.config); self.tn_gram=TNGram(self.config)
+        self.shake_memory=ShakeMemory(self.config); self.blocks=[Block(self.config) for _ in range(5)]
+        self.pass_embeddings=mx.zeros((3,self.config.d_model)); self.mod_routers=[nn.Linear(self.config.d_model,1) for _ in range(6)]
+        self.logit_scale=mx.zeros((self.config.vocab_size,)); self.misc_controls=mx.zeros((12,)); self._verify_ledger()
+    def _select(self,x,old,router,capacity,lam):
+        delta=1-mx.sum(unit_norm(x).astype(mx.float32)*unit_norm(old).astype(mx.float32),axis=-1)
+        score=router(x).squeeze(-1)+lam*delta
+        positions=mx.arange(x.shape[1])
+        mask=(positions%4!=3) if capacity==.75 else (positions%2==0)
+        return mx.broadcast_to(mask[None,:],score.shape),mx.sigmoid(score)
+    def __call__(self,ids):
+        x=self.tokens.encode(ids); x=unit_norm(x+mx.sigmoid(self.misc_controls[0])*self.tn_gram(ids))
+        x=unit_norm(x+mx.sigmoid(self.misc_controls[1])*self.shake_memory(x)); x,v=self.blocks[0](x); ri=0
+        for p in range(3):
+            old=x; x=unit_norm(x+self.pass_embeddings[p]); prior=None
+            for bi in range(1,4):
+                proposal,prior=self.blocks[bi](x,prior)
+                if p==0: x=proposal
+                else:
+                    mask,route_gate=self._select(x,old,self.mod_routers[ri],.75 if p==1 else .5,self.misc_controls[2+ri]); ri+=1
+                    routed=unit_norm(x+route_gate[:,:,None]*(proposal-x))
+                    x=mx.where(mask[:,:,None],routed,x)
+        x,_=self.blocks[4](x,v); return self.tokens.decode(x,self.logit_scale)
+    def parameter_ledger(self):
+        ledger=OrderedDict(token_embedding_tied_head=int(self.tokens.weight.size),attention_total=sum(_count(b.attention)-2 for b in self.blocks),
+          swiglu_total=sum(_count(b.gate)+_count(b.up)+_count(b.down) for b in self.blocks),tn_gram=_count(self.tn_gram),shake_memory=_count(self.shake_memory),
+          ngpt_steps=sum(int(b.attn_step.size+b.ffn_step.size) for b in self.blocks),pass_embeddings=int(self.pass_embeddings.size),
+          mod_routers=sum(_count(r) for r in self.mod_routers),logit_scale=int(self.logit_scale.size),value_residual=10,
+          misc_controls=int(self.misc_controls.size),xsa=0); ledger['deployed_total']=sum(ledger.values()); return dict(ledger)
+    def deployed_parameter_count(self): return _count(self)
+    def _verify_ledger(self):
+        expected={"token_embedding_tied_head":393216,"attention_total":491520,"swiglu_total":1474560,"tn_gram":573460,"shake_memory":51520,"ngpt_steps":1920,"pass_embeddings":576,"mod_routers":1158,"logit_scale":2048,"value_residual":10,"misc_controls":12,"xsa":0,"deployed_total":DEPLOYED_PARAMS}
+        if self.parameter_ledger()!=expected or self.deployed_parameter_count()!=DEPLOYED_PARAMS: raise RuntimeError(f"parameter ledger mismatch: {self.parameter_ledger()}, actual={self.deployed_parameter_count()}")
+
+def future_token_causality_error(model,sequence_length=35):
+    ids=mx.arange(sequence_length,dtype=mx.int32)[None]%model.config.vocab_size; boundary=sequence_length//2
+    changed=mx.concatenate((ids[:,:boundary],(ids[:,boundary:]+17)%model.config.vocab_size),axis=1)
+    a,b=model(ids),model(changed); mx.eval(a,b); return float(mx.max(mx.abs(a[:,:boundary]-b[:,:boundary])))
+
+def load_model(model_dir):
+    """Load the standalone MLX model and frozen tokenizer from a snapshot."""
+    from pathlib import Path
+    from tokenizers import Tokenizer
+
+    model_dir = Path(model_dir)
+    model = SolMilkshake(SolConfig())
+    model.load_weights(str(model_dir / "model.npz"))
+    mx.eval(model.parameters())
+    tokenizer = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
+    return model, tokenizer
+
+def generate(model, tokenizer, prompt, max_new_tokens=64, temperature=0.0, seed=0):
+    """Generate a single continuation with greedy or temperature sampling."""
+    mx.random.seed(seed)
+    tokens = tokenizer.encode(prompt).ids
+    if not tokens:
+        tokens = [0]
+    for _ in range(max_new_tokens):
+        context = tokens[-model.config.max_context:]
+        logits = model(mx.array([context], dtype=mx.int32))[0, -1]
+        if temperature and temperature > 0:
+            next_token = int(mx.random.categorical(logits / temperature).item())
+        else:
+            next_token = int(mx.argmax(logits).item())
+        tokens.append(next_token)
+        if next_token == 1:
+            break
+    return tokenizer.decode(tokens)

requirements.txt

@@ -0,0 +1,4 @@
+mlx>=0.29
+tokenizers>=0.22
+numpy>=2.0
+huggingface_hub>=1.0

sol_config.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+
+MODEL_NAME = "Sol Milkshake"
+ORGANIZATION = "Sol Labs"
+DEPLOYED_PARAMS = 2_990_000
+VOCAB_SIZE = 2_048
+D_MODEL = 192
+MAX_CONTEXT = 2_048
+N_PHYSICAL_BLOCKS = 5
+N_Q_HEADS = 6
+N_KV_HEADS = 2
+HEAD_DIM = 32
+FFN_HIDDEN = 512
+TN_RANK = 26
+TIED_EMBEDDING_HEAD = True
+XSA_PARAMETER_COUNT = 0
+TOKENS_PER_PARAMETER = 544
+GLOBAL_TOKENS = 1_626_560_000
+GLOBAL_BATCH = 32_768
+COMPLETE_UPDATES = 49_638
+FINAL_PARTIAL_TOKENS = 22_016
+
+assert DEPLOYED_PARAMS * TOKENS_PER_PARAMETER == GLOBAL_TOKENS
+assert COMPLETE_UPDATES * GLOBAL_BATCH + FINAL_PARTIAL_TOKENS == GLOBAL_TOKENS
+
+
+@dataclass(frozen=True)
+class SolConfig:
+    model_name: str = MODEL_NAME
+    vocab_size: int = VOCAB_SIZE
+    d_model: int = D_MODEL
+    max_context: int = MAX_CONTEXT
+    n_blocks: int = N_PHYSICAL_BLOCKS
+    n_q_heads: int = N_Q_HEADS
+    n_kv_heads: int = N_KV_HEADS
+    head_dim: int = HEAD_DIM
+    ffn_hidden: int = FFN_HIDDEN
+    tn_rank: int = TN_RANK
+    rope_theta: float = 20_000.0
+    memory_width: int = 64
+    memory_slots: int = 32
+    chunk_size: int = 32
+    passes: int = 3
+    dropout: float = 0.0
+
+    def as_dict(self) -> dict:
+        return asdict(self)

tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "model_max_length": 1000000000000000019884624838656,
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}

training_state.json

@@ -0,0 +1,21 @@
+{
+  "release_checkpoint": "checkpoint_000900005888",
+  "base_pretraining_tokens": 1626560000,
+  "recovery_tokens": 900005888,
+  "total_token_exposures": 2526565888,
+  "recovery_peak_learning_rate": 0.000005,
+  "recovery_terminal_learning_rate": 0.0000005,
+  "recovery_mixture": {
+    "original_curriculum": 0.70,
+    "cosmopedia_v2_english": 0.15,
+    "finephrase": 0.15
+  },
+  "recovery_context": {
+    "initial_phase": "75% 512 / 25% 1024",
+    "final_400982016_tokens": "512 only"
+  },
+  "optimizer": "nGPT-aware AdamW",
+  "seed": 20260920,
+  "weight_format": "MLX NPZ",
+  "selection_note": "Selected from recovery checkpoints using the reported evaluation suite."
+}