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

  1. KD distillation from the original model on FineWeb-Edu (earlier HRR stage, 2K sequences).
  2. 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).
  3. 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, giving student4).
  4. 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-attention chunk_gated_delta_rule kernel. 4K sequences, 70% recall data, answer-loss weight 2, 35 min (train4.py, giving student5).
  5. 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 with MEM=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=delta2 requires pip 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.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Model tree for staccs/Qwen2.5-0.5B-lecore-infmem

Finetuned
(762)
this model