jrz97619761/test-model-thing
A small proof-of-concept language model (not an LLM) that aims to incorporate latent-space prediction, internal state using recurrent trace units, and byte-by-byte output, built with MLX.
Test‑Model‑Thing (TMT) – a tiny proof‑of‑concept language model
What it is
- A research‑grade language model (≈4.5 M parameters) built with MLX, a library that runs efficiently on Apple Silicon and also works on Linux (CUDA/CPU). It is not a large‑scale LLM; it is a small experimental model meant to explore several ideas:
- Latent‑space prediction – the model tries to predict hidden representations rather than just the next token.
- Recurrent Trace Units (RTUs) – internal state that carries information across time steps, giving the model a memory beyond the usual transformer context.
- Byte‑level I/O – inputs and outputs are raw bytes, so the model can generate any character sequence (hence the occasional misspellings).
- Continuous data streaming – the architecture is designed to handle a never‑ending stream of data rather than fixed‑size batches.
- Test‑time training – you can keep learning while the model is being used for inference.
Why it matters
- Demonstrates that sophisticated mechanisms (latent prediction, recurrent state) can be implemented in a tiny network that runs on a laptop.
- Provides a playground for researchers who want to experiment with byte‑level models, online learning, or custom memory modules without needing massive GPU clusters.
How to get started
- Clone the repo and create a Python virtual environment.
- Install MLX:
# macOS (Apple Silicon) pip install mlx # Linux with GPU pip install "mlx[cuda]" # Linux without GPU pip install "mlx[cpu]" - Train (or continue training) a model:
python main.py <path-to-dataset> train- The script saves weights automatically and can be interrupted with
Ctrl‑C.
- The script saves weights automatically and can be interrupted with
- Chat / inference:
python main.py <path-to-weights> chat- Options include
--ce-only(cross‑entropy loss only),--no-save, and--frozen(run without updating weights).
- Options include
- Benchmark on the CoLA linguistic acceptability task:
python benchmark.py <path-to-weights> <epochs> <slice>
Dataset
- The author trained the demo model on a small Wikipedia dump (
simplewiki-20260801-pages-articles.xml.bz2). You can use any text dataset; the code expects a path to the raw files.
Current state
- The model is a proof‑of‑concept: 4.5 M parameters, 512‑dimensional hidden size, 16 layers, trained for only a few hours. It can generate coherent byte streams, close brackets/quotes, and shows modest performance on CoLA (far below GPT‑1’s 45.4 score).
- No pre‑trained weights are shipped (GitHub file‑size limits). Users must train their own model or obtain weights elsewhere.
Roadmap / opportunities
- Scale up the parameter count and training time to improve language quality.
- Replace the byte‑level decoder with a tokeniser for cleaner text.
- Explore larger datasets or domain‑specific corpora.
- Contribute improvements to the RTU implementation or add other memory mechanisms.
- Port the code to Windows‑native environments (currently only WSL works).
Who might find this useful
- Students or hobbyists learning about language model internals.
- Researchers interested in online learning, latent‑space objectives, or memory‑augmented transformers.
- Developers who need a tiny, self‑contained model that runs on a MacBook without a GPU.
Links
- YouTube demo video
- Architecture sketch (Freeform diagram) is included in the repo for reference.
All information above is taken directly from the repository’s README.
Related
- Project
- Project
- Project
- Project
- Project