Technical Summary: Energy-Gated Attention and Wavelet Positional Encoding

1. Problem Statement

Standard Transformer attention mechanisms compute pairwise token similarity via dot-product attention but suffer from two structural limitations regarding inductive biases:

1. Lack of Spectral Salience: Standard attention treats all tokens as equally salient, weighting them solely by content similarity to the query. It lacks a mechanism to distinguish tokens with high intrinsic informational density (e.g., morphological boundaries, syntactic heads) from background tokens (e.g., function words, repeated patterns).
2. Lack of Adaptive Locality: Standard positional encodings (e.g., fixed sinusoidal PE) assign fixed frequencies to embedding dimensions without spatial envelopes. Consequently, every position contributes equally at every scale, failing to adapt the extent of positional influence to the natural scale of linguistic phenomena or context length. While Rotary Positional Embeddings (ROPE) encode relative position, they still utilize fixed frequencies without Gaussian locality.

The paper posits that these two gaps—what to attend to (salience) and where attention operates at each scale (locality)—are complementary and orthogonal.

2. Methodology

The authors propose two simple, learnable components to address these gaps: Energy-Gated Attention (EGA) and Morlet Positional Encoding (MOPE).

2.1 Energy-Gated Attention (EGA)

EGA introduces a salience mechanism by gating value aggregation based on a learned energy estimate of key token embeddings.

• Mechanism: A single linear projection ($w_{proj}^T x_j$) estimates the "energy" of a token. This scalar is z-normalized and passed through a sigmoid gate ($g_j$).
• Function: The gate suppresses low-energy tokens (background/filler) and amplifies high-energy tokens (information-dense).
• Efficiency: It adds minimal parameters ($d+2$ per head, <0.3% overhead) and no computational cost.
• Theoretical Basis: The projection is interpreted as estimating spectrally-weighted energy via Parseval's identity, though the authors note the learned signal may represent general informational salience (e.g., rarity, boundaries) rather than strict spectral energy.

2.2 Morlet Positional Encoding (MOPE)

MOPE replaces fixed sinusoidal positional encodings with learned Gaussian-windowed wavelets to provide adaptive time-frequency localization.

• Formulation: For embedding dimension $i$ and position $b$:
  $$MOPE(b, 2i) = \cos(\omega_i b) \cdot e^{-b^2/2\sigma_i^2}$$
  $$MOPE(b, 2i + 1) = \sin(\omega_i b) \cdot e^{-b^2/2\sigma_i^2}$$
  where $\omega_i$ (center frequency) and $\sigma_i$ (locality bandwidth) are learned parameters.
• Generalization: MOPE strictly generalizes standard sinusoidal PE (recovered as $\sigma_i \to \infty$) and connects to ROPE (phase structure in the limit) and ALiBi (zero-frequency locality).
• Cross-Correlation Structure: Substituting MOPE into the attention cross-correlation yields a Morlet kernel in lag space, $C_i(\tau) \propto e^{-\tau^2/4\sigma_i^2} \cdot \cos(\omega_i \tau)$. This allows dimensions to capture different scales of linguistic structure (e.g., character-level vs. sentence-level) based on their learned $\sigma_i$.

2.3 Combined Model

The EGA-MORLET model applies EGA gating to attention weights computed using MOPE. The authors hypothesize superadditivity: since EGA controls which tokens are attended to and MOPE controls where influence extends at each scale, their combination should yield performance greater than the sum of their individual parts.

3. Key Contributions

1. Cross-Correlation Interpretation: The paper reframes dot-product attention as the zero-lag value ($C_{ij}(0)$) of a richer cross-correlation function, identifying that standard attention discards scale selectivity, lag structure, and spectral salience.
2. MOPE Architecture: A localized wavelet positional encoding that strictly generalizes sin/cos PE and provides a theoretical bridge to ROPE and ALiBi, offering adaptive frequency and locality.
3. Empirical Complementarity: Demonstration that EGA and MOPE are complementary inductive biases. Their combination achieves superadditive results, whereas individual components show mixed or neutral performance in isolation.
4. Structured vs. Learned Priors: A structured ablation showing that unconstrained learned projections consistently outperform structured spectral priors (e.g., Daubechies filters, fixed scale initialization) in language models, except when a structured component (MOPE) provides genuinely complementary information (adaptive locality) not discoverable by standard gradient descent on dot-products.

4. Experimental Results

Experiments were conducted on TinyShakespeare (character-level, $L=6, H=8, d=256$, $\le 6M$ parameters) over 5,000 steps.

• Main Finding (Superadditivity):

  • Baseline (Standard Attention): Validation Loss = 1.4742.
  • EGA Alone: +0.092 improvement (Loss = 1.3821). Salience is a strong inductive bias.
  • MOPE Alone: -0.032 degradation (Loss = 1.5060). Adaptive locality without salience is insufficient and slightly harmful.
  • EGA + MOPE: +0.119 improvement (Loss = 1.3550).
  • Analysis: The combined gain (+0.119) exceeds the sum of individual gains (+0.092 - 0.032 = +0.060) by +0.059. This supports the complementarity hypothesis.

• Ablation Studies:

  • Convolution Attention: Extending attention to non-zero lags (+0.007) confirms lag structure carries information.
  • Scale Initialization: Initializing heads at specific frequencies provided no benefit (-0.007), suggesting gradient descent discovers optimal scales from random initialization.
  • Spectral Flux: Using energy change rates as a gating signal provided a small positive gain (+0.012), suggesting boundary detection is useful.
  • Phase Coherence: Using phase information resulted in a slight negative impact (-0.007).

• Learned Parameters: MOPE parameters clustered into four distinct regions corresponding to character, word, clause, and sentence scales, suggesting the model learns linguistically natural temporal scales from data.

5. Significance and Claims

The paper claims that spectral salience and scale-selective locality are complementary inductive biases that standard attention lacks.

• Complementarity: Neither bias is sufficient alone at this scale; their interaction is superadditive.
• Learning vs. Priors: The consistent failure of structured priors (like fixed wavelet gates) suggests that the optimal basis for character-level embeddings is non-sinusoidal and corpus-specific. Unconstrained learning is preferred unless a structured component (like MOPE's adaptive locality) offers information orthogonal to what gradient descent naturally discovers.
• Modesty: The authors explicitly state that all experiments are small-scale (single seed, character-level). The primary result is considered credible at this scale, but secondary effects require multi-seed validation. The paper does not claim these methods solve long-context problems yet but identifies them as a high-priority direction for future work (specifically for $T \ge 4096$).

The work is presented as part of a series on spectral methods in transformer attention, with code made available for reproducibility.