Sparse Query Attention (SQA): A Computationally Efficient Attention Mechanism with Query Heads Reduction

1.1 The Computational Burden of Self-Attention

• For long sequences, MHA’s quadratic attention cost becomes dominant and limits the context length models can feasibly process.The attention computation scales quadratically with sequence length, N.Source: §1.1 · pp. 1–2

1.2 A Tale of Two Bottlenecks: Computation vs. Memory Bandwidth

• Training, encoding, and prompt processing are compute-bound, whereas autoregressive decoding is often memory-bandwidth-bound by loading the growing KV cache.The paper argues that attention optimization should distinguish these two bottlenecks.Source: §1.2 · p. 2

1.3 Existing Paradigms: Optimizing for the Memory Bottleneck

• MQA, GQA, and MLA reduce KV-cache data movement, but keep the query-head count unchanged and therefore do not reduce QK^T FLOPs.The paper notes that these methods do not accelerate compute-bound training or full-sequence encoding.Source: §1.3 · p. 3

1.4 The Broader Landscape of Efficiency

• Efficiency research also explores approximating full attention and replacing the Transformer architecture entirely.These are the two broader directions identified beyond memory-centric optimizations.Source: §1.4 · p. 3

1.4.1 Approximating Full Attention: Sliding Window Mechanisms

• Sliding Window Attention reduces complexity from O(N^2) to O(N · k), but cannot directly capture dependencies beyond its local window.Longformer combines local attention with global tokens to address that limitation.Source: §1.4.1 · p. 3

1.4.2 Architectural Alternatives to the Transformer

• SSMs, RetNet, and hierarchical-memory approaches offer alternatives to Transformers, while the paper contrasts their challenges with SQA’s preservation of attention-based parallelism.The cited passage identifies limitations including training stability, locality bias, compression overhead, and hardware scalability.Source: §1.4.2 · pp. 3–4

1.5 Our Contribution: Sparse Query Attention (SQA)

• SQA reduces query heads from H to Hq, lowering attention-layer computation by H/Hq; preliminary small-scale experiments report only minor model-quality impact.The paper also presents mathematical complexity analysis and variants sSQA and xSQA to explore efficiency–capacity trade-offs.Source: §1.5 · p. 4
• SQA accelerates compute-bound training and encoding on long sequences by up to 3x, where MQA and GQA offer no speed advantage.The reported benchmarks cover sequences of 32k–200k tokens.Source: §1.5 · p. 4

2.1 The Foundation: Multi-Head Attention (MHA)

• MHA runs parallel scaled dot-product heads, concatenates their outputs, and has O(N^2·dmodel) attention-score complexity for long sequences.Each head receives projected queries, keys, and values; the concatenated output passes through a final linear projection.Source: §2.1 · p. 5

2.2 Multi-Query Attention (MQA): A Radical Solution for Memory Bandwidth

• MQA shares one Key and Value head across all Query heads, reducing KV-cache size and per-step HBM data by a factor of h, with reported inference speedups up to 12x in certain configurations.This efficiency can reduce model capacity, degrade quality, and sometimes cause training instability.Source: §2.2 · p. 6

2.3 Grouped-Query Attention (GQA): The Balanced Interpolation

• GQA divides h Query heads into g groups that share Key and Value heads, interpolating between MHA at g = h and MQA at g = 1.An intermediate g can substantially reduce KV-cache size and memory bandwidth while bringing model quality closer to MHA than MQA does.Source: §2.3 · pp. 6–7

2.4 Multi-head Latent Attention (MLA)

• MLA compresses Key and Value tensors into a learned low-rank latent representation before caching, yielding a smaller memory footprint than GQA for the same number of query heads.The trade-off is a slight increase in computation during projection.Source: §2.4 · pp. 7–8

2.5 Sliding Window Attention (SWA)

• SWA limits each token’s attention to a fixed local window, reducing complexity to O(N·k·dmodel) but restricting direct access to distant tokens.This locality bias can hurt tasks requiring long-range reasoning; Longformer combines local windows with global-attention tokens.Source: §2.5 · p. 8

3 Sparse Query Attention (SQA)

• SQA targets attention-score computation by reducing query heads, complementing MQA and GQA’s memory-focused optimizations while retaining key/value-head reduction options.Some configurations can match GQA’s memory footprint; others increase K/V heads to improve quality when cache size matters less.Source: §3 · p. 8

3.1 The Core Concept: Reversing the GQA Paradigm

• SQA shrinks Q by setting Hq < H, reducing the attention-weight matrices, value-head aggregation, and total floating-point operations.Unlike MQA and GQA, which repeat or broadcast keys and values to match query heads without reducing FLOPs.Source: §3.1 · p. 9

3.2 Mathematical Formulation

• SQA projects the input into Q, K, and V, repeats K/V to Hq heads, computes attention across those query heads, then concatenates and projects outputs; Hq < H and Hkv ≤ Hq.Source: §3.2 · pp. 9–10

3.2.1 Complexity Analysis

• A 50% reduction in query heads yields a 2x reduction in attention computational cost relative to MHA.SQA’s attention complexity is proportional to Hq · N^2 · dhead, compared with H · N^2 · dhead for MHA.Source: §3.2.1 · p. 10

3.3 Architectural Variants

• Standard SQA sets Hq and Hkv independently; sSQA sets both to H/2 for a 2x theoretical speed-up, while xSQA at H/8 offers a theoretical 8x speed-up.Source: §3.3 · pp. 10–11

3.4 Synergy and Composability with Other Mechanisms

• Hybrid SW-SQA layers apply a local window and then reduce query heads within it, combining SWA’s linear scaling with SQA’s constant-factor FLOP reduction.The combination can also use longer sliding windows with the same efficiency.Source: §3.4 · p. 11

3.5 Comparative Positioning

• SQA reduces computation while GQA/MLA reduce memory bandwidth; it can complement SWA, modifies an attention layer, and differs in abstraction from SSMs, RetNet, and HMT.Source: §3.5 · pp. 11–12

4.1 Experimental Setup

• The experiments used small dense and MoE models, with one epoch on a Wikipedia subset for dense models and five TinyStories epochs for MoE models.Dense models had ∼10–12M parameters; MoE models had ∼8.5M parameters, and evaluation was described as preliminary and constrained by budget.Source: §4.1 · p. 12, §4.2 · p. 13

4.2.1 Dense Models (∼10-12M parameters)

• In dense models, sSQA and SQA had validation loss very close to GQA while training significantly faster; xSQA was fastest and slightly better than MQA.Source: Table 1 · p. 13, §4.2.1 · p. 13

4.2.2 Micro Mixture-of-Experts Models (∼8.5M parameters)

• At a 256-token context, MoE SQA variants trained ∼3–4% faster with minimal validation-loss differences; sSQA was 2% faster with a ∼0.3% loss difference from GQA.Source: Table 2 · p. 13, §4.2.2 · p. 13

4.3 Computational Performance Benchmarks

• At 200k tokens, standard SQA is over 2x faster than GQA (1.41s vs 2.86s), and xSQA is over 3.4x faster (0.82s vs 2.86s).Across the benchmark, all SQA variants are significantly faster than MHA, MQA, and GQA, with the gap widening as sequence length increases.Source: Table 3 · p. 14, §4.3 · p. 14
• For Sliding Window Attention with 128-token windows, sliding-loop overhead dominated attention calculations.Source: §4.3 · p. 14
• The authors hypothesize that SQA’s throughput advantage grows with model and dataset scale while its validation-loss gap remains minimal, and identify validating this scaling hypothesis as future work.Source: §4.3 · p. 14

5.1 The Performance Profile of SQA: When and Why it Excels

• SQA complements sliding-window attention by reducing the cost of attention within each local window, potentially enabling longer windows.Source: §5.1 · p. 15
• sSQA and xSQA demonstrate a 2-3x speed-up in the most expensive model component, reducing training time and financial cost in pre-training and supervised fine-tuning.These processes are described as limited by computational throughput.Source: §5.1 · p. 15
• For a 100k-token context, 2-3x faster prompt processing can noticeably improve the user’s time to first token.Prompt processing is a parallel, non-autoregressive computation over the input sequence.Source: §5.1 · p. 15
• During autoregressive generation, SQA’s computational advantage is less impactful because performance is primarily determined by KV-cache size, and a larger cache could slow generation.An sSQA configuration with Hkv = 16 has a larger cache than comparable GQA with Hkv = 8; xSQA with Hq = 8 and Hkv = 8 can match GQA’s memory footprint and performance in memory-bound scenarios.Source: §5.1 · p. 15

5.2 Trade-offs and Broader Implications for LLMs

• An xSQA configuration with Hq = 8 and Hkv = 8 matches the KV-cache size of GQA with 8 KV heads while offering a theoretical 4x computational speed-up in training and prompt processing.Source: §5.2 · p. 16

Future validation

• The proposed next step is to replace GQA layers with sSQA and xSQA in a pretrained open-source model such as Qwen3-0.6B to evaluate quality and performance trade-offs at scale.Source: §6 · p. 16

Light SQA (lSQA)

• Light SQA would test a more modest query reduction, such as Hq = 0.75 · H, for a possible 25% computational speed-up and a different quality-performance trade-off.Source: §6 · p. 16

Reverse SQA (rSQA)

• Reverse SQA would use fewer query heads than key/value heads, repeating query heads to match; its computational complexity would scale with Hkv and may offer no direct performance benefit.Source: §6 · p. 16

Flex-SQA

• Flex-SQA would combine SQA with local and global sparse patterns; symmetric head configurations could simplify implementation and potentially support efficient sequences of 1M+ tokens.Source: §6 · p. 16

SW-SQA (Sliding-Window SQA)

• Sliding-Window SQA would combine SQA’s FLOPs reduction with sliding-window attention’s linear complexity for tasks where locality is a strong prior.Source: §6 · p. 16

Future directions

• The proposed extensions treat SQA as a building block for designing more efficient and scalable Transformer models, rather than as an endpoint.Source: §6 · p. 16

• Reducing query heads lowers attention-layer computation by a factor of H/Hq; long-sequence benchmarks report up to 3x acceleration for compute-bound training, fine-tuning, and encoding.Preliminary experiments indicate a modest trade-off in model quality.Source: §7 · p. 17
• SQA is particularly suited to architectures that prioritize full-sequence computational throughput over minimizing the autoregressive KV cache.Source: §7 · p. 17
• Variants such as xSQA can potentially speed training and prompt processing while matching the inference memory footprint of state-of-the-art GQA models.Source: §7 · p. 17
• The paper reports that the SQA implementation and variants are publicly available in the RxNN-Attention library.The reported experiments used an internal RxNN library.Source: §7 · p. 17