README.md

---
base_model:
- meta-llama/Llama-3.1-8B
- meta-llama/Llama-3.2-1B
- openai-community/gpt2-xl
library_name: pytorch
tags:
- umim
- token-merging
- kv-cache
- language-model-compression
- pytorch
---

# 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](https://github.com/Zesearch/Umim-LLM)
- [Project page](https://zesearch.github.io/Umim-LLM/)
- 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:

```text
<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:

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

To download one backbone only:

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

## Inference

Clone the UMIM codebase and install its runtime dependencies:

```bash
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:

```bash
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](https://github.com/Zesearch/Umim-LLM/issues)
for questions about the code or released artifacts.