Naturemort/Gemma-3-270m-it-GroomAttention

Gemma 3 270M GroomAttention

This is an experimental MLX checkpoint derived from google/gemma-3-270m-it.

The model keeps the original Gemma 3 270M architecture and dense tensor shapes, but replaces the first transformer layer's self-attention projection weights with a block-modified variant selected by evolutionary search.

Main Result

GroomAttention is not a new attention kernel and not a structurally pruned Transformer. It is a dense-compatible modification of the learned weights in the first self-attention block.

The attention computation itself remains standard Gemma 3 attention: the model still forms queries, keys, and values, applies normalization and RoPE, computes masked attention weights, mixes values, and returns the result through the output projection.

What changes is the learned geometry inside that first attention block. The query and key projections change which token-to-token relations receive high attention scores. The value and output projections change what information is carried and how the attention result is remixed back into the residual stream.

What Changed

Only the first layer self-attention projections were modified:

• query projection,
• key projection,
• value projection,
• output projection.

All other layers, tokenizer behavior, normalization weights, MLP weights, embedding weights, output head weights, and architecture settings are inherited from the base model.

The internal search reported size_ratio = 0.75263916015625, corresponding to an estimated block-representation compression gain of about 24.74% for the searched attention-module representation. The released checkpoint remains a standard dense MLX-compatible model.

Architecture

The underlying model configuration is unchanged:

Property	Value
Architecture	Gemma3ForCausalLM
Layers	18
Hidden size	640
MLP size	2048
Attention heads	4
KV heads	1
Head dim	256
Layer 0 attention type	sliding attention
Sliding window	512
Context length	32768
Dtype	bfloat16

Layer 0 uses grouped-query attention: 4 query heads share 1 KV head.

Mechanism Analysis

Literal Before / After Weight Map

The figure below shows the same four projection matrices before and after GroomAttention as block-level mean absolute weights. The right column shows where the real weight values changed most strongly.

• Each small square is a block of real model weights, not a single weight.
• A block is a rectangular slice of one projection matrix. In this visualization each block summarizes a 32 x 32 region of the original tensor.
• The x-axis is the block column inside the projection matrix.
• The y-axis is the block row inside the projection matrix.
• The left column shows the base model before modification.
• The middle column shows the same projection after GroomAttention.
• The right column shows the absolute change between the two.
• Brighter color means a larger mean absolute value in that block.
• Darker color means a smaller mean absolute value in that block.
• In the change column, bright blocks are the places where GroomAttention changed the original weights most strongly.
• Black or nearly black blocks in the after column usually mean the block was heavily suppressed by the block transformation.

The four rows correspond to the four modified attention projections:

• Query projection: changes how tokens form search queries.
• Key projection: changes how tokens expose matchable features.
• Value projection: changes what information is carried by attended tokens.
• Output projection: changes how the attention result is mixed back into the residual stream.

What the transformation does:

Before GroomAttention, each projection was a normal dense Gemma weight matrix. Every cell in that matrix was a learned floating-point number, and the model used the whole matrix directly during attention.

GroomAttention keeps the same matrix shape, but treats the matrix as a grid of blocks. Each block is processed independently. The search process selected a different operation for each block:

• some blocks are preserved nearly unchanged;
• some blocks are rounded to a lower numerical precision;
• some blocks are strongly suppressed, which is why black regions appear in the after column;
• some blocks are rescaled so their magnitude changes but their location in the matrix stays the same;
• in the output projection, an additional compensation step adjusts the returned attention signal so the modified module remains usable by the rest of the model.

So the transformation is not:

• deleting the attention layer,
• changing the number of heads,
• changing the hidden size,
• changing RoPE or the attention mask,
• replacing Gemma attention with a new algorithm.

The transformation is:

• start with the original dense projection matrix;
• split it into local rectangular blocks;
• apply a selected numeric operation to each block;
• write the transformed blocks back into a matrix with the exact same shape;
• use that dense matrix in the original Gemma forward pass.

In the figure, the before column is the original dense Gemma weight structure. The after column is the transformed dense matrix after block operations. The change column is the absolute block-level difference between those two matrices. Bright regions in the change column are where the block transformation had the largest effect.

What this produced:

• the query/key projections now produce a different attention routing pattern;
• the value/output projections now carry and remix a different payload signal;
• the output projection became the main behavioral bottleneck in later ablation tests;
• the final checkpoint remains loadable as a normal dense MLX model.

The important visual pattern is that the architecture stays rectangular and shape-compatible, while many local weight blocks are suppressed, rescaled, or shifted. This is why GroomAttention is a dense-compatible weight modification, not a new attention architecture.

Weight-Space Change

The modified checkpoint stays tensor-shape compatible with the base model, but the first attention block changes substantially:

Projection	Shape	Changed ratio	Relative RMS delta	Cosine similarity
Query	[1024, 640]	0.6232	0.4258	0.9057
Key	[256, 640]	0.5799	0.4299	0.9035
Value	[256, 640]	0.6301	0.5103	0.8615
Output	[640, 1024]	0.6917	0.5096	0.8622

The largest relative RMS changes are in the value and output projections. This means the experiment changes both where attention points and what information is returned from the attention block.

Attention Entropy

The diagnostic prompts show non-uniform entropy changes across heads:

Head	Entropy before	Entropy after	Delta	Interpretation
0	0.4229	0.5361	+0.1132	More distributed attention
1	0.3189	0.2834	-0.0355	More concentrated attention
2	0.4762	0.4601	-0.0160	Slightly more concentrated
3	0.3286	0.3955	+0.0669	More distributed attention

This is the central behavioral difference: GroomAttention keeps the same attention algorithm, but the learned projections redistribute probability mass differently across token relations.

Token-to-Token Relation Changes

The following sparse heatmap shows the strongest changed attention edges for a diagnostic prompt about a wet road, rain, a driver, and a bridge.

Rows are query tokens. Columns are key tokens. Red means the token relation became stronger in GroomAttention; blue means it became weaker.

Example shifts in this case:

• the driver-to-article relation before driver weakens by about -0.19,
• the bridge-to-article relation before bridge weakens by about -0.19,
• several tokens reduce attention to the beginning-of-sequence token,
• self/local relations such as driver to itself and wet to itself strengthen.

Prompt-Case Diagnostics

Case	Tokens	Mean JS divergence	Entropy delta	Layer-0 output cosine
Commonsense causal relation	21	0.0562	+0.1176	0.7785
Coreference / pronoun resolution	21	0.0681	+0.1586	0.7640
Multi-clause dependency	24	0.0533	+0.1687	0.8111
Multiple-choice style continuation	13	0.0398	-0.0556	0.7936
Short arithmetic context	17	0.0465	+0.1091	0.7448

Aggregate over the diagnostic cases:

Diagnostic	Value
Mean JS divergence	0.0528
Mean entropy delta	+0.0997
Mean layer-0 output cosine	0.7784

These diagnostics are local to layer 0 and to the prompt set above. They show a mechanism shift, not a general proof of quality improvement.

Evaluation

The table below compares the original base model and this selected checkpoint on the same full benchmark splits using 0-shot accuracy.

Task	Original	GroomAttention	Delta	n
HellaSwag	0.3225	0.3137	-0.0089	10042
PIQA	0.6556	0.6436	-0.0120	1838
ARC Challenge	0.2308	0.2207	-0.0100	299
WinoGrande	0.5012	0.5043	+0.0032	1267
Mean	0.4275	0.4206	-0.0069

Summary: the selected model preserves most of the 0-shot benchmark behavior of the original model, with an average accuracy drop of 0.0069 across the four benchmarks. WinoGrande slightly improves in this run, while HellaSwag, PIQA, and ARC Challenge decrease slightly.

Important: these numbers should not be compared directly to official Google Gemma benchmark numbers. They were produced with a local custom MLX multiple-choice continuation scorer, not with Google's official evaluation pipeline. Prompt formatting, tokenizer handling, chat-template usage, and benchmark harness details can materially change reported scores.

Evaluation Protocol

The evaluation uses 0-shot multiple-choice continuation scoring. For each task, the model scores each answer option as a continuation of the prompt and selects the option with the highest summed continuation log-probability. No free-form generation is used for the benchmark table.

Setting	Value
Evaluation type	0-shot multiple-choice continuation scoring
Generation	Not used
Batch size	1
Evaluation limit	Full validation split
Seed	Not used by the deterministic scoring loop
Chat template	Not applied during benchmark scoring
Choice selection	Highest summed continuation log-probability
Metric	Accuracy

Datasets:

Task	Dataset	Config	Split
HellaSwag	Rowan/hellaswag	default	validation
PIQA	ybisk/piqa	plain_text	validation
ARC Challenge	allenai/ai2_arc	ARC-Challenge	validation
WinoGrande	allenai/winogrande	winogrande_xl	validation

The exact dataset revision hashes were not pinned in the original run. For a stricter reproduction, pin dataset revisions and record the resolved dataset versions before publishing final benchmark claims.

Limitations

• This is an experimental checkpoint.
• Only the first self-attention layer was modified.
• The released model is dense; it does not provide a smaller physical model.
• The compression gain refers to the internal searched block representation, not to released model size.
• The mechanism diagnostics are local to layer 0 and a small prompt set.
• The reported evaluation uses a custom local continuation-scoring harness and is not directly comparable to official Google Gemma benchmark numbers.
• Results may differ with other prompts, decoding settings, benchmark harnesses, or dataset revisions.

Citation

If you use this checkpoint, cite the original Gemma model according to Google's model card and documentation. This checkpoint is a derivative research artifact and inherits the original model license terms.