LongChenTong/LSTMem-Qwen3-4B

LSTMem Adapter for Qwen3-4B-Instruct-2507

This repository contains the LSTMem adapter for Qwen/Qwen3-4B-Instruct-2507, as presented in the paper LSTMem: Hierarchical Long Short-Term Online Memory for Large Language Models.

LSTMem is an LSTM-inspired online memory for a frozen large language model. Each attention layer gets two matrix-valued states: a cell state that accumulates the history and a hidden state whose readouts correct the backbone's attention. Input and forget gates control what the cell stores, while an output gate separately controls what the cell exposes through the hidden state. Memory is also connected across depth: hidden states propagate forward from layer to layer, and at the end of every history block, block-end feedback uses the reconstruction gradients of higher layers to refine the cell states of lower layers.

• Paper: arXiv:2609.33268
• Code: https://github.com/Longchentong/LSTMem

Model Details

Item	Value
Base model	Qwen/Qwen3-4B-Instruct-2507 (frozen)
Adapter type	LSTMem
Memory layers	all 36 attention layers
Memory heads x rank	4 x 8, scaling alpha = 16
Attention corrections	query and output
Block-end feedback	32 sampled positions per history block, step size 0.01, 5% cap on the relative cell update
Parameters	20.2M (468 BF16 tensors, 40.5 MB)

Files:

• lstmem_config.json: memory configuration (heads, rank, gate initialization, feedback settings)
• lstmem_adapter.safetensors: adapter weights

Results

MemoryAgentBench	LoCoMo	HotpotQA EM / F1	IFEval	GPQA-Diamond
45.61	54.29	53.86 / 68.22	81.70	42.42

Usage

To use this adapter, first install the LSTMem codebase:

git clone https://github.com/Longchentong/LSTMem.git
cd LSTMem
bash scripts/setup_uv_env.sh
source .venv/bin/activate

Minimal Loading Example

import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer

from lstmem import load_lstmem_adapter
from lstmem.runtime import Decoding, generate_answer

base_model = "Qwen/Qwen3-4B-Instruct-2507"
adapter_dir = snapshot_download("LongChenTong/LSTMem-Qwen3-4B")

tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    device_map={"": "cuda:0"},
).eval()
load_lstmem_adapter(model, adapter_dir)  # attaches the memory modules and loads the adapter weights

prompt = open("report.txt").read() + "\n\nQuestion: Which site was adopted as the backup?"
print(generate_answer(model, tokenizer, "cuda:0", prompt, Decoding(max_new_tokens=64)))

Citation

@article{shi2026lstmem,
  title   = {{LSTMem}: Hierarchical Long Short-Term Online Memory for Large Language Models},
  author  = {Shi, Xianglong and Yang, Ruijie and Zhao, Sirui and Yin, Shukang and Bian, Zihao and Yi, Tinghao and Chen, Enhong},
  journal = {arXiv preprint arXiv:2609.33268},
  year    = {2026}
}