Flow-based Spectral Kernel Learning for Nonstationary Attention
Authors: Zi Yang, Ying Li, Alejandro Lancho, Darío C. Larese, Pablo M. Olmos, Michael Minyi Zhang
Organizations: School of Computing and Data Science, The University of Hong Kong · Department of Statistics and Data Science, National University of Singapore · Department of Signal Theory and Communications, Universidad Carlos III de Madrid
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.
Figures & tables
Figure 1: Visualization of kernels and spectral densities p(ω1,ω2) under NF-based and NN-based implicit parameterizations. NF-based spectral modeling recovers more structured spectra and induces more faithful kernels.
Model
QK Proj.
Learnable V
Flexible Kernel
Positive RF
Acc ↑
Transformer
✓
✓
✗
–
0.8556
Transformer w/o QKV
✗
✗
✗
–
0.5267
IKAN
✓
✓
✓
✗
0.8656
Ours (FSKA)
✗
✗
✓
✓
0.8722
+ Direct RFF
✗
✗
✓
✗
0.3956
+ Learnable V
✗
✓
✓
✓
0.8725
Table 1: Ablation study on architectural design choices for attention mechanisms on MNIST classification. We report test accuracy (Acc). ✓ indicates that the corresponding component is used, ✗ indicates that it is removed, and – indicates that the component is not applicable.
MODEL
Venue
DBPedia
AGNews
TREC
RT-Polarity
SST-2
SST-5
METRIC
F1-macro ( ↑ )
Transformer
NeurIPS 2017
98.13 ± 0.15
90.43 ± 0.07
86.14 ± 0.93
70.52 ± 0.45
76.07 ± 1.10
35.11 ± 1.46
MIKAN
AAAI 2021
97.74 ± 0.01
91.15 ± 0.24
85.47 ± 1.07
71.30 ± 0.92
76.04 ± 0.86
32.84 ± 0.39
Performer
ICLR 2021
98.12 ± 0.09
90.23 ± 0.31
85.70 ± 0.57
70.88 ± 0.62
77.24 ± 0.41
35.92 ± 0.96
RKA
ICLR 2021
98.26 ± 0.01
90.80 ± 0.28
85.56 ± 0.75
71.15 ± 0.72
77.08 ± 1.46
36.10 ± 0.49
GMM-RKS
TMLR 2022
96.76 ± 0.21
90.58 ± 0.41
85.36 ± 0.91
71.05 ± 1.02
76.81 ± 0.89
36.42 ± 0.84
Table 2: F1-macro (%) on text classification benchmark datasets. Mean and standard deviation are computed over five runs. The best result per column is bolded and the second-best is underlined.
Model
SCOTUS
BBC News
20NG
arXiv
Acc.
F1
Acc.
F1
Acc.
F1
Acc.
F1
Transformer
55.12 ± 1.24
35.76 ± 0.91
88.43 ± 1.52
88.01 ± 1.57
24.91 ± 0.73
24.18 ± 0.69
75.62 ± 1.61
74.90 ± 1.66
FlashAttention
55.68 ± 0.79
36.21 ± 0.87
88.91 ± 0.49
88.47 ± 0.54
25.46 ± 0.70
24.92 ± 0.67
76.18 ± 0.58
75.41 ± 0.63
Performer
52.38 ± 0.77
31.05 ± 0.88
83.74 ± 1.26
82.96 ± 1.21
25.33 ± 0.81
24.97 ± 0.78
69.81 ± 1.04
68.93 ± 1.09
Mamba
59.84 ± 0.68
37.91 ± 0.78
89.34 ± 0.46
88.96 ± 0.51
28.91 ± 0.63
28.47 ± 0.66
76.92 ± 0.54
76.24 ± 0.59
RKA
53.06 ± 0.99
31.72 ± 1.26
85.18 ± 0.48
84.71 ± 0.55
28.62 ± 0.64
28.34 ± 0.67
70.15 ± 1.58
69.42 ± 1.62
Table 3: Full-document classification results (%). Values are mean ± standard deviation over three independent runs. Bold denotes the highest mean within each column. F1 denotes macro-F1. Documents are processed in full; BBC News and 20NG are shorter than SCOTUS and arXiv.
Figure 2: Parameter efficiency and memory scalability under the LRA-style benchmark. Left: different methods allocate trainable parameters differently across projection, spectral, and feed-forward components. Right: training memory growth with increasing sequence length.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Additional stability analysis of RFFs and PRFs.
Frequency pairs m
Correlation ρ
Absolute discrepancy Δ
16
0.921
0.794
32
0.945
0.794
64
0.963
0.800
128
0.978
0.798
256
0.989
0.806
512
0.994
0.802
Appendix
Table 4: ELU+1 similarity gap from the author response. Metrics use off-diagonal token pairs after cosine normalization. The frequency-pair count m is denoted by M in the response and reference analysis.
Figure 4: Regression plots on the toy nonstationary example across baselines.
Model
Acc
Runtime (s)
Ours (FSKA: Q=10 , M=96 , Nf=5 )
0.8743
85.32
(a) NGSM Mixture Count Q
Q=1
0.8483
76.52
Q=3
0.8543
76.30
Q=5
0.8612
76.39
Q=20
0.8431
75.32
Appendix
Table 5: Parameter sensitivity analysis on MNIST.
MODEL
Venue
DBPedia
AGNews
TREC
RT-Polarity
SST-2
SST-5
METRIC
Accuracy ( ↑ )
Transformer
NeurIPS 2017
98.13 ± 0.15
90.45 ± 0.07
86.67 ± 0.90
70.93 ± 0.83
75.99 ± 1.06
38.16 ± 0.67
MIKAN
AAAI 2021
97.94 ± 0.01
91.15 ± 0.24
86.60 ± 0.16
71.57 ± 0.38
75.40 ± 0.40
37.33 ± 0.73
Performer
ICLR 2021
98.12 ± 0.09
90.27 ± 0.29
85.00 ± 0.86
70.31 ± 0.42
77.26 ± 0.31
39.44 ± 0.42
RKA
ICLR 2021
98.16 ± 0.01
90.82 ± 0.28
86.80 ± 0.82
72.65 ± 0.82
76.77 ± 0.88
39.87 ± 0.50
GMM-RKS
TMLR 2022
96.82 ± 0.18
90.61 ± 0.34
85.74 ± 0.88
71.42 ± 0.96
76.38 ± 0.73
38.72 ± 0.61
Appendix
Table 6: Accuracy (%) on text classification benchmark datasets. Mean and standard deviation are computed over five runs. The best result per column is bolded and the second-best is underlined.
Dataset
Classes
Train
Validation
Test
SCOTUS
13
3,600
400
1,400
BBC News
5
1,102
123
1,000
20NG
20
3,600
400
2,000
arXiv
11
3,600
400
2,000
Appendix
Table 7: Document classification splits produced by the data preparation pipeline. Training counts exclude the held-out validation subset.
BBC News
20NG
Model
Accuracy
Macro-F1
Accuracy
Macro-F1
FSKA
90.76
90.31
30.24
28.86
Fixed bivariate Gaussian
64.20
63.50
21.80
20.40
Appendix
Table 8: Learned versus fixed bivariate spectral density. Accuracy and macro-F1 are reported in percent.
Random-feature capacity sweep, fixed Nf=2
M
32
64
96
128
–
Signed gap
+0.90
+0.81
−0.50
−0.87
–
Flow-depth sweep, fixed M=32
Nf
2
3
5
7
9
Signed gap
+0.90
−0.51
−0.55
−0.51
−0.47
Appendix
Table 9: Capacity sweeps on 20NG. Entries are signed Q/K accuracy gaps in percentage points.
λspec
Accuracy
Macro-F1
0
27.42
25.83
10−4
29.58
28.03
10−3 (default)
30.24
28.86
10−2
29.91
28.41
10−1
28.76
27.19
Appendix
Table 10: Spectral regularization sensitivity on 20NG. Accuracy and macro-F1 are reported in percent.
Instituto Superior Técnico & Instituto de Telecomunicações, Universidade de Lisboa, Portugal · Carnegie Mellon University, Pittsburgh, USA · Champalimaud Research, Lisbon, Portugal +1