UMIM Merge Modules and Merge Rules

This repository contains the task-agnostic merge modules and tokenizer-specific 2-, 3-, and 4-gram merge rules released with UMIM: Distilling Sequential Computation in Transformer Language Models.

UMIM replaces matched contiguous token spans with surrogate embeddings produced by a lightweight merge module. The pretrained language-model backbone remains frozen. During inference, UMIM compresses both prompt prefill and the growing KV cache while preserving fully autoregressive token generation.

• Code
• Project page
• Paper: coming soon

Repository layout

Directory	Backbone	Merge-module dtype	Rules
llama-3.1-8b/	meta-llama/Llama-3.1-8B	bfloat16	2-, 3-, and 4-grams
llama-3.2-1b/	meta-llama/Llama-3.2-1B	float32	2-, 3-, and 4-grams
gpt2-xl/	openai-community/gpt2-xl	float32	2-, 3-, and 4-grams

Each directory contains:

<backbone>/
├── merge_module.pt
├── filtered_bigrams_tensor.pt
├── filtered_trigrams_tensor.pt
├── filtered_fourgrams_tensor.pt
└── runtime_config.json

merge_module.pt is a clean PyTorch state dictionary containing only the nine merge-module tensors. Training optimizer, scheduler, scaler, and progress states are intentionally excluded. manifest.json records file sizes, SHA-256 hashes, tensor metadata, and rule counts for reproducibility.

Download

Download the complete release with the Hugging Face CLI:

hf download Zesearch/UMIM --local-dir artifacts/umim

To download one backbone only:

hf download Zesearch/UMIM \
  --include "llama-3.1-8b/*" \
  --local-dir artifacts/umim

Inference

Clone the UMIM codebase and install its runtime dependencies:

git clone https://github.com/Zesearch/Umim-LLM.git
cd Umim-LLM
python -m pip install torch transformers accelerate

After downloading the release into artifacts/umim, run:

python -m generation.generate \
  --model meta-llama/Llama-3.1-8B \
  --config artifacts/umim/llama-3.1-8b/runtime_config.json \
  --prompt "Language models can" \
  --max-new-tokens 128 \
  --device cuda:0 \
  --dtype bfloat16

The paths inside each runtime_config.json are relative to that configuration file, so no manual path editing is needed when the directory is kept intact.

Artifact compatibility

Merge rules contain tokenizer token IDs and therefore must be paired with their listed backbone. The two Llama releases share the same rule tensors because both use the compatible Llama tokenizer vocabulary in these experiments. GPT-2 XL uses its own rule tensors.

These files are the WikiText-103-trained, task-agnostic base merge modules and rules. They are the artifacts used for direct transfer without downstream retraining. Task-specific SFT and preference-based RL update only the merge module while keeping the backbone frozen; task-adapted checkpoints are not part of this initial artifact bundle.

Citation

The arXiv entry and BibTeX citation will be added when the final paper version is available.

Questions

Please use the GitHub issue tracker for questions about the code or released artifacts.