TriSSM

Associative Recall Under Ternary Quantization: A Controlled Comparison of State Space Models and Transformers at 5M Parameters

This repository contains the code for the paper. TriSSM is a selective state space model whose input-dependent sequence-mixing projections (Δ, B, C) are BitNet b1.58 BitLinear layers: ternary weights in {-1, 0, +1} and int8 activations, trained from scratch with straight-through estimators. The code trains TriSSM and a set of parameter-matched baselines on TinyStories, then adapts and evaluates every model on key-value associative recall.

The BitSSM sequence mixer. An RMSNorm precedes each ternary projection, and the recurrence is evaluated with a parallel associative scan.

Overview

The study crosses architecture (SSM or Transformer) with precision (full precision or ternary) at about 5M parameters, under one data set, token budget, optimizer and learning-rate schedule.

Model	Registry key	Sequence mixing	Precision
TriSSM	trissm	selective SSM, 12 layers, d=256	ternary Δ, B, C
FP SSM	fp_ssm	same as TriSSM	full precision
BitNet Transformer + RoPE	bitnet_transformer_rope	attention, 4 layers, d=192	ternary attention and MLP
FP Transformer + RoPE	fp_transformer_rope	same	full precision
QKVO Transformer + RoPE	qkvo_transformer_rope	same	ternary Q/K/V/O, full-precision MLP

The embedding and LM head are full precision in every model. Each Transformer is also available without position encoding (drop the _rope suffix), and four reshaped ~5M-parameter Transformer configurations (v2t_{deep,mid,shallow,widehead}_{fp,tern}[_rope]) check that the results do not depend on one particular depth or width.

Key results

Strict recall accuracy with N = 10 key-value pairs and validation perplexity, mean ± std over seeds 42, 43 and 44.

Model	Recall (%)	Perplexity
TriSSM	18.64 ± 0.48	17.25 ± 0.04
FP SSM	18.94 ± 1.49	14.96 ± 0.01
BitNet Transformer + RoPE	14.54 ± 1.67	16.90 ± 0.11
QKVO Transformer + RoPE	14.71 ± 0.05	15.31 ± 0.10
FP Transformer + RoPE	12.45 ± 1.13	11.24 ± 0.10

• TriSSM recalls more than every Transformer, in every seed.
• The lead over the Transformers holds from 5 to 50 key-value pairs.
• Ternary weights raise TriSSM's perplexity by 15.3%, against 50.3% for the ternary Transformer.

 
Left: recall as the number of key-value pairs grows. Right: validation perplexity in full precision and ternary form.

Recall of the parameter-matched Transformer configurations against the two SSMs.

Repository structure

trissm/                    model package
  bitlinear.py             BitLinear: ternary weights, int8 activations, STE
  ssm.py                   BitSSM block and parallel associative scan
  model.py                 TriSSM language model
  transformer.py           Transformer baselines with optional RoPE
  registry.py              every architecture used in the paper, by name
  utils.py                 parameter accounting, batching, checkpoints, logs
scripts/
  prepare_data.py          download TinyStories, train the tokenizer, write train/val.bin
  train.py                 pretrain one (architecture, seed) cell
  recall.py                associative recall adaptation and evaluation
  eval_perplexity.py       validation perplexity across seeds
  toy_ablation.py          toy-scale ablation of the BitSSM block (CPU)
experiments/
  run_seeds.py             main comparison: training and N=10 recall over seeds
  run_capacity.py          recall with N = 5, 20, 50, 100 key-value pairs
  run_configs.py           parameter-matched Transformer configurations
tests/
  test_models.py           unit tests (scan, quantization, RoPE, causality, sizes)

Installation

git clone https://github.com/PrnvKK/TriSSM.git
cd TriSSM
pip install -e ".[dev]"

PyTorch 2.4 or newer is required. A CUDA GPU is used when available. Run the unit tests with:

pytest -q

Usage

1. Prepare the data

python scripts/prepare_data.py

This downloads TinyStories V2, trains an 8k SentencePiece BPE tokenizer and writes data/train.bin and data/val.bin (510,281,874 and 26,853,217 tokens).

2. Train

python scripts/train.py --arch trissm --seed 42
python scripts/train.py --arch bitnet_transformer_rope --seed 42

Each run uses 2,500 iterations of 64 × 256 tokens (40.96M tokens). Checkpoints are written to out/<arch>_s<seed>/model.pt; add --resume to continue an interrupted run.

3. Associative recall

python scripts/recall.py --model trissm --seed 42 --n_pairs 10

The checkpoint is adapted on the recall task for 3,000 steps and evaluated on 3,200 held-out prompts with a fixed evaluation seed, so every model sees the same prompts. --distance_mode stratified reports accuracy by the distance between the queried pair and the query.

4. Perplexity

python scripts/eval_perplexity.py --seeds 42,43,44

5. Full experiments

Each driver runs its cells in sequence, skips cells already recorded under logs/, and prints its commands with --dry_run.

python experiments/run_seeds.py              # main comparison
python experiments/run_capacity.py           # capacity sweep
python experiments/run_configs.py --pos rope # parameter-matched configurations
python scripts/toy_ablation.py --all         # toy ablation, 7 configurations x 3 seeds

Citation

@article{krishnakumar2026trissm,
  title  = {Associative Recall Under Ternary Quantization: A Controlled Comparison of State Space Models and Transformers at 5M Parameters},
  author = {Krishnakumar, Pranav and Pandey, Garima and Murthy Y, Vishnu Srinivasa and Koolagudi, Shashidhar G.},
  note   = {Under review},
  year   = {2026}
}

License

Released under the MIT License. See LICENSE.