Fourier Head: Helping Large Language Models Learn Complex Probability Distributions
Authors: Nate Gillman
Daksh Aggarwal
Michael Freeman
Chen Sun
Venue: ICLR 2025 Poster (2025)
Keywords: LLM
Fourier
smooth function
multi-class classification
Primary area: foundation or frontier models, including LLMs
TL;DR: Using Fourier series, we build a neural network layer which learns categorical distributions that have a continuous structure.

Abstract: As the quality of large language models has improved, there has been increased interest in using them to model non-linguistic tokens. For example, the Decision Transformer recasts agentic decision making as a sequence modeling problem, using a decoder-only LLM to model the distribution over the discrete action space for an Atari agent. However, when adapting LLMs to non-linguistic domains, it remains unclear if softmax over discrete bins captures the continuous structure of the tokens and the potentially complex distributions needed for high quality token generation. We introduce a neural network layer, constructed using Fourier series, which we can easily substitute for any linear layer if we want the outputs to have a more continuous structure. We perform extensive analysis on synthetic datasets, as well as on large-scale decision making and time series forecasting tasks. We also provide theoretical evidence that this layer can better learn signal from data while ignoring high-frequency noise. All of our results support the effectiveness of our proposed Fourier head in scenarios where the underlying data distribution has a natural continuous structure. For example, the Fourier head improves a Decision Transformer agent's returns across four benchmark Atari games by as much as 377\%, and increases a state-of-the-art times series foundation model's forecasting performance by 3.5\% across 20 benchmarks unseen during training.
We release our implementation at https://nategillman.com/fourier-head

Published as a conference paper at ICLR 2025
FOURIER HEAD:
HELPING LARGE LANGUAGE MODELS LEARN COMPLEX PROBABILITY DISTRIBUTIONS
Nate Gillman Daksh Aggarwal" Michael Freeman', Saurabh Singh?, Chen Sun 1.2 Brown University, GGoogle DeepMind
ABSTRACT
As the quality of large language models has improved. there has been increased interest in using them to model non-linguistic tokens. For example. the Decision Transformer recasts agentic decision making a sequence modeling problem, using decoder-only LLM to model the distribution over the discrete action space for an Atari agent. However, when adapting L.MMs to non-linguistic domains, it remains unclear if softmax over discrete bins captures the continuous structure of the tokens and the potentially complex distributions needed for high quality token generation. We introduce neural network layer, constructed using Fourier series, which we can easily substitute for any linear layer we want the outputs to have a more continuous structure. We perform extensive analysis synthetic datasets, well as on large-scale decision making and time series forecasting tasks. We also provide theoretical evidence that this layer better learn signal from data while ignoring high-frequency noise. All of our results support the effectiveness of our proposed Fourier head in scenarios where the underlying data distribution has a natural continuous structure. For example, the Fourier head improves Decision Transformer agent's returns across four benchmark Atari games by as much as 377%, and increases a state-of-the-an times series foundation model's forecasting performance by 3.5% across 20 benchmarks unseen during training. We release our implementation at httptt::////...
iameeegiilaaa.aa..mmm Iman www...com/..o......m....mm.o.o.. wtt./ooumfoooooeooo
Fourier Head Learns Higher Quality Densities
zamal
Scoothosss 0075
-1.75
-0.90
-3.25
6.11
11:2
0.50
0:26
100
-1.00
-0.25
-1.50
-025
0.50
0.21
3.50
0.5
1:00
Figure 1: We task an MLP with learning to approximate a continuous bimodal density using a categorical distribution and cross-entropy objective. We observe that standard linear head fails to distinguish between the two modes, and overfits to high-frequency noise in the training set. In contrast, our proposed Fourier head learns smoother, accurate categorical distribution.
INTRODUCTION
Human language can be viewed discretization for continuous, often probabilistic represen- tation of the world that is construed in our mind (Spivey 2008). The continuous structure can be partially captured by language models with their token embeddings, where "nearby" tokens are em- bedded to have latent representations with high cosine similarities. The embeddings themselves are acquired as a result of the data-driven learning process. Can we, based on rich prior knowledge
"qqual contribution. Correspondence to: nate gillman herown.edu chensun brown.cdu

---

Published as a conference paper at ICLR 2025
about the continuous world, inform the language model about the underlying continuity of its in- puts, like the fact that the word "emerald" is more similar to "shamrock" than "pine" when they are used to describe different shades of green? As large language models (LLMs) have evolved into "foundation models" that are adapted to a diverse range of tasks, tokens that are a priori continuous more essential than ever, for example for arithmetic computations (Liu et al. 2023). decision making with continuous or discrete actions (Chen et al. (02T). future anticipation and time-series forecasting (Ansari al. 2024), simply drawing random numbers given probability distribu- tion (Hopkins et al.) 2023).
We view the problem of informing LLMs to utilize the continuity prior from the perspective prob- ability density estimation. For simplicity. we adopt the standard next token prediction framework whose training objective is softmax cross-entropy. Assuming non-overlapping vocabulary, continu- ous values can be discretized via binning (Ansari et al. 2024). On one hand, the linear head adopted by LLMs independently projects each token into probabilities, and has the expressive power flex- ibly approximate arbitrary probability density functions subject to the "quantization" errors. The linear head however does not consider any continuous structure that resides among the tokens (i.e. random re-shuffle of the tokens in the vocabulary would not change the predictions). On the other hand, a head based on a parameterized distribution (e.g. Gaussian or Gaussian Mixtures) naturally incorporates the continuous structure, but is often too simple (and overly "smooth") account for multi-modal distributions for future prediction or decision making. Can we design a head that is both expressive and incorporates continuous structures?
We introduce the Fourier head, motivated by Fourier series as universal function approximators. The Fourier head learns continuous probability density function, and returns a discrete ap- proximation of it. Intuitively, returning a discretization of a continuous density in this way allows the classification head to better model the low- now-frequency signals from the training data, because the Fourier head is forced to approximate the categorical distributions using finite number of frequen- cies. a high level, the Fourier head inputs x E IR"", uses a linear layer to learn the coefficients for a Fourier series with N frequencies over [-1, and quantizes the interval [-1, 1] into m equal bins. Then, the Fourier head evaluates the learned Fourier PDF at those m bin center points, and returns those likelihoods Basis Density Model (De la Fuente et al. 2024).
categorical distribution. The Fourier head builds upon the Fourier
Our main contributions as follows.
Contribution #1: We reveal the underlying principle on the trade-off between the Fourier head's s expressive power and the "smoothness" of the predicted distributions. We prove theorem which demonstrates a scaling law for the Fourier head. Namely, as we increase the quantity of Fourier coefficients learned by the Fourier head, the layer model increasingly complicated distributions; however, the Fourier head will necessarily fit to more high-frequency noise, thereby outputting cat- egorical distributions which are less smooth.
Contribution #2: We propose a practical implementation of the Fourier head capable of sequential prediction tasks by modeling complex multi-modal distributions. Additionally, we propose strate- gies to improve the layer's performance, including Fourier coefficient norm regularization, weight initialization, and the choice of how many Fourier frequencies use.
Contribution #3: We demonstrate the effectiveness of the Fourier head on two large scale tasks, where intuitively a continuity inductive bias over the output dimensions ought help the model's generation performance. In the first task, an offline RL agent which uses decoder-only transformer to model the next-action distribution for an Atari agent, we improve returns across four benchmark games by as much as 377%. In the second, we outperform a state-of-the-art time series foundation model on zero-shot forecasting by 3.5% across benchmark of 20 datasets unseen during training.
2 FOURIER HEAD
2.1 FOURIER HEAD:
MOTIVATION
When practitioners apply LLMs model complex probability distributions non-linguistic to- kens, a standard technique is to quantize the latent space into m tokens and learn a conditional categorical distribution over those tokens. We share two examples here:
2

---

Published as a conference paper at ICLR 2025
The Decision Transformer (Chen et al. 2021 models an Atari agent's behavior in the Seaquest game by learning a categorical distribution over the 18 possible actions (move left, move right, shoot left, etc.). They use decoder-only transformer architecture.
Example 1:
The Chronos time series foundation model (Ansari et al. 2024) models the distribution of next numerical values by quantizing the closed interval | 15, 15 into 4096 bins, and learning categorical distribution over those bins. They use an encoder-decoder transformer.
Example 2:
In pure language modeling task. token ID 1000 and token ID 1001 likely represent unrelated words. However, in a task where the token IDs represent numerical values, the token ID 1000 and 1001 would represent numbers that close together.
The final layers of an LLM for such a task are generally a linear layer, followed by softmax, fol- lowed by cross-entropy loss. We hypothesize that in scenarios where nearby token IDs encode similar items, an inductive bias that encourages them to have similar probabilities will improve per- formance. A generic linear layer learns an unstructured categorical distribution and thereby allows more arbitrary probabilities. In this work, we propose to give the model this inductive bias by let- ting the classification head learn a categorical distribution as the discretization of a continuous learned function from suitably flexible class. We consider the very flexible class of truncated Fourier series with N frequencies. These are functions of the form
((x) ao + Σ (ax cos(k<u) + by sin(k==)))
(2.1)
Fourier series are a classical tool for solving quantitative problems (Stein & Shakarchi (003) be- cause functions like Equation 2.1 are universal function approximators, with the approximation improving as N increases.
2.2 FOURIER HEAD: DEFINITION
We now propose replacement for the generic linear layer token classification head, built using Fourier series. We call our replacement the Fourier Series Classification Head, or the Fourier head for short. The Fourier head inputs any vector x E TR.", and outputs a categorical distribution in RT. For a high level summary of how it works-the Fourier head inputs E R", uses linear layer to extract the coefficients for a Fourier series over [-1,1], quantizes the interval [-1, 1] into m equal bins, evaluates the learned Fourier PDF those m bin centerpoints, and returns those The m likelihoods as a categorical distribution. We formally define this layer in Algorithm Fourier head is constructed using the Fourier Basis Density Model from (De la Fuente al. 2024
For more details on the original method (e.g. justification for how learning the autocorrelation coefficients guarantees that the Fourier series has integral and justification for normalizing the Fourier coefficients by R(co)) we refer the author to (De la Fuente et al. 2024). We direct the curious reader to Appendix C.I for a low-dimensional demonstration the Fourier head in action.
2.3 FOURIER HEAD: CONSIDERATIONS FOR TRAINING
We highlight the main design choices of the Fourier head so that users may apply it most effectively.
Training objective: The Fourier head inputs a signal I and extracts from that signal an inter- mediate representation of probability distribution Pa (z) defined over z E [-1, 1]. This probability distribution has a closed formula equal Fourier series. In our experiments, we optimize the pa- rameters of the Fourier PDF by discretizing it over the latent space and training using cross-entropy loss. However, we should note that the Fourier layer allows MLE training directly on continuous values, by evaluating the Fourier PDF directly on the ground truth value in the latent space. But for consistency of comparison, and to demonstrate how easy it is to swap the Fourier head with linear layer, use softmax cross-entropy loss the objective.
Choice of hyperparameter N: The Fourier head has one crucial hyperparamettr--naelly the num- ber of frequencies. How should one choose this in practice? We offer Theorem 3.3 as guiding principle beyond simple trial and error. This result provides a scaling law which formalizes the smoothness-exprestee
power trade-off in choosing the number of frequencies. In general, using more frequencies leads to more expressive power, and generally better success metrics, but at the cost of learning less smooth densities, well more model parameters.
3

---

Published as a conference paper at ICLR 2025
Algorithm 1 Fourier head
Hyperparameters: the input dimension n, output dimension m, number of frequencies N Initialization: define a linear layer A R R2(N+1) // maps input to autocorrelation coefficients
INPUT (#1, ........ xn) IR" (oo: , BN) t Ax
ak an iBk E C, for every lo O,...,N
I/ compute autocorrelation coefficients
Ck + -i= acd C, for every K = ....., N
I/ compute Fourier coeflicients
p(2) 2 + +R EE Ríco) exp(ihrz))
define Fourier PDF over [-1,1]
by + - + 1+2k for every k: = ,,..., m
// define m bin centerpoints
yk ΣT p(b,)' for every kt ....,
plbe)
// evaluate PDF at m bin centerpoints
OUTPUT (V1.... .m) E Rm
// by design, Ex k= Uk = and each y/s 2 0
Fourier regularization: For a given number of frequencies N, there could be many learned Fourier models that fit the given data equally well. To encourage a smoother learned model and penalize unnecessary high frequency content, we follow (De la Fuente et al. 2024) and add a regularization term that measures the total squared variation for the Fourier model, to prevent higher order Fourier coefficients from growing too large during training. This helps ensure that the learned Fourier PDF doesn't 't overfit to noise in the data, and therefore has a bias towards learning smoother densities. In the notation from Algorithm I this means adding regularization term of 7 2x² Ex= k2 cxlc2 to the loss function, where 7 is a hyperparameter. When picking regularization strength, we find that in the low-frequency domain (e.g, frequencies in the single digits) using 7 = 0 works best, and in the high-frequency domain (c.g. greater than 10 frequencies), using y = 10 works best.
Binning strategy: The choice of data binning can impact performance. As discussed, the Fourier head should only be applied when nearby bins are 'similar' in some sense, requiring a semantically meaningful bin ordering. When bins represent quantized numerical values over continuous la- tent space, a mnixed-precisiog
binning strategy can improve performance. For example, to model values in [-15, 15] with most data in [-1, 10], allocating more bins the dense interval improves performance. Given m total bins, a hyperparameter d E [0, 1) controls allocation, with [d m] bins for the sparse interval and the rest for the dense range (estimated from training data). Fourier theory supports this approach, as increasing precision in dense regions de-localizes the quantized data distribution, localizing the Fourier spectrum. This accelerates higher frequency decay, enabling effective learning with lowerffreuency Fourier heads. Separately, we note that (De la Fuente et al., 2024] suggests re-parameterizing the periodic domain the real line, though we do not use this in our work.
Weight initialization: The learned parameters for the Fourier head consist of the learned línear layer which extracts autocorrelation parameters dg. In Py Torch, linear layers use He initialization (He et al. 2015) by default, which ensures that the linear layer outputs values close to zero in expectation. Similarly, initializing the Fourier densities to be uniform p(2) 1/2 improves learning dynamics. We accomplish this by dividing the weights and biases by a large number, such as 1000, after He initialization; this guarantees that the linear layer outputs very small values, so that Fourier coefficients output from the autocorrelation step are very small as well.
3 THEORETICAL ANALYSIS OF FOURIER HEAD
3.1 "SMOOTHNESS": A METRIC FOR HIGH FREQUENCY CONTENT
In this subsection we propose a smoothness metric which inputs a categorical distribution y (y1...., Bm)) RT and assigns numerical value depending on how smooth it is. The score will output if is the smoothest possible categorical distribution, and larger values if y is less smooth. We will first specify what we mean by "smooth":
4

---

Published as a conference paper at ICLR 2025
Heuristic 3.1. We say a function is smooth if contains very little high-frequency information.
For example, the uniform categorical distribution contains no high-frequency information, so it is the smoothest possible function, and should get a smoothness score of 0. In contrast, cafegorical distribution containing samples from sin (10000)z) contains lots of high frequency information, so should get a smoothness score greater than O. We seek to define a metric which measures smoothness according to Heuristic 3.1
We will first develop a smoothness metric in the general case of a function f : [a, b] + R. then specialize to case of the discrete categorical distribution that we consider in the paper. If we let O R be weights satisfying Jo a, 1, and D be some measure of discrepancy such as L?,
and let go (ar) * ((x) denote the convolution of (x) with a Gaussian kernel of standard deviation o, then it is reasonable to define the smoothness of f to be the quantity
s sn)= 10=I110 [ oa ((S()),), [xx)dxzd..
(3.1)
In this expression, the discrepancy D[f(x), (z) f()) Gaaussin-smoothhe version of itself. Because the Gaussian is a low-pass filter, we can interpret Equation 3.1 saying, at a high level, that a function is "smooth" if it doesn't change that much when you remove high frequency content from it.
measures how different fxx is from
In our experiments,
consider discrete categorical distributions, and wish to tractably quantify their smoothness. Accordingly, we define specific case of this as follows.
Definition 3.2 (Smoothness metric for categorical distributions). Suppose y = (yl..... ... Vm E R is a categorical distribution, every yk 2 0 and Σk=1 Vk = 1. Demote by ga E R2 2m the discrete Gaussian kernel of standard deviation a and radius m 1. Define the weights 0, = 6/ o²² Then we define the smoothness of y to be the constant
20
s(y) Σ
ao bly - go * vl12.
(3.2)
We direct the curious reader to Appendix B where we conduct additional experiments to justify this choice of smoothness metric for our experiments.
3.2 SCALING LAW FOR THE FOURIER HEAD, FREQUENCY-ASPECT
In this subsection, we share a theorem that analyzes the quality of the Fourier head as the quantity of frequencies changes. We refer to this as the Fourier head scaling law as it quantifies the trade-off between modeling capacity and smoothness the number of frequencies increases. On one hand, it a celebrated result from Fourier analysis that a Fourier series with greater number of frequencies models larger class of functions; but the other hand, show that increasing frequencies also incurs loss smoothness. This is to be expected, as we designed our smoothness metric with the intention of identifying a distribution as less smooth if contains more high- figh-frequency information.
Theorem 3.3. (Fourier head scaling law.) Consider Fourier head with inpul dimension n, output dimension m, and N frequencies. Suppose that 1 < <
m. Then the following are true:
2.
1. (Increasing N improves modeling power.) As N increases, the Fourier head is capable of learning a larger class of densities.
2. (Increasing N degrades smoothness.) Consider an input to the Fourier head x E R", and denote by fa [-1, R the optimal conditional distribution that we would like the Fourier head to approximate for this input. Suppose that there exists some 1, 2 2 such that the Fourier coefficients of fa decay on the order of 1/k". Denote by f2.N the truncation of fa to its first N frequencies, denote by b E Im the m bin centerpoints in [-1, 1}, and denote by y( = Na,N (b) (f (bo) .... + fr.N (... (bm--1) E Rum the discretization of fr.N into bins. Then, there exist constants C1, C2 > O such that
C2
M(y!N)) = C1 2 O(1/N2).
(3.3)
5

---

Published as a conference paper at ICLR 2025
Toy Example: Learned Conditional Distribution vs True Conditional Distribution
Guussian Drinect
Canseian Dansel
Guussiaa Ditrect
Probability Mass
6.1
Tra X31
Hill hanurle
0.
Tra FDR
Dae FEC
-0.5
0,5
-1.4
-0.5
00
-0.5
0.0
1/
Untasee
Heto Ditnaeet
Beta Dotaset
022
Livar 1ith Swarrras
1.0
0.6
1.0
Figure 2: Comparison between the PMFs learned by the linear head, GMM head, and the Fourier head, for two of the datasets in the toy example- Gaussian and Beta. (The GMM dataset in Figure We observe that the Fourier head learns smoother categorical distribution than the linear head over its predicted values. Furthermore, the Fourier head better fits the true conditional PDF; this is reflected in the KL divergence and smoothness metrics.
Note that the smoothness scaling law asymptotic in Equation 3.3 shows that as N increases, so does ()). Further, note that if the Fourier spectrum of the underlying distribution decays quicker (controlled by 1) then the rate at which smoothness degrades is slower; this is because what we learning has little high frequency content, then increasing the frequencies shouldn't affect the smoothness of the learned distribution very much. In part (2). our assumption that the Fourier coef- ficients decay at least quadratically is reasonable since if fz is at least twice continuously differen- tiable, we already know its Fourier coefficients corresponding to the k-th frequency are in ((1/A²) (Stein Shakarchi] 20033 Ch.2, Cor. 2.4). Our Fourier weight decay regularization helps toward ensuring that this condition is met in practice well. We include a full proof in Appendix
4 ToY EXAMPLES
4.1 LEARNING A CONTINUOUS CONDITIONAL DISTRIBUTION
We demonstrate the advantage of using the Fourier head to learn a probability distribution for a simple task: learning the conditional distribution of the third number in the sequence given the first two. Here we will use (22) to denote the quantization of z.
Dataset: We create 3 synthetic datasets, which we name Gaussian, GMM-2, and Beta. Each dataset consists of 5000 quantized triples {g((z), (()) -- 1,3 Crucially, 2 is sampled from distribution which is conditioned x and y, and have explicit closed formula for this distribution. By design, the Gaussian dataset is unimodal in 2, whereas the more challenging GMM- 2 and Beta datasets are not unimodal. Full details about the datasets can be found in Appendix C.2
Task: Predict the conditional distribution of g(z) given the quantized tuple (((r),(())
Model architecture: Our model is MLP with ReLU activations and one hidden layer, which maps R22 + 64 + 32 + 58 The output of the model has dimension 50 because we quantize the interval |-1, 1] into 50 bins. We consider two baselines alongside the Fourier model. For the first baseline, the classification head is a linear layer; for the second baseline, the classification head a Gaussian model mixture classification layer with two Gaussians, where the means and standard deviations are learned; for the Fourier model, the classification head is the Fourier head. We sweep over frequencies N = 2, 4, ..... 20, and consider regularization 7 (0, 10-9). We train those models via cross-entropy loss We also consider a cegression-basede model, trained using MSE.
Model evaluation: We use three metrics for evaluation. Our first metric is the average KI. di- vergence ( (PPP((x)))]]) 2))ll q(y))), where P(z, y) is the fixed conditional distribution of given (x, u): ((P(((P. (x, u)) is the quantized approximation of P(x, y). obtained by evaluating the den-
Note that we also demonstrate the possibility of training continuous version of the Fourier head, using siaiimum-iikelhhoo hased objective. Accordingly, carry out experiments in the continuous domain analogous to those did in the quantized domain; for more details, Appendix

---

Published as a conference paper at ICLR 2025
KL Divergence (4)
Smoothness (L)
Dataset
Linear
Fourier
Linear
Fourier
Gaussian
0.170 0.052
0.116 = 0.043
0.116 = 0.049
0.057 = 0.011
GMM-2
0.238 + 0.032
0.146 = 0.033
0.068 + 0.022
0.038 + 0.007
Beta
0.234 + 0.032
0.191 = 0.016
0.127 + 0.044
0.076 + 0.021
Table 1: We compare metrics between the linear head, and the Fourier head with 12 frequencies and no regularization, for every dataset in toy example. We observe that the Fourier head outperforms the linear head across all metrics. Notably, using Fourier head improves the KL divergence (the primary success metric) on average by approximately 40%. We aggregate metrics over different seeds and report the standard deviation.
Using Llama-3.1 1-----BBInntrue Simulate Gaussian Sampling
Llama
Liama LoRA
Liama LoRA + Fourier Head (6 freus)
Probubility Mass
Prodicted PMF
Proficted PMF
Prodicted PMF TVD 0.1188
TVD 0.8745
- 0.3248
Figure 3: We demonstrate that the baseline Llama model does a poor job simulating Gaussian sam- pling, measured by the Total Variation Distance between the ground truth quantized Gaussian histogram, and the empirical histogram of samples. We find that LoRA fine-tuning improves the results by factor of 2.07, and that using the Fourier head improves the output distribution by factor of ~ 4.86.
sity function of P(z, y) the bin centers, multiplying by the bin width, and finally scaling by the sum of the likelihoods; and M (g(x), (v)) denotes the predicted categorical conditional distribution of (2)). Our second metric is smoothness. And our third metric is MSE, where we consider the expected value of g(=) under the learned categorical distribution as a prediction for g(=).
The metries for the best performing model on each dataset are reported in Table Figure presents sample visualizations of the learned conditional distributions alongside the true densities. And in Appendix C.2 we present the results a study on the impact of number of frequencies and Fourier regularization. Notably, this study provides empirical evidence for the Fourier head scaling law in Theorem smoothness degrades, and model performance improves until it reaches saturation point. Crucially, we observe that the Fourier head flexibly learns all three distributions better than the linear baseline does. We note that the Fourier head outperforms the linear head MSE well; include complete comparison with both Linear and GMM head baselines in Appendix C.2 Additionally, in Figure IO (Appendix), we demonstrate that the regression model simply regresses to the mean of the conditional distribution. Accordingly, the regression model performs well for the unimodal Gaussian dataset, and it performs poorly for the bimodal datasets GMM-2 and Beta.
Results:
3.3 as it demonstrates that for all datasets, as frequency increases, the
4.2 ARE LLMS RANDOM NUMBER GENERATORS?
Suppose that we query an LLM with the following prompt, repeatedly: "The following a list of normally distributed random numbers in the interval [-1, 1] with mean 0.55 and std 0.10: 0.57, 0.36, " Would the model outputs be approximately Gaussian? In this empirical study, we simulate Gaussian sampling using the Llama-3.1 --8-1-8---nnntree model (Dubey et al. 2024). We demonstrate that this language model struggles to generate high quality numerical Gaussian samples. We con- sider two possible interventions: LoRA fine-tuning the base Llama model, and LoRA fine-tuning the base model while also replacing the linear classification head with a Fourier head. We find that LoRA tine-tuning improves the learned distribution significantly, and replacing the linear head with the Fourier head improves the distributions even further. We present an illustrative example this phenomenon in Figure 3 See Appendix C.4 for experiment details and some related works.

---

Published as a conference paper at ICLR 2025
5 LARGE-SCALE STUDY: OFFLINE REINFORCEMENT LEARNING
The Decision Transformer (Chen et al. 2021 reframes reinforcement learning as sequentially mod- eling rewards, states, and actions. We evaluate its performance on the Seaquest game from the Atari (Bellemare et al., 2013) benchmark. The Seaquest game contains 18 actions, with two groups of eight actions that have natural "closeness" metric defined on them: move left, up left, up, up right, right, down right, down, down left; as well shooting in those eight directions. The original chitecture uses decoder-only language model (Radford et al. 2018) to encode context and map through a linear layer, producing categorical distribution over actions. At test time, the agent samples from this distribution to select its next action. We replace the linear classification head with Fourier head, introducing prior that semantically similar actions (e.g., 'move left' and 'move up left') should have similar likelihoods. Our results show the Fourier head improves returns by as much as 46% in the
setting, using identical training hyperparameters.
Task: In the Seaquest game, the agent moves a submarine to avoid enemies, shoot at enemies, and rescue divers. We consider this task in the Offline RL setting. The agent observes the past states, actions, and rewards, as well as the return-to-go and attempts to predict the action that matches what an agent operating like the dataset would likely do. We also consider three other Atari games with the same action space: BankHeist, DoubleDunk, and Gravitar.
Dataset: We use the same dataset from the original Decision Transformer implementation (Chen al. 202Ib. This dataset consists of 500k transitions experienced by online deep Q-network agent Mnih et al., 2015) during training on the Seaquest game.
Model architecture: (Chen et al. 2021) used the GPT-1 model (Radford et al. 2018 to autoregres- sively encode the context, which is then fed through a linear layer dimension 18, and the model ultimately optimizes the cross entsoyy loss between the action logits and the ground truth action from the dataset. We refer to this model as the linear baseline. To create our Fourier-N version, we simply replace the linear head with Fourier head with N frequencies and Fourier regularization = 10-6. In our experiments we consider frequencies E ...,3,,,
Normalized Returns for Decision Transformer Agent
Atari Game
Classification Head
BankHeist
DoubleDunk
Gravitar
Seaquest
Linear head
0.09 + 0.05
772772 + 33.08
1.32 0.17
2.53
0.63
Fourier head
0.92 1 0.33
45.45 1 36.36
4.98 1 0.93
3.70 1 0.47
We present returns obtained by the Decision Transformer agent using the linear baseline, and the Fourier head, across the four Atari games. We compute the returns (mean and standard deviation) by averaging over four seeds. Across all these games, the Fourier head significantly improves the normalized returns obtained by the agent.
Table 2:
Impact of Varying Fourier Frequencies Decision Transformer
Normalized Returss
Smnoothhsss
[ore kad Taka: badd
Frerirhcal
Litereel
18
Fourier Froquencies
Fourier Frequencies
Figure 4: We present empirical results for how the quantity of Fourier frequencies impacts returns and smoothness for the imitation learning task. For normalized returns, higher better; for smooth- lower is better. We the linear baseline agent when sufficiently many Fourier frequencies are used, while still learning smoother next-action distributions.
that the Fourier agent achieves higher normalized returns than
8

---

Published as a conference paper at ICLR 2025
Model Evaluation: We present results for the linear baseline and Fourier-N head (N 2, 2,4, 6, ......... ,0, 32) across four Atari games, showing mean reward totals for rollouts at the best epoch across four seeds. Tabossddeossttratss significant return gains with the Fourier head. For ex- ample, Seaquest returns increase by up to 46.2%, while Gravitar sees as much as a 377% boost. Fig- ure/fsshow improved Seaquest performance the number of frequencies grows, with learned PMFs becoming less smooth, aligning with Theorem 3.3 Qualitative results in Figure (3 high- light the smoother PMFs produced by the Fourier head. Additional results for BankHeist, Double- Dunk, and Gravitar in Figure T6 (Appendix) confirm that the Fourier agent consistently outperforms the linear baseline while maintaining smoother next-action distributions.
Ablations: We analyze whether model size has any effect on the relative performance of the Linear head and the Fourier head. The results in Figure 14 (Appendix) demonstrate that, across model sizes, the Decision Transformer with a Fourier head is better at learning high-quality next action distributions than the Decision Transformer with a Linear head. We also analyze whether dataset size has any effect on the relative performance of the Linear head and the Fourier head, and obtain similar result. In Figure [3 Appendix) we show that, across dataset sizes, the Decision Transformer agent with the Fourier head achieves larger returns than the agent with linear head.
6 LARGE-SCALE STUDY: PROBABILISTIC TIME SERIES FORECASTING
The Chronos time series foundation models (Ansari et al. 2024) "learn the language of time se- ries". They do this by approaching time series forecasting as language modeling, by tokenizing the quantized number line, learning token embeddings for each of those quantized values, and finally learning a categorical distribution to decide what the next value ought to be. This model is built on top of the encoder-decoder T5 model (Raffel et al.) 2020 In particular, this model normalizes time series values to the range [-15, 15] and quantizes this interval into 4096 tokens. As usual for language modeling, the final layer is linear map which learns a categorical distribution over next tokens. In particular, we observe that token represents a number very close to tokens - 1 and ++1. However, we note that there is no inductive bias in the T5 architecture which pushes their likelihoods to be similar. This is not hypothetical problem; in Figure [7 (Appendix), we can see that the linear token prediction PMFs fit the noise, and appear very jagged.
The motivation for replacing the linear head with the Fourier head is to "smooth" out the distribution in the left side of Figure [7] to help the forecasting model better learn the signal, and ignore the noise. In this figure, we can see that the Fourier head accomplishes this successfully.
In this section, we study how the performance of the Chronos time series foundation model changes when we pre-train using the Fourier head, instead of the linear head. For all of the frequencies that we consider, the Fourier head outperforms the Chronos linear baseline on the MASE metric, while learning next token multinomials which are at least 8x smoother, with fewer parameters than the baseline.
Dataset: We use the same training dataset for large-scale pretraining that (Anaai et al. (2024) used. We gather an evaluation benchmark of 20 time series datasets which were not seen during train- ing. These 20 come from the zero-shot eval from (Ansari et al. [024]. The reader can check Appendix D.2 for details on the training and evaluation datasets we used.
Model architecture: We use the Chronos model, which is built using the T5 architecture (Raffel ct al.) 2020). The original model has a linear classification head. For our study, we will replace this with Fourier head with frequencies N 64, 128, 256, 550. We use mixed precision binning; this is informed by an analysis of the Fourier spectrum of the next-token distribution, as described in Section 2.3 We also use Fourier weight decay regularization with 7 = 10-6 For the task, the model
Chronos Time Series Model
MASE (L)
WQL (1)
Smoothness (0)
Linear
0.883
0.750
0.1689 + 0.1087
F5u
0.852
0.749
0.0283 + 0.0224
Table 3: We present large-scale experiments on Chronos time series forecasting. The best- performing Fourier model outperforms the linear baseline both terms of the continuity of the learned probability mass functions (smoothness) for the quality of the forecasts (MASE, WQL)
9

---

Published as a conference paper at ICLR 2025
learns to input time series context of length 512, and output probabilistic forecast of length 64. At test time, the model chooses the next numerical token by sampling from the next-token distribution
Model evaluation: We have two sets of metries: model performance from (Ansari al. 2024) (MASE measures the accuracy of median forecast, and WQL measures the quality of the proba- bilistic forecast), as well as smoothness metric. Our Fourier metries in Table demoostttttttttttee that every Fourier model outperforms the linear baseline for MASE and smoothness. Furthermore, for the largest Fourier model that we consider, Fourier outperforms linear on WQL well.
The results in Table g (Appendix) show that mixed precision binning and regularization improve the MASE and smoothness for the Fourier head, and that using more Fourier frequencies improves MASE and WQL. Additionally, we show that the Fourier head yields more accurate fore- casts than the linear head across dataset sizes and model sizes (Figures 18 and 19 Appendix).
Ablations:
RELATED WORK
LLMs outside of natural language domains: LLMs are often adapted to domains beyond natural
language, as general purpose sequence models. For example. they have been used in protein syn- thesis (Madani et al.) 2023), time series forecasting (Ansari et al.) 2024 Das al.
2024 Jin al.)
2024
Nate Gruver & Wilson 2023 Requeima et al. 2024 Jia et al. 2024
et al. 2023 Wang
ct al. 2024). music generation (Dhariwal et al. [0200] A Agostinelli et al. 2023 Copet et al. 2023
Yuan et al. 2024), and as well in decision making (Li et al. 2022 Chen et al. 2021).
We consider three categories to adapt LLMs to non-language domains: when the output of a language-trained LLM is used as a feature for some out-of-domain task; when a language-pretraing iLM fine-tuned on a domain sppecff task; and when an LLM architecture is trained on a domain- specific dataset from scratch. Our work directly considers the latter method of LLM adaptation, particularly in settings where the outputs approximate continuous values. We note that using I L.MMs to model numerical functions has seen success in continuing sequences (Mirchandani et al. 2023)
but has been challenging for modeling samplers for probability distributions (Hopkins et al. 2023
In a related direction, Razeghi et al. (2022) found that model performance on numerical reason- ing tasks is correlated with the frequency of specific numbers in its corpus. Further, some have re-framed continuous regression as a descretized classification problem leverage LLMs in numer- ical modeling contexts (Song et al.) 2024) or RL contexts (Farebrother et al.) 2024). While even frozen L.Ms with no further training show interesting empirical results as regressors (Vacareanu et al., 2024), there is conceptual mismatch between the downstream task and model construc- tion because tokenized numerical values trained using cross-entropy loss does not explicitly enforce numerical relationships between the tokens.
Fourier series in neural networks: Many works leverage the Fourier transform as a data pre- processing step deterministic transformation within the network,
Fourier analysis to motivate design choices. It is far less common to learn the Fourier series directly. De la Fuente al. 2024) learned marginal univariate densities parameterized using Fourier basis; our work extends their Fourier Basis Density model to multivariate settings with an autoregressive scheme. Our method learns conditional univariate densities using Fourier basis, where the coefficients of the Fourier density model are input dependent. SSizzmann et al. (2020) proposed sinusoidal activation a Fourier series; in contrast, we fix
functions, which can be seen as learning the frequencies the frequencies to the canonoical choice (1,2..... N}. and learn the amplitudes. This allows the Fourier head to more directly benefit from approximation results from Fourier analysis.
CONCLUSION
We propose the Fourier head and demonstrate its positive impact on performance on several tasks. We prove a scaling law that characterizes the trade-off between the model's expressivity and the smoothness of its output distribution. The Fourier head is a modular architecture that can be easily added to existing models that would benefit from the continuity inductive bias that the head imparts. The Fourier head extends the already extensive reach of LLMs into more diverse, numerical, and probabilistic domains. Future work includes exploring alternative training objectives that do not depend on discretizing probability density functions, and incorporating the Fourier head in general- purpose L.LM training, where the head can be adaptively employed when needed.
10

---

Published as a conference paper at ICLR 2025
9 REPRODUCIBILITY STATEMENT
We have made efforts to ensure reproducibility. In Algorithm | we provide all the mathematical details that one needs to reproduce the Fourier head. In Appendix D.2 we prove our scaling law, Theorem 3.3 in full detail, and we list all assumptions in the statement of the theorem. Additionally, we have released the research code on GitHub: htttps://qithii com/ nate lee--glllmarr fourier-head
ACKNOWLEDGMENTS
We would like thank Jona Balle, Alfredo De la Fuente, Calvin Luo, Singh Saluja, Matthew Schoenbauer, and Megan Wei for the useful discussions. This work is supported by the Samsung Advanced Institute of Technology, NASA, and a Richard B. Salomon Award for Chen Sun. Our research was conducted using computational resources at the Center for Computation and Visual- ization at Brown University. Chen would like to thank Mia for inspiration.
REFERENCES
Andrea Agostinelli, Timo Denk, 7alán Borsos, Jesse Engel, Mauro Verzetti, Antoine Caillon, Qingqing Huang, Aren Jansen, Adam Roberts, Marco Tagliasacchi, et al. Musiclm: Generating music from text. arXiv preprint arXiv:2301. 1335, 2023.
Abdul Fatir Ansari, Lorenzo Stella, Caner Turkmen, Xiyuan Zhang, Pedro Mercado, Huibin Shen, Oleksandr Shchur, Syama Syndar Rangapuram, Sebastian Pineda Arango, Shubham Kapoor, Jasper Zschiegner, Danielle C. Maddix, Michael W. Mahoney, Kari Torkkola, Andrew Gor- don Wilson, Michael Seankee-ehhnider, and Yuyang Wang. Chronos: Learning the language of time series. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL
net/ http://wwwee//oouuteemmm Doorumiid=ggrNNcggtt.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling, The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Peter Brockwell and Richard Davis. Time series: theory and methods. Springer science & business media, 1991.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, arXiv preprint oxXiv22106001465 2021.
Jade Copet, Felix Kreuk, Itai Gat, Tal Remez, David Kant, Gabriel Synnaeve, Yossi Adi, and Alexan- dre Défossez. Simple and controllable music generation. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023.
Abhimanyu Das, Weihao Kong, Rajat Sen, and Yichen Zhou. A decoder-only foundation model for time-series forecasting. In International Conference on Machine Learning, 2024.
Alfredo De la Fuente, Saurabh Singh, and Johannes Ballé. Fourier basis density model. arXiv preprint 1.Xi::2402.15345 2024.
Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Ilya Sutskever. Jukebox: A generative model for music. arXiv preprint arXiv:2005.00341, 2020.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint urXiv:2407.21788, 2024.
Jesse Farebrother, Jordi Orbay. Quan Vuong, Adrien Ali Taiga, Yevgen Chebotar, Ted Xiao, Alex Irpan, Sergey Levine, Pablo Samuel Castro, Aleksandra Faust, et al. Stop regressing: Training value functions via classification for scalable deep rl. arXiv preprint arXiv:2403.03955) 2024.
11

---

Published as a conference paper at ICLR 2025
Jort F. Gemmeke, Daniel P. W. Ellis, Dylan Freedman, Aren Jansen, Wade Lawrence, R. Channing Moore, Manoj Plakal, and Marvin Ritter. Audio set: An ontology and human-labeles dataset for audio events. In Proc. IEEE ICASSP 2017, New Orleans, LA, 2017.
Yuan Gong, Yu-An Chung, and James Glass. Psla: Improving audio tagging with pretraining, sampling. labeling, and aggregation. JEEE/ACM Transactions on Audio. Speech, and Language Processing, 2021. doi: 10.1109/AALLP22211112633
Jia Gu, Liang Pang, Huawei Shen, and Xueqi Cheng- Do llms play dice? exploring probabil- ity distribution sampling in large language models for behavioral simulation. arXiv preprint arXiv:2404.09043 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Aspen K Hopkins, Alex Renda, and Michael Carbin. Can Ilms generate random numbers? evaluat- ing IIm sampling in controlled domains. In ICML 2023 Workshop: Sampling and Optimization in Discrete Space, 2023.
Tsuyoshi Inouye, Kazuhiro Shinosaki, H. Sakamoto, Seigo Toi, Satoshi Ukai, Akinori Iyama, Y Katsuda, and Makiko Hirano. Quantification of eeg irregularity by use of the entropy of the power spectrum. Fllectrenncepallograhhh and Clinical Neurophysiology, 29(3)2204-100 1991. ISSN 0013-4694. doi:
URL https:
www. www. www.sciencedirect.oo
Furong Jia, Kevin Wang, Yixiang Zheng. Defu Cao, and Yan Liu. Gpt4mts: Prompt- large language model for multimodal time-series forecasting. Proceedings of the AAAI Conference on Mar. 2024, doi: 10. (0.116//aaiivv88/// v38121.30383, URL httpss///jjs aaai org index. http://AA//arrieeeeeeee/////////
Artificial Intelligence, 38(2))
Ming Jin, Shiyu Wang, Lintao Ma, Zhixuan Chu, James Y Zhang, Xiaoming Shi, Pin- Yu Chen, Yux- uan Liang, Yuan-Fang Li, Shirui Pan, and Qingsong Wen. Time-LLM: Time series forecasting by reprogramming large language models. In International Conference on Learning Representations (ICLR), 2024.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836 2020.
Shuang Li, Xavier Puig, Chris Paxton, Yilun Du Clinton Wang, Linxi Fan, Tao Chen, De-An Huang, Ekin Akyürek, Anima Anandkumar, et al. Pre-trained language models for interactive decision- making. Advances in Neural Information Processing Systems, 35:31 199-31212, 2022.
Yixin Liu, Avi Singh, C Daniel Freeman, John D Co-Reyes, and Peter Liu. Improving large language model fine-tuning for solving math problems. arXiv preprint arXiv:2310. /0047, 2023.
Madani, Krause, and et al. Greene. Large language models generate functional protein se- quences across diverse families. Nature Biotechnology, 41:1099-1106, 2023, doi: 10.
$415877002-01661-22
Suvir Mirchandani, Fei Xia, Pete Florence, Brian Ichter, Danny Driess, Montserrat Gonzalez Are- nas, Kanishka Rao, Dorsa Sadigh, and Andy Zeng- Large language models as general pattern machines. In Proceedings of the 7th Conference on Robot Learning (CoRL), 2023.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518 (77754)) 2015.
Shikai Qiu Nate Gruver, Mare Finzi and Andrew Gordon Wilson. Large Language Models Are Zero Shot Time Series Forecasters. In Advances in Neural Information Processing Systems, 2023.
12

---

Published as a conference paper at ICLR 2025
Akshay Paruchuri, Jake Garrison, Shun Liao, John Hernandez, Jacob Sunshine, Tim Althoff, Xin Liu, and Daniel McDuff. What are the odds? language models are capable of probabilistic reasoning. arXiv preprint 2tti::20061123300 2024.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. OpenAl website, 2018.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
Yasaman Razeghi, Robert L Logan IV, Matt Gardner, and Sameer Singh. Impact of pretraining term frequencies on few-shot numerical reasoning. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.). Findings of the Association for Computational Linguistics: EMNLP 2022, pp- 840-854, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.18653/v 105553//122222.ndiinngsss findinng-eemll..59
James Requeima, John Bronskill, Dami Choi, Richard E Turner, and David Duvenaud. Llm processes: Numerical predictive distributions conditioned on natural language. arXiv preprint 1282v 2024.
Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. Advances in neural information processing systems, 33:7462-7473, 2020.
Xingyou Song, Oscar Li, Chansoo Lee, Bangding Yang, Daiyi Peng, Sagi Perel, and Yutian Chen. Omnipred: Language models as universal regressors. CoRR, abs/2402.114444 2024. doi: 10. 48550/ARXIV.2402.4455 URL Inttp:: http//ooiooorgo/////// 0e.e://1..4855000 www.w.wwwww........m.m 14547
Michael Spivey. The continuity of mind. Oxford University Press, 2008
Elias M Stein and Rami Shakarchi. Fourier analysis: an introduction, volume 1. Princeton Univer- sity Press, 2003.
Elias M Stein and Rami Shakarchi. Real Analysis: Measure Theory, Integration, and Hilbert Spaces,
volume 3. Princeton University Press, 2005.
Terry Tao, Dec 2014. URL http://s//errrraaoooorrereeeeeea -ote- http://wwe...ooooooaoooaaaaaoaaa
Robert Vacareanu, Vlad Andrei Negru, Vasile Suciu, and Mihai Surdeanu. From words to numbers: Your large language model is secretly a capable regressor when given in-context examples. In First Conference Language Modeling, 2024. URL wwwpreview....c... wwt. forum?
id=ZzpaUxcNFK
Shiyu Wang, Haixu Wu, Xiaoming Shi, Tengge Hu, Huakun Luo, Lintao Ma, James Y Zhang, and JUN ZHOU, Timemixer: Decomposable multiscale mixing for time series forecasting. In International Conference on Learning Representations (ICLR), 2024.
Megan Wei, Michael Freeman, Chris Donahue, and Chen Sun. Do music generation models encode music theory? International Society for Music Information Retrieval, 2024.
Eric W. Weisstein. Square wave. From Math World-A Wolfram Web Resource, 2024. URL https: Accessed: September 16, 2024.
mathwwre wol fram. com Squa reWave. html
Ruibin Yuan, Hanfeng Lin, Yi Wang, Zeyue Tian, Shangda Wu, Tianhao Shen, Ge Zhang, Yuhang Wu, Cong Liu, Ziya Zhou, Ziyang Ma, Liumeng Xue, Ziyu Wang, Qin Liu, Tianyu Zheng, Yizhi Li, Yinghao Ma, Yiming Liang, Xiaowei Chi, Ruibo Liu, Zili Wang, Pengfei Li, Jingcheng Wu, Chenghua Lin, Qifeng Liu, Tao Jiang, Wenhao Huang, Wenhu Chen, Emmanouil Benetos, Jie Fu, Gus Xia, Roger Dannenberg, Wei Xue, Shiyin Kang, and Yike Guo. Chatmusician: Understand- ing and generating music intrinsically with Ilm. arXiv preprint arXiv:2307.07443, 2024.
Tian Zhou, Peisong Niu, Xue Wang. Liang Sun, and Rong Jin. One Fits All: Power general time series analysis by pretrained Im. In NewrlPS, 2023.
13

---

Published as a conference paper at ICLR 2025
Appendix
Table of Contents
Proof of Fourier Head Scaling Law, Theorem 3.3
14
A.
Definitions
14
Overview of Proof
15
Proving Theorem 3.3 Using the Lemmala
A.4 Proving the Lemmata
B Smoothness Metric
22
Additional Experiment Details, Toy Examples
25
Motivating Example: Audio Spectrogram Transtormer
25
Learning Continuous Density
25
MI F. hased Fourier Head
26
C.4
Are LLMs Random Number Generators?
28
Additional Experiment Details, Large-Scale Examples
30
D.I Decision Transtormer
30
D22 Chronos
30
A PROOF OF FOURIER HEAD SCALING LAW, THEOREM 3.3
In this section we prove Theorem 3.3 the Fourier head scaling law. To do this, we must first discuss the Ayquist-Shannor Sampling Theorem. This result states that order to avoid distortion of signal (such as aliasing) the sampling rate must be at least twice the bandwidth of the signal. In the setting of the Fourier head, our sampling rate is m/2 because we have m bins uniformly spaced in (-1, 1), and the bandwidth is N/2 because the frequency of sin(TNas is N/2. Thus the Nyquist Theorem requires us to have
m/2 2 2 (N/2) N
in order for the higher order frequency content learned by our model to not be fallacious when we are learning from only m bins. This justifies why only theoretically study the case 1 << N < m/2 the scaling law.
A.1 DEFINITIONS
Consider input R the Fourier head, and denote by fa 1-1.]] R the optimal conditional distribution that we would like the Fourier head to approximate for this input. We will assume that fa is periodic, since the Fourier head learns 2-periodic Fourier density. We denote by S2,N the truncation of the Fourier series of fa to its first N frequencies. Note that fa, also integrates to over [-1, 1] since its first Fourier coefficient the same that fz. Further, fa.N is non-negative on |-1. 1] since its Fourier coefficients, being subsequence of the coefficients of fa, are non-negative definite; a periodic function with non -egaaivee definite Fourier coefficients non- negative by Herglotz's Theorem (Brockwell & Davis 1991| Corollary 4.3.2). For completeness, we will recall the convolution formulas, specialized to the cases we consider in our argument.
Definition A.1 (Discrete convolution). Let by := + 25+1 0 < < m the center points of 2222
the m bins in 1), and let us denote 5 = (bo,...,bm Denote by Go(z) V/20 the
Gaussian PDF with standard deviation Then the discrete Gaussian convolution filter of radius m iss
90 = Go((1 m. m, m. m 1))
E F 2m
(A.1)
S(m, o)
14

---

Published as a conference paper at ICLR 2025
where the normalization constant is
m-1
0(m,o) =
Σ (k(.).
(A.2)
k=1-m
The discrete convolution of 9a E 122-1 and (2, E R the vector (go S2,N) E R whose j 'th coordinate is given by
(go (b,))(b;) 5(m, σ) k=1-m
1
Σ Go Sa.N (bz-z).
(A.3)
Definition A.2 (Continuous convolution). The continuous Gaussian convolution filter Do 2, 2] + R>0 is
2
S(m, σ) 25( Go (m))
9- (2)
(A.4)
This function go (2) is a normalized truncation of a Gaussian PDF with mean 0 and standard devia- tion 20 m.. The continuous convolution of go 1, 2] and the periodic function f2.N [-1, R
go * fa,N(2) == [ go (u)f2,N(2 u) du.
(A.5)
A.2 OVERVIEW OF PROOF
In this subsection, we provide an overview of the proof of Theorem 3.3 by presenting the statements of the lemmata that will need, and connecting each one the overall argument. In the next subsection, we rigorously prove the scaling law by careful applications of these lemmata. And in the following subsection, we will rigorously prove each of the lemmata.
This first lemma allows us to replace the discrete Gaussian convolution in the definition with continuous Gaussian convolution.
Lemma A.3. (Discrete convolution is close to continuous convolution) If we define the constant
B1 (m, o) == 1 + [mm,o) then we have that
(m))
I12,, (6) Do (,)N( | [Bi a) Go (0,N / (A.6) Furthermore, B, (m, σ) satisfies the following bound, uniformly in o,
B₁ (m,, σ) 1+ 2m
(A.7)
This next lemma, standard result from analytic number theory allows us to upper bound the sums of the norms of the Fourier series coefficients. This is proved in various places, see e.g. (Tao 2014 Equation 21).
Lemma A.4 (Asymptotic expansion of Riemann zeta function). Consider the Riemann zeta function <(t) Σm Ift 2, then
N
c(t) 1 K
1 1- 11-ING 1N-I + +O1/N')
(A.8)
This next lemma allows us to extract the main asymptotic behavior the scaling law. Lemma A.5. (Main term asymptotic) Denote by au(x) the constant coefficient of f2.N- Let us suppose that the Fourier coefficients of f2,N decay like Bz(x) and define the constant
B2(a,m,at
+ 28 (2))))
(A.9)
Then we know that Bi ff, go fz,N
(A.10)
= BB(a, m,))
B1 (m, Bs(z))
N + ((1/N22)
1
(A.11)
2t -
Furthermore, B2(a, is bounded from above and below function of a and m.
15

---

Published as a conference paper at ICLR 2025
This final lemma allows to relate the continuous case, where analysis works out easier, the discrete case, where our smoothness metric is actually defined.
Lemma A.6. (The average value of the truncated Fourier PDF is 1/2) N m/2, then
Ja,N(b))=
m
(A.12)
2
1=0
A.3 PROVING THEOREM 3.
USING THE LEMMATA
We now prove the theorem that provides scaling law for the Fourier head. This result quantifies the trade-off between modeling capacity and smoothness as the number of frequencies increases. In order prove this, must assume that fa. the conditional distribution being learned by the Fourier head, is sufficiently smooth. For example, if fo is twice continuously differentiable, then the Fourier coefficients corresponding to the k-th frequency of fo are in O(1/k²) (Stein Shakarchi, 2003 Ch.2, Cor. 2.4). Thus, our assumption that the Fourier coefficients decay quadratically is reasonable, and our Fourier weight decay regularization helps ensure that this condition is in practice as well. In our theorem, we generalize this hypothesis to the cases where the Fourier coefficients corresponding to the k-th frequency fa are O(1/k)).
Theorem 3.3. (Fourier head scaling law.) Consider Fourier head with input dimension n, output dimension m, and N frequencies. Suppose that 1 & N < m.
2. Then the following are true:
(Increasing N improves modeling power.) As N increases, the Fourier head capable of learning larger class of densities.
2. (Increasing N degrades smoothness.) Consider an input to the Fourier head I E RR"" and denote by fa [-1,1] - R the optimal conditional distribution that we would like the Fourier head to approximate for this input. Suppose that there exists some t 2 2 such that the Fourier coefficients of fa decay on the order of 1/k'. Denote by f2,N the truncation of fa to its first N frequencies, denote by b E Tem the m bin centerpoints in [-1,1] 1), and denote by y( fu,N /fo (bo) +... + fa,N (bm-1)) m the discretization of fr.N into bins. Then, there exist constants C1, C2 > O such that
C₂ N + +
((y(N)) = C1
(3.3)
Proof of Claim 2 of Theorem
3.3
We can estimate that
((y
1 m- Σ == (b,)
aollfe.N( 9o * fa.N (b)(
(Definition 3.2)
0o (IB - (āo * (2)
(Lemma A.3
Ejo Sa.N(bj)
0=1
+ +/m((1/22411)
Vm
Σ a (B(a,,,,z) Bi B3(x22
(Lemma A.5
Enzo (b,)(b,)
(21 1)N
+ + 0(1/N2+1)
+0(1/N29
2 m
(Ca
C4 N 0(1/N2))
(Lemmata A.5
A.6
In the last step we used the convergence of the respective series (which follows from boundedness of B2(a, m) and B1 (m,,o) in σ) and we assigned C3 and C4 to be those sums. This completes the proof.
Proof of Claim I of Theorem 3.3 The proof of this claim is more straightforward. For any function on [-1, 1] that is least twice continuously differentiable, know that the Fourier series of
16

---

Published as a conference paper at ICLR 2025
converges uniformly and absolutely to f (Stein & Shakarchi 2003 Ch. 2, Cor. 2.4). In other words, the function fN being learnt by the Fourier head converges uniformly and absolutely to f.
A.4 PROVING THE LEMMATA
In this subsection, we will restate and prove Lemmata A.3
A.5 and A.6
Lemma A.3. (Discrete convolution is close to continuous convolution) If we define the constant B1 (m, σ) == 1 3(m,o)
(m) then we have that
(b) 9o * fa.N (b)l2 = [Bi(m, (b) - go * fu.N(6)l22 VmO(1/N24H) (A.6)
Furthermore, (m, σ) satisfies the following bound, uniformly in o,
(m,,, σ) S 1 2m
(A.7)
Proof of Lemma A.3 Extending fa.N periodically to [-2, we can compute that the continuous convolution Go (2)N(2) is
[ Do Bo(u)aa,N( u) du
(āa [2,N)(2)
(A.13)
m
Go ((2)
1788
[a - du
(A.14)
25(m. σ)
-2
m
Go(s)fz,N (z
28
ds.
(A.15)
S(m,o) 6)
-m
where in the third step applied the change of variables a continuous approximation of the discrete convolution in Definition 3.2 To see this, we will apply the Euler-Maclaurin formula. This formula says that the integral in Equation A.15 is a Riemann sum over rectangles of width 1 evaluated at the right endpoints of each interval, minus an error term E(m, o), follows:
We claim that this is precisely
m
2k
(ā. Sb,,)(by)) o(m, S(m, σ) k=
Σ
Go(k)
LE,N (b
(A.16)
m
Σ
m
Go(l)) fz.NN
2ÿ 2k
(A.17)
S(m, σ)
m
m
m
-1 2 k) +1
Go(kk) S2,N
(A.18)
S(m,o) σ)
m
k=1-m
1
Σ
(k) S2.N (by-n) Ga
0(m,o)
kk=l-m
(A.19)
1
= 0(m,
(S(m, (go * fb.N)(b,) (m
(A.20)
1
(go * fb,n) s(m,o) (b)))
(A.21)
where the error term defined as
1
d (Go (z 22)) (s)
E(m, o) 5(m, o)
(A.22)
ds
(Go(m)faNNz -2) 22) GG((-m)faN((= + +2)),
(A.23)
25(m,0)
17

---

Published as a conference paper at ICLR 2025
where (s(s) [-] 1/2 is the periodized Bernoulli polynomial. We will now estimate this error term. Note that since G is an even function and fz.N is periodic with period 2, the difference compute that
in A.23 is O. Therefore,
O(m.o)= = S(m, σ) J m
1
2s (=-
ds
(A.24)
mS(m, [ Go())
2s
ds.
(A.25)
Using the triangle inequality, we can bound E(m, a) terms of convolutions with Jr:
())))) m
(=-2)
2s m
(8m,,o)
ds
(A.26)
8(m,0)
28
(G(s))
ds
(A.27)
((,,o)
m
m
1
(x-2)
2s
Go(s)feNN
ds+
(A.28)
0(m,o)
20
m
1
2s
Go(d) Sa.N
ds
(A.29)
m
m
m
(go * (z)
(āo * (()))))(()
(A.30)
20
where in Equation A.29 we used that [P(()) < 1/2 and that [G) = |s| Gx)/o2 <
mGa(s)/o for I-m, m].
Note that since Ôo is a truncated Gaussian on [-2, 2]. it is infinitely differentiable on the open set (-2,2), however, it is not differentiable at the endpoints -2 and 2 when treated as 4-periodic function. This technical difficulty can be resolved using mollifiers: we can replace Jo with Zo * Pe. where { is a family of mollifiers indexed by c <> O. The key properties of a mollifier are that To fP is infinitely differentiable as 4-periodic function for all c > 0 and lim 9o (P. Ôo (Stein Shakarchi] 2005 Ch. 3). We are ultimately interested in only bounds on absolute values of fr convolved with various functions, and since absolute values are continuous and inequalities are preserved under taking limits, all our bounds are still true. particular, this shows that the k'th Fourier coefficients of ÿo decay faster than any polynomial. And on the other hand, by assumption we know that the Fourier coefficients of fa.N decay on the order of 1/k'; and we know that I is continuous and 2 periodic, its Fourier coefficients converge. So by the convolution theorem, can deduce that the Fourier coefficients of ğ fr.N and g- * *1f2,x1 decay faster than any polynomial.
Summed over the N frequencies, this shows that [g. (2)) and Iğo than any polynomial as well. Since m is fixed and σ 2 1, this implics that
| (2)l decay faster
|m Ø) = 011/22+1)
(A.31)
Using Definition A.1 and Equation A.16 we have that
9o * fz,N 5(m, σ) k=1-m
1
Σ
m
(k) fa,N (b
2k
1
Go (o(m)aa)(b)
(A.32)
S(m,
m
= fe,N)(b,) E(m,o) S(m, σ)
Go (m.
(A.33)
If we define (i = Gm(m) - then Equation A.33 combined with A.31 together imply that
9o (b,) 3o (b,) + (,(,,,) 0(1/N2(+1).
(A.34)
18

---

Published as a conference paper at ICLR 2025
Finally, estimate that
(b)) - 90 * fa,N (b)ll2
(A.35)
(b),, Do * f2, fa
(A.36)
+ +19 * f2,N(0) do * (a) + (im,
(A.37)
(l) -o (2) Ci o) )
(A.38)
= =I[(1 + C1 (1 f.. (6) - 9o (b)l2 + VmO
(A.39)
This completes the first part of the proof. For the second part of the proof, since S(m, o) 2 (2m 1)G. (m 1), we estimate that
(m)) Bi (m, o) 1 (2m I)G. (m 1) 2mn
11+
(2m 1 1)/20"
(A.40)
2m
This completes the proof.
Lemma A.5. (Main term asymptotic) Denote by ag(x) the constant coefficient of Jz.N- Let us suppose that the Fourier coefficients of f2,N decay like B2(z)/K² and define the constant
Bzlo, m) v (oll))BB((0)) 0)2 + 2B, (m, (
(A.9)
Then we know that
[B)(m,)))a.NN( Do (b)l
(A.10)
= =mm m2,,, m. Bi (m, Br())
1
N + O(1/N2))
(A.11)
Furthermore, B2(a, m, is bounded from above and below function of a and m.
Proof of Lemma A.5 We will first argue that B2(a, m, r),, as function of and m, bounded from above and below. Indeed, from its definition, (2 m, x) 2 / // (m or 2 [=o())] But the Fourier PDF has integral 1 over [-1, 1], its constant term is an(x) 1/2. This implies that (2 m, x) 2 1/2. To see (2 m, z) is bounded from above, we simply recall that in Lemmal A.3 we showed that [Bi (m, o) < 2, which implies that [2 m)) < /4ao(z)² + (z)(2))((2))) This shows that B2(a, m, x) is bounded above and below function of and a, as claimed.
Now, let (do(z),... i E R be the discrete Fourier transform of (Bi (m,, JJ,,N Ja (b) E R For notational simplicity, will write B1 (1 o) as long as is fixed. By Plancherel's Theorem, we have
(Bifa.N 3o * fb,)) m A=
(x)
(A.41)
Let h. be the Fourier coefficients of 9.. treated periodic function with period 4, and defined over |-2 2]:
(2)
Σ
ha /2/e/2
(A.42)
Since fr.N is defined over [-1, 1] and is periodic with period 2, can likewise treat it function over 2, 2] with period 4, in which case we can rewrite its Fourier series as
[a,(()) Σ
(A.43)
JN-2N
where
(1)/(())
if 0
(mod 2)
(z(2)
(A.44)
0
else.
19

---

Published as a conference paper at ICLR 2025
Then, by the Convolution Theorem, we have
2N
(B1f2, - Do * f2,N)(bi) = Σ
(B(x)(B - hhos) emihb//
(A.45)
j=-2N
N
Σ (2)(2)(Bi - -ha22)) e
(A.46)
N
=
ag (2) - lo 20,22)) e e
(A.47)
where the second equality we used the fact that 0j is 0 for odd j and therefore re-indexed using k = j/2.. Thus, using the definition of DFT along with Equation A.47 we get
(x)x) = Σ
(BIJ.NN - Do * (ai,))(i)) e
- 22rkk/mm
(A.48)
1=0
Σ
Σ ax(x)(B1 - -h, ) 22rikl/m
(A.49)
1=0
N
m-1 xij(-1+ 2111
Σ
(B(x)(Bi - -ho,23)
2rikt/m m
(A.50)
N
1=0
N
=
Σ
(B(x - -ho2zs) eij(
2ni(j
(A.51)
First, we claim that at most a single summand (in j) is represented. Towards this, we note that
E (mod m) k (mod m).
2mili-k)1/m
(A.52)
1=0
Then, we note that since 0 < N m/2, 2, for each 0 < k < m, there is at most one j (-N, -NN + 1,....N] N). such that & (mod m.) This shows that there is at most single summand. We will now find the exact formula for each summand. We consider three cases.
Case 1: k < N. In this case, j k satisfies j = k (mod m)) so this index gives the exponential sum m.
Case k m N. this case, k is too large be an index in the sum, so we can't 't choose j = k; the next smallest equivalent value is j = k - m, which satisfies = (mod m). But N m < j -N in this case, so le is too small be an index in the sum; therefore, every exponential sum is zero in this range.
Case 3: m N < k < m - 1. In this case, j = k m satisfies j = k (mod m). We have -N < 1, this is valid index in the sum.
This gives the following closed formula:
m an (B) ho 2e
0 S k S N N < < m N N N 5-mm-
de (x)
(A.53)
[B: 2 2m)-m) (
20

---

Published as a conference paper at ICLR 2025
Using this closed formula in A.41 obtain
(BiJa,N Do (b)))
(A.54)
N
m (z) - )e
(A.55)
m k=0
m ag-m (BI hn. jort--mmmmm
+
(A.56)
k= N
m-1 1
m (ax(t) - 10,22)1 m Σ
0k ax= - 200 (A.57)
A m
where in the last step we used that [etck(1-1/m) 1 |c / since they both complex exponentials. Now, since Do is real and even function, we know that its Fourier coeffi- cients h are real. Further, since Ga is infinitely differentiable, we also know that ha. O(1/k) (in fact they decay faster than 1/k for any p 2 1). Thus, using that [ax(2)) decays like (s(r)/k² we see
1ox (B) ho,211) = |ax(x)) (B) - 2h),2((x)) + 112222)
(A.58)
B3 1:21
+0
+
(A.59)
(.2((2))²
From A.59 it is clear that since we are interested in only the dominant asymptotic, we can safely ignore the higher order terms coming from the Na As result, we can estimate that
N
BB h-2t
mh
B3B (K - m2
[BBJ2,N - 3o Ixx))(,111 ~ mB)ao(2)
+ m
MI
Σ kmm
(A.60)
B
B
-ma mma(()) (x)² +
(A.61)
k
1:21
k=
N
B²B
N
B3B
(x:)² mma0(x)2 + m
-ma
771
(A.62)
1
k=1
k=1
B3
mag (aa(())23) (x) B + 2m Σ
(A.63)
1:21
k=1
Next, we note that our asymptotic in Lemma A.4
applied at 2f, yields
N
1
((2))
21 N + )
(A.64)
Substituting this into A.63 we obtain
m-1
(Bif..N go * fb,))
(A.65)
0=
m + 2m B3 (((2t) 21 - N21-1 + 0(1/N21)
(A.66)
B3 2/
BB
N B3B3O(1/N22)
1
(A.67)
B2
BB
1
0(1/22)).
(A.68)
2
2
21

---

Published as a conference paper at ICLR 2025
where defined B₂ B2 m,) V (z) 22 (21 as in the statement of the
Lemma, and in the third line we applied Lemma A.3 to estimate that B1 I 2, and we used that B3 (x(x) only depends on æ. Then, using the Taylor expansion (1 + x)/// 1 + + (x²) about 0, estimate that
1/2
(BiJa,N Do
* 0
(x,N))(;)1
(A.69)
-MBB3
m2B B3 1
N + mO(1/N2))
1/2
(A.70)
2t
2B
1/2
VmB2
1
(2t 1)B3 N + 01//N2)
(A.71)
B
.mB2
1
BB
1) N + O(L/N21)
1
(A.72)
(2t
B B3
mm
B2
N +(1/N22)
(A.73)
2t
To justify our application of the Taylor expansion, note that N > 1, and B2 B2(a, m, x) is bounded below as a function of a and m. This completes the proof.
Lemma A.6. (The average value of the truncated Fourier PDF is 1/2) If m/2, then
-ma(())= 2
(A.12)
Proof of Lemma A.6
Denote by ag the Fourier coefficients of fa.N. We can compute that
m
m-1
N
Σ Ja,((b)==
=
Σ
Σ
Oge imk( ane
-14 ( 2
(A.74)
1=0
1=0
A=-N
Σ
age tnn
2néjk
(A.75)
Note that by hypothesis, N m/2, which implies that [k] m for every outer sum index We 2mij consider two cases; if k 0, then the innermost summand is e m = 1; and if k * 0, then the
innermost sum is truncated geometric series with first term 1, common ratio e and m terms. In summary, the innermost summand is
2xsih
m h 0 x O,
(A.76)
0
m- Sa,N (by) man- But an 1/2 because Sa,N is PDF implies it has which implies that ン average value 1/2 over [-1, 1]. This completes the proof.
B SMOOTHNESS METRIC
We will examine how the proposed smoothness metric Equation 3.1 behaves in a toy example setting to gain intuition for its behavior. Consider a square wave, which can be expressed as an infinite sum of odd integer harmonics that decay in amplitude proportional to their frequency:
f(x)
sin
(B.1)
n
L
n= 1,3.5....
Here, the wavelength is 27. (Weisstein 2024)
22

---

Published as a conference paper at ICLR 2025
We construct a truncated version of the square wave with finite and fixed number of frequencies. The waveform will slowly approach its jagged. square shape as more sine waves are added. We frame these increasingly jagged waves discretized multinomial densities to simulate the output of the Fourier head. To do this, we simply set the height to zero when the wave crest becomes negative and normalize the sum 1. The output of this transformation for few representative waveforms is pictured in Figure 5
I Size; smouthness 0.00004
2 Sines; smmottinese 0. 0.0000
(.000)
00001
0.0005
Sines, smoothness 0.00011
4 Sines, smoothness 0000014
0.0000
0.0015
0.0000
1100015
110000
Sines; smoothness 0.00027
20 Sines; smoothness 0.00010
110000
100000
0.0000
Figure 5: Truncated square waves framed as densities and their smoothness.
Intuitively, the truncated square wave with a single sine wave ought be the smoothest. Thus our metric in this context should be smallest at that point, and increase monotonically as we add more sine waves. The plot in demonnttraes that this indeed the case.
Choice of L2 Distance over L1 Distance: The proposed smoothness metric Equation 3.1 permits
a general measure of discrepancy D, and we've chosen D to be L distance as indicated in 3.2 We empirically observe that 1,2 distance better preserves monotonicity than the L for higher frequency content, thus motivating this choice. With a sample rate 2048Hz, the L distance exhibits some undesirable warping when square-wave multinomial uses over 80 sine waves (see Figure 6). A Fourier head in a practical setting may possess several more than 80 frequencies; accordingly, we favor the 2 distance our diserepancy measure.
Alternative Notions of Smoothness: In validating our choice smoothness metric, we compare to the spectral entropy (Inouye al. [991). which has similar purpose in quantifying the "smooth- ness" of the frequency content of a signal. Spectral entropy is defined as the Shannon entropy of the power spectral density of sampled signal which is defined as follows:
(11)=-- -
Sn loga
Sn Scotal
H(f:N) Ep(n) loga (p(n)
Σ NEN
(B.2)
Stotal
DEN
Here, N' is the number of Fourier frequencies and S is the power of a frequency n E N: Sn is the power spectrum of the nith frequency, and Stotal is the power of the signal using all N frequencies. For some frequency at index n, S,, is called its relative power and EnEN Stote! us consider each frequency's power probability.
Sp_ enables
In the discrete case, the maximum entropy distribution is the uniform distribution. Thus, white noise will have the highest spectral entropy. This has the consequence that power spectral densities have more high frequency information will have lower entropy than that of white noise, provided that there relationship between amplitude and frequency. More concretely, blue noise, which is
23

---

Published as a conference paper at ICLR 2025
Smoothness Metric on Square Wave with increusing Number Waves
Distance
D- Distrnee
01011
Doenjency Morene
Docrspancy Moourz
Dunnce
010022
000000
0 0.0000
0.007
0.000
00074
0.0000
08
80
100
120
1400
160
180
200
30
120 Narcher of Sine Waves
200
Mattber of Sine Waves
Figure 6: Values of the smoothness metric 3.2 on our square-wave-like multinomials as we increase the number of sine waves. We desire the value of this metric to be close to zero when there are few sine waves, and be monotonically increasing with each additional wave, indicating that adding more high frequency content results in a less smooth distribution. On the right, we can see that T,² discrepancy measure leads to non-monotonicity. motivating our choice of T,? distance in measuring results.
defined by the amplitude increasing proportionally to the frequency, will have lower spectral entropy than white noise. We sought metric that always quantified 'sharper' signals like blue noise less smooth. In Table [4] we frame sampled noises different types as multinomial distributions to match our model setting by normalizing their amplitudes to be in [0, 1] and normalizing their 1. Our noise types are defined before normalization, in order of smoothest to sharpest:
1 Brown: S F2
Pink: S x
White: S ~ 10,1
Blue: S F
where is the power density and F the frequency. To obtain samples of each type, we first generate white noise. We do this by sampling Gaussian with mean 0 and standard deviation to obtain amplitudes for samples. We then apply the Fourier transform, and multiply (or divide) the amplitudes of each component by their frequency, and apply the inverse Fourier transform to recover the waveform. Finally we adjust the range of amplitudes of the signal to be within [0, 1] and
normalize the sum to
Discrepancy
Noise
Mean + Std. Deviation
Diff
Delta
Desired Delta
L2
Brown
0.0003 0.0001
n/a
n/a
n/a
1.2
Pink
0.0017 + 0.0002
0.0014
+
L2
White
0.0034 0.0003
0.0016
+
+
L2
Blue
0.0038 0.0003
0.0005
+
Spectral Entropy
Brown
0.4516 0.0894
n/a
n/a
n/a
Spectral Entropy
Pink
0.3878 0.0603
-0.0638
+
Spectral Entropy White
0.4266 + 0.0614
0.0388
+
Spectral Entropy
Blue
0.4191 + 0.0583
-0.0076
+
Table 4: Smoothness measurements for four types of noise bootstrap aggregated over 1,000 trials. The color red emphasizes how the value of Spectral Entropy undesirably not monotonic increasing for what we consider increasingly "sharp" noise types.
24

---

Published as a conference paper at ICLR 2025
C ADDITIONAL EXPERIMENT DETAILS, TOY EXAMPLES
C.I MOTIVATING EXAMPLE: AUDIO SPECTROGRAM TRANSFORMER
To illustrate simple problem setting where the design of the Fourier head is appropriate. it as a drop-in replacement for a linear classification head the Audio Spectrogram Transformer (Gong al. 2021). We consider the task of beats per minute (BPM) classification for metronome- like audio samples (Wei et al. 2024) within the tempo range (50,51. 51,...... 210}. While this task is can difficult, use this audio classification task to illustrate some of the design choices make when using the Fourier head. In this case, it is natural to group the BPMs into contiguous bins [50, 54]. [55, 59..... and use the Fourier head classify them. These bins have a natural contin- uous structure, which is where the Fourier head performs well. We also expect that the categorical distribution over possible BPMs for given audio clip ought to unimodal and therefore require few frequencies to approximate. In fact, our best performing model for this example uses only one frequency.
We initialize the Audio Spectrogram Transformer with pretrained weights from AudioSet (Gem- meke et al. 2017), and we train two different models-one with standard linear classification head, and one with the Fourier head. The Fourier head outperforms the linear classification head by an Fl score improvement of +118%. We attribute this success to the inductive bias continuity that the Fourier head imparts. In Figure 7 present the learned probability masses of both heads the same input sample. This graph illustrates that the Fourier head learns smoother PMFs than the linear head, a concept which will later formalize and explore.
Audio Classification Task: Learned Linear vs. Fourier PMFs
True Label
Truc Label Predicted Pourier PMF Smaothness
Predieted Linear PMF
Mass
BPM
BPM
Figure 7: Comparison between the PMF learned by the linear head, and the Fourier head with 2 frequencies, for the toy BPM classification task, on a single audio example. We observe that the Fourier head learns smoother categorical distribution over its predicted values, and is better centered around the ground truth label. We also note the small mini-sine wave artifacting on the left side of the Fourier model, which tends to occur when using few frequencies.
C.2 LEARNING CONTINUOUS DENSITY
Here provide full details of the datasets used in our toy example of learning a known conditional distribution.
Dataset: We create synthetic dataset D ((g(x)) ((v). R follows. Fix probability distribution Pi = (x) that is parameterized by one variable and second distribution Pa Pa(2,y) parameterized by two variables. Fix an interval I R. Sample uniformly from I. sample ~ Pi (z:), and finally sample ~ P2( y). We can repeat this sampling procedure N times obtain set of N triples for which we know the conditional distribution of given x and Finally, we quantize this set to a fixed number of uniformly spaced bins in the range [-1, 1] to 50 bins
obtain the dataset Dp.,Pu. We will denote the quantization of by q(2). We quantize and our dataset has size 5000, with 80-20 split between the train and test set. We describe three choices for the distributions we used to create our datasets. We fix I 1-0.8, 0.8] and o² 0.01 in all of them.
1. Gaussian dataset: Pi (2x) (x, o²), and P2(x, y) N(y,a2)
2. GMM-2 dataset: P1 = Uniform[]), and Pz(x, y) is a GMM centered at x and y with variance o².
25

---

Published as a conference paper at ICLR 2025
3. Beta dataset: Pi(x) Mxxx, and P2(x.y) U({ x Beta(100) |æ| 100 (y/),
where U({+1}) denotes the Rademacher distribution supported on {+1} with probability 1/2 each.
Additional results: In Figure g we present results from training over range of frequencies, and for each frequency we ran experiments with and without Fourier regularization. Table 6 we present results on the MSE metric, that show that the Fourier head outperforms the linear classification head.
Impacct Varying Fourier Frequencies Coussian datasel
Linar harditz
GMM heal
Spodtaes
Foorle: regalattotion
0.00
Fousie with regelarizzzien
0.05
10
12 Progaancics
12
14
Proquencics
Impact of Varying Fourier Frequencies on GiMM2 dataset
0.4
0.06
GMM leal
Spoottnss
Framder ingutaltation
regtlatizziio
0.02
Fourict Proquancics
Founler Proqucncics
Impact Varving Fourier Frequencies dataset
Lince baclit:
0.15
KJ Diarggeze
GMM beal
Sooobass
0.6
Pourier rigalaication
Foonle with regeladzzziinn
10
12
20
12
14
Froguamcics
Froquencics
Figure 8: We study how the quantity Fourier frequencies impacts KL divergence and smoothness for the toy example each dataset. For both KL divergence and smoothness, lower is better. We observe that the Fourier models with and without regularization performed similarly to each other, and outperformed the linear baseline. We also note that the 50% error bars larger for the linear baseline model; this indicates that the Fourier models (both with and without regularization) are in general more stable. This is in contrast to our large scale time series forecasting experiments, where we find that regularization helps; this is likely because those experiments use order of magnitude more frequencies, and their conditional distributions are more complicated. While the GMM head has better KI. divergence on the Gaussian and GMM-2 datasets, which is to be expected, the Fourier model (both with and without regularization) eventually has the best KL divergence on the Beta dataset, since is non-Gaussian Notice also how each of the datasets, the smoothness degrades as frequency increases, in a fashion that follows the asymptotic from our Theorem
C.3 MLE- BASEDD FOURIER HEAD
We carry out experiments in the continuous domain analogous to those we did in the quantized domain from the toy example.
Dataset: We use the same synthetic datasets-Gaussae GMM-2, and Beta-as in the previous sub- section, except we do not quantize the data into bins.
Task: Predict the conditional distribution of given (z,y).
Model architecture: Our model is an MLP with one hidden layer and the final layer is an MLE- based Fourier head which returns the 2NV + learned real coefficients of the Fourier series, mapping I² - 04 + 32 + 2 Alongside the Fourier-MLE model, we consider a baseline where the final layer is Gaussian model mixture whose means and standard deviations are learned using an MLE objective. For the MLE-Fourier model, we sweep over frequencies 2, 2...... regularization 7 E (0, 10-63
.,......5 20 and
26

---

Published as a conference paper at ICLR 2025
Impact of Varying Fourier Frequencies Gaussinn dataset
0.7
KJ Uuveegncce
0.4
Fornee with reghlettoiinn
resubiizziion
0.2
Erequascies
Fousier Frequences
Impact Varying Fourier Frequennciss on GMM-2 dataset
GNN beod
81. 1..00egnec
regalariorion
Founier 2o regulaziaaion
02-
10
15
18
14
16
Fouule: E'equunciess
Foureer Frequenssss
Impact of Varving Fourier Frequenciiss on Beta dataset
0.8
Diviggtte
Fourio, with mpulaaiera
1.0
Fouriar regubeloiionn
0.2
0.6
10
Fourset Propucncice
Finurier Frequenciss
Figure 9: We study how the quantity of Fourier frequencies impacts KL divergence and perplex- ity for the toy example on each dataset for the MLE experiments. For both KL divergence and perplexity, lower is better. We observe that the Fourier models with and without regularization per- formed similarly to each other. While the GMM head has better KL divergence on the Gaussian and GMM-2 datasets, which be expected, the Fourier model (both with and without regularization) has the best KI. divergence on the Beta dataset for sufficiently large Fourier frequencies, since is non-Gaussian.
KL Divergence (L) GMM
Dataset
Linear
Fourier
Gaussian
0.170 + 0.052
0.026 + 0.011
0.116 0.043
GMM-2
0.238 + 0.032
0.030 + 0.006
0.146 = 0.033
Beta
0.234 + 0.032
0.407 = 0.012
0.191 = 0.016
Smoothness (L)
Dataset
Linear
GMM
Fourier
Gaussian
0.116 0.049
0.068
0.012
0.057 + 0.011
GMM-2
0.068 - 0.022
0.043 = 0.009
0.038 + 0.007
Beta
0.127 = 0.044
0.061 0.003
0.076 + 0.021
Table 5: KL divergence and Smoothness for the three classification heads (Linear, GMM, and Fourier) each of the three synthetic datasets (Gaussian, GMM-2, Beta), As expected, the GMM head achieves the best KL divergence on the Gaussian and GMM-2 datasets, as their conditional dis- tributions Gaussian. However, the Fourier head has the best KL divergence the Beta dataset. This demonstrates the flexibility of the Fourier head in modeling non-Gaussian distributions as well.
Model evaluation: We use two metries for evaluation. The first metric the average KI. divergence DKL (P(x, M) (z, y)), where P(x, v) is the fixed conditional distribution of given (z, y) and M (x, y) denotes the predicted probability density function of z. Our second metric is perplexity, which is the exponential of the average negative log likelihood of the test set.
Results: The metrics for the best performing model on each dataset are reported Table Figure I1 presents sample visualizations of the learned conditional distributions alongside the true densities. While, as expected, the GMM-MLE head outperforms the Fourier-MLE head on the Gaussian and GMM-2 datasets due to the Gaussian nature of the datasets, the Fourier-MI. ML.E head outperforms
27

---

Published as a conference paper at ICLR 2025
Toy Example: MSE (4)
Classification Head
Dataset
Pointwise Regression
Linear
GMM
Fourier
Gaussian
0.010 = 0.001
0.013 = 0.001
0.010 0.001
0.012 + 0.001
GMM-2
0.121 + 0.004
0.126 = 0.004
0.120 + 0.004
0.123 + 0.005
Beta
0.275 0.009
0.276 = 0.008
0.273 0.009
0.275 + 0.008
Table 6: We compare the MSE between the linear head, GMM head, and the Fourier head with 12 frequencies and no regularization, for every dataset in the toy example. We also include a Pointwise Regression model baseline, whose base architecture is same as the classification heads, except the last classification layer is replaced with a dense layer having output dimension 1. We train the Pointwise Regression model using MSE. For a given dataset, the MSE values across all of the models is roughly similar. This is because the pointwise regression model tends regress to the mean, as does the expected value of each of the classification heads.
Toy Example: Ground Truth Conditional Distribution vs. Pointwise Regression Output
Gaussas Ditaset
Gaussian Dataset
Carssian Dutaset
Aui
Poutride Argini
Probability Mass
50
GMMM2 Dataset
GMM Dataset
GMM-22 Douseet
0.2
Falitsse Rogoniia
60
50
aD
Beta Dutasct
Betu Dusset
Beta Dauset
Poe aus Eepecan
00
50
0.0
-10
Figure 10: We present some examples of the ground truth conditional distribution versus the point predicted by the Pointwise Regression model. The regression model simply regresses the mean of the conditional distribution. Accordingly, the regression model performs extremely well for the unimodal Gaussian dataset, and it performs poorly for the bimodal datasets GMM-2 and Beta.
the GMM-MLE head on the Beta dataset, highlighting the flexibility of the Fourier- MLE head in learning large variety of probability distributions. In Appendix C.2) present the results of study on the impact of number of frequencies and Fourier regularization in the MIE setting.
KL Divergence (L)
Perplexity (L)
Dataset
GMM-MLE
Mo.
GMM-MLE
Mourier-ML.E
Gaussian
0.012 + 0.002
0.034 = 0.003
0.410 + 0.014
0.422 + 0.019
GMM-2
0.018 0.001
0.072 = 0.005
0.702
0.015
0.740 + 0.014
Beta
0.257 + 0.03
0.130 = 0.005
0.623 + 0.035
0.542 + 0.017
Table 7: We compare metrics between the GMM-MLE head, and the Fourier-MLE head with 12 frequencies and no regularization, for every dataset in our toy example. We aggregate metries over 4 different seeds and report the standard deviation.
C.4 ARE LLMS RANDOM NUMBER GENERATORS?
Dataset: We create training dataset using the prompt template: "The following is a list of normally distributed random numbers in the interval [-1, 1J with mean μ and std a: I1, " and response template: "xz". ", where (H,,a) a) {(-0.55, 0.10), (-0.03, 0.24), (0.42, 0.16), (0.55, 0.10)). and each
28

---

Published as a conference paper at ICLR 2025
Toy Example: Learned Conditional Distribution True Conditional Distribution
Caussian Danset
Goussian Dunaeet
(MM me
25
Tre
Prghabjlity Density
0.0
GMM-2 Dalnest
GMM-22 Dateel
0V9 ML.
Teals VLT
-1.00
-0.35
-0.30
-0.25
0.000
0.23
0.30
0.75
1.00
-10
-0.75
-0.30
-0.25
5000
9.2
9.75
110
Dataset
Beta Donasct
Teaten
-1000
-0.75
-0.30
-0.25
0.00
0.25
0.50
0.75
1.00
-1.00
-0.25
-0.50
-0.25
5.00
925
0.30
0.75
1.0
Figure 11: Comparison between the PDFs learned by the GMM-MLE head and the Fourier-MLEE head for each of the datasets the toy example. While GMM-MLF outperforms Fourier-MLEE on the Gaussian and GMM-2 datasets, Fourier-MLE performs better the Beta dataset.
J; (), We write each number using decimal places. Our training dataset consists of 256 such (prompt, response) pairs, divided evenly among the four distributions.
Model Architecture: We consider three different models: the original Llama-3. 1-8--Instruuc model Dubey et al.] (2024), the original model after LoRA fine-tuning, and the original model where we replace the linear classification head with the Fourier head and perform LoRA fine-tuning. For the Fourier head, we use output dimension of 200 and the original latent space [-1, 1] because of our chosen decimal precision. We conduct LoRA fine-tuning for 16 epochs with a learning rate of 3 10-4 and a linear decay schedule, and a batch size of 64. We release all of our training code on our project page.
Model Evaluation: We compute two metries: the first is Total Variation Distance; we define this to be one half the L distance between ground truth quantized Gaussian histogram, and empirical histogram of samples, quantized into 20 bins. Our second metric is the Quantity of Unique Samples. In Figure 3 we present example histograms for each of the three models that we consider for the median TVD in each class. Those results demonstrate that the Fourier head learns more accurate PMF. And in Figure [2] we demonstrate that the Fourier head model consistently obtains lower TVD and greater diversity samples. We hypothesize that the LoRA-finetueed baseline model has fewer diverse samples because it memorizes training data instead of learning the actual distribution. In contrast, the Fourier head is forced to learn a continuous distribution, and samples directly from that distribution.
Related works which consider LLMs random number generators: Gu et (2024) explores probability distribution sampling in LLMs in the context of behavioral simu lationtion. which demon- strates an application of using L.LMs as random number generators. And Paruchuri et al. (2024) explores the tasks of using LLMS for estimating percentiles, drawing samples, and calculating prob- abilities. Lastly, Hopkins et al. (2023) explores sampling from numerical probability distributions using LLMs, but they only consider single non-uniform density, namely the normal distribution [0.55, 0.2887). They find that LLMs don't perform well at this task, but they don't thoroughly investigate model-basee interventions. In contrast, thoroughly investigate the impact of fine- tuning, and replacing the token classification head one with a better inductive bias for modeling complex probability distributions.
29

---

Published as a conference paper at ICLR 2025
Distance Between Predicted Distribution and True Distribution
Llama- 3.1 in-instrcct
Llama-3. 1-mma---.--B-instrar LoRA
TVD Emor
Llamma-3. 1-8B-instract
Head
0.4
3
5
7
10
12
Fourier Frequencies
Quantity of Samples that are Unique Values
100
Llama-33J 3.1maa-3.11-BBB 30-iinslrcct
Quantity of Samples
Llamaa-3.1 1eaee-3.------steee + LoRA
80
IJaraa333 1.0--..1----imme -SB--instracct LoRA Fourier Head
60
40
2
4
7
9
10
11
12
Fourier Frequencies
We demonstrate that the Fourier head model consistently obtains lower total variation distance, as well as greater diversity of samples. For TVD (top), lower is better, because lower values indicate that the learned distribution closer to the ground truth distribution. And for quantity of samples (bottom), higher is better, because lower values indicate that the LM has just memorizes specific numbers instead performing sampling. We present here the mean values across the four distributions, for all ten seeds. We can see that the Fourier head obtains more diverse samples, and learns distribution closer the ground truth.
Figure 12:
D ADDITIONAL EXPERIMENT DETAILS, LARGE-SCALE EXAMPLES
D.1 DECISION TRANSFORMER
Following the original Decision Transformer implementation, trained on 500k transitions ob- served by a DQN agent during training, for epochs. trained on the same model size as original implementation (a GPT-1 model with approximately 2.01 2M parameters) which takes about hours on a single GPU, We can see that in Figure 13 that the PMFs learned by the Fourier head smoother. In Figure 16 we present results for more Atari games. In Figure 14 present results from an ablation study of the model size. The results demonstrate that, across model sizes, the Deci- sion Transformer with Fourier head is better learning high-quality next action distributions than the Decision Transformer with linear head. And in Figure 15 we present results from an ablation study of the dataset size, which show that the Fourier head obtains larger returns than the Linear classification head across dataset sizes.
D.2 CHRONOS
In Figure 17 present learned next-token PMF from a linear Chronos model, and a next-token PMF from Chronos model which uses the linear head. The Fourier head is about 4x smoother. In Table g we present results from an ablation study on the quantity of Fourier frequencies, choice of regularization, and binning strategy. We followed the original Chronos implementation, keeping all hyperparameters the same. In particular, original implementation (the T5 model with approximately 20M parameters) and this takes about 48 hours on 8 GPUs. See Table for the datasets used to train and evaluate Chronos.
trained for 200k steps, the same model size as the
30

---

Published as a conference paper at ICLR 2025
Decision Transformer Example: Learned Next Action Distributions
Linear head: smoothaess 0,2800
Fourier hend; smoothness 0.2581
Graand
Proboblliry Mass
0.3-
0.3
Probibiliyy
0.2-
0.2
4
Figure 13: We present example next action distributions for single step in the Decision Transformer test split. The Fourier agent with 8 frequencies produces a "clump" of actions that is semantically meaningful. Namely, this agent almost certainly wants to shoot in the down right or right direction, presumably because there is a submarine in that direction. In contrast, the linear agent's next-action distribution doesn't t clearly depict strategy, and incorrectly assigns higher likelihoods to incorrect actions. Because the Fourier head outputs smoother PMF, it learns to concentrate more probability mass near the correct action.
Impact of Model Size on Decision Transformer Returns: Seaquest
Average Normalized Returns
3.5
Fouricr bead Linese head
0.5
Quantity Model Parameters (Millions)
Figure 14: We present an ablation study on the effect of the model size on the relative performance of the Fourier head and the Linear head. The results demonstrate that, across model sizes, the Decision Transformer with Fourier head better at learning high-quality next action distributions than the Decision Transformer with linear head.
31

---

Published as a conference paper at ICLR 2025
Fourier Head Ablation Study: Dataset Size
3.5
Fourier head
Normalized returns
3.0
Linear head
2.5
2.0
1.5
1.0
0.5
5x102
5x103
5x
5x103
5X10
Dataset size
Figure 15: In this ablation study, we analyze whether dataset size has any effect on the relative performance of the Linear head and the Fourier head. Our results show that, across dataset sizes, the Decision Transformer agent with Fourier head achieves larger returns than the linear head on the Seaquest game.
Chronos Time Series Model
MASE (C)
WQL (1)
Smoothness (J)
Linear
0.383
0.750
0.1689 + 0.1087
Fourier-64
0.875
0.798
0.0032
0.0012
Fourier- 128
0.872
0.767
0.0068
1 0.0035
Fourier-256
0.859
0.755
0.0139 1 0.0087
Fourier-550
0.852
0.749
0.0283
0.0224
Fourier-550 (no regularization)
0.861
0.753
0.0286 + 0.0219
Fourier-550 (uniform precision binning)
0.873
0.747
0.0395 + 0.0252
Table 8: We present large-scale experiments on Chronos time series forecasting. Notably, every Fourier model outperforms the linear baseline on MASE and smoothness metrics. We can see that within the Fourier model class, decreasing the number of frequencies lets you trade off the continuity of the learned probability mass functions (smoothness) for the quality of the forecasts (MASE, WQL). In the bottom two rows, we present an ablation for our large-scale experiments on Chronos time series forecasting. The best overall performing Fourier-550 model uses Fourier regularization and mixed precision binning, which are both techniques informed by Fourier analysis. We observe that both of these interventions improve the MASE, but have minimal effect the WQL. Note that the choice of binning strategy doesn't 't affect the performance of the linear baseline.
32

---

Published as a conference paper at ICLR 2025
Impact of Varying Fourier Frequencies on Decision Transformer: BankHHiist
Normalized Returns
1.5
Smoothness
44
Frodeer
Fratle head
-10
10
10
32
30
12
14
15
26
Fourier Frequencies
Fourier Frequencies
Impact of Varying Fourier Frequencies Decision Transformer: DoubleDunk
Normalized Retumns
106
50-
hhess
Fooder lcad Lind lea)
Fonterheal
Lier heartea
Fourier Frequencies
Fourier Frequencies
Impact of Varying Fourier Frequencies on Decision Transformer: Gravitar
Nornalized Returns
Smnothocss
0.6
Fourier Proquencies
Fourier Frequencies
Figure 16: We present empirical results for how the quantity of Fourier frequencies impacts returns and smoothness for additional imitation learning games. For normalized returns, higher is better; for smoothness. lower is better. We can see that for the BankHeist, DoubleDunk, and Gravitar games, the Fourier agent consistently achieves higher normalized returns than the linear baseline agent, while still learning smoother next-action distributions.
Chronos Example: Learned Token Distribution
Lincer
Fourier head;
0.000
0.012
0.010
0.008
Paduahiitt
0.010
Prdbability
0.006
0.005
0.00
0.000
0.13
0.14
0.15
0.16
0.17
0.18
-0.22
-0.21
-0.20
-0.19
-0.15 -0.17
Figure 17: We present the next token value distribution for single forecasted timestep the Tourism Monthly dataset. We observe that the Fourier head's learned conditional distribution is smoother, fitting signal more robustly, whereas the linear head overfits to the noise, and is therefore more jagged. We note that the z-axis represents the bins in the latent space [-1, 1]: the a-axis values for the Fourier head are lower because the linear head uses uniform binning, and the Fourier head uses mixed precision binning.
33

---

Published as a conference paper at ICLR 2025
Fourier Head Ablation Study: Chronos Dataset Size
2.0
Lineur head
18
Fourier head
16
MASE
1.2
1.
0.8
1.1X10
1.1x10
1.1x107
1.1x10
Dataset size
Figure 18: In this ablation study, we analyze whether dataset size has any effect on the relative performance of the linear head and the Fourier head for the probabilistic time series task. Our results show that, across dataset sizes, the Fourier head yields more accurate forecasts than the linear head. For the dataset sizes 1.1 x 10º, 1.1 x 10 and 1.1 X 10², we report the average MASE across four seeds; for the dataset size 1.1 x 10 we report the MASE from Table 3 We generate the plot following (Kaplan et al. (020) and observe a similar power-law scaling behavior for both methods, with the Fourier head consistently outperforming the linear head.
Fourier Head Ablation Study: Chronos Model Size
1.2
Lineur head
Fourier head
MASE
1.0
0.8
1.25M
2.5M
5M
10M
20M
Quantity of Model Parameters
Figure 19: In this ablation study, analyze whether model size has any effect the relative performance of the linear head and the Fourier head for the probabilistic time series forecasting task. Our results show that, across model sizes, the Fourier head yields more accurate forecasts than the linear head. For the model sizes 1.25M, 2.5M, 5M, and 10M, we report the average MASE across three seeds; for the model size 20M we report the MASE from Table We generate the plot following (Kaplan et al./ 2020) and observe a similar law scaling behavior for both methods, with the Fourier head consistently outperforming the linear head.
34

---

Published as a conference paper at ICLR 2025
Table 9: All datasets that are used for our time series forecasting experiments. We built our time series forecasting experiments on top of Chronos (Ansari al. 2024), and this table is mostly copied from their paper. The datasets are partitioned according to how they are used for training and evaluation of models: pretraining-one data only used for training; evaluation data is not used in training models, but only for evaluation (final H observations). All of our evaluation datasets came from the zero-shot evaluation set from Chronos.
Series Length
Prediction
Dataset
Domain Freq. # Series
min
avg
max Length (H)
Pretraining
Brazilian Cities Temperature
nature
M
12
492
757
1320
MexIco City Bikes
transport
IH
494
780
78313
104449
Solar (5 Mm..)
energy
5min
5166
105120
105120
105120
Solar (Hourly)
energy
IH
5166
8760
8760
8760
Spanish Energy and Weather
energy
IH
66
35064
35064
35064
Taxi (Hourly)
transport
IH
2428
734
739
744
USHCN
nature
ID
6090
5906
38653
59283
Weatherbench (Daily)
nature
ID
225280
14609
14609
14610
Weatherbench (Hourly)
nature
IH
225280
350633
350639
350640
Weatherbench (Weekly)
nature
IW
225280
2087
2087
2087
Wiki Daily (TOOK)
web
ID
100000
2741
2741
2741
Wind Farms (Daily)
energy
ID
337
71
354
366
Wind Farms (Hourly)
energy
IH
337
1715
8514
8784
Evaluation
Australian Electricity
energy
30min
5 230736 231052
232272
48
CIF 2016
banking
1M
72
28
98
120
12
Car Parts
retail
IM
2674
51
51
51
12
Hospital
healthcare
IM
767
84
84
84
12
MI (Monthly)
various
IM
617
48
90
150
18
M1 (Quarterly)
various
3M
203
18
48
114
8
MI Yearlly)
various
IY
181
15
24
58
6
M3 (Monthly)
various
IM
1428
66
117
144
M3 (Qaateerly)
various
3M
756
24
48
72
8
M31 Yearly)
various
IY
645
20
28
47
6
M4 (Quarterly)
various
3M
24000
24
100
874
8
M4 (Yearly)
various
IY
23000
37
841
6
M5
retail
ID
30490
124
1562
1969
28
NN5 (Daily)
finance
ID
111
791
791
791
56
NN5 (Weklly)
finance
IW
111
113
113
113
8
Tourism (Monthly)
various
IM
366
91
298
333
24
Tourism Quarterly)
various
IQ
427
30
99
130
8
Tounsm Ycaarly)
various
IY
518
11
24
47
4
Trathe
transport
IH
862
17544
17544
17544
24
Weather
nature
ID
3010
1332
14296
65981
30
35


Reviews:
[{"review_id": "BonP0FB7Oo", "reviewer": "Reviewer_qzaH", "summary": "The paper proposes a Fourier Head based on the Fourier series as a replacement for the usual linear classification head to induce continuous densities across the class IDs. It presents a theoretical analysis of the expressiveness and smoothness trade-off as the number of frequencies increases and empirically shows the advantage of the Fourier Head over the conventional linear head in tasks with output classes with a continuous structure.", "questions": "- Can you provide more details on the Decision Transformer and Chronos experiments? How did you choose the size of the models, and how long to train?\n- Can you show how the benefit of using the Fourier head varies as the model size or amount of training data increases? That is, does the benefit of the Fourier head persist at scale or does it vanish?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The proposed Fourier Head layer is well-motivated in domains where classes have a continuous structure\n- The method is straightforward and clearly explained\n- Visualizations clearly demonstrate the advantage of the Fourier head on toy problems in learing continuous densities\n- Experiments in RL and time-series show the Fourier head can improve performance in non-toy settings", "weaknesses": "- As the paper is focused on improving LLM's ability to model numerical values, there are important related works that explore alternative ways of extracting continuous probabilistic predictions from LLMs over numerical data [1, 2], which are worth discussing. These methods use a hierarchical representation of the numerical values, encouraging nearby values to have similar densities, as they do not correspond to independently predicted classes. These methods therefore do not have the limitations of \"not consider any continuous structure that resides among the tokens\", which the Fourier head claims to address.\n- Similar to methods based on classification over binned values, the Fourier head can only represent a finite range of values. Methods like [1, 2] in principle do not have this issue.\n- The advantage of using the Fourier head seems most significant with small models trained on limited data. At a large scale, the model should be able to learn the continuous structure in the output classes, diminishing the benefit of using the Fourier head. It would be useful to show how the benefit of replacing the linear head with the Fourier scale with training data and model size, such as for Chronos models of different sizes.\n\n[1] Gruver et al. 2023. Large Language Models Are Zero-Shot Time Series Forecasters\n\n[2] Requeima et al. 2024. LLM Processes: Numerical Predictive Distributions Conditioned on Natural Language", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FHbXbrqimB", "reviewer": "Reviewer_7iR1", "summary": "The paper introduces a novel Fourier head for large language models (LLMs), designed to improve the modeling of continuous structures in non-linguistic tokens, such as decision-making sequences in games and time series forecasting", "questions": "I am unclear about the organization of the paper, such as why the related work is placed in the latter half.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Fourier head allows LLMs to better capture continuous structures in non-linguistic tokens, addressing the limitation in traditional models that use softmax over discrete bins.  The authors provide both theoretical justifications and empirical analysis.", "weaknesses": "1. The author posits that Fourier Head can endow the model with a continuity prior, , which can be described as semantic adjacency. However, since LLMs inherently incorporate attention mechanisms that aggregate tokens with higher similarities, the contribution of the Fourier Head seems incremental.\n\n2. Regarding the time series prediction section, the author has employed the T5 architecture, yet the baseline comprises only this architecture, which is overly simplistic. There is a significant body of work on time series LLMs currently, with most eventually employing a linear layer (could also be replaced with a Fourier head), such as TimeLLM GPT4TS[1,2]. I believe the author needs to further supplement the experiments.\n\n3. Additionally, I think the effectiveness of the Fourier Head may stem from its ability to analyze input frequency and amplitude through Fourier series. The author should consider comparing methods that are based on decoupling[3].\n\n[1]Jin M, Wang S, Ma L, et al. Time-llm: Time series forecasting by reprogramming large language models[J]. arXiv preprint arXiv:2310.01728, 2023.\n\n[2]Zhou T, Niu P, Sun L, et al. One fits all: Power general time series analysis by pretrained lm[J]. Advances in neural information processing systems, 2023, 36: 43322-43355.\n\n[3]Cao D, Jia F, Arik S O, et al. TEMPO: Prompt-based Generative Pre-trained Transformer for Time Series Forecasting[C]//The Twelfth International Conference on Learning Representations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "u6GJbrgQo8", "reviewer": "Reviewer_aPiu", "summary": "The authors argue that current methods for parameterizing a discrete distribution over numerical data suffer from ignoring ordinal structure, which should imply that adjacent discrete buckets will have similar density and therefore \"smoothness\" in the probability mass function. To fix this oversight, the authors propose a new parameterization on the coefficients of a Fourier series, leading to a smooth function on the interval [-1,1], which is then quantized. The new parameterization is therefore a drop-in alternative to a uniform bucketing of the interval. The method is evaluated on toy univariate densities as well as on an offline reinforcement learning problem and in time-series forecasting, and the result indicate that using Fourier head leads to lower errors in density estimation and higher returns in reinforcement learning.", "questions": "I included a few questions in my \"Weaknesses\" response. I've also included a few below:\n\n1. How were the frequencies used in the time series experiments chosen? Were they chosen a priori or through a cross-validation procedure? If cross-validation, how were the splits constructed?\n\n2. Why not explore other bases besides the Fourier basis? Is there something intrinsically better about that basis? Alternatively, there are many other parameterizations that would encourage smoothness. For example, one could parameterize only the differences between buckets and regularize these differences to be small. The final probability mass function would be calculated by integrating the differences. Is there a reason to believe a priori that this approach might perform worse?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed method is simple and outlined with clarity in the paper. While the method is not complex, it is relatively novel to my knowledge. The significance is also reasonably large because modeling continuous numerical values using discrete tokens is increasingly popular.", "weaknesses": "To my understanding, the main goal of the paper is to propose a practical method, and given this goal, the empirical evaluation is not very impressive. I'll break this criticism down into a few subcategories\n\n1. Emphasis on smoothness: The authors devote a lot of space and attention to the notion of \"smoothness\", proposing a new metric to measure it and including this metric in all the evaluations. However, from a practical standpoint it's not clear why we should care about smoothness independent of its effect on metrics like MAE or RMSE. In fact, it's possible to contrive examples where we want less smoothness (related to the square wave examples in the appendix), and it's not clear a priori that the marginal distributions for a particular downstream application will be \"smooth\". The \"smoothness\" numbers therefore feel like a distraction from what really matters, which is whether this ordinal inductive bias actually helps the model fit the data distribution. In many cases, the method seems to improve smoothness without affecting reward/loss or vice versa. \n\n2. Limited empirical impact: while Fourier head does seem to yield significant benefits in offline RL, it doesn't seem to have a significant effect on time series modeling. The benefit in terms of MASE and especially in terms of WQL is very marginal, and if I were looking for ways to improve my model, I might not adopt the additional complexity needed for such a small improvement, which is probably on par with tuning other hyperparameters or making small architectural changes. It might be helpful to identify possible explanations for why the effect is relatively minor in time series but more pronounced in offline RL. For example, are the next-action distributions significantly different in their multimodality? It might be much more compelling to replace the time series experiments with additional offline RL experiments if that application happens to be the ideal use case for this method. \n\n3. Limited baselines: Fourier head is only compared to the most naive possible baseline, uniform binning on [-1, 1]. In practice, there are more widely-used alternatives, such as Gaussian mixture models (GMMs) and quantile regression. Both of these techniques have an ordinal bias and should learn solutions that are much more smooth. I don't know if these methods are viewed as out-of-scope in this paper because they are not learned with cross-entropy loss. From one perspective, it might be reasonable to limit the investigation to discrete tokenization methods and discrete loss functions, but it does make the practical impact lower, as it's hard to tell whether this method is actually the best among all simple options or just an improvement upon simple uniform binning. This particular subcategory of criticism feels especially pertinent given the toy experiments in Section 4, where Fourier head is shown to approximate a GMM. It seems reasonable to conclude that in many cases a GMM should also therefore be able to approximate Fourier head. Is the converse not true and how important are the case where GMMs might not be able to match the performance of Fourier head?\n\nBeyond empirical evaluation I think there are also other potential weaknesses:\n\n1. Limited expressiveness: Presumably this method only works for bounded sequences. In the case of RL this might be reasonable if the state and action spaces are constrained. In the case of time series, this limits applications to series without a significant trend component, which would eventually cause the values to exit the range of past observations. \n\n2. Additional hyper-parameters in the form of chosen Fourier series frequencies and regularization strength.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "THlYet6LnH", "reviewer": "Reviewer_WB6H", "summary": "The paper proposes the \"Fourier head\" as an alternative to linear classification heads for tasks where the output can be thought of as a quantization of a continuous space. The Fourier head learns the Fourier coefficients of the target function and then quantizes it instead of learning the quantized values directly. The authors theoretically and empirically demonstrate that varying the number of frequencies in the Fourier head trades off between smoothness and modelling performance. The Fourier head is shown to improve over the baseline linear head on a toy task, an agentic decision-making task, and for time-series modelling.", "questions": "- The justification of the Fourier regularization term as imposing that the coefficients decay as $1/n^2$ is a little strange to me -- this is an asymptotic condition and, in practice, there are a finite number of coefficients, so isn't the condition always vacuously met?\n- For the Decision Transformer, the output space as described in the paper is more naturally a quantization of $S^1 \\sqcup S^1\\sqcup \\\\{0,1\\\\}$ instead of $[-1, 1]$. (Either a shooting direction or a moving direction, each of which takes eight different values arranged on the circle $S^1$. Also two actions without an associated direction.) It would be interesting to see if the Fourier head can be generalized to output spaces that are not naturally interval-shaped.\n- Actually, if I remember correctly, functions can only be approximated by Fourier series if they are periodic, i.e. functions on $S^1$. I suppose this does not affect the toy example and the time-series modelling, since the interval is chosen to be large enough that the distribution is near zero at the boundaries and so is approximately periodic. But I wonder if this is a limitation in other settings.\n- Often, for tasks with continuous-valued target output (e.g. the toy example and time-series example), only a point estimate is necessary, not the full distribution. Hence a good baseline to include for the toy example is an MLP model with only one output dimension (possibly with atan nonlinearity to map outputs to the interval), evaluated on MSE. Likewise for the time-series example, but with MASE.\n\nMinor typos:\n- Line 193: \"hyperparamter\" -> \"hyperparameter\"\n- Line 205: $c_n$ should be $c_k$ \n- Line 244: \"and $D$ be some measure of discrepancy such that $L^2$,\" should be \"such as\"?\n- Line 257: \"Denote by $g_\\sigma(x)$ is\" delete \"is\"\n- Line 515: \"descretized\" -> \"discretized\"", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper provides evidence that the Fourier head is an improvement over the baseline linear head in a wide variety of settings (toy example, agentic decision making, time-series modeling).\n\t- The Fourier head improves both smoothness and accuracy (MLE, MASE, WQL).\n- The exposition of the Fourier head is clear and easy to understand.\n- Various practical details (training objective, hyperparameter choice, regularization, binning strategy) are provided. This is helpful for reproducibility and to those who wish to apply the Fourier head to other tasks.", "weaknesses": "- For regression tasks with continuous-valued target output, it is not clear to me the practical motivation for outputting an entire probability distribution, instead of just a point estimate. Thus one of the main advantages of the Fourier head, that its outputs are more smooth, feels somewhat unmotivated to me. I would like to see more discussion of why exactly the smoothness is beneficial in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)