Prism Transformer: Progressive Head Schedules for Hierarchical Attention Processing

Abstract: Multi-head attention conventionally partitions the hidden dimension equally across all heads at every layer, enforcing an identical representational subspace dimension (dh = dmodel/h) throughout the models depth. In this work, we identify this uniform allocation as a fundamental structural bottleneck: due to their restricted dimensional space, early-layer heads are unable to faithfully capture complex, high-dimensional contextual patterns. To resolve this, we introduce the Prism Transformer, a novel architectural paradigm that replaces the static, uniform head configuration with a progressive head schedule. By monotonically increasing the head count across layers, the Prism Transformer naturally establishes a local-to-global representational hierarchy: early layers leverage fewer, exceptionally wide heads to capture complex, local compositional patterns, while deep layers deploy many, narrow heads to decompose these patterns into specialized linguistic features. Crucially, this structural shift is parameter-neutral, compute-neutral, and introduces zero training or inference overhead, preserving identical weight matrices and FLOP budgets as the standard Transformer. Across three model scales (124M, 354M, and 757M), the Prism Transformer consistently outperforms uniform baselines, achieving consistent reductions in validation loss alongside consistent gains on downstream zero-shot benchmarks (including PIQA, HellaSwag, ARC-Easy, and WinoGrande). Our findings demonstrate that non-uniform subspace allocation unlocks latent capacity within the standard Transformer budget, enabling more effective use of model capacity.

• The paper introduces a progressive attention head scheduling mechanism that adjusts head counts per layer to match specific representational needs.
• Prism consistently improves validation loss and zero-shot benchmark performance over uniform-head Transformers across multiple scales.
• The approach remains parameter- and FLOP-neutral while accelerating convergence, demonstrating hardware-aligned efficiency improvements.

Prism Transformer: Progressive Head Schedules for Hierarchical Attention Processing

Introduction

The Prism Transformer introduces a progressive attention head scheduling mechanism for decoder-only Transformers, challenging the architectural convention of fixed uniform head count across layers. The key hypothesis is that the representational and contextual demands at different network depths are not uniform; therefore, the standard MHA mechanism—which divides $d_{\text{model}}$ evenly among a fixed head count $h$ in all layers—imposes a bottleneck. Prism progressively increases the number of attention heads with depth, starting with a small number of wide heads at shallow layers and ramping to a higher count of narrow heads in deeper layers. This approach induces a local-to-global hierarchy, enabling early layers to capture rich local patterns and deeper layers to refine and integrate information globally, all while remaining strictly parameter- and compute-neutral relative to baseline architectures.

Figure 1: Head allocation comparison: uniform baseline (left) vs. Prism progressive ramp (right), resulting in a monotonically decreasing per-head dimension $d_h$ across depth.

Architectural Overview

Progressive Head Schedules

Let $d_{\text{model}}$ denote the model hidden size and $L$ the number of layers. The Prism schedule defines a non-decreasing sequence $(h_1, h_2, ..., h_L)$ of head counts with three constraints:

• Divisibility: $h_l$ divides $d_{\text{model}}$ in each layer.
• Monotonicity: $h_1 \leq h_2 \leq \ldots \leq h_L$.
• Boundary Matching: $h_L$ equals the standard baseline head count, ensuring compatible late-layer behavior.

The per-head subspace width $h$0 decreases with depth, forming a narrowing “prism” through the stack.

Parameter and Compute Neutrality

Despite modifying the intra-layer head allocation, the total parameter count is unchanged since all projection matrices ($h$1, $h$2, $h$3, $h$4) remain size $h$5. Computational cost in FLOPs per token is strictly invariant, as all operations exploit the entire hidden dimension, with only the split/merge across heads changing per layer. Hardware-aligned schedules—where early layers have head dimensions divisible by two—are prioritized for GPU efficiency.

Empirical Results

Pretraining Performance

Across three model scales (124M, 354M, 757M parameters), the Prism Transformer consistently achieves lower validation losses than uniform head baseline architectures, given matched datasets, training budgets, and hyperparameters. Performance improvements appear early during training and persist throughout the optimization trajectory.

Figure 2: (a) Dynamics of validation loss gap for large model: Prism's advantage widens and stabilizes over training; (b) Mean loss delta at convergence across all tested model scales, showing consistent gains for Prism.

Zero-Shot Benchmark Evaluation

Prism consistently matches or exceeds baseline accuracy on a suite of zero-shot downstream benchmarks: PIQA, HellaSwag, ARC-Easy, WinoGrande, BLiMP, and WikiText. Gains are most pronounced in structurally complex tasks requiring deep semantic integration and robust context modeling.

Figure 3: Zero-shot task-level benchmark accuracy comparisons across model sizes, showing Prism at parity or better across all tasks and sizes.

Scaling Behavior

The architectural advantage of progressive scheduling does not diminish with scale. For all evaluated parameter regimes, Prism yields a persistent validation loss gap. The best case ablation at small scale (Config-7) further shows that validation loss improvements can be increased—at the cost of hardware throughput—when hardware-alignment constraints are relaxed.

Attention Profile Analysis

A distinctive mechanistic effect is observed in per-layer attention distance analyses. Early wide-head layers in Prism attend more locally than baseline, resulting in lower average token-to-token distances, while deeper layers—post head ramp—exhibit more global attention than baseline. Importantly, both models converge to similar attention distance in their terminal layers.

Figure 2: Per-layer mean attention distance gap: Prism produces sharper local attention in early (wide head) layers, then transitions to more global focus mid-stack, with a sign change at the schedule boundary. The effect is consistent across model scales.

Comparative Perspective

The Prism approach addresses two deficiencies of uniform MHA scheduling. First, it directly mitigates functional head redundancy and low-rank bottlenecks in standard Transformers—identified in previous pruning and rank studies (Michel et al., 2019, Bhojanapalli et al., 2020, Voita et al., 2019)—by maximizing early-layer head expressivity and smoothly staging the transition to higher parallelizeability in deeper layers. Second, it distinguishes itself from parameter-variant approaches such as LayerDrop or variable-width transformers (Fan et al., 2019, Wu et al., 16 Jun 2026): all layers remain fully active under rigid parameter/FLOP constraints.

Prism is orthogonal and composable with inference-oriented head repartitioning methods such as Multi-Query Attention (MQA) and Grouped-Query Attention (GQA), which optimize KV caches rather than representational expressivity.

Practical Implications and Future Work

The Prism Transformer demonstrates a significant underexplored architectural axis: moving beyond uniform attention head allocations to depth-aware, hardware-aligned head scheduling. It offers consistent, scale-invariant representational efficiency enhancements at no sacrifice to training FLOPs, wall-clock time, or parameter envelope. These findings have direct implications for budget-constrained and efficiency-focused LLM pretraining regimes, especially where increased model quality must be achieved within rigid computational budgets.

Notably, Prism's benefits extend to convergence speed—minimal validation loss milestones are consistently achieved earlier. This could translate into meaningful reductions in overall pretraining resource requirements and operational costs for large-scale foundation models.

Future directions include:

• Extending progressive head schedules to encoder-only and encoder-decoder architectures (especially those with bidirectional attention).
• Exploring dynamic, learned, or continuously parameterized schedules rather than fixed, hand-tuned ramps.
• Evaluating very long-context and streaming inference regimes, particularly for applications needing robust local/global token mixing.

Conclusion

The Prism Transformer reveals that architectural inductive biases, when aligned with functional depth specialization, provide systematic representational improvements without increasing model size or compute. Progressive head schedules capture the natural local-to-global processing trajectory across Transformer depth, yielding robust efficiency and accuracy gains. Prism's parameter- and FLOP-invariance, hardware-aligned execution, and empirical improvements make it a highly relevant baseline for future LLM research and deployment strategies.