This project implements the Memory Decoder (MemDec) approach to extend a base LLM (Large Language Model) with domain-specific knowledge through a modular use-case system under MemDec-Universal/config, where each domain is defined by its own config files — demonstrated on the German legal domain using the open-source Open Legal Data (a non-profit platform providing free API access to German court decisions, laws, and legal texts). The pipeline enables plug-and-play domain adaptation across lightweight, instruction-tuned SLMs (Small Language Models), including Gemma 3-270m, Qwen 3.5-0.8B, and SmolLM3-3B, with the active use-case selected via the environment variable MEMDEC_USECASE and all pipeline behavior controlled by config files alone. To maximize efficiency, this repository integrates official Unsloth models. By replacing standard PyTorch layers with custom-written NVIDIA Triton kernels, it achieves 2x to 3x faster training speeds and significantly reduces memory overhead.
MemDec offers a parametric alternative to Retrieval-Augmented Generation (RAG) for domain adaptation in LLMs. While RAG retrieves relevant documents from an external vector database at inference time to augment the prompt, it introduces substantial runtime overhead, increased latency, and scaling costs as the datastore grows. In contrast, the MemDec plug-and-play approach eliminates external retrieval by internalizing domain knowledge directly within a lightweight, pretrained transformer decoder. Once trained on a target domain, a single MemDec can be flexibly integrated into any base LLM sharing the same tokenizer with low latency and improved accuracy, as shown in the paper's results by the reduced perplexity.
Caution
For readability, the linked Features-Table above is shown in a condensed version. The full Terminal Commands for each script can be copied from here:
utils.py
TEST
python -m unittest test.test_utils -v
RUN
1. Set Default LLM Model (default "gemma3"):
python -m src.utils <model_name>
ℹ️ python -m src.utils qwen3.5
ℹ️ python -m src.utils smollm3
2. Show Current Default Model:
ℹ️ python -m src.utils --show
step0_dataset.py
TEST
python -m unittest test.test_step0_dataset -v
RUN
Download & Clean File-Based Dataset:
python -m src.step0_dataset <dataset_name>
ℹ️ URL: python -m src.step0_dataset bverfg_decisions
ℹ️ Local: python -m src.step0_dataset gerdalir_queries
step0_data_management.py
TEST
python -m unittest test.test_step0_data_management -v
RUN
Combine Cleaned Datasets with Corpus Name:
python -m src.step0_data_management --corpus-name <corpus_name> <dataset1> <dataset2> <datasetN>
ℹ️python -m src.step0_data_management --corpus-name test_corpus bverfg_decisions gerdalir_queries
step1_cleaning.py
TEST
python -m unittest test.test_step1_cleaning -v
RUN
1a. Clean HuggingFace Dataset:
python -m src.step1_cleaning --hf-dataset <dataset_name>
ℹ️ python -m src.step1_cleaning --hf-dataset DomainLLM/german-law-qa
1b. With Specific Config (Subsets):
python -m src.step1_cleaning --hf-config <config>
ℹ️ python -m src.step1_cleaning --hf-config dump-20221018-10k or
… --hf-config cases2022_10k
1c. With Duplicated Config:
python -m src.step1_cleaning --hf-dataset <dataset_name> --hf-config <config>
ℹ️ python -m src.step1_cleaning --hf-dataset openlegaldata/court-decisions-germany --hf-config dump-20260520-10k or
… --hf-dataset openlegaldata/court-decisions-germany --hf-config cases2026_10k
ℹ️ python -m src.step1_cleaning --hf-dataset openlegaldata/laws-germany --hf-config dump-20260520-10k or … --hf-dataset openlegaldata/laws-germany --hf-config laws_germany_10k
📌Only via pipeline_config.yaml:
ℹ️ min_text_length, min_alpha_ratio, chunk_size, check_duplicates
step2_tokenization.py
TEST
python -m unittest test.test_step2_tokenization -v
RUN
1a. Tokenize Cleaned Dataset:
python -m src.step2_tokenization <dataset_cleaned>
ℹ️ python -m src.step2_tokenization cases2026_10k
1b. Override Base LLM:
python -m src.step2_tokenization <dataset_cleaned> --model <model_name>
ℹ️ python -m src.step2_tokenization cases2026_10k --model gemma3-1b
2. Custom Params:
… --train-test-split <ratio> --num-workers <n>
ℹ️ … --train-test-split 0.9 --num-workers 4
step3_pretraining.py
TEST
python -m unittest test.test_step3_pretraining -v
RUN
1a. Knowledge Base Creation:
python -m src.step3_pretraining <dataset_tokenized>
ℹ️ python -m src.step3_pretraining cases2026_10k_tokenized-gemma3
1b. Override Base LLM:
python -m src.step3_pretraining <dataset_tokenized> --model <model_name>
ℹ️ python -m src.step3_pretraining cases2026_10k_tokenized-gemma3 --model gemma3-1b
2. Custom Params:
… --batch-size <n> --seed <X> --ncentroids <n> --num-keys-to-add-at-a-time <int>
ℹ️ … --batch-size 32 --seed 777 --ncentroids 4096 --num-keys-to-add-at-a-time 1_000_000
step4_training.py
TEST
python -m unittest test.test_step4_training -v
RUN
1a. MemDec Training:
python -m src.step4_training <dataset_tokenized>
ℹ️ python -m src.step4_training cases2026_10k_tokenized-gemma3
1b. Without Unsloth (plain HF/PyTorch):
python -m src.step4_training <dataset_tokenized> --no-unsloth
ℹ️ python -m src.step4_training cases2026_10k_tokenized-gemma3 --no-unsloth
1c. Override Base LLM:
python -m src.step4_training <dataset_tokenized> --model <model_name>
ℹ️ python -m src.step4_training cases2026_10k_tokenized-gemma3 --model gemma3-1b
1d. With Specific Checkpoint:
python -m src.step4_training <dataset_tokenized> --checkpoint <step_X>
ℹ️ python -m src.step4_training cases2026_10k_tokenized-gemma3 --checkpoint step_500
2. MemDec Params:
… --k-neighbors <k> --alpha <float> --lmbda <float>
ℹ️ … --k-neighbors 8 --alpha 0.7 --lmbda 0.5
3. Training Params:
… --max-steps <X> --checkpointing-steps <interval> --num-train-epochs <X> --batch-size <n> --per-device-train-batch-size <n> --gradient-accumulation-steps <n> --learning-rate <rate> --seed <X>
ℹ️ … --max-steps 1500 --checkpointing-steps 100 --num-train-epochs 23 --batch-size 4 --per-device-train-batch-size 4 --gradient-accumulation-steps 4 --learning-rate 5e-5 --seed 777
step5_testing.py
TEST
python -m unittest test.test_step5_testing -v
RUN
1a. Test with Checkpoint:
python -m src.step5_testing --checkpoint <step_X>
ℹ️ python -m src.step5_testing --checkpoint step_1000
1b. Override Base LLM (uses latest Checkpoint-Save):
python -m src.step5_testing --model <model_name>
ℹ️ python -m src.step5_testing --model gemma3-1b
2. MemDec Params:
… --knn-temp <float> --lmbda <float>
ℹ️ … --knn-temp 0.8 --lmbda 0.5
3. Generation Params:
… --max-new-tokens <X> --repetition-penalty <float> --do-sample --temperature <float> --top-p <float> --top-k <n>
ℹ️ … --max-new-tokens 200 --repetition-penalty 1.5
ℹ️ … --do-sample --temperature 0.8 --top-p 0.9 --top-k 40 (sampling instead of greedy)
4. Custom Params:
… --scenarios <scenario> --tasks <task> --save-results --compare-base
ℹ️ … --scenarios withdrawals employment --tasks question_answering summarization
--save-results / --compare-base invert the config default in pipeline_config.yaml, leading to disabled result saving & base-model comparison.
step5_evaluation.py
TEST
python -m unittest test.test_step5_evaluation -v
RUN
1a. PPL Evaluation:
python -m src.step5_evaluation <dataset_tokenized> --checkpoint <step_X>
ℹ️ python -m src.step5_evaluation cases2026_10k_tokenized-gemma3 --checkpoint step_1000
1b. Override Base LLM:
python -m src.step5_evaluation <dataset_tokenized> --model <model_name>
ℹ️ python -m src.step5_evaluation cases2026_10k_tokenized-gemma3 --model gemma3-1b
2. MemDec Params:
… --lmbda <float> --knn-temp <float>
ℹ️ … --lmbda 0.3 --knn-temp 1.0
3. Evaluation Params:
… --split {train|validation|test} --max-examples <X> --batch-size <n> --seed <X>
ℹ️ … --split test --max-examples 100 --batch-size 8 --seed 777
4. Custom Params:
… --save-results --compare-base
--save-results / --compare-base invert the config default in pipeline_config.yaml, leading to disabled result saving & base-model comparison.
Important
All CLI flags listed above are optional overrides: arguments not passed on the command line resolve to their configured values in pipeline_config.yaml and dataset_config.yaml, so each step can equally be invoked as python -m src.<step> without flags. Exception: for base_model, the global .env setting MEMDEC_MODEL outranks even the --model flag (remove it from .env to let --model take over). The complete YAML schemas, parameter descriptions and instructions for defining custom domains are documented in config/HowTo.md.
MemDec-Universal/
├── test/ # Test Scripts (Unit Tests)
├── src/ # Source Code
│ ├── utils.py # Utility functions
│ ├── step0_dataset.py # Data Preparation & Cleaning
│ ├── step0_data_management.py # Data Management
│ ├── step1_cleaning.py # HF-Dataset Preparation & Cleaning
│ ├── step2_tokenization.py # Data Tokenization
│ ├── step3_pretraining.py # MemDec's Knowledge Base Creation
│ ├── step4_training.py # MemDec-Model Training
│ ├── step5_testing.py # MemDec-Model Testing
│ └── step5_evaluation.py # MemDec-Tests Evaluation
├── MemoryDecoder/ # MemDec's Training Scripts (used & modified)
│ ├── knn_utils/saveEmbedMulti.py # Step 3
│ ├── utils/cal_loss.py # Step 4
│ ├── train_memdec.py # Step 4
│ ├── train_base.py # Step 4
│ └── demo/memDec.py # Step 5 (Testing)
├── config/ # Config Files
│ ├── custom_usecase/ # Custom Usecase (Template)
│ ├── legal_usecase/ # German Legal Usecase (Project Example)
│ ├── model_config.json # LLM Metadata
│ └── HowTo.md # Config-Guide
├── dataset/ # Datasets (cleaned & tokenized)
├── knowledge_base/ # KNN Datastore & FAISS Index
├── outputs/ # Output Files
│ ├── logs/ # Script Logs (while running "stepX".py)
│ ├── test_results/ # Test & Evaluation Results
│ └── step_5000/ # Trained Model Weights
├── assets/ # README Media & Tables
├── .env.example # Environment Variables (Template)
├── .gitattributes # Git Attributes (LFS for large files)
├── .gitignore # Git Ignore (files to exclude from commit)
├── requirements.txt # Project Base Dependencies
├── THIRD-PARTY-NOTICES.txt # 3rd-Party Software & Dataset Notices
├── LICENSE # Project License
└── README.md # This fileThe pipeline runs on Python 3.11 with PyTorch (CUDA 12.8) and Unsloth for efficient model loading. All pinned dependencies are listed in requirements.txt.
-
Set Environment Variables →
.env.example -
Create & Activate Virtual Environment
a) Using Python:# Python v3.11 must be pre-installed: https://www.python.org/downloads/ py -3.11 -m venv memdec-env .\memdec-env\Scripts\Activate.ps1
b) Using Conda:
# Anaconda or Miniconda (lightweight ver.) must be pre-installed: # https://www.anaconda.com/docs/getting-started/installation conda create -n memdec-env python=3.11 conda activate memdec-env
-
Install Project Dependencies
# 1) Project base packages pip install -r requirements.txt --no-cache-dir # 2) CUDA 12.8 toolkit libraries required by PyTorch pip install nvidia-cuda-nvcc-cu12 nvidia-cuda-runtime-cu12 --no-cache-dir # For Conda: conda install nvidia::cuda-nvcc nvidia::cuda-runtime -c nvidia # 3) PyTorch with CUDA 12.8 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 --no-cache-dir # 4) Unsloth last, so it patches the final PyTorch build pip install unsloth
Caution
On Windows, the install order is critical: pip re-resolves shared dependencies, so installing PyTorch or Unsloth first can silently upgrade/downgrade the pinned packages from requirements.txt, resulting in CUDA DLL conflicts and memory-isolation crashes (Hypervisor Errors).
-
Authenticate with Hugging Face
Setup Steps:- Login to Hugging Face to request access on the Open Legal Data datasets which are used in this project: openlegaldata/court-decisions-germany and openlegaldata/laws-germany.
- Create an User Access Token at huggingface.co/settings/
tokensto get yourHF_TOKEN. - Run the following command in the terminal to authenticate:
huggingface-cli login # Paste your Access Token when prompted
Important
Authentication via huggingface-cli login is required to access the gated Open Legal Data dataset. However, as this pipeline uses Unsloth-optimized model weights (including Gemma 3), there is no need to accept separate model license agreements on HuggingFace.
To run all test scripts in the MemDec-Universal/test directory:
python -m unittest discover -s test -vImportant
To run individual test scripts, refer to the Terminal Commands listed in Main Components.
This project is a config-driven prototype for domain adaptation. While ongoing maintenance is limited, contributions via Pull Requests are highly appreciated! 🤝
Key areas for contribution include:
- New Use-Cases: Add ready-made
<your_domain>_usecasefolders with config files, data sources, and test scenarios for new domains or languages. - Config Enhancements: Extend
dataset_config.yamlfield mappings orpipeline_config.yamloptions to cover more source types and pipeline behaviors. - Framework Robustness: Fix bugs, refine pipeline steps, and expand test coverage or evaluation benchmarks.
For inquiries, technical feedback, or potential collaborations, feel free to reach out via my website's contact form.
This is an open-source reimplementation of the research paper
Memory Decoder: A Pretrained, Plug-and-Play Memory for Large Language Models by Jiaqi Cao, Jiarui Wang, Rubin Wei, Qipeng Guo, Kai Chen, Bowen Zhou, and Zhouhan Lin (2025)
trained on the Open Legal Data 
datasets for research, testing, and educational purposes. It is built on the open-weights foundation models Gemma 3 , Qwen 3.5 , and SmolLM3 . This repository does not host or redistribute model weights or datasets. The base models provided via Unsloth 
and the gated Open Legal Data datasets must be accessed directly from Hugging Face, subject to their respective licenses and terms of use.
This project is licensed under the 
This repository includes and modifies code from the official Memory Decoder (MemDec) repository, vendored under MemDec-Universal/MemoryDecoder and licensed under the 
