File size: 4,985 Bytes
18eb577 677c01c 276bc2b 677c01c 5211579 677c01c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
---
license: apache-2.0
language:
- pyt
base_model:
- openai-community/gpt2
---
# MemoryDecoder-GPT2-Small
## Model Description
Memory Decoder is a pretrained, plug-and-play memory component designed for efficient domain adaptation of large language models. This checkpoint contains the GPT2-small Memory Decoder trained on WikiText-103, as described in our NeurIPS 2025 paper.
- **Paper:** [Memory Decoder: A Pretrained, Plug-and-Play Memory for Large Language Models](https://www.arxiv.org/abs/2508.09874)
- **GitHub:** [https://github.com/LUMIA-Group/MemoryDecoder](https://github.com/LUMIA-Group/MemoryDecoder/tree/main)
- **Conference:** NeurIPS 2025 (Poster)
- **Model Size:** 124M parameters
- **Base Architecture:** GPT2-small transformer decoder
## Overview
Memory Decoder bridges the gap between non-parametric retrieval methods and parametric fine-tuning approaches. By pre-training a compact transformer decoder to internalize retrieval patterns, it provides:
- **Plug-and-Play Integration:** Works with any GPT2 model variant without modifying original parameters
- **Efficient Inference:** No retrieval overhead - just parallel forward passes
- **Domain Expertise:** Captures long-tail knowledge like kNN-LM but with parametric efficiency
- **Preserved Capabilities:** Original model remains unchanged
## Quick Start
### Step 1: Import Libraries and Initialize Models
```python
from memDec import MemoryDecoder
import transformers
from transformers import AutoModelForCausalLM
from loguru import logger
# Define paths to your models
base_lm_path = "gpt2-xl" # or any GPT2 variant
knn_generator_path = "Clover-Hill/MemoryDecoder-gpt2-small"
# Load tokenizer and models
tokenizer = transformers.AutoTokenizer.from_pretrained(base_lm_path)
base_lm = AutoModelForCausalLM.from_pretrained(base_lm_path)
knn_generator = AutoModelForCausalLM.from_pretrained(knn_generator_path)
```
### Step 2: Prepare Models and Create Joint Model
```python
# Resize embeddings and set to evaluation mode
base_lm.eval()
knn_generator.eval()
# Create the joint Memory Decoder model
joint = MemoryDecoder(base_lm, knn_generator, lmbda=0.55, knn_temp=1.0).to("cuda")
```
### Step 3: Generate Text and Compare Results
```python
# Prepare input prompt
prompt = "As with previous Valkyira Chronicles games , Valkyria Chronicles III is"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
# Generate with Memory Decoder
out_ids = joint.generate(**inputs, max_new_tokens=20, do_sample=False)
logger.info(f"Memory Decoder output: {tokenizer.decode(out_ids[0], skip_special_tokens=True)}")
# Generate with base model for comparison
out_ids = base_lm.generate(**inputs, max_new_tokens=20, do_sample=False)
logger.info(f"Base Model output: {tokenizer.decode(out_ids[0], skip_special_tokens=True)}")
```
**📊 Generation Results Comparison:**
| Model | Generated Continuation |
|-------|------------------------|
| **Base Model** | *"...is a turn-based strategy game. The player takes control of a squad of Valkyria soldiers..."* |
| **+Memory Decoder** | *"...is a **role-playing** video game developed by Sega and published by Sega for the PlayStation 2."* |
> [!NOTE]
> Memory Decoder correctly identifies Valkyria Chronicles III as a **role-playing game** (factually accurate), while the base model incorrectly predicts it as a strategy game.
## Performance on WikiText-103
| Model Configuration | Perplexity | Improvement |
|:-------------------|:----------:|:-----------:|
| GPT2-small (baseline) | 24.89 | - |
| GPT2-small + MemoryDecoder | **13.36** | -11.53 |
| GPT2-medium (baseline) | 18.29 | - |
| GPT2-medium + MemoryDecoder | **12.25** | -6.04 |
| GPT2-large (baseline) | 15.80 | - |
| GPT2-large + MemoryDecoder | **11.53** | -4.27 |
| GPT2-xl (baseline) | 14.39 | - |
| GPT2-xl + MemoryDecoder | **10.93** | -3.46 |
## Key Features
- **Universal Compatibility:** Works with all GPT2 model sizes (small, medium, large, xl)
- **Parameter Efficient:** Only 124M additional parameters enhance models up to 1.5B
- **Domain Adaptation:** Trained to capture WikiText-103 domain knowledge
- **Inference Speed:** Minimal overhead compared to retrieval-based methods
## Training Details
- **Training Data:** WikiText-103
- **Training Objective:** Hybrid KL divergence and language modeling loss
- **Supervision Signal:** kNN distributions from GPT2-xl, it is suggested to use the finetuned version of GPT2-xl [here](https://huggingface.co/Clover-Hill/gpt2-xl-finetuned-wikitext103).
- **Hyperparameters:**
- Learning rate: 1e-3
- Beta (loss balance): 0.5
- Training Epoch: 70
## Citation
```bibtex
@article{cao2025memory,
title={Memory decoder: A pretrained, plug-and-play memory for large language models},
author={Cao, Jiaqi and Wang, Jiarui and Wei, Rubin and Guo, Qipeng and Chen, Kai and Zhou, Bowen and Lin, Zhouhan},
journal={arXiv preprint arXiv:2508.09874},
year={2025}
}
```
## Contact
For questions and support: maximus.cao@outlook.com |