We’re questioning something that’s been baked into transformer design since the original paper: why does every layer use the same number of attention heads?
The uniform allocation implicitly assumes every layer benefits equally from the same representational granularity — but early layers and late layers have fundamentally different structural needs.
We propose the Prism Transformer: a progressive head schedule that starts with fewer, wider heads in early layers and monotonically increases head count with depth. Early layers get a wide per-head subspace (dh=256) for rich local pattern capture. Late layers get many narrow heads for specialized decomposition.
Crucially — this requires zero architectural overhead:
- Weight matrices retain their standard dmodel×dmodel shapes (parameter neutral)
- Total FLOPs are mathematically invariant to head count (compute neutral)
- Power-of-2 head dims (dh ∈ {256, 128}) preserve Tensor Core alignment (throughput neutral)
Results:
→ Lower validation loss at every scale vs uniform baseline
→ Identical tokens/sec and wall-clock time
→ Gains or parity across PIQA, WinoGrande, HellaSwag, ARC-Easy
Per-layer attention-distance analysis confirms the gain is structural — early Prism layers attend significantly more locally than baseline, flipping to more global integration precisely where the head schedule completes.
Implementation is a one-line change per attention layer — head count becomes layer-dependent, everything else stays identical.
Paper: [2606.27449] Prism Transformer: Progressive Head Schedules for Hierarchical Attention Processing
This is a really elegant result and it aligns closely with something we’ve been observing empirically from the other direction.
We’ve been measuring the geometric trajectory of hidden states layer-by-layer across multiple open-weight models (GPT-2, OPT, Qwen, TinyLlama, Phi, Llama 3.2) using a framework called LIMEN. One of the consistent findings is that functional zones are not uniformly distributed across depth:
-
early layers (~0.54 relative depth) show strong local structure (syntax / lexical)
-
mid-to-late layers (~0.78) show global integration and decision-like regimes
-
the transition between these regimes is marked by measurable curvature and entropy shifts
Your Prism Transformer is basically the architectural consequence of that geometry: early layers need fewer, wider heads to capture local patterns; late layers need more, narrower heads to handle decomposition. The fact that this improves validation loss without changing FLOPs or parameters suggests that the standard uniform head allocation is structurally under-optimized relative to the depth-dependent geometry the model actually uses.
One thing I’d be curious about: your head schedule increases monotonically with depth have you tested whether the peak of the head-count curve aligns with the decision / integration zone we observe around 0.78–0.85 depth? In our measurements, that zone is where curvature and functional capacity peak, and where transitions between regimes are most frequent.
Also, your observation that early Prism layers attend more locally while late layers attend globally matches the pattern we see in ct_t (trajectory curvature): low curvature in early layers (smooth local motion), high curvature in late layers (sharp directional shifts toward the final token).
This seems like two independent lines of evidence converging on the same structural insight one from architecture design, one from latent geometry. Would be interesting to see if the head schedule could be predicted from the geometry alone, rather than tuned.
Hi,
Thank you so much for reaching out and sharing these insights! This is an incredibly elegant connection.
To answer your question about the alignment:
In our current optimal schedules (detailed in Appendix B, Table 4), the progressive ramp for our 24-layer models (Medium and Large) completes relatively early, reaching the baseline head count of 16 heads by layer 7. This corresponds to a relative depth of roughly 0.29. Empirically, we found that allocating the remaining ~70% of the network to a “baseline preservation phase” consistently produced the strongest semantic convergence.
As we note in Section 4.1, the attention distance delta reaches its maximum around the middle of the network, closely coinciding with the point where the progressive ramp stabilizes and the model transitions from predominantly local aggregation to long-range global integration.
One interesting observation is that our optimized schedules settle into their steady-state configuration well before the late-stage transition zone (0.78–0.85 relative depth) identified by LIMEN. While our search space was designed around empirically effective schedules, the correspondence with LIMEN suggests that there may be additional opportunities to explore schedules that explicitly incorporate this later geometric transition. Whether doing so yields further improvements remains an interesting open question.
Hi Shubham,
Thank you for this detailed breakdown. The fact that your optimal schedule stabilizes at layer 7 (~0.29 depth) while LIMEN identifies a critical geometric transition zone much later (0.78–0.85) is a fascinating discrepancy.
It suggests two possibilities:
- Semantic Convergence vs. Dynamic Stability: Your schedule optimizes for semantic loss (which stabilizes early), but might be missing dynamic instabilities that emerge in the deeper layers during complex reasoning or edge-case inputs.
- Untapped Potential: The “baseline preservation phase” (layers 7–24) might be under-utilized. If the model is still undergoing significant geometric reorganization (as LIMEN detects) in the 0.78–0.85 zone, forcing a static head count there might constrain its expressive power.
Proposal for a Quick Test:
I can run a targeted analysis on your Prism Transformer checkpoints (if available) or a similar architecture. I will specifically measure ct_t (trajectory curvature) and entropy in that late-stage zone (layers 18–22) during:
Standard inference.
High-complexity reasoning tasks (e.g., multi-step logic or ambiguous prompts).
Hypothesis: We might see spikes in geometric instability in that late zone where your schedule is static. If so, a second, smaller adaptive ramp in the deep layers could potentially reduce this friction and improve robustness without adding compute.
Would you be open to sharing a set of hidden states from a Prism model run? Or I can replicate the architecture locally and run the LIMEN audit myself.
Best,
Happy to find a team of the same interest.
I also did a similar research. In your proposal your model starts with fewer and ends with wider, but in my research my model is almost the same as normal transformers, but some layers with different heads for some special tasks.
Would love to learn more about it!!
Hi,
Feel free to replicate the architecture locally and run the LIMEN audit yourself. I’d be very interested to see the results.
Looking forward to your findings !!
