Release UMIM merge modules and merge rules

README.md

@@ -1,3 +1,116 @@
 ---
-license: apache-2.0
+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.

gpt2-xl/filtered_bigrams_tensor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5314ce0d98784afbafd7530a1d6d99c4bcc3cbb47ef207551f76d66340cca112
+size 1748652

gpt2-xl/filtered_fourgrams_tensor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24feb7320011daf547cd242a8265304cd112ca19f2b1031669913f843dc6f619
+size 19867830

gpt2-xl/filtered_trigrams_tensor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7da885e42794dfe91ff1a74dec58aeb1b2aea7928a6345ed838fa3a34941af52
+size 8432177

gpt2-xl/merge_module.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:15357651b23f6ab78e2fdfa64dd4622b0714b813a2f9f3f61a0cae8dc0ac4706
+size 20529236

gpt2-xl/runtime_config.json

@@ -0,0 +1,11 @@
+{
+  "merge_weights": "merge_module.pt",
+  "rule_files": {
+    "2": "filtered_bigrams_tensor.pt",
+    "3": "filtered_trigrams_tensor.pt",
+    "4": "filtered_fourgrams_tensor.pt"
+  },
+  "num_heads": 4,
+  "surrogate_cache_size": 4096,
+  "local_files_only": false
+}

llama-3.1-8b/filtered_bigrams_tensor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75a202e6f29be5d4bc8137790f35b80af68c53bd30e06b918a42177c4db57a5c
+size 1734892

llama-3.1-8b/filtered_fourgrams_tensor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:636e9e275573e7a55a55b36974447ca6a75e85c38bc777b1f7d3e4d086a62a44
+size 20258614

llama-3.1-8b/filtered_trigrams_tensor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5129bf14b7542e7fde1e7e51721cc0779c4dd39959c6383860b4d6efab10e938
+size 8218801

llama-3.1-8b/merge_module.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f221a3f87a457b3e74e20e68413f6c21dfb0fa7873ffff298733a21bcb7ce14
+size 67168588

llama-3.1-8b/runtime_config.json

@@ -0,0 +1,11 @@
+{
+  "merge_weights": "merge_module.pt",
+  "rule_files": {
+    "2": "filtered_bigrams_tensor.pt",
+    "3": "filtered_trigrams_tensor.pt",
+    "4": "filtered_fourgrams_tensor.pt"
+  },
+  "num_heads": 4,
+  "surrogate_cache_size": 4096,
+  "local_files_only": false
+}

llama-3.2-1b/filtered_bigrams_tensor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75a202e6f29be5d4bc8137790f35b80af68c53bd30e06b918a42177c4db57a5c
+size 1734892

llama-3.2-1b/filtered_fourgrams_tensor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:636e9e275573e7a55a55b36974447ca6a75e85c38bc777b1f7d3e4d086a62a44
+size 20258614

llama-3.2-1b/filtered_trigrams_tensor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5129bf14b7542e7fde1e7e51721cc0779c4dd39959c6383860b4d6efab10e938
+size 8218801

llama-3.2-1b/merge_module.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85f7d800d64130f8278480d24ced3752c929c64ff4249dd50647e8b870d2e0c7
+size 33616212

llama-3.2-1b/runtime_config.json

@@ -0,0 +1,11 @@
+{
+  "merge_weights": "merge_module.pt",
+  "rule_files": {
+    "2": "filtered_bigrams_tensor.pt",
+    "3": "filtered_trigrams_tensor.pt",
+    "4": "filtered_fourgrams_tensor.pt"
+  },
+  "num_heads": 4,
+  "surrogate_cache_size": 4096,
+  "local_files_only": false
+}

manifest.json

@@ -0,0 +1,142 @@
+{
+  "format": "UMIM release manifest",
+  "models": {
+    "gpt2-xl": {
+      "base_model": "openai-community/gpt2-xl",
+      "hidden_size": 1600,
+      "merge_module": {
+        "dtype": "float32",
+        "file": "merge_module.pt",
+        "sha256": "15357651b23f6ab78e2fdfa64dd4622b0714b813a2f9f3f61a0cae8dc0ac4706",
+        "size_bytes": 20529236,
+        "tensor_count": 9
+      },
+      "num_attention_heads": 4,
+      "rules": {
+        "2": {
+          "count": 218415,
+          "dtype": "int32",
+          "file": "filtered_bigrams_tensor.pt",
+          "max_token_id": 50255,
+          "min_token_id": 12,
+          "sha256": "5314ce0d98784afbafd7530a1d6d99c4bcc3cbb47ef207551f76d66340cca112",
+          "size_bytes": 1748652
+        },
+        "3": {
+          "count": 702570,
+          "dtype": "int32",
+          "file": "filtered_trigrams_tensor.pt",
+          "max_token_id": 50254,
+          "min_token_id": 6,
+          "sha256": "7da885e42794dfe91ff1a74dec58aeb1b2aea7928a6345ed838fa3a34941af52",
+          "size_bytes": 8432177
+        },
+        "4": {
+          "count": 1241657,
+          "dtype": "int32",
+          "file": "filtered_fourgrams_tensor.pt",
+          "max_token_id": 50253,
+          "min_token_id": 6,
+          "sha256": "24feb7320011daf547cd242a8265304cd112ca19f2b1031669913f843dc6f619",
+          "size_bytes": 19867830
+        }
+      },
+      "training_checkpoint_metadata": {
+        "best_val_loss": 3.5166141991783126,
+        "epoch": 14,
+        "global_step": 123390,
+        "val_loss": 3.5166141991783126
+      }
+    },
+    "llama-3.1-8b": {
+      "base_model": "meta-llama/Llama-3.1-8B",
+      "hidden_size": 4096,
+      "merge_module": {
+        "dtype": "bfloat16",
+        "file": "merge_module.pt",
+        "sha256": "1f221a3f87a457b3e74e20e68413f6c21dfb0fa7873ffff298733a21bcb7ce14",
+        "size_bytes": 67168588,
+        "tensor_count": 9
+      },
+      "num_attention_heads": 4,
+      "rules": {
+        "2": {
+          "count": 216700,
+          "dtype": "int32",
+          "file": "filtered_bigrams_tensor.pt",
+          "max_token_id": 127935,
+          "min_token_id": 12,
+          "sha256": "75a202e6f29be5d4bc8137790f35b80af68c53bd30e06b918a42177c4db57a5c",
+          "size_bytes": 1734892
+        },
+        "3": {
+          "count": 684793,
+          "dtype": "int32",
+          "file": "filtered_trigrams_tensor.pt",
+          "max_token_id": 128000,
+          "min_token_id": 11,
+          "sha256": "5129bf14b7542e7fde1e7e51721cc0779c4dd39959c6383860b4d6efab10e938",
+          "size_bytes": 8218801
+        },
+        "4": {
+          "count": 1266080,
+          "dtype": "int32",
+          "file": "filtered_fourgrams_tensor.pt",
+          "max_token_id": 128000,
+          "min_token_id": 6,
+          "sha256": "636e9e275573e7a55a55b36974447ca6a75e85c38bc777b1f7d3e4d086a62a44",
+          "size_bytes": 20258614
+        }
+      },
+      "training_checkpoint_metadata": {}
+    },
+    "llama-3.2-1b": {
+      "base_model": "meta-llama/Llama-3.2-1B",
+      "hidden_size": 2048,
+      "merge_module": {
+        "dtype": "float32",
+        "file": "merge_module.pt",
+        "sha256": "85f7d800d64130f8278480d24ced3752c929c64ff4249dd50647e8b870d2e0c7",
+        "size_bytes": 33616212,
+        "tensor_count": 9
+      },
+      "num_attention_heads": 4,
+      "rules": {
+        "2": {
+          "count": 216700,
+          "dtype": "int32",
+          "file": "filtered_bigrams_tensor.pt",
+          "max_token_id": 127935,
+          "min_token_id": 12,
+          "sha256": "75a202e6f29be5d4bc8137790f35b80af68c53bd30e06b918a42177c4db57a5c",
+          "size_bytes": 1734892
+        },
+        "3": {
+          "count": 684793,
+          "dtype": "int32",
+          "file": "filtered_trigrams_tensor.pt",
+          "max_token_id": 128000,
+          "min_token_id": 11,
+          "sha256": "5129bf14b7542e7fde1e7e51721cc0779c4dd39959c6383860b4d6efab10e938",
+          "size_bytes": 8218801
+        },
+        "4": {
+          "count": 1266080,
+          "dtype": "int32",
+          "file": "filtered_fourgrams_tensor.pt",
+          "max_token_id": 128000,
+          "min_token_id": 6,
+          "sha256": "636e9e275573e7a55a55b36974447ca6a75e85c38bc777b1f7d3e4d086a62a44",
+          "size_bytes": 20258614
+        }
+      },
+      "training_checkpoint_metadata": {
+        "best_val_loss": 3.3169725206163196,
+        "epoch": 14,
+        "global_step": 111600,
+        "val_loss": 3.3169725206163196
+      }
+    }
+  },
+  "version": 1
+}