From Attention to Atoms: Spectral Dictionary Learning for Fast, Interpretable Language Models

TL;DR

The paper tackles the quadratic cost of self-attention in Transformers by replacing it with a learnable spectral dictionary of $K$ Fourier atoms and per-token mixing coefficients, enabling linear-time language modeling with complexity $O(KL)$. It introduces a dual-domain training objective that combines time-domain embedding reconstruction, frequency-domain STFT magnitude matching, and a standard language modeling loss, plus a GMM prior learned from mixing coefficients for generation. Empirically, SDGM achieves competitive perplexities on WikiText-2 and Penn Treebank with far fewer parameters ($22.8$M) and lower memory usage, while delivering higher inference throughput compared to several baselines. The work demonstrates that spectral dictionary learning can provide a scalable, interpretable alternative to self-attention, with practical benefits for long-context language modeling and resource-constrained deployment, and suggests avenues for enhancing generation coherence and coefficient modeling.

Abstract

We propose a novel spectral generative modeling framework for natural language processing that jointly learns a global time varying Fourier dictionary and per token mixing coefficients, replacing the ubiquitous self attention mechanism in transformer architectures. By enforcing reconstruction losses in both the time domain (embedding reconstruction) and the frequency domain (via Short Time Fourier Transform magnitude matching) alongside a standard language modeling objective, and fitting a Gaussian Mixture Model (GMM) prior over the learned mixing vectors, our approach achieves competitive perplexity and generation quality on standard benchmarks such as WikiText2 and Penn Treebank. In contrast to the quadratic computation complexity of self attention, our method operates with linear complexity, delivering substantial efficiency gains. We demonstrate that spectral dictionary models can achieve competitive performance compared to transformer baselines while significantly reducing inference latency and memory footprint, offering a compelling alternative for scalable language modeling.

Figures (1)

• Figure 1: Architecture of the Spectral Dictionary Generative Model. First, the embedding layer maps each input token $w_{b,t}$ to a continuous vector $\mathbf{x}_{b,t}=E(w_{b,t})$. Next, the mixing encoder applies a one‐dimensional convolution to produce soft coefficients $C_{b,t,k}$. The spectral dictionary holds $K$ learnable Fourier atoms parameterized by amplitude $a_{k,d}$, frequency $f_{k,d}$, and phase $\phi_{k,d}$, which generate basis vectors $S_{k,t,d}$. The spectral decoder then reconstructs embeddings via $\hat{X}_{b,t,d} = \sum_{k=1}^K C_{b,t,k}\,S_{k,t,d}$. Finally, the pointer‐generator head combines each reconstructed vector $\hat{\mathbf{x}}_{b,t}$ with a context vector $\mathbf{c}_{b,t}$ to compute a mixture of vocabulary and copy distributions for token prediction.