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

  1. Clone the repo and create a Python virtual environment.
  2. Install MLX:
    # macOS (Apple Silicon)
    pip install mlx
    # Linux with GPU
    pip install "mlx[cuda]"
    # Linux without GPU
    pip install "mlx[cpu]"
    
  3. 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.
  4. 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).
  5. 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