TinyEngram

[!NOTE] TL;DR: TinyEngram demonstrates that Engram-based memory injection outperforms LoRA in both parameter efficiency and catastrophic forgetting resistance—and extends seamlessly to vision (e.g., Stable Diffusion) for lightweight, composable concept injection. All code, logs, and experiments are open!

If you find TinyEngram useful, a ⭐ helps support the project.

⚙️ Quick Environment Setup

TinyEngram provides a clean, pinned direct-dependency file for training and vision reproduction:

conda create -n tinyengram python=3.10 -y
conda activate tinyengram
pip install --upgrade pip
pip install -r requirements.txt

For CUDA notes and optional evaluation dependencies, see doc/reproduction/environment.md.

Introduction

TinyEngram is an open research project exploring the Engram architecture—an LLM enhancement that boosts phrase-level understanding by integrating a compact N-gram memory module and a gated retrieval mechanism into key transformer layers.

Built on Qwen, TinyEngram provides a lightweight, ready-to-train codebase for anyone to reproduce, experiment with, or extend Engram-style models. We actively share new experiments, training logs, and findings right here—making this repo both a toolkit and a living research notebook.

Beyond LLMs, we propose a modality-agnostic memory architecture. Engram in Stable Diffusion serves as one instantiation, proving that non-textual concepts can be retrieved and integrated as efficiently as language.

[!TIP] Join the Research You are welcome to propose any questions in the Issues. We will burn our own GPUs to research on any interesting questions. Join us in evolving how LLMs remember what matters! ✨

️ TinyEngram-Vision: Engram Goes Multimodal

Technical Report: We have organized the TinyEngram-Vision findings into a complete technical report. Read the report for the full methodology, experiments, and conclusions, or check out the arXiv version if you prefer.

Can Engram's memory mechanism work beyond text?

We extended TinyEngram to Stable Diffusion, treating visual concepts as "memories" capable of being injected into the Text Encoder. By simply recognizing specific N-grams in the prompt, we inject learned embeddings that guide the generation—all without fine-tuning the massive U-Net or DiT backbone.

It's a lightweight, composable way to "teach" the model new subjects (like specific characters) while keeping the original weights frozen.

Why is this cool?

• Minimal & Surgical: We construct a minimal Engram vocabulary specifically for your target phrase.
• Infinite Composability: Since Engram relies on exact N-gram matching (hard hash collisions), memories strictly do not interfere with each other. You can stack thousands of different character/style engrams together, and they will only trigger when their exact name is called—zero degradation to the base model's general capabilities.

Click here to view the full report (SD1.5 & SD3.5 experiments)

Reproduce our experiments

Interested in injecting your own concepts (or even your cat ) into Stable Diffusion? Check out the reproduction guide here

We provide everything needed to get started:

• Training scripts for SD1.5 & SD3.5
• Pre-processed datasets
• Inference demos

Have Fun!

Key Finding 1. Engram as Parameter Efficient Fine-Tuning Method

1. Engram works as Parameter Efficient Fine-Tuning Method

  Training Setup

We insert several Engram modules into decoder layers of Qwen. We fine-tune the Engram module on a subset of the Biomed-Enriched dataset. Only added parameters are trainable during the fine-tuning.

The train and eval loss demonstrate robust convergence. This confirms that the Engram module effectively learns specialized biomedical knowledge while preserving the stability of the underlying pre-trained knowledge base.

  Training Loss


  
  

  Validation Performance

2. Catastrophic Forgetting

• Objective: Verify if integrating Engram memory harms the model's pre-trained general capabilities while adapting to new domains.
• Methodology: We fine-tune the Biomed-Enriched on Qwen, and evaluate the checkpoint on general benchmarks (We evaluate on MMLU, excluding all biomedical-related subtasks).
• Full Results: Click here to view detailed results

** Update (2026.01.30):** We have added a new set of experiments comparing Engram and LoRA on catastrophic forgetting. Please refer to Engram vs LoRA Catastrophic Forgetting Experiment for details.

3. Vocabulary Scalability Analysis

• Objective: Investigate the relationship between Engram memory size (vocabulary size) and performance gains.
• Methodology: Train multiple models with varying engram_vocab_size (e.g., 2k vs 10k vs 20k vs 100k) and observe the impact on biomedical validation loss.
• Full Results: Larger representation capacities do not necessarily translate into better performance. In our experiments, we observe an apparent trade-off: smaller capacities may suffer from semantic collisions, while larger ones can become difficult to fully utilize given limited data. Click here to view detailed results

** Update (2026.01.30):**
We have expanded our study with a comprehensive ablation of TinyEngram’s configurable hyperparameters. Please refer to Engram Systematic Hyperparameter Tuning Experiment for details.

Reproduce our experiments

To reproduce the experiments conducted in Key Finging 1, please refer to this guide.

Key Finding 2. Engram Outperforms LoRA in Catastrophic Forgetting

LoRA is the de-facto PEFT method, So how does Engram compare? We also conduct systematic hyperparameter tuning to understand Engram better.

1. Engram vs LoRA Catastrophic Forgetting Experiment

Preliminary observation: In our experiments, Engram shows noticeably better resistance to catastrophic forgetting than LoRA.

It is worth noting that LoRA generally converges faster. In our experiments, LoRA could reach an even lower loss (0.1458) quickly, but the trade-off was severe: catastrophic forgetting worsened significantly ($MC1: 0.2472$, $MC2: 0.3993$). Engram provides a safer learning path.

We fine-tune models on "poisoned" function-call-style data (see processing script) based on the glaive-function-calling-v2 dataset, which encourages a strong bias toward structured function-call outputs. We then evaluate both LoRA and Engram on TruthfulQA, a natural language QA benchmark, to examine how well they retain general-language capabilities under this distribution shift. Click here to view detailed results.

2. Engram Systematic Hyperparameter Tuning Experiment

During initial trials, we observed that LoRA converges faster than the default Engram configuration. To enable a scientifically sound comparison, we conducted a systematic hyperparameter study to calibrate Engram such that it reaches evaluation loss levels comparable to LoRA on the same training data.

Using the small-scale, filtered glaive-function-calling-v2 dataset, we ablated key Engram parameters beyond vocabulary size, including:

• N-gram order
• Vocabulary size
• Embedding dimension per n-gram
• Number of hash heads per n-gram
• Target layer(s) for Engram injection

We hope this experiment can serve as a solid starting point for parameter selection in similar small-scale supervised fine-tuning (SFT) scenarios. Click here to view detailed results.

Reproduce our experiments

Reproduction details of experiments conducted in Key Finging 2: please refer to this guide.

️ More Research is on the way!

Acknowledgements

We borrowed a lot of code from the following excellent projects:

• Engram
• Qwen
• lm-evaluation-harness
• Diffusers

We thank the authors of training datasets that help our research:

• Biomed-Enriched
• glaive-function-calling-v2

Citation

If you find TinyEngram useful for your research or