Research map / Time series / Foundation models for time series
Fourier Head: Helping Large Language Models Learn Complex Probability Distributions
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes the Fourier head, a drop-in replacement for the linear softmax output layer that learns a continuous Fourier-series density discretized over bins, improving Decision Transformer Atari returns and Chronos time-series forecasting.
Problem
Standard linear classification heads over discretized continuous tokens ignore the ordering/continuity among bins, leading to jagged distributions that overfit noise when LLMs model non-linguistic, numerical, or probabilistic outputs.
Approach
A linear layer predicts autocorrelation coefficients of a truncated Fourier series; the resulting continuous PDF is evaluated at bin centers to yield a categorical distribution, imposing a smoothness inductive bias over ordered tokens, with a scaling law tying number of frequencies to expressivity vs. smoothness.
Key results
- Fourier head improves Decision Transformer normalized returns on Gravitar by up to 377% and Seaquest by up to 46.2% over the linear head.
- On Chronos zero-shot forecasting over 20 datasets, Fourier-550 head reduces MASE from 0.883 to 0.852 (about 3.5%) and smoothness score from 0.1689 to 0.0283.
- On toy synthetic datasets, Fourier head with 12 frequencies reduces average KL divergence by roughly 40% vs. the linear head across Gaussian, GMM-2 and Beta.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the linear classification head with a Fourier head (12 frequencies, no regularization) | MLP · Gaussian | KL divergence: 0.170 ± 0.052 -> 0.116 ± 0.043 | Helped |
| Replace the linear classification head with a Fourier head (12 frequencies, no regularization) | MLP · GMM-2 | KL divergence: 0.238 ± 0.032 -> 0.146 ± 0.033 | Helped |
| Replace the linear classification head with a Fourier head (12 frequencies, no regularization) | MLP · Beta | KL divergence: 0.234 ± 0.032 -> 0.191 ± 0.016 | Helped |
| Replace the linear classification head with a Fourier head | Decision Transformer · BankHeist | normalized returns: -0.09 ± 0.05 -> 0.92 ± 0.33 | Helped |
| Replace the linear classification head with a Fourier head | Decision Transformer · DoubleDunk | normalized returns: -72.72 ± 33.08 -> 45.45 ± 36.36 | Helped |
| Replace the linear classification head with a Fourier head | Decision Transformer · Gravitar | normalized returns: 1.32 ± 0.17 -> 4.98 ± 0.93 | Helped |
| Replace the linear classification head with a Fourier head | Decision Transformer · Seaquest | normalized returns: 2.53 ± 0.63 -> 3.70 ± 0.47 | Helped |
| Replace the linear classification head with a Fourier head using 550 frequencies | Chronos · 20 zero-shot time-series datasets | MASE: 0.883 -> 0.852 | Helped |
| Replace the linear classification head with a Fourier head using 550 frequencies | Chronos · 20 zero-shot time-series datasets | WQL: 0.750 -> 0.749 | Helped |
| Replace the linear classification head with a Fourier head using 550 frequencies | Chronos · 20 zero-shot time-series datasets | smoothness: 0.1689 ± 0.1087 -> 0.0283 ± 0.0224 | Helped |
| Use mixed-precision binning (ablation) | Fourier head in Chronos · 20 zero-shot time-series datasets | MASE | Helped |
| Add Fourier weight decay regularization (γ=10⁻⁶) (ablation) | Fourier head in Chronos · 20 zero-shot time-series datasets | MASE | Helped |
Surprising findings
- Regression (MSE) model simply regresses to the conditional mean and performs poorly on bimodal datasets, showing that distributional heads matter.
- Fourier head with ~550 frequencies gives only a 0.749 vs 0.750 WQL gain, i.e. nearly no improvement on the probabilistic metric despite large MASE and smoothness gains.
Limitations
- Fourier head only appropriate when bins have a meaningful ordering/continuity; not for pure language tokens.
- Theory requires assumptions (periodic, Fourier coefficients decay at rate 1/k^t, N much larger than 1 and below m/2) and the smoothness metric is a chosen proxy.
- Training uses discretized cross-entropy rather than direct continuous likelihood, so the continuous objective is not the main experimental setting.
Experimental setup
Datasets: Atari Seaquest offline DQN-replay dataset (Decision Transformer) Chronos pretraining corpus Synthetic Gaussian, GMM-2, Beta datasets Chronos zero-shot benchmark (20 datasets)
Benchmarks: Atari BankHeist Atari DoubleDunk Atari Gravitar Atari Seaquest Chronos zero-shot benchmark (20 datasets)
Metrics: normalized return KL divergence smoothness MSE MASE WQL total variation distance
Baselines: Linear head GMM head Chronos (linear head) Decision Transformer (linear head)
Builds on: Fourier Basis Density Model Decision Transformer Chronos
Closest papers by meaning
- XTC: Head-Aware Sampling by Excluding Top Choices (NeurIPS 2026 · similarity 0.86)
- Enhancing Decision-Making of Large Language Models via Actor-Critic (ICML 2025 · similarity 0.85)
- FFN Fusion: Rethinking Sequential Computation in Large Language Models (NeurIPS 2025 · similarity 0.85)
- Large Language Models to Diffusion Finetuning (ICML 2025 · similarity 0.85)
- Flexible Language Modeling in Continuous Space with Transformer-based Autoregressive Flows (NeurIPS 2025 · similarity 0.85)
- Reasoning is Periodicity? Improving Large Language Models Through Effective Periodicity Modeling (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
