Replacing Gemma-3 Global Attention with Concepts (Toward Longer Context)

TL;DR #

What you will learn:

• How I replace Gemma-3’s global-attention layers with a recurrent bank of 128 concepts on frozen Gemma-3 1B + LoRA
• Why short-context CE left that replacement unused — local windows still did the job
• What changed at 4K / 1B tokens under Muon: beyond-local ΔCE rose from ~0.003 to ~2.35 nats
• Why this is a step toward longer context, not a claim that 10M context is solved

This is a progress report from MrCogito — Concept Reasoning Model. After Phase 1 hit concept collapse, I stopped paying the from-scratch language-acquisition tax on every memory test. The Gemma track asks a narrower architectural question:

If I remove Gemma’s global attention and put concepts in its place, can that concept state actually carry long-range information?

The honest answer so far: not under short-context CE — local windows are too strong. Under a 4K long-document Muon run, yes: beyond-local next-token prediction started to depend on the concepts. That is not a 10M-token demo. It is evidence that the replacement can work, which is why I am still training in this direction.

Why replace global attention at all? #

Agents need state that survives tools, turns, and long documents. Today that usually means stuffing more tokens into a window, retrieving chunks, or bolting on external memory. Those help. They still inherit expensive attention and a narrow token interface.

Gemma-3 already took one smart step toward efficiency: most layers are local sliding-window attention; only every sixth layer is global (Gemma 3 Technical Report). That cuts KV-cache cost for long context. But the global layers still do full attention over N tokens. At million-token scale, that remains the expensive part.

My working bet in MrCogito is to go one step further:

1. Keep local attention for fluent token↔token work
2. Replace global attention with reads/writes over C concept vectors (C ≪ N)
3. Carry those concepts recurrently across blocks, so sequence length grows by looping, not by growing a full global KV cache

If that works, three things become more tractable:

• Longer context as an architectural consequence. Attention becomes O(N·(K+C)) with fixed block size K, not O(N²) global attention. Multi-million-token windows become a compute and state-quality problem. The ~10M ambition on the project page is still ambition — this post trains at 4K — but replacing global attention is the mechanism that makes the ambition coherent.
• Higher reasoning bandwidth. Discrete CoT writes roughly O(log |V|) bits per step (Chen et al.). Continuous latent work such as COCONUT and Soft Thinking (Zhang et al., 2025) points at wider internal channels. Kuratov et al. show a vector can hold far more than one token’s worth of information when used well. Concepts are my trainable version of that wider channel.
• Reusable state for agents. A concept workspace is a candidate bus for memory and tools — but only if the model actually depends on it.

Phase 1 trained encoders from scratch and watched concepts collapse. Useful. Expensive. For this track I refuse to re-learn English every time I test the long-range channel. Start from a pretrained decoder. Replace its global path. Measure whether concepts take over that job.

Why Gemma-3 1B is the right place to do the surgery #

I chose google/gemma-3-1b-pt because the hybrid pattern is explicit and the model still fits evening GPUs:

• iterable on 3× RTX 3090
• strong enough that LoRA fine-tunes are informative
• 5 local + 1 global already marks where long-range routing lives

On the 1B text model (num_hidden_layers=26, hidden_size=1152) there are four global layers at 0-indexed indices 5, 11, 17, 23. Local window is 512 tokens.

Those four layers are the surgical site. Gemma was pretrained to send long-range information through them. Design C keeps that inductive bias and changes the substrate:

I am not adding concepts on top of global attention. I am replacing the global-attention channel with concepts, then asking whether the model learns to use that channel.

What I built #

BackboneConceptLM in MrCogito: frozen Gemma-3 1B + LoRA, plus concept I/O. Related to Perceiver-style latent bottlenecks (Jaegle et al., 2021), but grafted into the pretrained hybrid stack rather than trained as a separate encoder.

Working configuration:

Read — concepts stand in for global attention #

Only the former global layers may attend concepts. Everywhere else, tokens see a local window. If information from far earlier still matters, it cannot come from full global attention anymore. It has to live in z.

Write — concepts are recurrent, not a static prompt #

After each concept-reading depth, a tied BiXT write updates the same bank:

z ← z + tanh(α) · RMSNorm(BiXT_lat←tok(z, block_hiddens))

α is a scalar gate. Gate init 0.01 keeps the path live without immediately dominating the pretrained backbone.

Depth-wise updates — refine the long-range state inside the block #

The first graft wrote once after all 26 layers, so every concept-reading layer saw the same stale z. The better schedule writes after each global layer:

layer 6  reads z     → tied write → z^(6)
layer 12 reads z^(6) → tied write → z^(12)
layer 18 reads z^(12)→ tied write → z^(18)
layer 24 reads z^(18)→ tied write → z^(24)
next block starts from z^(24)

That is not full recursive reasoning yet. It is the minimum recursive property the long-context bet needs: one workspace refined as depth progresses, then carried to the next token block.

How I measure whether the replacement worked #

If concepts replace global attention, the decisive question is not “do concepts look diverse?” It is: at positions that used to need global reach, does damaging concepts hurt next-token CE?

Those are concept ablation ΔCE metrics on beyond-local positions:

“Beyond-local” means past what local windows + one-block carry can explain — roughly ≥1024 on 2K runs, with longer buckets on 4K runs. That is exactly the region where stock Gemma would have used global attention.

Pre-registered gate for the long-context run:

min(Δstatic_beyond, Δshuffle_beyond) ≥ 0.01 nats

RankMe tracks geometry. Gates track whether the path is open. Neither substitutes for causal use.

Healthy concept geometry ≠ a working global-attention replacement.

Short context: local windows still win #

Concepts present, long-range channel unused #

First 100M-token pilot:

• RankMe 77.1 — not collapsed
• beyond-local recurrent ablations < 0.001 nats
• WandB: first Gemma concept pilot

I had replaced global attention in the graph. The model still solved short-context CE mostly from local windows. The replacement existed; the job had not transferred.

Calibration opened the path, not the transfer #

I tried the usability knobs:

• RMSNorm on concept reads
• nonzero gate init (0.01)
• higher LR on concept-memory params
• then the combination at 100M tokens

CE could improve. Beyond-local Δ stayed near zero. Under 2K plain CE, calibration was not enough.

Forced-memory diagnostics #

Delayed-recall tasks, where only earlier-block state identifies a later answer, died early or stayed at chance on the hard block. Useful negative: the issue was not only “missing a synthetic task.” At short length, local Gemma context still carried too much of the prediction.

Depth-wise writes at 2K #

Better update schedule, same regime limit at 50M / 2K:

• RankMe 62.2
• Δstatic / Δshuffle beyond ≈ 0.0005 / 0.0010
• WandB: shared-depth run

Adam vs Muon #

Matched 100M / 2K optimizer A/B on the shared-depth architecture:

Muon orthogonalizes matrix updates; Liu et al., 2025 show weight decay and update-scale care matter at LLM scale. Muon won CE and RankMe. Both arms still missed the 0.01 causal-use gate.

Conclusion: the optimizer was not the main blocker. The short-context regime still let local attention finish too much of the work that global attention / concepts were supposed to own. Keep Muon. Make long-range dependence harder to avoid.

Promising result: when local context is not enough #

I kept the replacement architecture fixed and changed the operating regime:

At 2K / K=512 there are ~4 concept blocks. At 4K there are 8, and longer documents reward cross-block state — the job global attention used to do.

Training finished stably after ~38 wall hours on Odra (3× 3090). Beyond-local dependence grew with scale:

Offline Tier-1 on the best checkpoint:

• within-sample RankMe 101.0 / 128
• Δshuffle ≥1024 2.47
• Δstatic ≥1024 2.35
• Δone-block ≥1024 0.58
• ≥3/4 read and write gates above 0.005 magnitude

WandB: long-context Muon run

This is the promising signal for the title claim:

concepts started doing the long-range job that global attention used to do.

Not STS-B success. Not recursive reasoning. Causal use of the replacement channel for beyond-local CE. Δone-block at 0.58 further suggests multi-block accumulation, not only a one-block carry.

Why this may have worked #

1. Pressure on the replaced channel. Longer docs make local windows insufficient; the model has to use whatever long-range path remains — here, concepts.
2. Scale curve. Short-context Muon at 0.0028 looks like an early point on a curve 2K runs never climbed.
3. Compound regime change. Architecture fixed; length + data + budget + Muon moved together. Factor isolation remains open.
4. The graft was viable earlier than it looked. Short-context near-nulls were mostly regime evidence.

What this does not prove yet #

• Not 10M context. Toward longer context means the mechanism is pointed that way. This run trained at 4K.
• Not parity with stock Gemma global attention on every long-range behavior. I measured causal concept use, not a full quality bake-off against unmodified Gemma at matched length.
• Not semantic / reasoning success. Probes are next.
• Not a single-variable proof that Muon, 4K, or the long-doc mix alone is enough.
• Not a bake-off against other read designs (memory tokens, per-layer KV prefixes).
• Causal use ≠ desired abstractions. The channel is used. Whether concepts are the right latent factors remains open.

Why I continue from here #

I keep training because this is one of the few bets where architecture, long context, and agent state point the same way. Gemma already showed hybrid local/global attention is useful. The next question is whether concepts can own the global role more scalably than full attention over N.

Phase 1 showed concepts can collapse. The early Gemma runs showed the opposite failure: concepts present while local windows still carry the load. The 4K result is the first time this replacement cleared a pre-registered beyond-local causal-use gate by a large margin.

Next experiments, in order:

1. semantic probes on this checkpoint
2. longer sequences and C-scaling once semantics look real
3. recursive refinement in concept space — the actual reasoning bet on top of the long-range channel

Practical lessons if you are attempting a similar graft:

1. Measure beyond-local causal use. That is where the replaced global path must prove itself.
2. Pretrained hybrid models are excellent surgical targets: the global layers are already labeled.
3. If local context can solve the objective, your global replacement will stay ornamental.
4. Optimizers matter; making the long-range job unavoidable matters more.

I am not claiming a finished long-context concept model. I am claiming a clearer architectural step: replace Gemma-3 global attention with concepts, and show the model can learn to use that channel. That is enough reason to keep going — carefully, and toward longer context rather than announcing it early.

References #

1. Chen, L. et al. (2025). The Information Bottleneck of Chain-of-Thought and How Latent CoT Overcomes It. ICLR 2026 submission.
2. Kuratov, Y. et al. (Feb 2025). Cramming 1568 Tokens into a Single Vector and Back Again: Exploring the Limits of Embedding Space Capacity. ACL 2025. arXiv:2502.13063
3. Hao, S. et al. (Dec 2024). Training Large Language Models to Reason in a Continuous Latent Space. arXiv:2412.06769
4. Zhang, Z. et al. (May 2025). Soft Thinking: Unlocking the Reasoning Potential of LLMs in Continuous Concept Space. NeurIPS 2025. arXiv:2505.15778
5. Gemma Team et al. (Mar 2025). Gemma 3 Technical Report. Google DeepMind. arXiv:2503.19786
6. Jaegle, A. et al. (2021). Perceiver: General Perception with Iterative Attention. ICML 2021. arXiv:2103.03206
7. Jordan, K. et al. (2024). Muon: An optimizer for hidden layers in neural networks
8. Liu, J. et al. (Feb 2025). Muon is Scalable for LLM Training. arXiv:2502.16982
9. Sopyla, K. (Mar 2026). Quicker Failures Lead to Better Questions. ai.ksopyla.com
10. Sopyla, K. MrCogito — Concept Reasoning Model. Project page.
11. Experiment sources: long-context run report · Gemma graft spec · shared-depth spec · master experiment log

Metrics are from my own training and offline Tier-1 evaluation on Odra (3× RTX 3090). Independent reproduction has not been published.