infocusp/stage1-VQ-VAE

Assorted Latents — Stage 1 (VQ-VAE)

Stage 1 of an unofficial implementation of Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning (Su et al., ICML 2025). It is a VQ-VAE that compresses chain-of-thought (CoT) text into a short sequence of discrete latent codes (16 text tokens → 1 code, codebook of 1024 entries) and decodes the codes back to text.

This is a component, not a chat model. Its job is to turn MetaMathQA reasoning traces into latent codes, which are used to fine-tune an LLM in Stage 2: Assorted-Latents-Stage2.

• Collection: assorted-latents (both stages)
• Code: https://github.com/infocusp/assorted_latents
• Paper: https://arxiv.org/abs/2502.03275

Model details

Architecture	Token embedding → 2-layer encoder → Conv1d down-sampler (k = stride = 16) → EMA vector quantizer → learned ConvTranspose1d up-sampler → 2-layer parallel (non-autoregressive) decoder
Size	d_model 512, 4 heads, ~89M parameters (mostly the 128k-token embedding table)
Codebook	1024 codes, Euclidean, EMA (decay 0.99), k-means init, dead-code restart
Tokenizer	unsloth/Llama-3.2-1B (ungated mirror of Llama-3.2). Stage 1 and Stage 2 must use the same one
Sequence length	256 tokens → 16 codes
Training data	meta-math/MetaMathQA train split, last 2,000 rows held out for validation
Training	Adam, lr 1e-5 (constant), batch 32, 100k steps, bf16, deterministic. About 3 h on one RTX PRO 6000 (96 GB)
Checkpoint	Seed 0, step 100000, val recon_acc 0.4747

Results

Our own measurements. The paper does not publish Stage 1 results. Mean ± std over seeds 0, 1, 2; recon_* and global_* are on the 2,000-row validation split, round-trip metrics on 100 held-out validation examples.

Metric	Value
recon_ce ↓	2.163 ± 0.007
recon_acc ↑	0.4735 ± 0.0010
global_active_frac ↑ (codebook usage)	0.874 ± 0.012
global_perplexity ↑	785.9 ± 14.3
round-trip token_f1 ↑	0.612 ± 0.003
round-trip number_preserved ↑	0.467 ± 0.017

The decoder sees only the codes and a pooled prompt vector, with no ground-truth tokens. So recon_acc measures what the codes carry, and reconstruction is lossy by design: less than half the numbers in a trace survive the round trip. The codes are meant to be read by an LLM, not to be a faithful text encoding.

Usage

The weights are in model.safetensors; config.json holds the training config and tokenizer name. Loading needs the model class from the GitHub repo:

import json
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
# from a clone of https://github.com/infocusp/assorted_latents
# Run `pip install -r requirements.txt` from cloned repo root before running
from src.utils import Config
from src.data import load_tokenizer, token_info
from src.model import TokenVQVAE

repo = "infocusp/stage1-VQ-VAE"
cfg_d = json.load(open(hf_hub_download(repo, "config.json")))
cfg = Config.from_dict({k: v for k, v in cfg_d.items() if not k.startswith("_")})
tok = load_tokenizer(cfg_d["tokenizer_name"])
model = TokenVQVAE.build_from_config(cfg, token_info(tok))
model.load_state_dict(load_file(hf_hub_download(repo, "model.safetensors")))
model.eval()

Limitations

• Trained only on MetaMathQA (English math word problems), so other domains are untested.
• Needs the Llama-3.2 tokenizer. Codes from different tokenizers are not interchangeable.
• Code differs from the paper in the codebook update (EMA instead of gradient descent) and the dataset (original MetaMathQA instead of the paper's Llama-3.1-405B-regenerated version). See the GitHub README for the full list.

Citation

@inproceedings{su2025token,
  title     = {Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning},
  author    = {Su, DiJia and Zhu, Hanlin and Xu, Yingchen and Jiao, Jiantao and Tian, Yuandong and Zheng, Qinqing},
  booktitle = {Forty-second International Conference on Machine Learning},
  year      = {2025},
  url       = {https://openreview.net/forum?id=hYfOPXrbUr}
}