Fourier Head: Helping Large Language Models Learn Complex Probability Distributions
Under review as a conference paper at ICLR 2025

FOURIER HEAD: HELPING LARGE LANGUAGE MODELS
LEARN COMPLEX PROBABILITY DISTRIBUTIONS

Anonymous authors
Paper under double-blind review

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 by 46% on the Atari Seaquest game, and increases a
state-of-the-art times series foundation model’s forecasting performance by 3.5%
across 20 benchmarks unseen during training.

Fourier Head Learns Higher Quality Densities

Figure 1: We task an MLP with learning to approximate a continuous bimodal density using a
categorical distribution and a cross entropy objective. We observe that a standard linear classification
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 a smoother, more accurate categorical distribution.

1

INTRODUCTION

Human language can be viewed as a discretization for a 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
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
are more essential than ever, for example for arithmetic computations (Liu et al., 2023), decision
making with continuous or discrete actions (Chen et al., 2021), future anticipation and time-series
forecasting (Ansari et al., 2024), or simply drawing random numbers given a probability distribu-
tion (Hopkins et al., 2023).

1

000
001

002
003
004
005
006
007

008
009
010
011
012
013

014
015
016
017
018

019
020
021
022
023
024

025
026
027
028
029
030

031
032
033
034
035
036

037
038
039
040
041

042
043
044
045
046
047

048
049
050
051
052
053

Under review as a conference paper at ICLR 2025

054
055

056
057
058
059
060
061

062
063
064
065
066
067

068
069
070
071
072

073
074
075
076
077
078

079
080
081
082
083
084

085
086
087
088
089
090

091
092
093
094
095

096
097
098
099
100
101

102
103
104
105
106
107

We view the problem of informing LLMs to utilize the continuity prior from the perspective of 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 to 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.
a 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”) to 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 a continuous probability density function, and returns a discrete approx-
imation of it. Intuitively, returning a discretization of a continuous density in this way allows the
classification head to better model the low-frequency signals from the training data, because overfit-
ting to high-frequency noise is explicitly penalized by the Fourier head’s built-in regularization. At
a high level, the Fourier head inputs x ∈ Rn, uses a linear layer to learn the coefficients for a Fourier
series with N frequencies over [−1, 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
m likelihoods as a categorical distribution. The Fourier head is constructed using the Fourier Basis
Density Model (De la Fuente et al., 2024).

Our first contribution is to reveal the underlying principle on the trade-off between the Fourier head’s
expressive power and the “smoothness” of the predicted distributions. We have proven a theorem
which demonstrates a scaling law for the Fourier head. Namely, as we increase the quantity of
Fourier coefficients that the Fourier head learns, the layer is able to model increasingly more com-
plicated distributions; however, the Fourier head will necessarily fit to more high-frequency noise,
thereby outputting categorical distributions which are less smooth.

Our second contribution is to propose a practical implementation of the Fourier head that allows us
to handle sequential prediction tasks by modeling complex multi-modal distributions. Alongside our
implementation, we propose strategies to improve the layer’s performance, including Fourier coef-
ficient norm regularization, weight initialization, and the choice of how many Fourier frequencies
to use. We demonstrate the effectiveness of the Fourier head on two large scale tasks, where intu-
itively a continuity inductive bias over the output dimensions ought to help the model’s generation
performance. In the first task, an offline RL agent which uses a decoder-only transformer to model
the next-action distribution for an Atari game, we improve returns by 46%. And in the second, we
outperform a state-of-the-art time series foundation model on zero-shot forecasting by 3.5% across
a benchmark of 20 datasets unseen during training. We commit to open source our models and code.

2 FOURIER HEAD

2.1 FOURIER HEAD: MOTIVATION

When practitioners apply LLMs to model complex probability distributions, 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:

• 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 an encoder-only transformer architecture.

• 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 learn-
ing a categorical distribution over those bins. They use an encoder-decoder transformer.

In a 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 are close together.

2

Under review as a conference paper at ICLR 2025

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
letting the classification head learn a categorical distribution as the discretization of a contin-
uous learned function from a suitably flexible class. In this paper, we consider the very flexible
class of truncated Fourier series with N frequencies. These are functions of the form

f (x) = a0 +

N
(cid:88)

k=1

(cid:0)ak cos(kπx) + bk sin(kπx)(cid:1).

(2.1)

Fourier series are a classical tool for solving quantitative problems (Stein & Shakarchi, 2003) 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 a 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 ∈ Rn, and outputs a categorical distribution
in Rm. For a high level summary of how it works–the Fourier head inputs x ∈ Rm, uses a 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 at those m bin centerpoints, and returns those
m likelihoods as a categorical distribution. We formally define this layer in Algorithm 1, and we
present a concrete low-dimensional demonstration of the Fourier head in action in Section 2.3. The
Fourier head is constructed using the Fourier Basis Density Model from (De la Fuente et 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 1, and justification for normalizing the
Fourier coefficients by ℜ(c0)) we refer the author to (De la Fuente et al., 2024).

Algorithm 1 Fourier head
Hyperparameters: the input dimension n, output dimension m, number of frequencies N
Initialization: define a linear layer A : Rn → R2(N +1) // maps input to autocorrelation coefficients

Step 1: INPUT x = (x1, . . . , xn) ∈ Rn
Step 2: (α0, β0, . . . , αN , βN ) ← Ax
Step 3: ak ← αk + iβk ∈ C, for every k = 0, . . . , N // compute autocorrelation coefficients
Step 4: ck ← (cid:80)N −k
ℓ=0 aℓa∗
Step 5: p(z) = 1
2 + ℜ
Step 6: ωk ← (−m + 1 + 2k)/m, for every k = 0, . . . , m − 1 // define m bin centerpoints
Step 7: yk ← p(ωk)
Step 8: OUTPUT (y1, . . . ym) ∈ Rm // by design, we know that each yi ≥ 0, and (cid:80)m

ℓ+k ∈ C, for every k = 0, . . . , N // compute Fourier coefficients

, for every k = 0, . . . , m − 1 // evaluate PDF at m bin centerpoints

// define Fourier PDF over [−1, 1]

ck
ℜ(c0) exp(ikπz)

j=0 p(ωj )

(cid:16)(cid:80)N

(cid:80)m−1

k=1

(cid:17)

k=1 yk = 1

2.3 FOURIER HEAD: MOTIVATING EXAMPLE

To illustrate a simple problem setting where the design of the Fourier head is appropriate, we use it as
a drop in replacement for a linear classification head in the Audio Spectrogram Transformer (Gong
et al., 2021). We consider the task of beats per minute (BPM) classification for metronome-like
audio samples within the tempo range {50, 51, . . . , 210}. While this task is not difficult, we use this
audio classification task to illustrate some of the design choices one can 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 to classify them. These bins have a natural continuous structure, which is where
the Fourier head performs well. We also expect that the categorical distribution over possible BPMs
for a given audio clip ought to be unimodal and therefore require few frequencies to approximate.
In fact, our best performing model for this example uses only one frequency.

3

108
109

110
111
112
113
114
115

116
117
118
119
120
121

122
123
124
125
126

127
128
129
130
131
132

133
134
135
136
137
138

139
140
141
142
143
144

145
146
147
148
149

150
151
152
153
154
155

156
157
158
159
160
161

Under review as a conference paper at ICLR 2025

We initialize the Audio Spectrogram Transformer with pretrained weights from AudioSet (Gem-
meke et al., 2017), and we train two different models–one with a standard linear classification head,
and one with the Fourier head. The Fourier head outperforms the linear classification head by an
F1 score improvement of +118%. We attribute this success to inductive bias of continuity that the
Fourier head imparts. In Figure 2 we present the learned probability masses of both heads on the
same input sample. This graph illustrates that the Fourier head learns smoother PMFs than the linear
head, a concept which we will later formalize and explore.

Audio Classification Task: Learned Linear vs. Fourier PMFs

Figure 2: 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 a 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.

2.4 FOURIER HEAD: DETAILS FOR USING IT DURING TRAINING

We highlight the main design choices for a user when applying the Fourier head in practice.
Training objective: The Fourier head inputs a signal x ∈ Rn and extracts from that signal an
intermediate representation of a probability distribution p(x) defined over [−1, 1]. This probability
distribution has a closed formula equal to a Fourier series. In our experiments, we optimize the
parameters 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 a linear
layer, we use softmax cross-entropy loss as the objective.

Choice of hyperparameter N : The Fourier head has one crucial hyperparameter–namely, 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-expressive 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 a learning less smooth densities, as well as more model parameters.

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 overfit to noise in the data, and therefore has a bias towards learning smoother densities. In
the notation from Algorithm 1, this means adding a regularization term of γ · 2π2
k=1 k2|ck|2 to
m
the loss function, where γ is a hyperparameter. When picking regularization strength, we find that
in the low-frequency domain (e.g. frequencies in the single digits) using γ = 0 works best, and in
the high-frequency domain (e.g. greater than 10 frequencies), using γ = 10−6 works best.

(cid:80)m

Binning strategy: The choice of how we bin the data can affect performance significantly. As we
already discussed, we should only apply the Fourier head when nearby bins are “similar” in some
sense. This means we should order our bins in a semantically meaningful ordering. Further, in the
case where the bins represent quantized numerical values over a continuous latent space, it can be
helpful to use a “mixed-precision” binning strategy. For instance, if we want to model all values

4

162
163

164
165
166
167
168
169

170
171
172
173
174
175

176
177
178
179
180

181
182
183
184
185
186

187
188
189
190
191
192

193
194
195
196
197
198

199
200
201
202
203

204
205
206
207
208
209

210
211
212
213
214
215

Under review as a conference paper at ICLR 2025

216
217

218
219
220
221
222
223

224
225
226
227
228
229

230
231
232
233
234

235
236
237
238
239
240

241
242
243
244
245
246

247
248
249
250
251
252

253
254
255
256
257

258
259
260
261
262
263

264
265
266
267
268
269

from [−15, 15], but we find that most values lie in the range [−1, 10], then we should allocate a
higher proportion of bins to the dense data interval. Specifically, if we would like to use m total
bins to quantize the data, then we control the allocation of bins using a hyperparameter d ∈ [0, 1),
where ⌊d · m⌋ uniformly spaced bins are allocated to the sparse data interval while the remaining
m − ⌊d · m⌋ bins are allocated to the dense range (estimated from training data). This is motivated
and supported by the Fourier theory as well, since by increasing precision in the dense data range we
are effectively de-localizing the quantized data distribution, which leads to a more localized Fourier
spectrum. This lets us obtain a quicker decay of higher frequency content, which ensures that we can
more effectively learn the same distribution with lower-frequency Fourier heads. We also note that
(De la Fuente et al., 2024) proposes an optional re-parameterization step that replaces the periodic
domain [−1, 1] with the real line, although we don’t use that in this work.

Weight initialization: The learned parameters for the Fourier head consist of the learned linear layer
which extracts autocorrelation parameters. In PyTorch, the linear layers uses the He initialization
(He et al., 2015) by default, which ensures that the linear layer outputs values close to zero in
expectation. Similarly, it’s better for the learning dynamics for the Fourier densities to be initialized
to uniform p(z) ≈ 1/2. 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 THEORY

3.1

“SMOOTHNESS”: A METRIC FOR HIGH FREQUENCY CONTENT

In this subsection, we propose a smoothness metric which inputs a categorical disribution y =
(y1, . . . , ym) ∈ Rm, and assigns a numerical value depending on how smooth it is. The score will
output 0 if y is the smoothest possible categorical distribution, and larger values if y is less smooth.
We will first specify what we mean by “smooth”:
Heuristic 3.1. We say a function is smooth if it 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, a categorical
distribution containing samples from sin(100πx) contains lots of high frequency information, so it
should get a smoothness score greater than 0. 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
ασ ∈ R be weights satisfying (cid:82) ∞
0 ασdσ = 1, and D be some measure of discrepancy such as L2,
and let gσ(x) ∗ f (x) denote the convolution of f (x) with a Gaussian kernel of standard deviation σ,
then it is reasonable to define the smoothness of f to be the quantity

s(f ) :=

(cid:90) ∞

(cid:90) b

0

a

ασD[f (x), gσ(x) ∗ f (x)]dxdσ.

(3.1)

In this expression, the discrepancy D[f (x), gσ(x) ∗ f (x)] measures how different f (x) is from a
Gaussian-smoothed version of itself. Because the Gaussian is a low-pass filter, we can interpret
Equation 3.1 as 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.

In our experiments, we consider discrete categorical distributions, and wish to evaluate how smooth
they are in a numerically tractable way. Accordingly, we define a specific case of this as follows.
Definition 3.2 (Smoothness metric for categorical distributions). Suppose y = (y1, . . . , ym) ∈ Rn
is a categorical distribution, so every yk ≥ 0 and (cid:80)m
k=1 yk = 1. Denote by gσ(x) the discrete
Gaussian kernel with zero-padding of standard deviation σ. Define the weights ασ = 6/π2σ2. Then
we define the smoothness of y to be the constant

s(y) :=

∞
(cid:88)

σ=1

ασ∥f (x) − (gσ ∗ f )(x)∥2

(3.2)

5

Under review as a conference paper at ICLR 2025

270
271

272
273
274
275
276
277

278
279
280
281
282
283

284
285
286
287
288

289
290
291
292
293
294

295
296
297
298
299
300

301
302
303
304
305
306

307
308
309
310
311

312
313
314
315
316
317

318
319
320
321
322
323

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 A SCALING LAW FOR THE FOURIER HEAD, IN 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 as the number of frequencies increases. On one hand, it
is a celebrated result from Fourier analysis that a Fourier series with a greater number of frequencies
models a larger class of functions; but on the other hand, we show that increasing frequencies also
incurs loss in smoothness. This is to be expected, as we designed our smoothness metric with the
intention of identifying a distribution as less smooth if it contains more high-frequency information.
Theorem 3.3. (Fourier head scaling law.) Consider a Fourier head with input dimension n, output
dimension m, and N frequencies. Suppose that 1 ≪ N < m

2 . Then the following are true:

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 ∈ Rn,
and denote by fx : [−1, 1] → R the optimal conditional distribution that we would like
the Fourier head to approximate for this input. We assume that fx is twice continuously
differentiable, and that its Fourier coefficients decay on the order of 1/k2. Denote by fx,N
the truncation of fx to its first N frequencies, and denote by y(N ) a discretized version of
fx,N into m bins. Then, there exist constants C1, C2 > 0 such that

s(y(N )) ≈

√

mC1 −

√
mC2
N 3 + O(1/N 4).

(3.3)

Note that the smoothness scaling law asymptotic in Equation 3.3 shows that as N increases, so
does s(y(N )). In part (2), since fx is at least twice continuously differentiable, we already know
its Fourier coefficients corresponding to the k-th frequency are in O(1/k2) (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 toward ensuring that this condition is
met in practice as well. We include a full proof of this result in Appendix A.

4 TOY EXAMPLE: 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 q(z) 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 {(q(x), q(y), q(z))} ⊆ [ − 1, 1]3. Crucially, z is sampled
from a distribution which is conditioned on x and y, and we have an explicit closed formula for this
distribution. By design, the Gaussian dataset is unimodal in z, whereas the more challenging GMM-
2 and Beta datasets are not unimodal. Full details about the datasets can be found in Appendix C.

Task: Predict the conditional distribution of q(z) given the quantized tuple (q(x), q(y)).

Model architecture: Our model is an MLP with ReLU activations and one hidden layer, which
maps R2 → R64 → R32 → R50. The output of the model has dimension 50 because we quantize
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 is 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 γ ∈ {0, 10−6}. We train those models
via cross entropy loss. We also consider a regression-based model, trained using MSE.

Model evaluation: We use three metrics for evaluation. Our first metric is the average KL di-
vergence DKL(q(P(x, y))||M (q(x), q(y))), where P(x, y) is the fixed conditional distribution of z

6

Under review as a conference paper at ICLR 2025

given (x, y); q(P(x, y)) is the quantized approximation of P(x, y), obtained by evaluating the den-
sity function of P(x, y) at the bin centers, multiplying by the bin width, and finally scaling by the
sum of the likelihoods; and M (q(x), q(y)) denotes the predicted categorical conditional distribution
of q(z). Our second metric is smoothness. And our third metric is MSE, where we consider the
expected value of q(z) under the learned categorical distribution as a prediction for q(z).

Results: The metrics for the best performing model on each dataset are reported in Table 1. Figure
3 presents sample visualizations of the learned conditional distributions alongside the true densities.
And in Appendix C, we present the results of 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 3.3, as it demonstrates that for all datasets, as frequency decreases, the smoothness
degrades, and model performance improves until it reaches a 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 on MSE as well; we include a complete
comparison with both Linear and GMM head baselines in Appendix C.

Figure 3: Comparison between the PMFs learned by the linear head, GMM head, and the Fourier
head, for each of the datasets in the toy example. We observe that the Fourier head learns a 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.

KL Divergence (↓)

Smoothness (↓)

Dataset
Gaussian
GMM-2
Beta

Linear
0.170 ± 0.052
0.238 ± 0.032
0.234 ± 0.032

Fourier
0.116 ± 0.043
0.146 ± 0.033
0.191 ± 0.016

Linear
0.116 ± 0.049
0.068 ± 0.022
0.127 ± 0.044

Fourier
0.057 ± 0.011
0.038 ± 0.007
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 our 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 4 different
seeds and report the standard deviation.

5 LARGE-SCALE STUDY: OFFLINE REINFORCEMENT LEARNING

The Decision Transformer (Chen et al., 2021) casts the problem of reinforcement learning as se-
quentially modeling rewards, states, and actions. Here, we study the performance of the Decision
Transformer on the Seaquest game in the Atari (Bellemare et al., 2013) benchmark. The Seaquest
game contains 18 actions, with two groups of eight actions that have a natural “closeness” metric
defined on them: move left, up left, up, up right, right, down right, down, down left; as well as
shooting in those eight directions. In their architecture, a decoder-only language model (Radford

7

324
325

326
327
328
329
330
331

332
333
334
335
336
337

338
339
340
341
342

343
344
345
346
347
348

349
350
351
352
353
354

355
356
357
358
359
360

361
362
363
364
365

366
367
368
369
370
371

372
373
374
375
376
377

Under review as a conference paper at ICLR 2025

et al., 2018) encodes the context and then maps it through a linear layer, outputting a categorical
distribution over the 18 possible actions. At test time, the agent chooses its next action by sampling
from the learned next-action distribution. In our study, we replace that linear classification head with
a Fourier head. Intuitively, this ought to give the model the prior that actions like “move left” and
“move up left” are semantically similar, and therefore should have similar likelihoods. Our study
confirms that the Fourier head outperforms the linear head in returns obtained by as much as 46%,
in the reward conditioned setting considered in the paper, using identical training hyperparameters.

Task: In the Seaquest game, the agent moves a submarine to avoid enemies, shoot at enemies,
and rescue divers. The Seaquest game contains 18 actions: move left, up left, up, up right, right,
down right, down, down left; as well as shooting in those eight directions; as well as no move, and a
generic fire move. 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
et al., 2021). This dataset consists of 500k transitions experienced by an online deep Q-network
agent (Mnih et al., 2015) during training on each of the games.

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 of dimension 18, and the model
ultimately optimizes the cross entropy 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 a Fourier head.

Normalized Returns for Decision Transformer Agent

Atari Game

Classification Head
Linear head
Fourier head

BankHeist

DoubleDunk
−0.09 ± 0.05 −72.72 ± 33.08
45.45 ± 36.36
0.92 ± 0.33

Gravitar
1.32 ± 0.17
4.98 ± 0.93

Seaquest
2.53 ± 0.63
3.70 ± 0.47

Table 2: 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.

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 is better; for smooth-
ness, lower is better. We can see that the Fourier agent achieves higher normalized returns than
the linear baseline agent when sufficiently many Fourier frequencies are used, while still learning
smoother next-action distributions.

Model evaluation: We present results for the linear baseline, as well as the Fourier-n head, for
n = {2, 4, 6, 8, . . . , 30, 32}, across the four Atari games. We present mean reward totals for rollouts
for each of these for the best epoch across 4 seeds. In Table 2, our results demonstrate that the
Fourier head increases agent returns significantly. For example, for the Seaquest game, normalized
returns increase by as much as 46.2%, and for Gravitar, normalized returns increase by as much as
300%. In Figure 4, we can see that the Fourier head is able to perform better in the Seaquest game

8

378
379

380
381
382
383
384
385

386
387
388
389
390
391

392
393
394
395
396

397
398
399
400
401
402

403
404
405
406
407
408

409
410
411
412
413
414

415
416
417
418
419

420
421
422
423
424
425

426
427
428
429
430
431

Under review as a conference paper at ICLR 2025

432
433

434
435
436
437
438
439

440
441
442
443
444
445

446
447
448
449
450

451
452
453
454
455
456

457
458
459
460
461
462

463
464
465
466
467
468

469
470
471
472
473

474
475
476
477
478
479

480
481
482
483
484
485

as the number of frequencies grows. We can also see that as we increase the quantity of frequencies,
learned PMFs become less smooth, in accordance with Theorem 3.3. Qualitatively, we can also
see that in Figure 12 (Appendix) the PMFs learned by the Fourier head are smoother. In Figure 13
(Appendix), we also include results for the remaining three games, BankHeist, DoubleDunk, and
Gravitar. Across these games, the results show that the Fourier agent consistently achieves higher
normalized returns than the linear baseline agent, while still learning smoother next-action distribu-
tions. And in Figure 9 (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.

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 10 (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 a
similar result. In Figure 11 (Appendix) we show that, across dataset sizes, the Decision Transformer
agent with the Fourier head achieves larger returns than the agent with a 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 series”.
They do this by approaching time series forecasting as language modeling–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 a linear map which learns a categorical distribution over next tokens. In
particular, we observe that token i represents a number very close to tokens i−1 and i+1. However,
we note that there is no inductive bias in the T5 architecture which pushes their likelihoods to be
similar. This is not a hypothetical problem; in Figure 14 (Appendix), we can see that the linear
next-token prediction PMFs fit to the noise, and appear very jagged. Here, the motivation for
replacing the linear head with the Fourier head is to “smooth” out this distribution, to help
the forecasting model better learn the signal, and ignore the noise. In Figure 14, we can see that
the Fourier head accomplishes this successfully.

In this section, we study how 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 outperform the Chronos linear baseline on the MASE metric, while
learning next token multinomials which are 8x smoother, with fewer parameters than the baseline.

Dataset: We use the same training dataset for large-scale pretraining that Ansari et al. (2024) used.
We gather an evaluation benchmark of 20 time series datasets which were not seen during training.
These 20 come from the zero-shot eval from (Ansari et al., 2024). The reader can check Appendix E
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
et al., 2020). The original model has a linear classification head. For our study, we will replace
this with a 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.4). We also use Fourier weight decay regularization. For the task, the model learns to
input time series context of length 512, and output a 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 metrics: model performance from (Ansari et al., 2024)
(MASE measures the accuracy of median forecast, and WQL measures the quality of the proba-
bilistic forecast), as well as our smoothness metric. Our Fourier metrics in Table 3 demonstrate 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 as well.

Ablations: The results in Table 7 (Appendix) show that mixed precision binning and regularization
improve the MASE and smoothness for the Fourier head. And in Figure 15 (Appendix) we demon-
strate that, across dataset sizes, the Fourier head yields more accurate forecasts than the linear head.

9

Under review as a conference paper at ICLR 2025

486
487

488
489
490
491
492
493

494
495
496
497
498
499

500
501
502
503
504

505
506
507
508
509
510

511
512
513
514
515
516

517
518
519
520
521
522

523
524
525
526
527

528
529
530
531
532
533

534
535
536
537
538
539

Chronos Time Series Model MASE ↓ WQL ↓
0.750
Linear
0.798
Fourier-64
0.767
Fourier-128
0.755
Fourier-256
0.749
Fourier-550

0.883
0.875
0.872
0.859
0.852

Smoothness ↓
0.1236 ± 0.0712
0.0027 ± 0.0012
0.0053 ± 0.0030
0.0101 ± 0.0072
0.0203 ± 0.0176

Table 3: We present large-scale experiments on Chronos time series forecasting. Notably, every
Fourier model outperforms the linear baseline on MASE and smoothness metrics. 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).

7 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 et al., 2024; Jin et al.,
2024; Nate Gruver & Wilson, 2023; Requeima et al., 2024; Jia et al., 2024; Zhou et al., 2023; Wang
et al., 2024), music generation (Dhariwal et al., 2020; Agostinelli et al., 2023; Copet et al., 2023;
Yuan et al., 2024), and as well as 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-pretrained
LLM is fine-tuned on a domain-specific 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 LLMs
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 to leverage LLMs in numer-
ical modeling contexts (Song et al., 2024). While even frozen LLMs with no further training show
interesting empirical results as regressors (Vacareanu et al., 2024), there is a conceptual mismatch
between the downstream task and model construction 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 or a deterministic transformation within the network, or use Fourier analysis to
motivate design choices. It is far less common to learn the Fourier series directly. De la Fuente
et al. (2024) learned marginal univariate densities parameterized using a Fourier basis; our work
extends their Fourier Basis Density model to multivariate settings with an autoregressive scheme.
Our method learns conditional univariate densities using a Fourier basis, where the coefficients of
the Fourier density model are input dependent. Sitzmann et al. (2020) proposed sinusoidal activation
functions, which can be seen as learning the frequencies of a Fourier series; in contrast, we seek to
fix 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.

8 CONCLUSION

We propose the Fourier head and demonstrate its positive impact on performance on several tasks.
We prove scaling laws that characterize 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 LLM training, where the head can be adaptively employed when needed.

10

Under review as a conference paper at ICLR 2025

540
541

542
543
544
545
546
547

548
549
550
551
552
553

554
555
556
557
558

559
560
561
562
563
564

565
566
567
568
569
570

571
572
573
574
575
576

577
578
579
580
581

582
583
584
585
586
587

588
589
590
591
592
593

9 REPRODUCIBILITY STATEMENT

We have made efforts to ensure reproducibility. In Algorithm 1 we provide all the mathematical
details that one needs to reproduce the Fourier head.
In Appendix E we prove our scaling law,
Theorem 3.3, in full detail, and we list all assumptions in the statement of the theorem. Additionally,
we release the research code in the supplemental section on OpenReview.

REFERENCES

Andrea Agostinelli, Timo I Denk, Zal´an 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.11325, 2023.

Abdul Fatir Ansari, Lorenzo Stella, Caner Turkmen, Xiyuan Zhang, Pedro Mercado, Huibin Shen,
Oleksandr Shchur, Syama Sundar Rangapuram, Sebastian Pineda Arango, Shubham Kapoor, et al.
Chronos: Learning the language of time series. arXiv preprint arXiv:2403.07815, 2024.

Marc 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.

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 arXiv:2106.01345, 2021.

Jade Copet, Felix Kreuk, Itai Gat, Tal Remez, David Kant, Gabriel Synnaeve, Yossi Adi, and Alexan-
dre D´efossez. Simple and controllable music generation. In Thirty-seventh 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´e. Fourier basis density model. arXiv

preprint arXiv: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.

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-labeled 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. IEEE/ACM Transactions on Audio, Speech, and Language
Processing, 2021. doi: 10.1109/TASLP.2021.3120633.

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 llms generate random numbers? evaluat-
ing llm 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. Electroencephalography and Clinical Neurophysiology, 79(3):204–210,
ISSN 0013-4694. doi: https://doi.org/10.1016/0013-4694(91)90138-T. URL https:
1991.
//www.sciencedirect.com/science/article/pii/001346949190138T.

Furong Jia, Kevin Wang, Yixiang Zheng, Defu Cao, and Yan Liu. Gpt4mts: Prompt-based large
language model for multimodal time-series forecasting. Proceedings of the AAAI Conference on
Artificial Intelligence, 38(21):23343–23351, Mar. 2024. doi: 10.1609/aaai.v38i21.30383. URL
https://ojs.aaai.org/index.php/AAAI/article/view/30383.

11

Under review as a conference paper at ICLR 2025

594
595

596
597
598
599
600
601

602
603
604
605
606
607

608
609
610
611
612

613
614
615
616
617
618

619
620
621
622
623
624

625
626
627
628
629
630

631
632
633
634
635

636
637
638
639
640
641

642
643
644
645
646
647

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.08361, 2020.

Shuang Li, Xavier Puig, Chris Paxton, Yilun Du, Clinton Wang, Linxi Fan, Tao Chen, De-An Huang,
Ekin Aky¨urek, Anima Anandkumar, et al. Pre-trained language models for interactive decision-
making. Advances in Neural Information Processing Systems, 35:31199–31212, 2022.

Yixin Liu, Avi Singh, C Daniel Freeman, John D Co-Reyes, and Peter J Liu.

Improving large
language model fine-tuning for solving math problems. arXiv preprint arXiv:2310.10047, 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.1038/
s41587-022-01618-2.

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, Marc G Belle-
mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level
control through deep reinforcement learning. nature, 518(7540):529–533, 2015.

Shikai Qiu Nate Gruver, Marc Finzi and Andrew Gordon Wilson. Large Language Models Are Zero
Shot Time Series Forecasters. In Advances in Neural Information Processing Systems, 2023.

Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under-

standing by generative pre-training. OpenAI 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/v1/2022.findings-emnlp.59.

James Requeima, John Bronskill, Dami Choi, Richard E Turner, and David Duvenaud. Llm
processes: Numerical predictive distributions conditioned on natural language. arXiv preprint
arXiv:2405.12856, 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.14547, 2024. doi: 10.
48550/ARXIV.2402.14547. URL https://doi.org/10.48550/arXiv.2402.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.

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 on Language Modeling, 2024. URL https://openreview.net/forum?
id=LzpaUxcNFK.

12

Under review as a conference paper at ICLR 2025

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.

Eric W. Weisstein. Square wave. From MathWorld–A Wolfram Web Resource, 2024. URL https:

//mathworld.wolfram.com/SquareWave.html. Accessed: September 16, 2024.

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 llm. 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 lm. In NeurIPS, 2023.

13

648
649

650
651
652
653
654
655

656
657
658
659
660
661

662
663
664
665
666

667
668
669
670
671
672

673
674
675
676
677
678

679
680
681
682
683
684

685
686
687
688
689

690
691
692
693
694
695

696
697
698
699
700
701

Under review as a conference paper at ICLR 2025

A MATHEMATICAL DETAILS

In this section we prove Theorem 3.3, the Fourier head scaling law. To do this, we must first discuss
the Nyquist-Shannon Sampling Theorem. This result states that in order to avoid distortion of a
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(πN x) is N/2. Thus the Nyquist
Theorem requires us to have

m/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.

We now present a theorem that provides a scaling law for the Fourier head. This result quantifies the
trade-off between modeling capacity and smoothness as the number of frequencies increases. In or-
der to prove this, we assume that the underlying function being learned is at least twice continuously
differentiable, which implies that the Fourier coefficients corresponding to the k-th frequency of fx
are in O(1/k2) (Stein & Shakarchi, 2003, Ch.2, Cor. 2.4). Thus, our assumption that the Fourier
coefficients decay quadratically is reasonable, and our Fourier quadratic weight decay regularization
helps ensure that this condition is met in practice as well.
Theorem 3.3. (Fourier head scaling law.) Consider a Fourier head with input dimension n, output
dimension m, and N frequencies. Suppose that 1 ≪ N < m

2 . Then the following are true:

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 ∈ Rn,
and denote by fx : [−1, 1] → R the optimal conditional distribution that we would like
the Fourier head to approximate for this input. We assume that fx is twice continuously
differentiable, and that its Fourier coefficients decay on the order of 1/k2. Denote by fx,N
the truncation of fx to its first N frequencies, and denote by y(N ) a discretized version of
fx,N into m bins. Then, there exist constants C1, C2 > 0 such that
√
mC2
N 3 + O(1/N 4).

s(y(N )) ≈

mC1 −

(3.3)

√

Proof. We first prove part (2). Let bj = −1 + 2j+1
m , 0 ≤ j < m be the center points of the m bins
in (−1, 1). Let aj(x) ∈ C, −N ≤ j ≤ N denote the the Fourier coefficients of fx corresponding to
the first N frequencies. In other words,

fx,N (y) =

N
(cid:88)

aj(x)eπijy

(A.1)

j=−N
is the function the Fourier head is learning. By our assumption, there exists a constant cx such that
|ak(x)| ≈ cx/k2. We want to study

s(y(N )) =

∞
(cid:88)

σ=1

ασ





m−1
(cid:88)

|(fx,N − gσ ∗ fx,N )(bj)|2



.

(A.2)



1/2

j=0

Let dj(x) be the Discrete Fourier transform of (fx,N − gσ ∗ fx,N )(bj) for 0 ≤ j < m. By Parseval’s
Theorem for the DFT, we have

m−1
(cid:88)

j=0

|(fx,N − gσ ∗ fx,N )(bj)|2 =

1
m

m−1
(cid:88)

k=0

|dk(x)|2 .

(A.3)

Since fx is supported only on (−1, 1), the convolution gσ ∗ fx,N is the same as (gσI(−1,1)) ∗ fx,N ,
where I(−1,1) is the indicator function for (−1, 1). Treating gσI(−1,1) as a function on (−1, 1), let
hσ,n be the Fourier coefficients of gσI(−1,1). Note that by the Convolution Theorem, we have

(fx,N − gσ ∗ fx,N )(bn) =

N
(cid:88)

j=−N

14

aj(x)(1 − hσ,j) · eπijbn .

(A.4)

702
703

704
705
706
707
708
709

710
711
712
713
714
715

716
717
718
719
720

721
722
723
724
725
726

727
728
729
730
731
732

733
734
735
736
737
738

739
740
741
742
743

744
745
746
747
748
749

750
751
752
753
754
755

(A.5)

(A.6)

(A.7)

(A.8)

(A.9)

(A.10)

(A.11)

Under review as a conference paper at ICLR 2025

Thus, using the definition of DFT along with Equation A.4, we get

dk(x) =

m−1
(cid:88)

(fx,N − gσ ∗ fx,N )(bn) · e−2πikn/m

n=0

m−1
(cid:88)

N
(cid:88)

n=0

j=−N

aj(x)(1 − hσ,j)eπijbn · e−2πikn/m

N
(cid:88)

j=−N

N
(cid:88)

j=−N

aj(x)(1 − hσ,j)

m−1
(cid:88)

n=0

eπij(−1+ 2n+1

m ) · e−2πikn/m

aj(x)(1 − hσ,j)eπij(1−1/m)

m−1
(cid:88)

n=0

e2πi(j−k)n/m.

=

=

=

m−1
(cid:88)

n=0

e2πi(j−k)n/m =

(cid:26)0

if j ̸≡ k (mod m)

m else.

Note that

We therefore have



dk(x) =

m · ak(x)(1 − hσ,k)eπik(1−1/m)
m · ak−m(1 − hσ,k−m)eπi(k−m)(1−1/m)
0



if 0 ≤ k ≤ N
if m − N ≤ k ≤ m − 1
otherwise.

Using this in A.3, we obtain

m−1
(cid:88)

j=0

|(fx,N − gσ ∗ fx,N )(bj)|2 =

1
m

N
(cid:88)

k=0

+

1
m

(cid:12)
(cid:12)

(cid:12)m · ak(x)(1 − hσ,k)eπik(1−1/m)(cid:12)

2
(cid:12)
(cid:12)

m−1
(cid:88)

k=m−N

(cid:12)
(cid:12)

(cid:12)m · ak−m(1 − hσ,k−m)eπi(k−m)(1−1/m)(cid:12)

(cid:12)
(cid:12)

2

= m

N
(cid:88)

k=0

|ak(x)(1 − hσ,k)|2 + m

(A.12)

m−1
(cid:88)

k=m−N

|ak−m(1 − hσ,k−m)|2 ,

(A.13)

where in the last step we used that (cid:12)
(cid:12) since they are both
complex exponentials. Now, since gσI(−1,1) is a real and even function, we know that hσ,k is
real. Further, since the truncated Gaussian gσI(−1,1) is infinitely differentiable, we also know that
hσ,k = O(1/k2). Thus, using that |ak(x)| ∼ cx/k2, we see

(cid:12)eπi(k−m)(1−1/m)(cid:12)

(cid:12)eπik(1−1/m)(cid:12)

(cid:12) = 1 = (cid:12)

|ak(x)(1 − hσ,k)|2 = |ak(x)|2 (1 − 2hσ,k(x) + h2

σ,k) ∼

c2
x
k4 + O(1/k6) + O(1/k8).

(A.14)

From A.14, it is clear that since we are interested in only the dominant asymptotic, we can safely
ignore the higher order terms coming from the hσ,k. As a result,

m−1
(cid:88)

j=0

|(fx,N − gσ ∗ fx,N )(bj)|2 ≈ ma0(x)2 + m

= ma0(x)2 + m

= ma0(x)2 + m

N
(cid:88)

k=1

N
(cid:88)

k=1

N
(cid:88)

k=1

c2
x
k4 + m

m−1
(cid:88)

k=m−N

c2
x
(k − m)4

c2
x
k4 + m

−1
(cid:88)

k=−N

c2
x
k4

c2
x
k4 + m

N
(cid:88)

k=1

c2
x
k4

= ma0(x)2 + 2m

N
(cid:88)

k=1

c2
x
k4 .

15

(A.15)

(A.16)

(A.17)

(A.18)

756
757

758
759
760
761
762
763

764
765
766
767
768
769

770
771
772
773
774

775
776
777
778
779
780

781
782
783
784
785
786

787
788
789
790
791
792

793
794
795
796
797

798
799
800
801
802
803

804
805
806
807
808
809

Under review as a conference paper at ICLR 2025

We can approximate the dominant terms of the sum using an integral:

N
(cid:88)

k=1

1
k4 =

(cid:90) N

1

1
x4 dx +

1
2

(cid:18)

1 +

(cid:19)

1
N 4

+ O(1/N 4) =

(cid:18)

1
3

1 −

(cid:18)

(cid:19)

1
N 3

+

1
2

1 +

(cid:19)

1
N 4

+ O(1/N 4).

Substituting the estimate into A.18, we obtain

m−1
(cid:88)

|(fx,N − gσ ∗ fx,N )(bj)|2 ≈ m

(cid:18)

C 2(x) −

2c2
x
3N 3 + O(1/N 4)

(cid:19)

,

j=0
(cid:113)

where C(x) =

a0(x)2 + 5

3 c2

x is a constant depending upon x.

Using the Taylor expansion (1 + x)1/2 = 1 + x

(cid:18)

mC 2(x) −

2m
3N 3 + mO(1/N 4)

(cid:19)1/2

≈

2 + O(x2) about 0 and using that N ≫ 1,
(cid:19)1/2
2c2
√
x

(cid:18)

mC(x)

1 −

3C(x)N 3 + O(1/N 4)

(cid:18)

= mC(x)

1 −

√

=

(cid:18)

m

C(x) −

1
2

·

2c2
x

(cid:19)
3C(x)N 3 + O(1/N 4)
c2
x
3N 3 + O(1/N 4)

(cid:19)

.

Putting it all together in A.2, we get

s(y(N )) ≈

√

(cid:18)

m

C(x) −

√

=

(cid:18)

m

C(x) −

c2
x
3N 3 + O(1/N 4)
c2
x
3N 3 + O(1/N 4)

(cid:19) ∞
(cid:88)

σ=1

(cid:19)

,

ασ

(A.19)

(A.20)

(A.21)

(A.22)

(A.23)

(A.24)

(A.25)

as claimed. This completes the proof of part (2).

The proof of part (1) is more straightforward. For any function f on [−1, 1] that is at least twice
continuously differentiable, we know that the Fourier series of f 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 , which is precisely the statement of part
(1).

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) =

4
π

∞
(cid:88)

n=1,3,5,...

1
n

sin

(cid:16) nπx
L

(cid:17)

.

(B.1)

Here, the wavelength is 2L (Weisstein, 2024).

We construct a truncated version of the square wave with a 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 as 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 to 1. The output of this transformation for a few representative waveforms is
pictured in Figure 5.

Intuitively, the truncated square wave with a single sine wave ought to 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 6 demonstrates that this is indeed the case.

16

810
811

812
813
814
815
816
817

818
819
820
821
822
823

824
825
826
827
828

829
830
831
832
833
834

835
836
837
838
839
840

841
842
843
844
845
846

847
848
849
850
851

852
853
854
855
856
857

858
859
860
861
862
863

Under review as a conference paper at ICLR 2025

Figure 5: Truncated square waves framed as densities and their smoothness.

Figure 6: Values of the smoothness metric Equation 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.

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 L2 distance as indicated in 3.2. We
empirically observe that L2 distance better preserves monotonicity than the L1 for higher frequency
content, thus motivating this choice. With a sample rate of 2048Hz, the L1 distance exhibits some
undesirable warping when our square-wave multinomial uses over 80 sine waves (see Figure 7). A
Fourier head in a practical setting may possess several more than 80 frequencies; accordingly, we
favor the L2 distance as our discrepancy measure.

Alternative Notions of Smoothness: In validating our choice of smoothness metric, we compare it
to the spectral entropy (Inouye et al., 1991), which has a similar purpose in quantifying the “smooth-

17

864
865

866
867
868
869
870
871

872
873
874
875
876
877

878
879
880
881
882

883
884
885
886
887
888

889
890
891
892
893
894

895
896
897
898
899
900

901
902
903
904
905

906
907
908
909
910
911

912
913
914
915
916
917

Under review as a conference paper at ICLR 2025

Figure 7: Values of the smoothness metric 3.2 on our square-wave-like multinomials as we increase
the number of sine waves. On the right, we can see that L1 as a discrepancy measure leads to
non-monotonicity, motivating our choice of L2 distance in measuring our results.

ness” of the frequency content of a signal. Spectral entropy is defined as the Shannon entropy of the
power spectral density of a sampled signal f , which is defined as follows:

H(f ; N ) =

(cid:88)

n∈N

p(n) log2

(cid:18) 1

(cid:19)

p(n)

= −

(cid:88)

n∈N

Sn
Stotal

log2

(cid:19)

(cid:18) Sn
Stotal

(B.2)

Here, N is the number of Fourier frequencies and S is the power of a frequency n ∈ N ; Sn is the
power spectrum of the nth frequency, and Stotal is the power of the signal using all N frequencies.
For some frequency at index n, Sn/Stotal is called its relative power and (cid:80)
= 1 enables
us to consider each frequency’s power as a probability.

Sn
Stotal

n∈N

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 is a relationship between amplitude and frequency. More concretely, blue noise, which is
defined by the amplitude increasing proportionally to the frequency, will have lower spectral entropy
than white noise. We sought a metric that always quantified ‘sharper’ signals like blue noise as less
smooth. In Table 4, we frame sampled noises of different types as multinomial distributions to
match our model setting by normalizing their amplitudes to be in [0, 1] and normalizing their sum to
1. Our noise types are defined before normalization, in order of smoothest to sharpest:

• Brown: S ∝ 1
F 2
• Pink: S ∝ 1
F
• White: S ∼ N (0, 1)
• Blue: S ∝ F

where S is the power density and F is the frequency. To obtain samples of each type, we first
generate white noise. We do this by sampling a Gaussian with mean 0 and standard deviation 1
to obtain amplitudes for t 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 1.

C TOY EXAMPLE DETAILS

Here we provide full details of the datasets used in our toy example of learning a known conditional
distribution.

18

918
919

920
921
922
923
924
925

926
927
928
929
930
931

932
933
934
935
936

937
938
939
940
941
942

943
944
945
946
947
948

949
950
951
952
953
954

955
956
957
958
959

960
961
962
963
964
965

966
967
968
969
970
971

Under review as a conference paper at ICLR 2025

Noise Mean ± Std. Deviation
Discrepancy
0.0003 ± 0.0001
L2
Brown
0.0017 ± 0.0002
L2
Pink
0.0034 ± 0.0003
L2
White
0.0038 ± 0.0003
Blue
L2
0.4516 ± 0.0894
Spectral Entropy Brown
0.3878 ± 0.0603
Spectral Entropy
Pink
0.4266 ± 0.0614
Spectral Entropy White
0.4191 ± 0.0583
Blue
Spectral Entropy

Diff Delta Desired Delta
n/a
n/a
+
0.0014
+
0.0016
+
0.0005
n/a
n/a
+
-0.0638
+
0.0388
+
-0.0076

n/a
+
+
+
n/a
-
+
-

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 is undesirably not monotonic increasing
for what we consider increasingly “sharp” noise types.

Dataset: We create a synthetic dataset D = {(q(x), q(y), q(z))} ⊂ R3 as follows. Fix a probability
distribution P1 = P1(x) that is parameterized by one variable and a second distribution P2 =
P2(x, y) parameterized by two variables. Fix an interval I ⊂ R. Sample x uniformly from I,
sample y ∼ P1(x), and finally sample z ∼ P2(x, y). We can repeat this sampling procedure N
times to obtain a set of N triples for which we know the conditional distribution of z given x and
y. Finally, we quantize this set to a fixed number of uniformly spaced bins in the range [−1, 1] to
obtain the dataset DP1,P2 . We will denote the quantization of z by q(z). We quantize into 50 bins
and our dataset has size 5000, with a 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 = [−0.8, 0.8] and σ2 = 0.01
in all of them.

1. Gaussian dataset: P1(x) = N (x, σ2), and P2(x, y) = N (y, σ2).

2. GMM-2 dataset: P1 = Uniform(I), and P2(x, y) is a GMM centered at x and y with

variance σ2.

3. Beta dataset: P1(x) = N (x, σ2), and P2(x, y) ∼ U ({±1}) × Beta(100 |x| , 100 |y|),
where U ({±1}) denotes the Rademacher distribution supported on {±1} with probability
1/2 each.

Additional results: In Figure 8, we present results from training over a range of frequencies, and for
each frequency we ran experiments with and without Fourier regularization. In Table 6 we present
results on the MSE metric, that show that the Fourier head outperforms the linear classification head.

Dataset
Gaussian
GMM-2
Beta

Linear
0.170 ± 0.052
0.238 ± 0.032
0.234 ± 0.032

Dataset
Gaussian
GMM-2
Beta

Linear
0.116 ± 0.049
0.068 ± 0.022
0.127 ± 0.044

KL Divergence (↓)
GMM
0.026 ± 0.011
0.030 ± 0.006
0.407 ± 0.012

Fourier
0.116 ± 0.043
0.146 ± 0.033
0.191 ± 0.016

Smoothness (↓)
GMM
0.068 ± 0.012
0.043 ± 0.009
0.061 ± 0.003

Fourier
0.057 ± 0.011
0.038 ± 0.007
0.076 ± 0.021

Table 5: KL divergence and Smoothness for the three classification heads (Linear, GMM, and
Fourier) on 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 are Gaussian. However, the Fourier head has the best KL divergence on the Beta dataset.
This demonstrates the flexibility of the Fourier head in modeling non-Gaussian distributions as well.

19

972
973

974
975
976
977
978
979

980
981
982
983
984
985

986
987
988
989
990

991
992
993
994
995
996

997
998
999
1000
1001
1002

1003
1004
1005
1006
1007
1008

1009
1010
1011
1012
1013

1014
1015
1016
1017
1018
1019

1020
1021
1022
1023
1024
1025

Under review as a conference paper at ICLR 2025

1026
1027

1028
1029
1030
1031
1032
1033

1034
1035
1036
1037
1038
1039

1040
1041
1042
1043
1044

1045
1046
1047
1048
1049
1050

1051
1052
1053
1054
1055
1056

1057
1058
1059
1060
1061
1062

1063
1064
1065
1066
1067

1068
1069
1070
1071
1072
1073

1074
1075
1076
1077
1078
1079

Figure 8: We study how the quantity of Fourier frequencies impacts KL divergence and smoothness
for the toy example on 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 are 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 an order of magnitude
more frequencies, and their conditional distributions are more complicated. While the GMM head
has better KL 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 it is non-Gaussian. Notice also how on each of the datasets, the smoothness degrades
as frequency increases, in a fashion that follows the asymptotic from our Theorem 3.3.

Toy Example: MSE (↓)

Dataset
Gaussian
GMM-2
Beta

Pointwise Regression
0.010 ± 0.001
0.121 ± 0.004
0.275 ± 0.009

Linear
0.013 ± 0.001
0.126 ± 0.004
0.276 ± 0.008

Classification Head
GMM
0.010 ± 0.001
0.120 ± 0.004
0.273 ± 0.009

Fourier
0.012 ± 0.001
0.123 ± 0.005
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 to regress to the
mean, as does the expected value of each of the classification heads.

D ADDITIONAL DECISION TRANSFORMER EXPERIMENT DETAILS

Following the original Decision Transformer implementation, we trained on 500k transitions ob-
served by a DQN agent during training, for 5 epochs. We trained on the same model size as the
original implementation (a GPT-1 model with approximately 2.012M parameters) which takes about
4 hours on a single GPU. We can see that in Figure 12 that the PMFs learned by the Fourier head

20

Under review as a conference paper at ICLR 2025

Toy Example: Ground Truth Conditional Distribution vs. Pointwise Regression Output

Figure 9: We present some examples of the ground truth conditional distribution versus the point
predicted by the Pointwise Regression model. The regression model simply regresses to 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.

are smoother. In Figure 13 we present results for more Atari games. In Figure 10, we present results
from an ablation study of the model size. The results demonstrate that, across model sizes, the Deci-
sion Transformer with a Fourier head is better at learning high-quality next action distributions than
the Decision Transformer with a linear head. And in Figure 11, 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.

Figure 10: 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 a Fourier head is better at learning high-quality next action distributions than the
Decision Transformer with a linear head.

E ADDITIONAL CHRONOS EXPERIMENT DETAILS

In Figure 14 we present a learned next-token PMF from a linear Chronos model, and a next-token
PMF from a Chronos model which uses the linear head. The Fourier head is about 4x smoother.
In Table 7 we present results from an ablation study on the choice of regularization, and binning
strategy. We followed the original Chronos implementation, keeping all hyperparameters the same.
In particular, we trained for 200k steps, on the same model size as the original implementation

21

1080
1081

1082
1083
1084
1085
1086
1087

1088
1089
1090
1091
1092
1093

1094
1095
1096
1097
1098

1099
1100
1101
1102
1103
1104

1105
1106
1107
1108
1109
1110

1111
1112
1113
1114
1115
1116

1117
1118
1119
1120
1121

1122
1123
1124
1125
1126
1127

1128
1129
1130
1131
1132
1133

Under review as a conference paper at ICLR 2025

Fourier Head Ablation Study: Dataset Size

Figure 11: 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 a Fourier head achieves larger returns than the linear head on the
Seaquest game.

Figure 12: We present example next action distributions for a single step in the Decision Trans-
former. We can see that the Fourier agent with 14 frequenices produces “clumps” of actions that are
semantically meaningful; for example, this agent either wants to move down right, down, or down
left, or else the agent wants to shoot down right, down, or down left, presumably because they need
to move in that general direction to rescue a diver in that direction but there are submarines in the
way. In contrast, there is no indication that the linear agent has learned a cohestive strategy which
relates moving and shooting.

(the T5 model with approximately 20M parameters) and this takes about 48 hours on 8 GPUs. See
Table 8 for the datasets we used to train and evaluate Chronos.

Chronos Time Series Model
Fourier-550
Fourier-550 (no regularization)
Fourier-550 (uniform precision binning)

MASE ↓ WQL ↓
0.749
0.753
0.747

0.852
0.861
0.873

Smoothness ↓
0.0203 ± 0.0176
0.0204 ± 0.0172
0.0292 ± 0.0205

Table 7: We present large-scale ablations 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 on the WQL. We note that the choice of binning strategy doesn’t
affect the performance of the linear baseline.

22

1134
1135

1136
1137
1138
1139
1140
1141

1142
1143
1144
1145
1146
1147

1148
1149
1150
1151
1152

1153
1154
1155
1156
1157
1158

1159
1160
1161
1162
1163
1164

1165
1166
1167
1168
1169
1170

1171
1172
1173
1174
1175

1176
1177
1178
1179
1180
1181

1182
1183
1184
1185
1186
1187

Under review as a conference paper at ICLR 2025

1188
1189

1190
1191
1192
1193
1194
1195

1196
1197
1198
1199
1200
1201

1202
1203
1204
1205
1206

1207
1208
1209
1210
1211
1212

1213
1214
1215
1216
1217
1218

1219
1220
1221
1222
1223
1224

1225
1226
1227
1228
1229

1230
1231
1232
1233
1234
1235

1236
1237
1238
1239
1240
1241

Figure 13: 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.

Figure 14: We present the next token value distribution for a single forecasted timestep on 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 x-axis represents the bins in the latent space [−1, 1]; the x-axis values
for the Fourier head are lower because the linear head uses uniform binning, and the Fourier head
uses mixed precision binning.

23

Under review as a conference paper at ICLR 2025

Fourier Head Ablation Study: Chronos 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 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 × 105, 1.1 × 106, and 1.1 × 107, we report the average MASE across
four seeds; for the dataset size 1.1 × 108 we report the MASE from Table 3. We generate the plot
following (Kaplan et al., 2020) and observe a similar power-law scaling behavior for both methods,
with the Fourier head consistently outperforming the linear head.

24

1242
1243

1244
1245
1246
1247
1248
1249

1250
1251
1252
1253
1254
1255

1256
1257
1258
1259
1260

1261
1262
1263
1264
1265
1266

1267
1268
1269
1270
1271
1272

1273
1274
1275
1276
1277
1278

1279
1280
1281
1282
1283

1284
1285
1286
1287
1288
1289

1290
1291
1292
1293
1294
1295

Under review as a conference paper at ICLR 2025

Table 8: All datasets that are used for our time series forecasting experiments. We built our time
series forecasting experiments on top of Chronos (Ansari et 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-only data is 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.

Dataset

Domain Freq. # Series

Series Length

Prediction

min

avg max Length (H)

Pretraining

M
Brazilian Cities Temperature nature
transport 1H
Mexico City Bikes
energy
Solar (5 Min.)
Solar (Hourly)
energy
Spanish Energy and Weather energy
Taxi (Hourly)
USHCN
Weatherbench (Daily)
Weatherbench (Hourly)
Weatherbench (Weekly)
Wiki Daily (100k)
Wind Farms (Daily)
Wind Farms (Hourly)

5min
1H
1H
transport 1H
1D
nature
1D
nature
1H
nature
1W 225280
nature
100000
1D
web
337
1D
energy
337
1H
energy

1320
757
492
12
494
780 78313 104449
5166 105120 105120 105120
8760
5166
66 35064 35064 35064
744
739
734
5906 38653 59283
225280 14609 14609 14610
225280 350633 350639 350640
2087
2741
366
8784

2087
2741
71
1715

2087
2741
354
8514

2428
6090

8760

8760

5 230736 231052 232272
120
72
51
2674
84
767
150
617
114
203
58
181
144
1428
72
756
47
645
874
24000
841
23000
1969
30490
791
111
113
111
333
366
130
427
47
518
862 17544 17544 17544
1332 14296 65981
3010

98
51
84
90
48
24
117
48
28
100
37
1562
791
113
298
99
24

28
51
84
48
18
15
66
24
20
24
19
124
791
113
91
30
11

Evaluation

Australian Electricity
CIF 2016
Car Parts
Hospital
M1 (Monthly)
M1 (Quarterly)
M1 (Yearly)
M3 (Monthly)
M3 (Quarterly)
M3 (Yearly)
M4 (Quarterly)
M4 (Yearly)
M5
NN5 (Daily)
NN5 (Weekly)
Tourism (Monthly)
Tourism (Quarterly)
Tourism (Yearly)
Traffic
Weather

30min
energy
1M
banking
retail
1M
healthcare 1M
1M
various
3M
various
1Y
various
1M
various
3M
various
1Y
various
3M
various
1Y
various
1D
retail
1D
finance
1W
finance
1M
various
1Q
various
1Y
various
transport 1H
1D
nature

25

-
-
-
-
-
-
-
-
-
-
-
-
-

48
12
12
12
18
8
6
18
8
6
8
6
28
56
8
24
8
4
24
30

1296
1297

1298
1299
1300
1301
1302
1303

1304
1305
1306
1307
1308
1309

1310
1311
1312
1313
1314

1315
1316
1317
1318
1319
1320

1321
1322
1323
1324
1325
1326

1327
1328
1329
1330
1331
1332

1333
1334
1335
1336
1337

1338
1339
1340
1341
1342
1343

1344
1345
1346
1347
1348
1349

