MCR-Attention V3.0

Multi-Scale Compressed Recurrent Attention — an autoregressive language model architecture that replaces quadratic self-attention with K multi-scale linear recurrences and a tiny 1×K attention router, stacked into configurable depth.

Mission: keep using attention — but simple and linear everywhere. O(N) training, O(1) per-token streaming inference, GPT-style autoregressive generation, optimized.

Why MCR-Attention?

Transformers scale quadratically with sequence length. MCR-Attention keeps the idea of attention (a learned query attending over context) but operates over K compressed recurrent memory vectors instead of all N tokens. Each memory scale is an independent linear recurrence with a learned decay, so:

Measured on the bundled benchmark suite (CPU, d_model=128, 2 layers):

[1] Parallel training cost: time ratio 27x over 32x length   -> ~LINEAR  PASS
[2] Streaming latency: 1.49 -> 1.49 -> 1.49 -> 1.70 ms/token -> FLAT     PASS (O(1))
[3] Streaming memory: 1.0 KB constant                        -> FLAT     PASS (O(1))
[4] Train/inference identity: cosine 0.99999988              -> IDENTICAL PASS

What's new in V3

V3 keeps the philosophy and fixes the real issues of V2:

• Stacked blocks (depth) — repeat the {memory + router} + {FFN} block num_layers times. Streaming stays O(1); capacity goes up.
• Train == inference — the parallel training forward and the streaming inference step compute exactly the same math (V2 diverged for long-memory scales).
• Geometric decay init — each scale starts with its target memory horizon.
• O(N) scan — a single fused cumsum (with a stable fallback).
• Real HuggingFace integration — AutoModelForCausalLM compatible.
• V2 compatibility — old models still load and run, bit-exact.

See CHANGELOG.md for the full diff.

Installation

pip install mcr-attention

Optional extras:

pip install "mcr-attention[dev]"        # pytest, hypothesis
pip install "mcr-attention[hf]"         # transformers + huggingface-hub
pip install "mcr-attention[benchmark]"  # matplotlib
pip install "mcr-attention[all]"

From source:

git clone https://github.com/bueormnew/MCRA.git
cd MCRA
pip install -e ".[dev]"

Quick start

Create a model

from mcr_attention import MCRConfig, create_model

config = MCRConfig(vocab_size=50257, d_model=512, num_layers=6)
model = create_model(config)
print(f"{sum(p.numel() for p in model.parameters()):,} parameters")

Train (causal multi-position LM)

import torch
from mcr_attention import Trainer
from mcr_attention.api import TrainingConfig

model = create_model(vocab_size=50257, d_model=512, num_layers=6)

def data_fn():
    # Return a batch of token sequences [B, N]; the trainer auto-shifts targets
    # (predicts token t+1 from position t) so every position is a training signal.
    return torch.randint(0, 50257, (16, 512))

trainer = Trainer(model, TrainingConfig(
    num_steps=10000, learning_rate=3e-4, warmup_steps=500,
    lr_schedule="cosine", use_amp=True, save_dir="ckpt/", save_interval=1000,
))
stats = trainer.fit(data_fn)
print(f"Final loss: {stats['final_loss']:.4f}")

You may also return (input_ids, targets) — targets can be [B, N] (multi-position) or [B] (single last-token, V2-compatible).

Save / load

from mcr_attention import save_model, load_model

save_model(model, "my_model/")            # config.json + model.pt + metadata.json
model = load_model("my_model/", device="auto")  # auto-detects V2 vs V3

Generate (O(1) per token streaming)

from mcr_attention import generate, GenerationConfig

tokens = generate(model, prompt=[1, 2, 3], config=GenerationConfig(
    max_tokens=200, temperature=0.8, top_k=50, top_p=0.9,
    repetition_penalty=1.1, seed=42,
))
# Greedy: temperature=0.0

Architecture

Token IDs
   │
   ▼
[Embedding + RoPE]                 ── position-aware token representations
   │
   ▼
┌─────────────── repeated num_layers times (pre-norm residuals) ───────────┐
│                                                                          │
│   m = RMSNorm(x)                                                         │
│   C = MultiScaleMemory(m)        ── K linear recurrences:                │
│   │        h_t = a_k · h_{t-1} + B_k · m_t   (parallel scan / 1 step)    │
│   │        C ∈ R^{K × d}  (one compressed memory per scale)              │
│   ▼                                                                      │
│   z = AttentionRouter(C, m)      ── 1 query × K keys softmax attention   │
│   x = x + Dropout(z)              (residual)                             │
│   │                                                                      │
│   ▼                                                                      │
│   x = x + Dropout(FFN(RMSNorm(x)))  ── SwiGLU / GELU position-wise MLP   │
│                                                                          │
└──────────────────────────────────────────────────────────────────────────┘
   │
   ▼
RMSNorm → LanguageHead → logits [B, N, vocab]

Training runs the parallel scan over the whole sequence and produces logits at every position. Streaming runs one recurrence step per token. The two are mathematically identical at each position — there is no train/inference gap.

Multi-scale memory

Each scale k has a learned decay a_k = exp(-softplus(λ_k)) and is initialized so its effective horizon is roughly L_k = scale_min · scale_factor^(k-1) tokens (short scales forget fast, long scales retain). With the defaults this gives horizons of 256, 512, …, 32768 tokens out of the box.

Attention router

A 1-query × K-key attention: the query is formed from the current token and the longest-scale memory; keys/values are the K memory vectors. It is O(K) (cheap) regardless of sequence length.

Configuration

MCRConfig(
    vocab_size=50257,     # token vocabulary
    d_model=512,          # residual-stream width (must be even for RoPE)
    d_state=128,          # recurrent state width per scale
    num_scales=8,         # number of memory scales (K)
    num_layers=6,         # depth (stacked blocks)
    scale_min=256,        # nominal shortest horizon
    scale_factor=2,       # geometric ratio between horizons
    rope_base=10000.0,    # RoPE base frequency
    dropout=0.1,
    use_input_gate=True,  # learned sigmoid input gate per scale
    use_rope=True,
    init_scales_geometric=True,
    ffn_type="swiglu",    # "swiglu" or "gelu"
    ffn_hidden=0,         # 0 = auto (8/3 · d_model)
    tie_embeddings=False,
    mlp_hidden=0,         # head MLP width (0 = auto: 4 · d_model)
    max_seq_len=32768,
)

Fine-tuning

from mcr_attention import load_model
from mcr_attention.finetune import finetune, FinetuneConfig

model = load_model("my_model/")
stats = finetune(model, data_fn, FinetuneConfig(
    learning_rate=1e-4, num_steps=500, freeze_embeddings=True, freeze_memory=True,
))

Selective freezing: freeze_embeddings, freeze_memory, freeze_router, freeze_ffn, or freeze_all_except(model, ["head"]).

HuggingFace integration

Real transformers integration (requires mcr-attention[hf]):

from mcr_attention.hf import MCRForCausalLM, MCRPretrainedConfig, push_to_hub, from_pretrained

# AutoModel-compatible
model = MCRForCausalLM(MCRPretrainedConfig(vocab_size=50257, num_layers=6, d_model=512))

# Hub
push_to_hub(model, "username/my-mcr", tokenizer=tokenizer)
model = from_pretrained("username/my-mcr")
print(model.generate_text("Once upon a time", max_new_tokens=200))

MCRForCausalLM supports the standard HF model.generate(...) and works with AutoModelForCausalLM.from_pretrained.

Compatibility: keep using V2 models

V3 loads V2 checkpoints bit-exact via the legacy architecture, and can migrate them to V3:

from mcr_attention import load_model, convert_v2_to_v3

old = load_model("v2_model/")        # returns the legacy V2 model, identical behavior
convert_v2_to_v3("v2_model/", "v3_model/", num_layers=1)  # warm-start migration

Version is auto-detected from config.json (arch_version). Old checkpoints without a version load as V2.

CLI

mcr-attention info path/to/model/                       # show model info
mcr-attention benchmark --lengths 256 1024 4096 --layers 2   # speed/memory
mcr-attention generate path/to/model/ --prompt 1 2 3 --max-tokens 50 --temperature 0.8

Benchmarks

Run the mission-verification suite:

python benchmarks/run_benchmark.py --d-model 128 --layers 2

It checks the four pillars and reports PASS/FAIL:

1. Linear training cost — parallel forward time grows ~linearly with N.
2. O(1) streaming — per-token latency is flat as N grows.
3. O(1) streaming memory — peak memory is flat.
4. Train/inference identity — parallel logits == streaming logits.

Synthetic-quality tasks (copy / needle / associative recall) live in mcr_attention.benchmarks.

Development

pip install -e ".[dev]"
pytest tests/ -v                 # 100+ tests, all green
python benchmarks/run_benchmark.py

License

GNU AGPL-3.0-or-later. See LICENSE.

• Author: Gerson Fabian Buenahora Ormaza (BUEORM)
• GitHub: bueormnew/MCRA
• Contact: dalusx64@gmail.com

MCR-Attention is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. Network use (e.g. offering the model as a service) triggers the source-disclosure obligation — see section 13 of the AGPL.