PaperlayerView exact revision on arXiv

Source-linked AI summary

WavePrune: One period is often enough for RoPE

Guancheng Du, Luotian Huang, Shaowen Wang, Si Li, Kaifeng Lyu

arXiv:2610.06963v1cs.CLcs.SD

TL;DR

RoPE’s periodic rotations can make distinct relative positions hard to distinguish, a concern for long-context modeling. WavePrune restricts each channel to its first rotation period. It improves long-context performance across most tested models, while also supporting length extrapolation and faster attention kernels.

  • Problem

    RoPE’s periodicity can make distinct relative positions acquire similar phases and attention patterns, creating ambiguity that is especially concerning in long-context settings.

    Source: §1 · p. 2

  • Method

    WavePrune restricts each RoPE channel to its first rotation period, removing later rotations without extra training.

    Source: §1 · p. 2

  • Results

    WavePrune improves long-context performance on four of five models, raising Qwen3-8B’s HELMET score from 35.7 to 40.0.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results suggest that one period is often enough for RoPE, and its sparsity pattern enables practical speedups with FlashWavePrune.

    Source: §6 · pp. 11–12

    WavePrune’s training-free performance varies across models.

    Source: §6 · p. 12

Abstract

from arXiv · show

Rotary Position Embedding (RoPE) encodes token positions by rotating each two-dimensional channel of the query and key vectors at a channel-specific frequency, making the attention logits invariant to a common shift of positions. However, this rotation is periodic, and it leads to position aliasing where relative positions separated by a full rotation period become hard to tell apart. To address this, we propose WavePrune, which restricts each channel to its first rotation period. We show that it removes the distractions in attention maps created by position aliasing and improves overall long-context performance. Specifically, WavePrune raises the HELMET score on four of five models we test without any extra tuning (e.g., 35.7 -> 40.0 on Qwen3-8B). When pretraining models from scratch, WavePrune also achieves lower validation loss at extrapolated lengths than pretraining without it. Because WavePrune restricts each channel to a sliding window, it induces a fine-grained sparsity that our hardware-aligned CUDA kernels exploit for 1.15x prefill and 1.24x decoding speedups over FlashAttention-2 at 32K context. Together, these results show that RoPE's periodic structure, widely regarded as essential, is largely redundant beyond the first rotation period.

1 INTRODUCTION

RoPE gives attention translation invariance, but its periodic rotations can make distinct relative positions ambiguous. WavePrune restricts each channel to its first period, removing aliasing-related attention patterns while preserving or improving long-context performance.

1 INTRODUCTION

  • RoPE encodes token positions by applying position-dependent rotations to queries and keys, making attention logits invariant to a common position shift.Each channel’s contribution varies sinusoidally with relative displacement and has a period of 2π/f_c for fixed query and key channels.Source: §1 · pp. 1–2
  • RoPE periodicity can make different relative positions acquire similar rotary phases and attention patterns, an ambiguity called position aliasing.This ambiguity can be especially problematic in long-context settings, where models must distinguish positions over much longer distances.Source: §1 · p. 2
  • WavePrune removes “ghost” sub-diagonals in attention maps, patterns the authors identify as distractions caused by position aliasing.The authors argue that repeated rotations beyond the first period are unnecessary: one period is often enough for RoPE.Source: §1 · p. 2
  • WavePrune restricts each RoPE channel to its first rotation period and improves long-context performance, raising HELMET on four of five models, including Qwen3-8B from 35.7 to 40.0.The method is training-free.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries introduce the Transformer attention computations and notation used in the paper. They then explain how RoPE injects position through channel-wise rotations and a frequency schedule.

2 PRELIMINARIES

  • A Transformer projects each token representation into query, key, and value vectors; multi-head attention runs multiple heads in parallel and concatenates their outputs.For a single head, the learned projection matrices map the input matrix X to Q, K, and V.Source: §2 · p. 3
  • Under causal masking, each query compares against keys using scaled dot products, attending only to itself and preceding tokens.A row-wise softmax turns the logits into attention scores, which weight the value vectors to form the head output.Source: §2 · p. 3
  • RoPE injects token position by rotating each two-dimensional query and key channel before their dot product.The channel-wise rotations are concatenated into a position-dependent map.Source: §2 · p. 3
  • The frequency schedule controls how rapidly query and key channels rotate; high frequencies distinguish positions, while low frequencies track distant tokens.The standard schedule uses a geometric progression governed by a base θ.Source: §2 · p. 3

3 WAVEPRUNE

WavePrune limits each RoPE channel’s attention to its first rotation period. The paper links periodic position aliasing to distracting ghost sub-diagonals and approximates channel-wise pruning in grouped computations for efficient execution.

3.1 WAVEPRUNE: CHANNEL-WISE ATTENTION TRUNCATION

  • WavePrune restricts each channel’s attention window to its wavelength, producing a channel-dependent mask rather than a conventional token-level sliding-window mask.The group-wise approximation is used in the efficient implementation.Source: §3.1 · p. 4, Figure 2 · p. 4
  • On Qwen3-8B at 32K context, HELMET peaks slightly below α = 2 and declines beyond one period, while aggressive pruning at smaller α degrades performance.The default α = 2 corresponds to one period.Source: §3.1 · p. 4, Figure 3 · p. 4

3.2 ANALYSIS: PERIODICITY OF ROPE DISTRACTS ATTENTION

  • In attention maps averaged over 100 PG19 sequences at length 4096, attention leaks from the lag-1 sub-diagonal to larger-lag ghost sub-diagonals, which grow more dilutive as context length increases.Per-channel high-frequency logits oscillate with lag; summed and softmax-normalized, this structure appears as ghost sub-diagonals. WavePrune suppresses the oscillations and removes these distractions by preventing rotation-angle wraparound.Source: §3.2 · pp. 5–6
  • Position aliasing makes logits near-periodic across distinct lags, so the intended attention lag competes with repeated ghost sub-diagonals.Under the theorem’s identical-query-and-key setting, the largest logit is approximately replicated on infinitely many ghost sub-diagonals.Source: §3.2 · pp. 5–6

3.3 EFFICIENT IMPLEMENTATION: FLASHWAVEPRUNE

  • FlashWavePrune groups head dimensions in sets of 16 and assigns each group a shared pruning threshold, enabling standard GEMM computation on Tensor Cores.It skips fully masked groups, uses the standard FlashAttention path for fully retained groups, and masks only the critical groups; the approximation has minimal reported performance effect.Source: §3.3 · p. 6
  • Comparing each group’s threshold with the interval of pairwise distances between query and key blocks identifies fully masked, fully retained, and critical groups with distinct compute paths.Fully masked groups skip QK multiplications and, during decoding, key-cache loads; fully retained groups use the standard path.Source: §3.3 · pp. 6–7, Figure 5 · p. 7

4 EXPERIMENTS

Across training-free integration, pretraining, controlled experiments, and kernel tests, WavePrune generally improves long-context or extrapolation performance while enabling more efficient attention. Results vary by model, and the ablations examine pruning thresholds, channel counts, and kernel optimizations.

4.1 TRAINING-FREE INTEGRATION

  • 35.7 →40.0 on Qwen3-8B; WavePrune improves HELMET scores on four of five models without fine-tuning, while Llama3.1-8B degrades.The table reports results averaged across 64K and 128K context splits.Source: §4.1 · p. 7, Table 1 · p. 8

4.2 PRETRAINING EXPERIMENTS

  • Across GPT-2 base, medium, and large, WavePrune lowers validation loss at length 4096 after training at length 2048.The gains persist when the RoPE base is increased after training, while in-distribution validation loss remains comparable to baseline.Source: §4.2 · p. 8, Figure 6 · p. 9
  • At length 4096, pretraining with WavePrune beats training-free integration: 3.0953 vs. 3.3314 for GPT2-medium and 2.9536 vs. 3.1797 for GPT2-large.Training-free integration also improves over the baseline, but pretraining with WavePrune performs better than that plug-and-play setting.Source: §4.2 · p. 8

4.3 CONTROLLED EXPERIMENTS

  • 82.8% vs. 76.7% of nearby-token attention mass is retained under length extrapolation for WavePrune and baseline, respectively, in the Tetranacci experiment.WavePrune also significantly reduces extrapolated validation loss on this short-range task.Source: §4.3 · p. 9, Figure 7 · p. 10
  • WavePrune substantially reduces extrapolated validation loss on NIAH and makes retrieval heads attend to the magic number more consistently across samples.The task tests recall of a seven-digit number placed in a long passage of irrelevant text.Source: §4.3 · p. 9, Figure 7 · p. 10

4.4 SYSTEM EFFICIENCY

  • WavePrune consistently outperforms FlashAttention-2 in latency; skipping QK multiplications reduces prefill FLOPs, and bypassing key-cache loads reduces decoding memory traffic.Removing any of the tested kernel optimizations—in suffix-only QK GEMM, suffix-only K loading, or split-point specialization—increases latency.Source: §4.4 · pp. 9–10

4.5 ABLATION STUDIES

  • With Qwen3-8B on HELMET at 32K, FlashWavePrune peaks at α = 2; aggressive pruning at small α degrades performance, while larger values add no gains.At moderately large α, the gap between FlashWavePrune and exact pruning is marginal; α →∞ corresponds to no pruning.Source: §4.5 · p. 10
  • On the HELMET ICL subset, performance generally improves as more high-frequency channels are pruned, with jumps at n = 8, 16, and 24; pruning beyond n = 40 adds no benefit.The channels beyond n = 40 have wavelengths exceeding the context length.Source: §4.5 · p. 10

5 RELATED WORK

Related work explains RoPE’s frequency-dependent attention patterns, explores alternative approaches to length extrapolation and efficient attention, and describes concurrent channel-wise truncation. WavePrune is distinguished by its one-period window, its account of ghost sub-diagonals, and its use during pretraining.

Related work on RoPE

  • Low RoPE frequencies retrieve semantic context, while high frequencies produce position-based patterns; slash-dominant heads use high-frequency channels to form strong sub-diagonal lines at specific lags (Jin et al., 2025; Cheng et al., 2026).Prior work also observes periodically repeated peaks and links them to positional ambiguity and inversion (Yang et al., 2026; Du et al., 2026).Source: §5 · p. 10

Length extrapolation

Efficient attention

Concurrent work

  • ATFlash also applies training-free channel-wise truncation, but its default window spans two periods, whereas WavePrune finds quality peaks near one period and also evaluates pretraining (Hayashi et al., 2026).WavePrune identifies periodicity-induced ghost sub-diagonals as its pruning motivation; ATFlash describes the removed terms as noise.Source: §5 · p. 11, Figure 8 · p. 11

6 CONCLUSION AND FUTURE WORK

The conclusion attributes long-context distractions to RoPE periodicity and presents WavePrune as a way to remove them by limiting each channel to its first rotation period. It reports benefits for pretrained models, pretraining extrapolation, and practical kernel speedups, while identifying model variation and shared thresholds as limitations.

6 CONCLUSION AND FUTURE WORK

  • WavePrune restricts each channel to its first rotation period, removing periodicity-related ghost sub-diagonals and improving long-context performance for most evaluated LLMs without training.Its resulting sparsity pattern also enables practical speedups from the FlashWavePrune CUDA kernel.Source: §6 · pp. 11–12

Limitations and future work

  • WavePrune’s training-free performance varies across models, and it uses the same per-channel pruning thresholds for all attention heads despite their potentially different roles.The authors identify mechanistic analysis and head-specific thresholds as directions for future work.Source: §6 · p. 12
  • The fine-grained attention mask is based on RoPE; generalizing it to other positional encodings and model architectures remains future work.The authors suggest this could potentially improve long-context processing and accelerate inference.Source: §6 · p. 12

A A DIRECT PROOF OF THEOREM 1

The paper proves that RoPE can produce infinitely many lag positions with logits near the maximum, then extends the periodicity analysis to symmetric query and key distributions. These repeated logits create competing attention sub-diagonals, with their period bounded in relation to channel sensitivity sparsity.

A A DIRECT PROOF OF THEOREM 1

  • Under identical queries and keys across positions, infinitely many lags have attention logits within any tolerance of the largest logit.The proof finds infinitely many almost-periods by bounding phase shifts and applying the pigeonhole principle.Source: A A DIRECT PROOF OF THEOREM 1 · p. 18

B.1 SLASH-DOMINANT ATTENTION PATTERN

  • A slash-dominant pattern has large attention scores along one fixed offset sub-diagonal, a pattern observed in tasks such as in-context learning.Each token attends strongly to the key at that fixed lag, enabling position-wise information flow.Source: §B.1 · p. 19

B.2.1 ASSUMPTIONS AND DEFINITIONS

  • The analysis assumes queries and keys are i.i.d. from symmetric distributions and defines τ-truncation sparsity to measure the sensitivity measures’ sparse approximation.The symmetry is about the respective means; the sparsity definition uses an ℓ1 approximation threshold τ.Source: §B.2.1 · p. 19

B.2.2 MAIN THEOREM

  • Under those assumptions, attention logits are approximately periodic across lags, and the minimum period has an upper bound that depends on channel-sensitivity sparsity.The paper describes similar logits at separated lags as “periodic” in this result, rather than requiring strict mathematical periodicity.Source: §B.2.2 · pp. 19–20

B.2.3 IMPLICATION ON ATTENTION SCORES

  • This logit periodicity prevents attention from concentrating exclusively on the intended sub-diagonal; sufficiently strict slash-dominance eventually collapses as sequence length increases.The collapse statement applies for sufficiently small ε and η; prior observations used looser thresholds and contexts up to 6000 tokens.Source: §B.2.3 · p. 20

B.3 EMPIRICAL EVIDENCE AND ANALYSIS

The analysis finds channel sparsity and periodic ghost sub-diagonals in attention maps, while WavePrune suppresses those patterns by truncating attention beyond each channel’s wavelength. The appendix proves periodicity for point-mass queries and keys, extends the result to symmetric distributions, and bounds the minimum period using channel sensitivity sparsity.

B.3 EMPIRICAL EVIDENCE AND ANALYSIS

  • In Qwen3-8B layer 0 head 20, only 7 of 64 channels are actively utilized, illustrating the observed channel sparsity.The example is a slash-dominant head with lag Δ=1.Source: §B.3 · p. 20
  • Observed periodicity produces logit spikes at multiple ghost sub-diagonals, which the authors argue can prevent lossless signal passage along a desired lag in non-periodic tasks.The passage identifies these patterns in Figures 1a, 15a, and 16a.Source: §B.3 · p. 20
  • The mathematical τ-truncation sparsity may nevertheless be non-small, because inactive channels can retain sensitivity; the definition omits rotation frequencies for tractability.The authors leave tighter analysis to future work.Source: §B.3 · p. 20
  • WavePrune effectively eliminates ghost sub-diagonals, with channel-wise attention-logit decompositions showing periodic patterns suppressed across channels.The reported visual comparisons are in Figures 15, 16, 19, and 20.Source: §B.3 · p. 21
  • For a single channel, WavePrune sets logits to zero beyond the wavelength, preventing periodic spikes and invalidating the periodicity result in that example.At lag 1, the rotation cancels the initial relative angle and yields the maximum logit; beyond the wavelength, the logit is zero.Source: §B.3 · p. 21

B.4.1 STEP 1: PERIODICITY UNDER POINT MASS DISTRIBUTIONS

  • For point-mass queries and keys, the proof establishes infinitely many periods and bounds the minimum period by a quantity dependent on sensitivity sparsity.The proof uses a grid covering of the phase space and extends the argument through an almost-period construction.Source: §B.4.1 · p. 21, §B.4 · p. 21

B.4.2 STEP 2: EXTENSION TO GENERAL SYMMETRIC DISTRIBUTIONS

  • For symmetric query and key distributions, the proof extends the periodicity result using a probability bound on logit differences and symmetry-preserving transformations.The supporting lemmas establish symmetry for differences and linear transformations, and a probability bound for inner products.Source: §B.4 · p. 21, §B.4.2 · pp. 23–24

C KERNEL IMPLEMENTATION DETAILS

FlashWavePrune’s forward kernel specializes retained suffixes to improve instruction-level parallelism, while the backward kernel adapts the pruning scheme to different matrix axes. Backward overhead over FA2 shrinks with context length, reaching 0.18% at 32K.

C.1 FORWARD KERNEL

  • Compile-time specialization makes the forward kernel’s retained suffix a static GEMM, improving ILP while confining masking overhead to boundary regions.Enumerating split-point pairs as template arguments enables loop unrolling and efficient scheduling of shared-memory loads and tensor-core MMA instructions.Source: §C.1 · p. 25
  • Specializing every split-point pair would bloat code quadratically, so the forward kernel specializes only reachable branches for each model.This requires recompilation for each distinct model, trading one-time compilation cost for inlined hot paths and simpler control flow.Source: §C.1 · p. 25

C.2 BACKWARD KERNEL

  • Backward gradients use different MMA axis arrangements: the head dimension is the N axis for dQ and dK rather than the shared K axis used in the forward kernel.The backward kernel modifies dQ and dK, while dV and dP follow FA2’s standard path; it masks dP before the MMA.Source: §C.2 · pp. 25–26
  • Jointly computing dQ and dK lets the backward kernel mask each dP fragment once for both gradients, halving masking work.The same pruning predicate applies to both gradients; fully retained suffix tiles are consumed unmasked.Source: §C.2 · p. 26
  • Specializing backward GEMMs per split-point pair can exceed compiler budgets, moving state to local memory and adding control overhead that outweighs pruning savings.The kernel body includes score recomputation, dV and dP GEMMs, and joint dQ and dK GEMMs.Source: §C.2 · p. 26
  • 0.18% backward-kernel overhead over FA2 at 32K, down from 8.58% at 8K and 3.81% at 16K.The paper characterizes the overall training-time gap as negligible.Source: §C.2 · p. 27, Table 4 · p. 27

D TRAINING-FREE INTEGRATION

The evaluation uses HELMET for diverse long-context tasks and supplements it with general benchmarks. Protocol adjustments address truncation and scoring constraints, while model settings vary by context length.

D TRAINING-FREE INTEGRATION

  • HELMET covers diverse long-context tasks up to 128K, including recall, RAG, reranking, citation, LongQA, in-context learning, and summarization.The paper reports only six tasks after dropping Summ.Source: D TRAINING-FREE INTEGRATION · p. 27
  • Summ is excluded because F1 poorly captures long summaries, and NarrativeQA uses ROUGE-1 F1 because LLM-as-a-Judge evaluation is costly.For other subtasks not requiring LLM-as-a-Judge, the HELMET protocol is retained.Source: D TRAINING-FREE INTEGRATION · p. 27
  • RULER MV and JSON KV output limits are raised to 200 and 300 tokens, respectively, to address unstable results from early truncation.Source: D TRAINING-FREE INTEGRATION · p. 27
  • General benchmarks cover language modeling, reasoning, code, and math, and are run with default lm-evaluation-harness configurations.The selected tasks are MMLU and CMMLU; BBH and DROP; MBPP and HumanEval; GSM8K and MATH.Source: D TRAINING-FREE INTEGRATION · p. 27
  • Qwen3-8B and Qwen3-32B use YaRN above 32K context and have thinking mode disabled to reduce evaluation cost; native RoPE is used otherwise.Table 5 summarizes the evaluated models’ architectural properties.Source: D TRAINING-FREE INTEGRATION · p. 27, Table 5 · p. 28

E FAILED ATTEMPTS ON LLAMA3.1-8B

In Llama3.1-8B, a high-frequency RoPE channel carries information for a long-range retrieval task despite its short wavelength; truncating that channel makes retrieval fail. The tested remedies do not recover the unpruned baseline.

E FAILED ATTEMPTS ON LLAMA3.1-8B

  • At 32K context, a channel with an approximately 571-token wavelength carries the needle signal to final tokens, and WavePrune truncation causes retrieval failure.The authors report similar patterns in other heads, sometimes driven by different RoPE channels.Source: E FAILED ATTEMPTS ON LLAMA3.1-8B · p. 28
  • None of the tested remedies recovers the unpruned baseline’s performance.Source: E FAILED ATTEMPTS ON LLAMA3.1-8B · p. 28
  • The attempted remedies include increasing α, pruning only selected high-frequency channels or layers, and continuing long-context training with the ProLong recipe.The continued training was tried on both base and instruct models.Source: E FAILED ATTEMPTS ON LLAMA3.1-8B · p. 28

F.1 TETRANACCI DATA CONSTRUCTION

The appendix describes synthetic Tetranacci and needle-in-a-haystack datasets and collects supplementary figures and tables. These materials include reported comparisons of attention patterns and validation loss.

F.1 TETRANACCI DATA CONSTRUCTION

  • Tetranacci samples use a four-term recurrence modulo 1000, with random initial values and nonperiodic sequences; the dataset has 500,000 samples split 99:1.Training and validation lengths are 2,048 and 4,096 tokens; initial values, intervening EOT tokens, and padding are excluded from loss.Source: §F.1 · p. 28

F.2 NIAH DATA CONSTRUCTION

  • NIAH samples place a randomly positioned seven-digit magic number in irrelevant text, then supervise only the requested target number and final EOT token.The dataset has 500,000 samples split 99:1, with lengths of 2,048 and 4,096 tokens.Source: §F.2 · pp. 28–29

G PRETRAINING HYPERPARAMETERS

  • Table 6 reports training hyperparameters for the GPT-2 series pretraining experiments.Source: Table 6 · p. 29

H SUPPLEMENTARY FIGURES AND TABLES

  • Figure 10 contrasts attention-logit decay over distance: standard RoPE fluctuates chaotically, whereas WavePrune decays near-monotonically.Source: Figure 10 · p. 30
  • For Qwen3-8B on PG19, Figure 11 shows nearly coincident NLL curves at 32K without YaRN, but lower WavePrune NLL at 128K with YaRN.Source: Figure 11 · p. 30
  • Figure 12 reports lower WavePrune validation loss than baseline across all tested RoPE base frequencies.Source: Figure 12 · p. 30
  • Tables 13 and 14 report that WavePrune retains more neighboring-token attention at extrapolated length and that its retrieval heads activate consistently across NIAH samples.Table 13 covers GPT2-base Tetranacci attention; Table 14 measures attention to magic-number tokens for selected retrieval heads.Source: Table 13 · p. 33, Table 14 · p. 33
  • With constant queries and keys and intended lag zero, Figure 14 shows WavePrune suppressing RoPE’s ghost sub-diagonals and restoring a clean attention map.Source: Figure 14 · p. 34

I.1 EXPERIMENTAL DETAILS

Across Qwen3-8B and Ministral3-3B visualizations, WavePrune removes ghost sub-diagonals from attention maps averaged over 100 PG19 sequences truncated to 4096 tokens.

I.1 EXPERIMENTAL DETAILS

  • Attention maps average scores extracted from 100 PG19 sequences, each truncated to 4096 tokens, for Qwen3-8B and Ministral3-3B.The visualizations include four layer-wide comparisons and two head-level per-channel logit comparisons.Source: §I.1 · p. 34, §I.2 · p. 34

I.2 LAYER-WIDE AND HEAD-LEVEL COMPARISONS

  • Centered per-channel logit contributions are compared before and after WavePrune for Qwen3-8B layer 0 heads 20 and 6.The full-map panels show layer-wide scores, while cropped panels show selected heads in greater detail.Source: §I.2 · p. 34, Figure 19 · p. 39, Figure 20 · p. 39
  • Across four layer-wide comparisons, WavePrune eliminates ghost sub-diagonals and concentrates attention on intended lags.The examples cover Qwen3-8B layers 0 and 32 and Ministral3-3B layers 0 and 1.Source: Figure 15 · p. 35, Figure 16 · p. 36, Figure 17 · p. 37
Loading 2610.06963v1…