kye/latent-mamba

Latent Mamba

Checkpoints for Latent Mamba, adaptive continuous reasoning in constant memory, in Pytorch

A Mamba-2 language model learns to think silently in continuous space between reading a question and writing the answer. The final hidden state is projected back into input-embedding space by a small thought projector and fed in as the next input, Coconut-style, while a halting head trained with the PonderNet objective decides when to stop. Every thought only updates the fixed-size recurrent state, so memory stays constant no matter how long the model thinks.

question ... <bot>   thought  thought  thought   <eot>   ### answer
                        ^        ^        ^
                        |        |        |
              last hidden state --> thought projector --> next input
              halting head: keep thinking, or stop and answer

The checkpoints here are a pretrained Mamba-2 130M fine-tuned on ProsQA, a graph-search reasoning benchmark, with a Coconut curriculum followed by a halting stage.

Results

ProsQA test set, 500 questions, one seed.

Model	Thoughts	Test accuracy
Mamba-2 130M, answer directly (no CoT)	0	74.6
Mamba-2 130M, written chain of thought	0	70.8
Latent Mamba, learned halting	3.11 (mean)	98.6
Latent Mamba, fixed budget	6	99.0
ablation: no thoughts	0	73.6
ablation: thoughts set to zero	6	76.8
ablation: thoughts replaced by noise	6	75.8

• Removing or scrambling the thoughts drops accuracy back to the level of the no-CoT baseline, so the content of the thoughts carries the reasoning.
• One thought already gives 94.2%, accuracy saturates at 99.0% from five, and budgets up to 12 (twice the training budget) lose at most 0.6 points.
• Learned halting uses half the thoughts for 0.4 points. It matches a fixed budget of 3 but adapts only weakly to question difficulty (2.99 thoughts on 3-step questions, 3.42 on 6-step ones).
• The recurrent state stays at 18.2 MiB for budgets up to 4,096 thoughts; GPT-2 124M's KV cache grows by 36 KiB per thought and stops at its 1,024-token context. The reference PyTorch implementation is 3 to 5 times slower per thought than GPT-2.

Training curves are on Weights & Biases. Full details are in the paper.

Files

Path	Contents
prosqa-mamba2-130m-latent/ponder.pt	the final model: best epoch of the halting stage, used for every result above
prosqa-mamba2-130m-latent/final.pt	end of the curriculum, always 6 thoughts, before halting training
prosqa-mamba2-130m-latent/stage1.pt ... stage6.pt	curriculum stages, k thoughts replace the first k reasoning steps
prosqa-mamba2-130m-latent/cot.pt	stage 0, written chain of thought
prosqa-mamba2-130m-latent/last.pt	weights at the end of training (epoch 55), for resuming
prosqa-mamba2-130m-cot/cot.pt	baseline: written chain of thought
prosqa-mamba2-130m-nocot/nocot.pt	baseline: answers directly
*/eval_test_*.json	test results, per inference mode and by number of reasoning steps
*/metrics.jsonl, */run.json	training logs and the exact settings of each run
bench/bench-mamba2-130m.json	memory and latency against thinking budget, versus GPT-2
report.md	summary table of all runs

Each checkpoint is a torch.save dictionary (float32, 642 MiB) holding the weights under "model" plus the model name, phase and training settings. Load it with weights_only = True.

Install

The model code is in the GitHub repository.

$ git clone https://github.com/The-Swarm-Corporation/latent-mamba
$ cd latent-mamba
$ uv sync --extra train

Requires Python >= 3.10 and PyTorch >= 2.4. Without uv, pip install torch transformers huggingface_hub setuptools works too.

Usage

Run from the repository root.

import sys
import torch
from huggingface_hub import hf_hub_download

sys.path.insert(0, 'latent_mamba')
from train import load_checkpoint, decode_answer

path = hf_hub_download('kye/latent-mamba', 'prosqa-mamba2-130m-latent/ponder.pt')
bundle, ckpt = load_checkpoint(path)

model = bundle.model.eval()
tokenizer = bundle.tokenizer    # GPT-NeoX tokenizer, as in the base model

question = 'Every grimpus is a jelpus. Every jelpus is a shumpus. Bob is a storpus. Every storpus is a grimpus. Tom is a bompus. Every bompus is a gerpus. Is Bob a shumpus or gerpus?'

# prompts are the question, a newline, then the begin-of-thought token

prompt = tokenizer(question + '\n', add_special_tokens = False)['input_ids'] + [bundle.bot_id]

with torch.no_grad():
    out = model.generate(
        torch.tensor([prompt]),
        max_thoughts = 6,           # upper bound on thoughts
        halt_rule = 'learned',      # 'learned' | 'fixed' | 'convergence'
        max_new_tokens = 32,
        eos_id = bundle.eos_id
    )

out.n_thoughts                                      # thoughts taken, e.g. tensor([3])
decode_answer(out.tokens[0].tolist(), bundle)      # '### Bob is a shumpus.'

Use halt_rule = 'fixed' with any max_thoughts to set the thinking budget yourself, 0 included. To batch prompts of different lengths, left pad them with bundle.pad_id and pass prompt_mask (True on real tokens); padding is exact.

Evaluate

Reproduce the test numbers, including the ablations and the 0 to 12 thought sweep.

$ uv run latent_mamba/train.py data --dataset prosqa
$ hf download kye/latent-mamba prosqa-mamba2-130m-latent/ponder.pt --local-dir checkpoints
$ uv run latent_mamba/train.py eval --ckpt checkpoints/prosqa-mamba2-130m-latent/ponder.pt --no-hf-upload

Train

The full pipeline (data, the three runs, test evaluation, memory benchmark) took about 16 hours on one H100 NVL. It uploads results only when HF_TOKEN is set.

$ DATASET=prosqa MODEL=mamba2-130m WORKDIR=$PWD uv run bash scripts/run_pipeline.sh

Or one run at a time.

$ uv run latent_mamba/train.py train --model mamba2-130m --dataset prosqa --method latent --no-hf-upload

--method is latent, cot or nocot. See docs/RUNPOD.md in the repository for GPU choices and cloud setup.

Limitations

• One synthetic dataset, one seed, one model size. ProsQA is close to saturated, so it cannot separate strong methods.
• Not yet compared against Coconut GPT-2 trained in the same pipeline, or against a pause-token model trained the same way.
• These checkpoints only know ProsQA-style questions. They are research artifacts, not general-purpose language models.

Citations

@techreport{gomez2026latentmamba,
    title       = {Latent Mamba: Adaptive Continuous Reasoning in Constant Memory},
    author      = {Kye Gomez},
    institution = {Swarms Corp},
    year        = {2026},
    url         = {https://github.com/The-Swarm-Corporation/latent-mamba}
}

@article{hao2024coconut,
    title   = {Training Large Language Models to Reason in a Continuous Latent Space},
    author  = {Shibo Hao and Sainbayar Sukhbaatar and DiJia Su and Xian Li and Zhiting Hu and Jason Weston and Yuandong Tian},
    journal = {arXiv preprint arXiv:2412.06769},
    year    = {2024}
}

@inproceedings{dao2024ssd,
    title     = {Transformers are {SSM}s: Generalized Models and Efficient Algorithms Through Structured State Space Duality},
    author    = {Tri Dao and Albert Gu},
    booktitle = {International Conference on Machine Learning},
    year      = {2024}
}

@article{banino2021pondernet,
    title   = {PonderNet: Learning to Ponder},
    author  = {Andrea Banino and Jan Balaguer and Charles Blundell},
    journal = {arXiv preprint arXiv:2107.05407},
    year    = {2021}
}