WORK IN PROGRESS / research checkpoint. Single-passkey recall works; multi-fact recall on real text FAILS (11/108 against 95/108 for the original model; the target of 70/108 was not reached).
Qwen2.5-0.5B-lecore-infmem
Qwen2.5-0.5B converted to constant-memory attention: every attention layer keeps the original Q/K/V/O weights and uses a 256-token sliding-window softmax plus 4 attention-sink tokens, plus a gated delta-rule recurrent memory trained into the weights. There is no gateway, no retrieval, and no external store. The decode state is a fixed 8.9 MB (delta2; 8,945,664 bytes; the earlier delta version was 8,687,616 bytes, measured identically at 2K, 8K and 32K context).
Note: the memory is not pure HRR. An earlier holographic (HRR, circular-convolution) memory with forget gates reached
0/36 passkey recall. This checkpoint replaces it with an error-correcting delta-rule memory: per head a 64x64 fp32 matrix
S_t = a_t S_{t-1}(I - b_t k k^T) + b_t v k^T, with L2-normalised keys and queries (a small learned projection on each),
data-dependent decay a_t and write strength b_t, and an output gated in through a learned per-head mix.
The legacy HRR branch is still in the code (MEM=hrr).
Training
- KD distillation from the original model on FineWeb-Edu (earlier HRR stage, 2K sequences).
- Delta memory + KD with 50% synthetic single-passkey examples (needle always > 256 tokens before the question), 2K sequences, 24 min on one RTX PRO 5000 (
train3.py). - Same, but with 1-8 key:value facts, 5 varied templates, mostly real FineWeb text as filler, and supervised answers for every fact, 30 min (
train4.py, givingstudent4). MEM=delta2(stage 4): separate learned write projections for query/key/value (initialised from q/k/v_proj), a causal depthwise conv of width 4 on q/k/v (token-shift context), data-dependent write gate b_t and decay, and a fused flash-linear-attentionchunk_gated_delta_rulekernel. 4K sequences, 70% recall data, answer-loss weight 2, 35 min (train4.py, givingstudent5).- Stage 5: same architecture; 2-12 interleaved facts per sequence, wordings randomised over 5 frames x 18 relation words, answer-loss weight 3, 38 min (
train6.py, giving this checkpoint,student6.pt, run withMEM=delta2). Evaluation keys (12 held-out words) were never used in training. Stage 5's random frames include the eval's "The {relation} for {key} is" frame and the relation word "code", so stage 5 only tests generalisation to unseen keys, not to an unseen wording. Evaluation templates and keys ("The code for {word} is ...", 12 held-out words) were never used in training.
Results (greedy decoding, constant-state streaming)
| perplexity @2K / @8K (FineWeb-Edu held-out) | single passkey beyond window, 1K / 4K / 8K | 16K / 32K | |
|---|---|---|---|
| original Qwen2.5-0.5B (full attention, KV cache grows) | 13.6-15.2 | 12/12 at all lengths | - |
| HRR memory (previous) | 13.67 / 13.63 | 0/12, 0/12, 0/12 | - |
| delta memory, stage 2 | 13.76 / 13.71 | 12/12, 12/12, 12/12 | 12/12, 12/12 |
| delta memory, stage 3 | 13.81 / 13.77 | 12/12, 12/12, 12/12 | - |
| delta2, stage 4 | 13.78 / 13.72 | 12/12, 12/12, 12/12 | - |
| delta2, stage 5 (this) | 13.79 / 13.73 | 12/12, 12/12, 12/12 | - |
Harder eval (hardeval.py): real-text filler, 4 or 8 key:value facts at random positions (all > 256 tokens before the questions), each fact queried separately.
| context / number of facts | original (full attention) | stage 2 | stage 3 | stage 4 | stage 5 (this) |
|---|---|---|---|---|---|
| 4K / 4 | 10/12 | 0/12 | 2/12 | 2/12 | 3/12 |
| 4K / 8 | 21/24 | 0/24 | 0/24 | 0/24 | 2/24 |
| 16K / 4 | 11/12 | 0/12 | 2/12 | 2/12 | 3/12 |
| 16K / 8 | 21/24 | 1/24 | 1/24 | 1/24 | 1/24 |
| 32K / 4 | 10/12 | 0/12 | 2/12 | 1/12 | 1/12 |
| 32K / 8 | 22/24 | 0/24 | 0/24 | 1/24 | 1/24 |
| total | 95/108 | 1/108 | 7/108 | 7/108 | 11/108 |
Honest caveats
- Single-passkey success is largely format-specific. Stage 2 recalled the passkey template perfectly, even in real text, but got near misses (e.g. 41213 instead of 48213) with other phrasings. General associative recall is not solved: 7/108 against 95/108 for the original model.
- Tiny training budget (under 1 GPU-hour for the memory stages, 2K sequences). Recall extrapolating to 32K is promising, but needs to be shown on real tasks.
MEM=delta2requirespip install flash-linear-attention.- What to try next (untested): train the memory from scratch instead of on top of a KD-tuned model, use far more associative-recall tokens (this used about 15M), use more heads or a larger value dimension, or add a retrieval loss on intermediate memory reads.
Files in this repo
| file | size | what it is |
|---|---|---|
student6.pt |
1.22 GB | support file |
he.log |
29.2 KB | support file |
he4.log |
29.2 KB | support file |
he5.log |
29.2 KB | support file |
he6.log |
29.2 KB | support file |
hrr_swa.py |
12.9 KB | support file |
student6.pt.log.json |
12.5 KB | config / tokenizer |
train6.py |
5.6 KB | support file |
Usage
from transformers import AutoModelForCausalLM; import torch
from hrr_swa import convert, stream_logits_last
m = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B", dtype=torch.bfloat16).cuda()
convert(m, 256); m.load_state_dict(torch.load("student6.pt"), strict=False) # set env MEM=delta2 before importing hrr_swa
last_logits, states, _ = stream_logits_last(m, input_ids, chunk=2048) # states stay a fixed size at any length
License
Base model: Qwen/Qwen2.5-0.5B β Apache-2.0. Only adaptation was done here: no retraining of the base weights.
Model tree for staccs/Qwen2.5-0.5B-lecore-infmem
Base model
Qwen/Qwen2.5-0.5B