PaperlayerView exact revision on arXiv

Source-linked AI summary

Flow-based Spectral Kernel Learning for Nonstationary Attention

Zi Yang, Ying Li, Alejandro Lancho, Darío C. Larese, Pablo M. Olmos, Michael Minyi Zhang

arXiv:2610.30465v1stat.MLcs.LG

TL;DR

Fixed or stationary kernels limit kernelized attention’s ability to adapt to complex token interactions. FSKA learns nonstationary kernels by modeling spectral density with normalizing flows and regularizing it with a structured prior. It reports competitive or improved classification performance, approximately 40%–47% fewer projection-related parameters than most baselines, and favorable linear-attention memory scaling.

  • Problem

    Existing kernelized attention often relies on fixed or stationary kernels, limiting its ability to model input-dependent, nonstationary token interactions.

    Source: §1 · p. 1, §3.1 · p. 3

  • Method

    FSKA parameterizes a bivariate spectral density with normalizing flows and a structured spectral prior, enabling flexible kernel learning compatible with RFF-based linear attention.

    Source: §1 · p. 2, §3.2 · p. 5

  • Results

    FSKA achieves competitive or improved performance across regression and classification settings, reduces projection-related trainable parameters by approximately 40%–47% compared with most baselines, and maintains a low-growth memory profile close to linear-attention baselines.

    Source: §1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    Expressive kernel learning supports removing query–key projections and using fixed orthogonal value projections to reduce projection-related parameters.

    Source: §1 · p. 2

    With finite random features, near-zero normalization denominators may still occur, especially when the feature count is limited for computational efficiency.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

Kernelized attention provides a useful perspective for understanding self-attention as a similarity-based aggregation mechanism, but existing approaches often rely on fixed or stationary kernels, limiting their ability to adapt the similarity structure to complex token interactions. In this paper, we propose flow-based spectral kernel attention (FSKA), a framework for learning expressive nonstationary attention kernels. FSKA formulates kernelized attention through spectral density learning and parameterizes the resulting bivariate spectral density with normalizing flows. This design enables flexible and regularizable kernel learning while remaining compatible with scalable linear-attention computation. We further show that expressive kernel learning enables a simplified attention design with removed query--key projections and fixed orthogonal value projections, resulting in a more parameter-efficient architecture. Controlled experiments on classification, ablation, and scalability profiling show that FSKA achieves competitive or improved performance over existing attention mechanisms, while reducing projection-related trainable parameters by approximately $40\%$--$47\%$ compared with most baselines and maintaining favorable linear-attention memory scaling. Code Available: https://github.com/ziyang18/FSKA.

1 Introduction

Kernelized attention can be efficient, but fixed or stationary kernels limit adaptation to complex, context-dependent token interactions. FSKA learns nonstationary kernels through flow-based spectral modeling and uses this expressiveness to simplify projections.

1 Introduction

  • Fixed-kernel methods such as RKA and Performer enable linear-time attention but often cannot adapt to complex, context-dependent token interactions.Prior learnable-kernel approaches add flexibility, but Gaussian mixtures can be costly and implicit neural parameterizations lack an explicit density for regularization.Source: §1 · pp. 1–2
  • FSKA learns nonstationary attention kernels by parameterizing spectral density with normalizing flows and regularizing it with a structured spectral prior.The resulting kernel remains compatible with RFF-based linear attention.Source: §1 · p. 2
  • FSKA removes query–key projections and uses fixed orthogonal value projections, reducing projection-related trainable parameters by approximately 40%–47% compared with most baselines.The simplified design is enabled by expressive kernel learning.Source: §1 · p. 2
  • Across regression and classification tasks, ablations, and scalability studies, FSKA achieves competitive or improved performance while maintaining a low-growth memory profile close to linear-attention baselines.The reported parameter reduction is approximately 40%–47% compared with most baselines.Source: §1 · p. 2

2 Background

Self-attention aggregates token information using learned pairwise similarities, and its kernel interpretation connects that operation to kernel smoothing. Random Fourier features approximate shift-invariant kernels and enable linear-time attention.

2.1 Self-Attention

  • Self-attention projects tokens into query, key, and value representations, computes similarity-based attention weights, and aggregates the values across one or more heads.Multi-head attention concatenates outputs from heads with independent projections.Source: §2.1 · p. 2

2.2 Kernel View of Attention

  • The kernel view expresses attention as a kernel smoother over token representations; the softmax choice is one kernel, while naive computation requires O(L^2) time and memory.The kernel interpretation is attributed to Tsai et al. [2019].Source: §2.2 · pp. 2–3

2.3 Random Fourier Features for Kernel Attention

  • Random Fourier features approximate shift-invariant kernels using sampled frequencies from their spectral density, yielding feature inner products that support linear-time attention.For RKA and Performer, predetermined frequency distributions induce stationary kernels that cannot adapt to data-dependent patterns.Source: §2.3 · p. 3

3 Method

The method learns nonstationary attention kernels through bivariate spectral densities and random-feature approximations, then uses kernel expressiveness to simplify projections while retaining multi-head diversity.

3.1 Nonstationary Attention Kernel

  • A bivariate spectral density represents a continuous nonstationary kernel, and its random-feature approximation permits attention computation linear in sequence length.The method applies a positive feature transformation to prevent negative normalization values; near-zero denominators may still occur with finite, especially limited, feature counts.Source: §3.1 · p. 4

3.2 Flow-based Spectral Density Modeling

  • Normalizing flows parameterize the bivariate spectral density through invertible transformations, while a structured spectral prior regularizes learning without removing flow flexibility.The paper uses RealNVP transformations and Monte Carlo estimates for the regularizer.Source: §3.2 · p. 5

3.3 Simplifying Attention via Kernel Expressiveness

  • The method removes query–key projections because a sufficiently expressive kernel can model token similarities directly on the original input representations.The result is conditional on the kernel family being uniformly dense in C(X × X); the theorem does not establish that a finite implementation meets this assumption.Source: §3.3 · p. 5, §3.3.1 · p. 6

3.3.1 Removing Query and Key Projections

  • Under the theorem’s uniform-density assumption, projected similarities can be uniformly approximated by kernels on the original inputs, making explicit query–key projections unnecessary for expressive kernel attention.Removing them reduces projection-related parameters by shifting interaction modeling to the kernel.Source: §3.3.1 · p. 6

3.3.2 Value Representation with Fixed Projections

  • The value branch remains to define the aggregation feature space, but its projection can be fixed because the learned kernel determines adaptive weights over values.Directly using inputs as values would give every head the same representation space and may limit diversity or produce redundant outputs.Source: §3.3.2 · pp. 6–7

3.4 Multi-Head Kernel Attention

  • Each head uses a head-specific learned kernel and a fixed orthogonal value projection, preserving distinct input-dependent similarities and complementary value subspaces.Orthogonality encourages heads to aggregate in distinct subspaces and reduces redundancy among outputs.Source: §3.4 · p. 7

4 Experiments

The experiments assess FSKA through architectural ablations, sentence- and full-document classification, and parameter-efficiency and memory-scaling analyses. Across these settings, results indicate competitive or improved performance alongside reduced projection parameters and favorable memory growth.

4.1 Ablation Study

  • The MNIST ablation supports that flow-based spectral density learning improves performance and permits simpler projections while preserving expressive capacity.Removing query–key projections maintains strong performance with the learned kernel, while making the value projection learnable brings little additional benefit over the fixed-value variant.Source: §4.1 · p. 8, Table 1 · p. 8
  • With a limited number of random features, positive random features stabilize linearized attention by preventing signed-feature cancellation in the denominator.Direct RFF features can severely degrade performance under this condition.Source: §4.1 · p. 8

4.2 Sentence Classification

  • FSKA achieves the strongest or highly competitive performance across six sentence-classification datasets, improving Transformer F1-macro from 76.07 to 78.74 on SST-2 and from 35.11 to 38.16 on SST-5.Table 2 reports F1-macro over five runs.Source: Table 2 · p. 9, §4.2 · p. 9

4.4 Parameter Efficiency and Memory Scalability

  • FSKA uses roughly 40%–47% fewer trainable parameters than most baseline attention modules.The reduction mainly comes from eliminating learnable query and key projections and replacing the learnable value projection with a fixed mapping.Source: §4.4 · pp. 9–10, Figure 2 · p. 10

4.3 Document Classification with Full Sequences

  • FSKA achieves the highest mean macro-F1 on all four full-document datasets and the highest mean accuracy on BBC News, 20NG, and arXiv.On SCOTUS, Mamba has the highest accuracy at 59.84%, while FSKA’s macro-F1 is 39.68% versus 37.91%.Source: §4.3 · p. 9, Table 3 · p. 10

4.4 Parameter Efficiency and Memory Scalability

  • At sequence length 4096, FSKA keeps training memory below 10K MB, while quadratic-style attention mechanisms use around 80K–120K MB.FSKA’s memory growth is close to linear-attention baselines; its RFF-based design avoids explicitly forming the full quadratic attention matrix.Source: §4.4 · p. 10, Figure 2 · p. 10

5 Conclusion

The paper concludes that FSKA learns expressive nonstationary attention kernels through flow-parameterized spectral densities while remaining compatible with RFF-based linear attention. Experiments report competitive or improved classification performance, fewer parameters, and favorable memory scaling.

5 Conclusion

  • FSKA parameterizes spectral density with normalizing flows and regularizes it with a structured spectral prior for flexible kernel learning compatible with RFF-based linear attention.The framework targets expressive nonstationary attention kernels.Source: §5 · p. 10
  • Expressive kernel learning permits removing query–key projections and using fixed orthogonal value projections, reducing parameter sizes by approximately 40%–47% compared with most baselines.The paper reports competitive or improved performance on sentence- and full-document classification tasks.Source: §5 · p. 10
  • On long-sequence benchmarks, FSKA maintains favorable linear-attention memory scaling.The conclusion characterizes FSKA overall as an expressive and parameter-efficient attention mechanism.Source: §5 · p. 10

A.1 Proof of Theorem 3.2

The supplied passages cover a Monte Carlo spectral representation and its unbiased estimator, a projected-similarity approximation result, and experimental hardware details.

A.1 Proof of Theorem 3.2

  • The proof rewrites a spectral representation as a randomized feature map and establishes that its Monte Carlo estimator is unbiased.The spectral pairs are sampled independently from the bivariate density.Source: §A.1 · p. 13

A.2 Proof of Theorem 3.3

  • The proof establishes that a kernel defined directly on the original input space can uniformly approximate projected query–key similarity.This result follows under the stated uniform-density assumption on the kernel class.Source: §A.2 · pp. 13–14

B Experimental Details and Additional Experiments

B.1 Additional Analysis on Positive Random Features

With limited random features, signed RFF components can cancel and make attention denominators unstable; PRF avoids negative or relatively small denominators in the tested settings. The positive transformation also changes the finite-feature similarity, so its affinity should not be treated as an unbiased estimate of the original spectral kernel.

B.1 Additional Analysis on Positive Random Features

  • Signed sine and cosine components can cancel at finite feature counts, producing very small or unstable attention-normalization denominators.The analysis examines the denominator distribution under different random-feature dimensions.Source: §B.1 · p. 14
  • With small feature counts, direct RFF denominators have non-negligible mass near zero; increasing the count concentrates the distribution away from zero.The reported trend suggests more features reduce cancellation probability.Source: §B.1 · pp. 14–15
  • PRF keeps the fraction of negative or relatively small denominators at zero across tested feature counts, while direct RFF's fraction decreases overall but is not eliminated.The analysis motivates PRF as a stabilization mechanism when memory and runtime constraints prevent making the feature count arbitrarily large.Source: §B.1 · p. 15, Figure 3 · p. 15

B.1.1 Kernel Induced by the Positive Transformation

  • For each fixed shared spectral draw, the positive-transformation finite-feature kernel is positive semidefinite, as is its population counterpart.For a general RealNVP spectral distribution, the population expectation has no simpler closed-form expression.Source: §B.1.1 · pp. 15–16
  • Unbiasedness applies only to the signed trigonometric feature map before ELU+1, not to the transformed affinity or normalized attention weights.ELU+1 changes the similarity function rather than merely adding a small constant to the denominator.Source: §B.1.1 · p. 16

B.1.2 Similarity Gap After ELU+1

The appendix compares positive-transformed and signed feature similarities in a synthetic diagnostic, then evaluates FSKA on a toy nonstationary regression task. FSKA leads the regression comparison, while the similarity diagnostic retains a substantial absolute discrepancy.

B.1.2 Similarity Gap After ELU+1

  • The signed matrix is a finite-feature reference, so the discrepancy measures disagreement between normalized feature similarities, not error against the exact analytical kernel.The diagnostic uses synthetic segmented tokens with varying mean directions and within-segment drift; it is separate from trained text-classification experiments.Source: §B.1.2 · p. 16
  • In the synthetic feature-map diagnostic, correlation rises from 0.921 to 0.994, but mean absolute discrepancy remains approximately 0.8.The analysis compares cosine-normalized signed and ELU+1-transformed Gram matrices using off-diagonal Pearson correlation and mean absolute discrepancy.Source: §B.1.2 · p. 16

B.2 Toy Example for Nonstationary Regression

  • The toy benchmark tests whether kernels capture global trend consistency and local regime transitions under distribution shift; nine models are compared by test MSE on a fixed split.Figure 4 reports per-model prediction plots and MSE.Source: §B.2 · p. 17
  • FSKA achieves the best toy-regression test MSE (0.0334), ahead of IKAN-direct (0.0338) and substantially below Transformer (0.1219), MIKAN (0.1938), and most text-kernel baselines (≈0.18−0.23).Its predictions track regime boundaries and high-curvature regions more faithfully while preserving global phase alignment.Source: §B.2 · p. 17

B.3 Hyperparameter Sensitivity

The MNIST sensitivity study varies NGSM mixture count, random-feature dimension, and normalizing-flow depth. Each parameter presents a different performance–runtime trade-off.

B.3 Hyperparameter Sensitivity

  • Increasing NGSM mixture count Q from 1 to 10 improves performance (0.8483 → 0.8743), while Q=20 degrades it to 0.8431; runtime stays at 75–77s.The passage attributes the late degradation to possible mixture-component overfitting and says mixture computation is efficiently parallelized.Source: §B.3 · p. 17
  • Increasing RFF dimension M from 32 to 96 improves accuracy (0.8656 → 0.8743) with mild runtime growth; M=128 reaches 0.8822 but nearly doubles runtime to 149.20s.The default M=96 balances accuracy (0.8743) and efficiency (85.32s).Source: §B.3 · pp. 17–18
  • Increasing flow depth Nf from 2 to 10 raises accuracy (0.8503 → 0.8707), but runtime grows by approximately 13s per added layer.Beyond Nf=4, marginal accuracy gains diminish; the default Nf=5 is described as near-peak without excessive overhead.Source: §B.3 · pp. 17–18

B.4.1 Additional Results for Sentence Classification

FSKA performs competitively across six sentence-classification benchmarks, with the best accuracy on four and second-best on AGNews. The appendix also specifies the benchmark baselines and their attention implementations.

B.4.1 Additional Results for Sentence Classification

  • FSKA ranks best on DBPedia, TREC, RT-Polarity, and SST-2, and second on AGNews; it improves over Transformer accuracy from 70.93 to 73.18 on RT-Polarity and 75.99 to 77.58 on SST-2.Table 6 reports accuracy on six sentence-classification benchmarks.Source: §B.4.1 · p. 18, Table 6 · p. 19

B.4.2 Baselines and Their Parameter Settings

  • The Transformer baseline uses learned Q, K, V, and output projections, while MIKAN replaces softmax attention with implicit kernel attention and uses 64 random Fourier features per head.The Transformer configuration has hidden dimension 128, four heads, and two encoder layers.Source: §B.4.2 · p. 19
  • Performer uses positive orthogonal random features, RKA generates Fourier features with a small network, and GMM-RKS learns per-head Gaussian-mixture spectral parameters.Their implementations use 110, 64, and 64 features or dimensions, respectively.Source: §B.4.2 · pp. 19–20
  • FlashAttention is included as an efficient exact-attention baseline, while Mamba supplies a linearly scaling, non-attention comparison for full-document classification.FlashAttention uses tiling to reduce data transfers between GPU memory levels.Source: §B.4.2 · p. 20

B.4.3 Datasets, Sources, and Access Links

The appendix documents sentence and document classification data protocols, an LRA-style synthetic scaling benchmark, and a fixed-spectrum ablation. In the ablation, learning the bivariate spectrum improves accuracy and macro-F1 on BBC News and 20NG.

B.4.3 Datasets, Sources, and Access Links

  • The six sentence-classification benchmarks span sentiment, question, news-topic, ontology, and review-rating tasks; preprocessing lowercases and whitespace-tokenizes text, builds vocabulary from training data, and caps sequences at 256 tokens.The vocabulary allows up to 50,000 entries and uses <PAD> and <UNK> special tokens.Source: §B.4.3 · pp. 20–21

B.5 Additional Details of LRA Experiments

  • The LRA-style scalability benchmark varies synthetic sequence lengths from 256 to 4096 to measure memory and throughput, isolating computational scaling rather than downstream task accuracy.It is a systems benchmark, not an evaluation on official LRA datasets.Source: §B.5 · p. 21

B.6 Document Classification Protocol

  • Document classification covers SCOTUS, BBC News, 20NG, and arXiv, with full tokenized documents supplied without truncation, chunking, or sliding-window segmentation.The 256-token limit applies only to sentence-classification benchmarks.Source: §B.6 · p. 21
  • The document protocol compares nine models and reports accuracy and macro-F1 as mean and standard deviation across three runs; FSKA removes learned query–key projections and uses fixed orthogonal value projections.Training data are capped at 4,000 documents and test data at 2,000 where available, with 10% of the training pool reserved for validation.Source: §B.6 · pp. 21–22

B.7.1 Learning the Bivariate Spectral Density

  • Learning the flow-based spectrum improves BBC News accuracy by 26.56 points and macro-F1 by 26.81 points, and 20NG accuracy by 8.44 points and macro-F1 by 8.46 points.The comparison replaces the learned spectrum with a fixed bivariate Gaussian while retaining the bivariate feature construction, PRF transformation, projection choices, and other settings.Source: Table 8 · p. 22, §B.7.1 · p. 22

B.7.2 Finite-Capacity Query–Key Ablation

On 20NG, removing query–key projections caused accuracy gaps of at most 0.90 points across tested settings, while the default spectral regularization weight produced the highest reported scores among tested values.

B.7.2 Finite-Capacity Query–Key Ablation

  • The finite-capacity comparison varied random-feature capacity and flow depth on 20NG, using Nf = 3 as the default flow depth for document-classification experiments.The sweeps used M ∈ {32, 64, 96, 128} with Nf = 2, and Nf ∈ {2, 3, 5, 7, 9} with M = 32.Source: §B.7.2 · p. 22, Table 9 · p. 23
  • At most 0.90 percentage points separated configurations with and without query–key projections across tested settings, with gaps changing sign as capacity and flow depth varied.At Nf = 2, gaps were positive for M = 32 and 64 but negative for M = 96 and 128; at M = 32, gaps were negative for every tested depth of at least three layers.Source: §B.7.2 · p. 22, Table 9 · p. 23

B.7.3 Spectral Regularization Strength

  • The default spectral regularization weight, λspec = 10^-3, gave the highest reported accuracy and macro-F1 among tested values.Removing the penalty reduced accuracy from 30.24% to 27.42% and macro-F1 from 28.86% to 25.83%; a weight of 10^-1 also reduced both metrics.Source: Table 10 · p. 23, §B.7.3 · p. 23

Related papers

Loading 2610.30465v1…