soumil1/simamba-midpoint-10m-slimpajama-500m

Simamba Midpoint 10M SlimPajama 500M

This is the best validation-loss Simamba checkpoint from our HPML discretization experiments. It is a small causal language model trained from scratch on a SlimPajama token subset to compare Simpson-style Simamba dynamics against Mamba2/trapezoidal baselines.

The model is a research checkpoint, not an instruction-tuned assistant.

Checkpoint

Field	Value
Variant	Simamba, Simpson discretization, midpoint control
Parameters	9,524,544 unique parameters with tied input/output embeddings
Training data	SlimPajama subset, data/slimpajama_500m_50m/train.bin
Validation data	Held-out SlimPajama subset, data/slimpajama_500m_50m/val.bin
Tokenizer	EleutherAI/gpt-neox-20b tokenizer files are included
Context length	128 tokens during training
Best checkpoint step	71,000
Best validation loss	4.917760848999023
W&B run	https://wandb.ai/ssb2234-columbia/simamba/runs/juxcboyg

Architecture

The model uses the local mamba_ssm implementation with a Simamba mixer:

{
  "d_model": 160,
  "n_layer": 8,
  "vocab_size": 50280,
  "rms_norm": true,
  "residual_in_fp32": true,
  "fused_add_norm": true,
  "tie_embeddings": true,
  "ssm_cfg": {
    "layer": "Simamba",
    "d_state": 64,
    "expand": 2,
    "headdim": 32,
    "rope_fraction": 0.5,
    "chunk_size": 16,
    "recompute_chunk_size": 16,
    "dt_limit": [0.001, 0.1],
    "A_max": 4.0,
    "use_midpoint_control": true,
    "discretization": "simpson",
    "simamba_backend": "triton",
    "is_outproj_norm": true
  }
}

Usage

This checkpoint requires the Simamba code in this repository, since the Simamba mixer and Triton kernels are not part of upstream transformers.

vLLM

This is a custom architecture, so stock vLLM cannot serve it through the native Mamba2 path without changing the model. Use vLLM's Transformers backend with remote code:

PYTHONPATH=/path/to/simamba \
vllm serve soumil1/simamba-midpoint-10m-slimpajama-500m \
  --trust-remote-code \
  --model-impl transformers \
  --dtype float32 \
  --max-model-len 128

The PYTHONPATH or editable install must point to this Simamba repository so the remote wrapper can import mamba_ssm.

Transformers

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "soumil1/simamba-midpoint-10m-slimpajama-500m"

tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    trust_remote_code=True,
    torch_dtype=torch.float32,
).cuda()
model.eval()

prompt = "State space models are"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")

with torch.no_grad():
    logits = model(input_ids).logits
next_token = torch.argmax(logits[:, -1], dim=-1)
print(tokenizer.decode(next_token))

Limitations

This is a small pretraining checkpoint trained for architecture comparison, not a production language model. It was trained on a limited SlimPajama subset with a short sequence length, so generation quality and downstream task performance should be interpreted cautiously.