Fourier Head: Helping Large Language Models Learn Complex Probability Distributions

Thank you for the thoughtful feedback! We were able to produce the additional results that you asked for, summarized here:

• We introduce a new pointwise regression baseline model, and demonstrate that the Fourier classification head outperforms it.
• We provide more evidence that the Fourier head can model latent spaces which are not circle-shaped, by adding results for three more RL tasks where the latent space is a quantization of $S^1\sqcup S^1\sqcup \\{0,1\\}$. In all these tasks, the Fourier head agent obtains significantly larger returns than the baseline.

Below, we provide more details, as well as question-specific responses. We have also uploaded a revised manuscript, with changes written in blue.

Q2: why should we impose the $1/n^2$ Fourier regularization term during training, if it’s a truncated Fourier series?

If the amplitudes of a truncated Fourier series decay like $1/n^2$, this ensures that the function is smoother.

In more detail–in the presence of limited data, there are many possible choices of Fourier coefficients which may fit the data equally well. We can express a preference towards smoother densities by penalizing the higher-order Fourier coefficients more than the lower-order Fourier coefficients. Intuitively, this ensures that the model extracts low-frequency signal from the training data, while ignoring high-frequency noise. We’ve added a summary of this in Section 2.4 of the revision. Thanks for this question, it seems like we didn’t make this point clear enough in the submission.

Q4: Is it a limitation that the Fourier series can only model periodic functions?

In our large-scale experiments we find that this isn’t an issue, for the very reason that you mentioned–namely, for the time series task the distribution is close to zero near the boundaries. If needed, a tanh reparameterization can be used to map the domain $[-1,1]$ to the real line, and you can then truncate it. This would let the Fourier head learn non-periodic functions. We have added an explanation about this in Section 2.4 of the updated manuscript.

However, maintaining the original periodic formulation and using a Fourier head with a sufficiently large number of frequencies tends to solve the problem. This follows from the Fourier theory: if you have a smooth function on $[-1,1]$, but $f(-1)\neq f(1)$, then you can learn a Fourier series that does a very good job approximating $f$ over the open interval $(-1,1)$, and it will only have problems near the boundaries $\pm 1$. Using more frequencies guarantees that the approximation improves close to the boundaries.

Q1: for tasks with a continuous-valued target output, does it make more sense to output a point estimate?

We have conducted additional experiments, and our results show that the probabilistic nature of the tasks leads to a pointwise regression model “regressing to the mean” and performing poorly on the MSE metric.

In more detail: we have created another toy example baseline model with the same architecture as the Fourier head model, but where we substituted the last classification layer with a linear layer having a single output dimension, and where we train using MSE loss. We refer to this model as “Pointwise Regression”. When we evaluate each of Pointwise Regression, Fourier head, and Linear head on the toy experiment datasets, their average MSE performances across seeds are all approximately equal to each other, showing no observable advantage in terms of MSE for one model over another. This is to be expected, given the probabilistic nature of the datasets. We include here our new results:

To demonstrate why a simple pointwise estimate in a probabilistic setting might be inadequate, we also created a new dataset (the "Beta" dataset in the above table) for which the conditional distributions are symmetric about $0$ by construction. More precisely, the conditional distribution of $z$ given $(x,y)$ is $b * \mathrm{Beta}(100 |x|, 100|y|)$, where $b$ is a random variable that is $+1$ or $-1$ with equal probability. We have used the Beta distribution since it is a non-Gaussian distribution, so it illustrates the flexibility of the Fourier head in learning a wide variety of naturally occurring complicated distributions. In fact, on the Beta dataset, our results show that the Pointwise Regression, Linear classification head, and Fourier head models all have identical MSE very close to 0.275, which is also the MSE of a naive baseline model that always predicts 0 regardless of the input. In other words, the Pointwise Regression model collapses to always predicting values very close to 0 since it regresses to the mean of the underlying distribution, completely missing any information about higher moments of the distribution.

On the other hand, the Fourier head is able to learn a high quality reconstruction of the underlying conditional distribution, showing why one might not choose to rely on a pointwise estimate even in our toy setting. We have included the Pointwise Regression baseline MSE values in the latest draft. We also include visual samples of the model predictions alongside the true conditional distribution, illustrating this unwanted "regression to the mean" behavior that arises. (link to graph)

Given these results in the toy setting, we believe that using a classification head that can generate a probabilistic forecast rather than a pointwise regression in the Chronos experiments is desirable since the underlying next-token distributions are complicated and not well-captured by regressing to the mean.

(On a technical note: your suggestion prompted us to realize that for evaluating the MSE performance of the categorical distributions predicted by the Fourier and Linear head, one should consider their expected values as the model’s implied pointwise estimate, rather than just the bin of maximum probability. This is especially applicable since some of our toy distributions are bimodal. As a result, we have updated our MSE metrics in the paper. So, thank you for asking this question.)

Q5: the minor typos

Thanks for catching those, we’ve fixed them in the latest draft :)

Q3: Can the Fourier head be generalized to output spaces which are not interval shaped?

This is certainly possible, but beyond the scope of the paper. One way to do this would be to have the learned output categorical distribution be a mixture of Fourier heads and a linear classification head. For example, since the action space for our RL task is a quantized version of $S^1 \sqcup S^1 \sqcup \\{0, 1\\}$, this hypothetical architecture would learn two Fourier heads, each with output dimension 8, as well as a linear classification head, with output dimension 2. Such a model would also have a 3-dimensional classification layer which functions as a high level controller, choosing which of the three classification heads to route the model input to. It would need to be trained with masking to ensure that gradients only update for the correct head.

We don’t conduct experiments on this generalized classification head for two reasons: first, it is sufficiently complicated that it is beyond scope, and we hope that future works will be inspired to make this applicable to those other domains. And second: our Decision Transformer results show that a single Fourier classification head does a much better job than the linear baseline, even while incorrectly placing a continuity inductive bias between some of the token boundaries, as your review noted. We believe that the positive results, despite this limitation, underscores just how problematic it is that many models continue to use a linear classification head, while the next-token distribution seems to be “starving” for at least some continuity inductive bias. To demonstrate this more robustly, we have added results for three more Atari games. In addition to our previously reported results on Seaquest, the paper now includes results on BankHeist, DoubleDunk, and Gravitar, which demonstrate that the Fourier head significantly outperforms the baseline.

In this table, we report normalized returns (mean $\pm$ standard deviation, averaged over four seeds) for the Decision Transformer agent across the four Atari games. The results show that the Fourier agent obtains higher returns than the Linear agent across all games.

We wanted to follow up to make sure that your concerns are being properly addressed. Please let us know if there are additional questions that we can provide further information / experiments on. Thank you again for your time and effort reviewing our paper!

Thank you for your detailed response and additional experiments. I have some remaining questions/comments:

$1/n^2$ regularization

The explanation of the regularization is still not entirely clear to me. You write:

for the class of Fourier series which have continuous second derivatives, the Fourier coefficients decay on the order of $1/n^2$. To impose this regularity assumption on the learned Fourier densities, we [...] add a regularization term to the loss

On first reading, I understood this to mean that the $k^2$ factor in the regularization $\sum_k k^2 |c_k|^2$ is used to enforce that the coefficients decay like $1/n^2$ (and this causes the Fourier head output to have continuous second derivatives?). However, it's unclear to me what this could mean in the case of truncated Fourier series. For instance, the condition that $c_n\in O(1/n^2)$ is an asymptotic statement; when the Fourier series is truncated, it vacuously holds regardless of whether the regularization is present. If you had some non-asymptotic interpretation of $c_n\sim 1/n^2$ in mind, I don't know what it could be or why the $k^2$ regularization would cause it to hold.

I find the motivation in de la Fuente (2024) II.D to be much more clear -- the regularization term is equal to the total squared variation (Eq 10), which is a measure of smoothness.

Also: it seems that smoothness could be increased either by increasing regularization strength or decreasing the number of Fourier frequencies $N$. How do these relate? When is it appropriate to tune one vs the other?

Exposition of the Fourier head and the de la Fuente (2024) paper

Additionally, in my opinion de la Fuente et al (2024) does a better job of explaining some other details:

• The purpose of first learning autocorrelation coefficients then converting to Fourier coefficients (Alg. 1 Step 4) is explained in their section II.A
• The normalization by $1/c_0$ (Alg. 1 Step 5) is explained in their II.B.

It would be helpful to briefly mention these to help explain Algorithm 1. (Or maybe just refer the reader to the sections in their paper.)

Overall, it appears that your work takes significant inspiration from the de la Fuente paper. Your paper cites theirs at several places, which is good, but it might also be appropriate to mention them in the intro and/or when the Fourier head is introduced in Sec 2.1/2.2. My understanding is that the Fourier head is effectively the "Fourier basis density model" that they propose swapped in for the last layer of a deep neural model.

Apologies for not bringing this up earlier -- I hadn't looked at the de la Fuente paper at the time.

Point estimate vs probabilistic forecasting

Thank you for the pointwise regression experiments. Intuitively, the fact that the pointwise regression predicts the mean even for bimodal distributions is not ideal. However, if one only cares about MSE, then your experiments show that the point estimate is perfectly fine. The paper would be better motivated if you could mention some examples of when the MSE alone is insufficient, and predicting the full density is superior to just the point estimate. It looks like the RL setting is one such example, since the agent needs to explicitly sample from the output distribution. (Maybe this should be made more explicit in the paper.) Are there similar examples for the time-series setting?

We appreciate the reviewer’s engagement with us! In response to each point:

$1/n^2$ regularization

We agree that interpreting the regularization term as the total squared variation is much more natural given its relation to smoothness. We have modified the manuscript to reflect this change (see e.g. updated Section 2.4).

As for choosing whether to increase regularization strength or decrease the number of Fourier frequencies: tuning the regularization strength gives finer control on penalizing the model for high frequency content to increase smoothness, while still allowing some high frequency information. The number of frequencies is a more drastic handle on this tradeoff, as made explicit in our scaling law. Choosing which to tune needs to be determined for each application – for applications in which the underlying distributions are representable in a few number of frequencies $N_0 \ll m/2$, minimizing the number of frequencies until $N_0$ might be preferable in order to increase smoothness. On the other hand, in applications where a high number of frequencies are required to obtain a reasonable reconstruction of the distribution, we would prefer to maximize the number of frequencies while tuning the regularization strength in order to encourage smoothness while not sacrificing modeling capacity.

Exposition of the Fourier head and the de la Fuente (2024) paper

As you noticed, we only tried to provide a minimal exposition of the Fourier Basis Density Model paper. Going over it now, we agree with you that our exposition was too minimal, so we’ve added additional references and explanations--

• In the introduction, we’ve added the sentence: The Fourier head is constructed using the Fourier Basis Density Model (de la Fuente et al, 2024).
• When we define Fourier head algorithm, we’ve added the sentence: 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 $\mathrm{Re}(c_0)$), we refer the author to (De la Fuente et al., 2024).

Point estimate vs probabilistic forecasting

In the paper, we only consider scenarios where probabilistic modeling is necessary, and where a point estimate is insufficient. We constructed the toy example with this in mind–the main success metric for the toy example is the KL divergence between the quantized ground truth distribution, and the learned categorical distribution. Similarly, the large-scale examples in the paper (probabilistic agentic decision making, and probabilistic time series forecasting) require learning a probability distribution over the latent space, and sampling from it at test time to obtain the success metrics.

In particular, probabilistic time series forecasting is a useful tool for decision making because probabilistic forecasts allow us to precisely quantify future uncertainty. We note that there are tradeoffs when deciding whether to model time series probabilistically versus deterministically:

• Deterministically modeling time series (e.g. learning to regress, with an MSE loss) is generally simpler, especially during data preprocessing. For example, tokenization is not needed.

• Probabilistically modeling time series (e.g. learning a distribution over the next value of the time series, using cross entropy loss, as in Chronos) is more complicated, as it requires design choices such as tokenization. But the upshot of these methods is that probabilistic forecasts contain all the information from deterministic forecasts, plus more. For example, from a probabilistic time series model, you can choose to sample many possible futures. Computing the median of those futures allows you to compute an accuracy metric like MSE. Additionally, you can extract error bars using the quantiles from the possible futures, which is a clear advantage for practical applications.

And as you requested, we have added explicit descriptions for our large-scale tasks to make it clear that they involve probabilistic sampling--

• For the RL task, we added: “ At test time, the agent chooses its next action by sampling from the learned next-action distribution.”
• For the time series task, we added: “At test time, the model chooses the next numerical token by sampling from the next-token distribution.”

Please let us know if you have any further questions, or if there is anything that we can clarify!

We just wanted to follow up, and say thank you for your effort reviewing our paper, and for your engagement throughout the rebuttal period!

Q7: exploring bases other than Fourier basis

We want to acknowledge that there may be other bases that perform equally well to the Fourier basis. Our paper is a first step towards exploring this direction, and we expect that our results would inspire follow-up works that explore different bases. However, at a practical level, the Fourier basis is more computationally tractable (i.e. fewer FLOPS) and more stable to compute than many alternatives, such as the Chebyshev basis. Thanks for asking this question.

Q6: how were frequencies chosen for the time series experiments?

We didn’t use any clever procedure for selecting the Fourier frequencies for the time series experiments. We just started by sweeping over powers of $2$ to be efficient, because the Fourier head scaling law indicates that you might want to try up to $\mathrm{n_bins} / 2 = 2048.$

Q5: to what extent is it a weakness of the Fourier head that the user must select hyperparameters?

We have included additional experimentally-backed details to the manuscript to make it easier for the user to select hyperparameters. Thank you for giving us the opportunity to clarify this point. In short:

• When picking regularization strength–we find that in the low-frequency domain (i.e. frequencies in the single digits) using $\gamma=0$ works best, and in the high-frequency domain (i.e. greater than 10 frequencies) using $\gamma=10^{-6}$ works best.

• When picking Fourier frequencies–our Decision Transformer results show that the model is reasonably robust to the choice of the number of frequencies. For example, our latest Decision Transformer results (link to graph) show that the Fourier agent obtains higher returns than the Linear agent, irrespective of the quantity of Fourier frequencies. Meanwhile, for more complex problems such as zero-shot probabilistic time-series forecasting, one may prefer to use more frequencies as this results in a Fourier head with more modeling power. This is stated in our Theorem 3.3 (Fourier head scaling law) and demonstrated in Table 3.

Q4: does the Fourier head have limited expressiveness because it is bounded?

In practice, we find that this is not an issue. If needed, a tanh reparameterization can be used to map the domain $[-1,1]$ to the real line. This would let the Fourier head learn unbounded values. We have added an explanation about this in Section 2.4 of the updated manuscript.

However, even if we decide to use a bounded domain (as is the case in all the experiments in the manuscript), this still isn’t an issue, as evidenced by the original Chronos model’s SOTA time series forecasting accuracy, with datasets which are not a priori bounded within some finite range. Those authors modeled unbounded sequences using a combination of two techniques: 1) normalizing the time series so the individual values are small, and 2) ensuring that the next-token distribution is defined over a sufficiently wide interval.

In more detail: in the Chronos paper, the authors normalize the time series so that the mean of the absolute value in the historical context is equal to 1; this ensures that the time series values are small. Then, the architecture defines the range of tokens as equally spaced inside of $[-15, 15]$; this ensures that the model is capable of learning from the examples with significant trend components. With this normalization and tokenization strategy, the Chronos model is SOTA because it turns out that even time series with aggressive trend components are normalized to values that can be forecasted accurately by the model.

Q3: the toy experiment has limited baselines, and in particular would benefit from a comparison with a Gaussian Mixture Model

To address this concern we have implemented a GMM classification layer for which the means and standard deviations are learned, and the number of Gaussian components is a hyperparameter. As you predicted, our results show that substituting the GMM head in our toy experiment yields better performance than the Fourier head on the previous datasets. This is to be expected since the conditional distributions being learned in those datasets are precisely GMMs.

To highlight the flexibility of the Fourier head over the GMM head, we have added a more challenging dataset to the toy example, where the conditional distributions are based on a Beta distribution. On this new Beta dataset, the Fourier head achieves KL divergence 0.191 while the GMM head (with 2 Gaussians) achieves KL divergence 0.407, more than twice that of Fourier. Since the Beta distribution is a common and naturally occurring distribution, this comparison shows the advantage of using the Fourier head when the underlying next-token distribution might be complicated and unknown.

We include results with the GMM head and Beta dataset in our latest draft. In particular, our numerical results show that the Fourier head learns the Beta distribution more robustly than both the GMM head and the linear head:

We also analyze the effect of Fourier frequencies on learning this Beta dataset, and we find that the Fourier head indeed outperforms the other classification heads for sufficiently many frequencies. (link to graph)

Lastly, thank you very much for asking this question. Our choice of synthetic datasets in the toy experiment certainly makes it seem that the Fourier head could easily be replaced by a learned GMM.

Q2: how to demonstrate more empirical impact?

To strengthen our empirical contribution, we have added results for three more Atari games. In addition to our previously reported results on Seaquest, the paper now includes results on BankHeist, DoubleDunk, and Gravitar, which demonstrate that the Fourier head significantly outperforms the baseline.

In this table, we report normalized returns (mean $\pm$ standard deviation, averaged over four seeds) for the Decision Transformer agent across the four Atari games. The results show that the Fourier agent obtains higher returns than the Linear agent across all games.

Furthermore, we also demonstrate that the Fourier head consistently outperforms the Linear head, irrespective of the quantity of Fourier frequencies, for all these additional games. (link to graph)

Lastly, while we agree that the paper’s RL results are relatively stronger than the time series forecasting results, we want to underscore that a 3.5% accuracy improvement on a recent SOTA forecasting model is no easy feat. For comparison, in the original Chronos paper, the authors find that their novel architecture beats the next best task-specific model by only 0.9%. Our paper’s 3.5% forecasting improvement didn’t require any hyperparameter changes to the original configuration, making the Fourier head an easy “drop-in” replacement for a performance increase.

Q1: why all the emphasis on “smoothness”?

A primary goal of the paper was to improve performance on the original metrics from the tasks we considered (e.g. accuracy for time series, returns for RL), and we accomplished that using the Fourier head, while also showing that the learned categorical distributions are smoother for the Fourier head than the linear head.

And you’re absolutely right that, in the context of this paper, smoothness is not a metric that we should care about on its own. Our experiments show that the Fourier head yields smoother categorical distributions than the Linear classification head. But we also find that smoothness of the categorical distribution doesn’t necessarily lead to better downstream performance.

To provide additional evidence that the Fourier head provides concrete improvements for the original success metrics, we have added two more ablation studies to the paper:

• Ablation study #1: We show the Fourier head is better at learning high-quality next action distributions than the Decision Transformer with a Linear head across all dataset sizes. (link to graph)
• Ablation study #2: We show the Fourier head is better at learning high-quality next action distributions than the Decision Transformer with a Linear head across all model sizes. (link to graph)

In these ablations, we show that when the Fourier head learns a very smooth density, the generalization is better. Intuitively, this is because if the model needs to learn the likelihood of the $N$’th token, it can lean on the learned likelihood of its neighboring $N-1$ and $N+1$ tokens, even if it saw very few examples during training where the $N$’th token was the correct answer.

Thank you for the thoughtful feedback! We were able to produce the additional results that you asked for, summarized here:

• We demonstrate additional empirical impact by adding results for three more RL tasks. In all these tasks, the Fourier head agent obtains much larger returns than the baseline.
• We demonstrate additional empirical impact by including two more ablations: one which demonstrates that the benefit of the Fourier head persists as the dataset size scales, and one which demonstrates that the benefit of the Fourier head persists as the model size scales.
• We introduce a new GMM baseline model, as well as a more complicated synthetic dataset, and demonstrate that the Fourier head’s added flexibility is needed to model this more challenging dataset.

Below, we provide more details, as well as question-specific responses. We have also uploaded a revised manuscript, with changes written in blue.

We wanted to follow up to make sure that your concerns are being properly addressed. Please let us know if there are additional questions that we can provide further information / experiments on. Thank you again for your time and effort reviewing our paper!

We just wanted to follow up, and say thank you for your effort reviewing our paper, and for your engagement throughout the rebuttal period!

Q3: the paper should include a comparison with models using decoupling such as TEMPO

Thank you for bringing the TEMPO paper to our attention. Since their paper also studies how to perform time series forecasting using an LLM, we have added a citation to TEMPO in our related works section. However, we believe TEMPO is not directly comparable to the Fourier head, for the following reasons:

• TEMPO employs an STL decomposition as a preprocessing step for improved training on time series data using an LLM, arguing that a transformer’s self-attention mechanism is not guaranteed to be able to disentangle the trend and seasonality components (their Theorem 3.1).
• TEMPO proposes novel prompting techniques to improve performance.

Both of these contributions are not comparable to our Fourier head since we propose an alternative way to improve the classification layer in the transformer, which is not related to either preprocessing or prompting.

Q2: the paper should include experimental results for the TimeLLM and GPT4TS models

Thank you for highlighting the TimeLLM and GPT4TS papers. In the updated manuscript we have added citations to them in the related works section. However, we don’t believe that a direct comparison to those models is relevant, because those models require in-domain training, whereas Chronos does zero-shot domain transfer.

More precisely, both the models TimeLLM and GPT4TS require fine-tuning and testing on each dataset separately. In contrast, Chronos is pretrained on a single dataset consisting of many real and synthetic time series, and it is designed to be evaluated on domains unseen during training. In our paper, we evaluate Chronos on 20 benchmark datasets unseen during training, without any additional dataset specialization or fine-tuning.

Additionally, we would like to emphasize that the goal of the paper is not to present the Fourier head as a tool just for time series, but rather as a tool for any domain where a classification head can benefit from having more continuous structure. Towards this, we would like to highlight the diverse types of transformer-based models that we used to study the Fourier head–our time series experiments use an encoder-decoder T5 architecture, our reinforcement learning experiments use a decoder-only GPT-class architecture, and our audio toy example uses an audio-spectrogram transformer.

Q1: the Fourier head is incremental because LLMs inherently aggregate tokens with higher similarities

The Fourier head’s job is to ensure that neighboring tokens have similar likelihoods during the next-token prediction step, and this is not guaranteed by LLMs. In fact, in the paper we provide examples of cases where LLMs fail to have neighboring tokens having similar likelihoods.

More precisely: if you enumerate the $m$ tokens $t_1, t_2, …, t_m$, then the Fourier head ensures that the learned next-token softmax distribution satisfies the property that the likelihood of $t_i$ is close to the likelihood of $t_{i+1}$. We demonstrate in our experiments that a general attention-based transformer with a linear classification head doesn’t satisfy this same property. We illustrate this in a figure in the paper: (link to graph).

This experiment demonstrates that an LLM with a standard classification head outputs a “jagged” next-token distribution, whereas the Fourier classification head outputs a “smooth” next-token distribution.

Thank you for the thoughtful feedback! Below we provide detailed answers to your questions. We have also uploaded a revised manuscript, with changes written in blue.

Q4: why is the paper organized this way (specifically, with related works near the end)?

We ultimately decided that we wanted the paper to flow from:

1. Motivating the desire to put a continuous structure over the next-token distribution (section 1, introduction)
2. Introducing our proposed method for putting a continuous structure over the next-token distribution (section 2, Fourier head)
3. Providing theoretical evidence that the Fourier head is capable of modeling complicated densities, explain the tradeoffs involved (section 3, Theory)
4. Demonstrating the Fourier head’s modeling capability in a low-dimensional intuitive example (section 4, toy example)
5. Demonstrating the Fourier head’s modeling capability in a large scale example (section 5, offline RL)
6. Demonstrating the Fourier head’s modeling capability in another large scale example (section 6, zero-shot probabilistic time series forecasting)

We reasoned that putting the related works between the motivation section, and the introduction of the Fourier head method, would break up this flow, so we opted to put it in the last section before the conclusion. Not all of us liked this, so our compromise was to have a “minimal” related works explanation embedded in the introduction; this is when we discussed the Decision Transformer, Chronos, and other instances where it could be beneficial to learn next-token distributions with a continuous structure. We’re very open to alternative ways of framing the story though, did you have any suggestions?

In the abstract, the author claims that "we can easily substitute for any linear layer if we want the outputs to have a more continuous structure." However, the output linear layer of methods like TimeLLM and GPT4TS could also be replaced with a Fourier head. I believe that it is not feasible to simply substitute Fourier head in TimeLLM and GPT4TS because they do not reconstruct the continuous vocabulary. (It is worth noting that LLM-based methods like TimeLLM and GPT4TS are applicable to zero-shot forecasting tasks, as demonstrated in Section 4.8 of GPT4TS. Although the datasets used differ, these LLM-based methods can perform well in out-of-domain tasks.)

Besides LLM-based methods, other TS foundation models such as UNITS[1] and MOMENT[2] cannot replace their output layers with Fourier head either. The author does not adequately explain how to implement Fourier head in these scenarios, which I think limits their usability.

Chronos is also a TS foundation model, and compared to LLM-based methods, TS foundation models require significant computational resources as they need to trainning from scratch or fully tuning on large datasets. Their broad applicability remains unproven (whereas LLM-based methods could be quickly trained and validated). If constructing an ordered vocabulary is an essential part of the process, then this disadvantage cannot be ignored.

[1] Gao S, Koker T, Queen O, et al. Units: Building a unified time series mode. NeurIPS2024

[2] Goswami M, Szafer K, Choudhry A, et al. MOMENT: A Family of Open Time-series Foundation Models[C]//Forty-first International Conference on Machine Learning.

We appreciate the reviewer’s engagement with us and would like to put the contributions of our Fourier Head in proper perspective with respect to the time-series forecasting methods the reviewer brought up.

1. When to use Fourier Head: As illustrated in our new toy experiments (revised manuscript, Figure 3 and Table 6), it is desirable to use the Fourier Head when we aim to model complex probabilistic distributions without knowing a priori the family of distributions. Particularly, Figure 3 shows that when the underlying distribution follows a Beta distribution, both vanilla linear head (no inductive bias) or GMM head (wrong inductive bias) are significantly outperformed by Fourier Head. Table 6 and Figure 9 (link to figure) further show that when the target distribution has multiple modes (e.g. GMM or Beta), pointwise methods regress to the mean, and fail to capture any of the modes.

2. TS forecasting baselines: To the best of our knowledge, GPT4TS, UNITS, and MOMENT are all pointwise forecasting models; they are hence “incompatible” with Fourier Head, which is designed to model (discretized) distributions. However, we believe the quoted statement from us remains true: one just needs to change both the “head” and the training objective (away from pointwise estimation), which we believe may lead to improved empirical performance as discussed in our response above. As the reviewer could see, these improvements are orthogonal to the design choices made by GPT4TS, UNITS, or MOMENT, meaning our contributions are complementary, not competing against these prior works.

3. Validity of the Chronos implementation: Chronos is a state-of-the-art probabilistic TS forecasting model recently accepted by the TMLR journal. We integrated the Fourier Head implementation into Chronos for its competitive performance, as well as open-source data preparation and training pipelines. We acknowledge that compared with TimeLLM or other similar methods, Chronos requires an additional training / fine-tuning stage. We would also like to clarify that the vocabulary construction process is independent from the requirement of “training from scratch”: Chronos actually explored fine-tuning with pre-trained LLM weights but observed that it has marginal impact on the performance. Additionally, updating the “vocabulary”, or even incorporating continuous features, is a standard practice, as exemplified by numerous multimodal LLM works, such as LLAVA. We hence believe this should not be considered as a fundamental disadvantage.

4. Broader impact of Fourier Head: Finally, we would like to clarify that we have demonstrated broader applicability of the Fourier Head beyond the time-series forecasting applications. We invite the reviewer to check our decision transformer experiments, especially the newly added evaluations (link to figure) where Fourier Head demonstrates consistent and significant performance improvements.

Lastly–please let us know if there are additional questions that we can provide further information / experiments on. Thank you again for your time and effort reviewing our paper!

We wanted to follow up, and check if our answers have satisfied your concerns. Thanks again for your effort reviewing our paper, and for your engagement throughout the rebuttal period!

Q4: providing more details on the Decision Transformer and Chronos experiments

For both Decision Transformer and Chronos, we followed the same training recipes from their original implementations. At your suggestion, we’ve added brief summaries of them to the paper. We will share some details here as well.

• Our Decision Transformer experiments: following the original implementation, we trained on 500k transitions observed 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 approx. 2.012M parameters) which takes about four hours on a single GPU.

• Our Chronos experiments: following the original implementation, we trained for 200k steps, on the same model size as the original implementation (a T5 model with approx. 20M parameters) which takes just under 48 hours on 8 GPUs.

Q3: how does the Fourier head’s performance scale with training data and model size?

Thank you for suggesting adding these valuable experiments, which we have added to the revised draft. We summarize the results of our additional ablations here, and also include links to the graphs:

• Ablation study #1: We show the Fourier head is better at learning high-quality next action distributions than the Decision Transformer with a Linear head across all dataset sizes. (link to graph)
• Ablation study #2: We show the Fourier head is better at learning high-quality next action distributions than the Decision Transformer with a Linear head across all model sizes. (link to graph)

These results demonstrate that the benefits of the Fourier head indeed persist with scale.

Q2: the Fourier head can only represent a finite range of values

In practice, we find that this is not an issue. If needed, a tanh reparameterization can be used to map the domain $[-1,1]$ to the real line. This would let the Fourier head learn unbounded values. We have added an explanation about this in Section 2.4 of the updated manuscript.

However, even if we decide to use a bounded domain (as is the case in all the experiments in the manuscript), this still isn’t an issue, as evidenced by the original Chronos model’s SOTA time series forecasting accuracy, with datasets which are not a priori bounded within some finite range. Those authors modeled unbounded sequences using a combination of two techniques: 1) normalizing the time series so the individual values are small, and 2) ensuring that the next-token distribution is defined over a sufficiently wide interval.

In more detail: in the Chronos paper, the authors normalize the time series so that the mean of the absolute value in the historical context is equal to 1; this ensures that the time series values are small. Then, the architecture defines the range of tokens as equally spaced inside of $[-15, 15]$; this ensures that the model is capable of learning from the examples with significant trend components. With this normalization and tokenization strategy, the Chronos model achieved SOTA accuracy because it turns out that even time series with aggressive trend components are normalized to values that can be forecasted accurately by the model.

Q1: alternative ways of extracting continuous probabilistic predictions from LLMs over numerical data

Thank you for mentioning the papers [1] Gruver et al and [2] Requeima et al, we have added their citations to the discussion in the related works section. However, we note that the Fourier head avoids a practical limitation from both of those methods. In short, the Fourier head’s design ensures that the categorical distribution which is learned during training is exactly the same distribution which is sampled from during evaluation. This is not the case for the methods in [1] and [2].

In more detail–in the methods from [2], the authors state the following limitation of their work: It must be noted that this approach does not guarantee that P(12) yields the mass assigned by the LLM to values in the bin [12, 13). However, we note that our method defines a valid predictive distribution, and we empirically observed that our predictive distribution closely matches the sampling distribution. In the appendix, they share images that show how the distributions are visually similar, but not identical. In contrast: in the Fourier head, the mass assigned to any range of numbers is exactly equal to the integral. This is a theoretical and practical benefit of modeling tokens over numerical distributions continuously using a smooth density (as in the Fourier head) rather than learning a discretized hierarchical softmax (as in [1] and [2]).

Thank you for the thoughtful feedback! We were able to produce the additional results that you asked for, summarized here:

• We conduct an additional ablation which demonstrates that the benefit of the Fourier head persists as the dataset size scales.
• We conduct an additional ablation which demonstrates that the benefit of the Fourier head persists as the model size scales.

Below, we provide more details, as well as question-specific responses. We have also uploaded a revised manuscript, with changes written in blue.

We wanted to follow up to make sure that your concerns are being properly addressed. Please let us know if there are additional questions that we can provide further information / experiments on. Thank you again for your time and effort reviewing our paper!

In the original Fourier Basis Density Model paper, the authors showed that the Fourier inductive bias is indeed natural for representing an (a priori) unbounded range of values once a tanh transformation is involved. In more detail–one of their modeling tasks involves learning a mixture of 25 Gaussians using a Fourier density with the tanh reparameterization, where most of the probability density is concentrated inside the range $[-10, 10]$. They find that the reparameterized Fourier density model can accurately learn the density, while capturing more modes and using fewer parameters than alternative models. (In case you’re curious, here is a link to an illustration of this example from their paper.)

We agree with the reviewer that results on scaling to even larger data size and model size would be interesting. We would like to clarify that our Chronos experiments, which demonstrate the effectiveness of Fourier Head, are already “large-scale”. The model has 20M parameters, and is pre-trained on over 10 million data points (the same setup as used by the Chronos paper). Further scaling up the Chronos model size and data size by orders of magnitude would be beyond the scope of our academic training size. Additionally, the scaling behavior of time-series forecasting models is on its own an active research area, we refer the reviewer to recent papers such as (Shi et al., NeurIPS 2024) on this topic.

We have conducted an additional study to analyze whether dataset size has any effect on the relative performance of the Linear head and the Fourier head for Chronos. Our results show that, across dataset sizes, the Fourier head indeed yields more accurate forecasts than the Linear head. We have updated the manuscript with these changes. (And here is a link to the figure that we added to the paper.) Thanks for the suggestion!

Additionally, we have started running experiments on scaling model size for Chronos, and we are hoping to finish them before the end of the extended rebuttal period.

As you requested, we have conducted an additional ablation study to analyze whether model size has any effect on the relative performance of the Linear head and the Fourier head for the Chronos time series forecasting task. Our results show that, across model sizes, the Fourier head indeed yields more accurate forecasts than the Linear head. It looks like we are not allowed to update the manuscript in OpenReview anymore, so here is a link to the paper updated with these changes. (And here is a link to just the figure that we added to the paper.)

Thanks for your patience during this rebuttal period while we ran our experiments, and thanks for suggesting that we run this experiment. Please let us know if you have any other questions!

We wanted to follow up, and check if the additional ablation studies on scaling model size and dataset size for Chronos satisfied your concerns. Thanks again for your effort reviewing our paper, and for your engagement throughout the rebuttal period!