uattn 0.1.0

Unified Attention: single projection replacing Q/K/V with 67% fewer parameters

Unified Attention

One projection replaces Q, K, V. 67% fewer attention parameters.

Standard self-attention uses three separate weight matrices (W_Q, W_K, W_V) to project the same input. Unified Attention uses one. The output splits into three functional bands that serve as query, key, and value. The bands form naturally during training through amplitude and phase differentiation in the weight matrix.

from uattn import UnifiedAttention

attn = UnifiedAttention(dim=528, num_heads=4)
y = attn(x)  # [batch, seq, 528] -> [batch, seq, 528]

Install

pip install uattn

From source:

pip install git+https://github.com/ReinforceAI/unified-attention.git

For FA3 (Hopper, 1.57x faster):

pip install uattn[fa3]

Why

Attention routes tokens to each other. The MLP transforms them. In parameter-constrained models, every byte spent on routing is a byte not spent on computation. Unified Attention cuts the routing budget:

	Standard Q/K/V	Unified
Attention projection params	3d²	d²
Share of layer params	~33%	~18%
Compressed attention bytes (16MB model)	5.10 MB	2.82 MB
Compressed MLP bytes (16MB model)	10.21 MB	12.70 MB

Results

Validated on OpenAI Parameter Golf (16MB artifact, FineWeb):

Setting	BPB	vs Standard SOTA
10-min record (8xH100)	1.1412	Rank 7, only novel architecture in top 10
1-hr unlimited (8xH100)	1.1088	Beats all submissions on both leaderboards

Also tested at small scale (484K params, TinyStories): 9.58 perplexity, matching models 2-4x larger.

Usage

Standalone (owns its own weights)

from uattn import UnifiedAttention

attn = UnifiedAttention(
    dim=528,              # must be divisible by 3
    num_heads=4,          # heads for multi-head attention
    rope_base=10000.0,    # RoPE base frequency
    seeking_gain_init=0.5,# per-head query scaling
    rope_fraction=1.0,    # fraction of head_dim for RoPE
)

x = torch.randn(batch, seq_len, 528)
y = attn(x)  # uses internal W_unified and W_out

With external weight banks (for Muon optimizer)

# Weights stored in contiguous 3D banks for batched optimization
unified_bank = nn.Parameter(torch.randn(K, d, d))
output_bank = nn.Parameter(torch.randn(K, d, d//3))

attn = UnifiedAttention(dim=d, num_heads=4, output_proj=False)
y = attn(x, unified_w=unified_bank[k], output_w=output_bank[k])

Check backend

from uattn import backend
print(backend())  # 'fa3', 'fa2', or 'sdpa'

How It Works

# Standard: 3 separate projections
q = W_q @ x;  k = W_k @ x;  v = W_v @ x

# Unified: 1 projection, 3 bands
unified = W_unified @ x
seeking, offering, content = unified.split(d//3, dim=-1)

RoPE is applied to seeking and offering bands only. Position affects routing, not content.

For FA3 (Hopper): head_dim is zero-padded to the nearest multiple of 8 before the kernel call and sliced back after. Padded zeros contribute nothing to dot products. Mathematically lossless.

Citation

@misc{deshwal2026yocto,
  title={Attention Fields: Unified Projections for Efficient Language Models},
  author={Deshwal, Viraj},
  year={2026},
  url={https://github.com/ReinforceAI/yocto},
  howpublished={\url{https://github.com/ReinforceAI/unified-attention}}
}

License

MIT