ArXiv: 2501.19383
🎯 Pitch
Autoregressive decoders trained with cross-entropy loss can match pointwise heads on standard regression, while also flexibly modeling complex numeric distributions for density estimation—unlocking a unified architecture for both tasks.
1. Executive Summary
This paper introduces decoding-based regression heads—autoregressive Transformer decoders that represent numeric outputs as token sequences—and analyzes their effectiveness when attached to any feature encoder for standard regression tasks. Across tabular benchmarks (OpenML-CTR23, AMLB, UCI) and synthetic BBOB functions using an MLP encoder, the authors compare these heads against traditional pointwise heads, Riemann (histogram) distribution heads, and Mixture Density Networks. The key finding is that properly tuned decoding-based heads are competitive with pointwise heads for pointwise estimation while offering superior flexibility for density estimation, with the decoder outperforming pointwise heads on the majority of real-world tasks and achieving reliable negative log-likelihood (<0.7) across all UCI datasets where Gaussian mixtures show high variance (e.g., 7.49 ± 0.73 on Kin 40K). The paper also provides theoretical risk guarantees under a K-bit universality assumption, establishing that the decoder head can approximate smooth densities with bias decaying as 2^−2k and variance growing as 2^k/N, with practical gains over histogram heads emerging in low-data regimes due to implicit regularization.
2. Context and Motivation
The Core Problem: Representing Numeric Outputs in Neural Networks
The fundamental question this paper tackles is deceptively simple: when a neural network needs to output a real number, should that number come from the network's final layer as a scalar value, or can it usefully be represented as a sequence of discrete tokens that get decoded into a number? This question sits at the intersection of architectural design, representation learning, and probabilistic modeling, and the paper argues that the prevailing answer in most applications—just use a pointwise head—leaves significant flexibility on the table.
To understand why this question even arises, we need to recognize that every regression model has two distinct components that jointly determine its capabilities:
-
The encoder : a neural network (MLP, CNN, Transformer, etc.) that processes the input and produces a feature representation . This component absorbs domain-specific structure—spatial patterns in images, sequential dependencies in text, or tabular features in structured data.
-
The regression head : the final component that maps from the feature representation to a predictive distribution over the output . This is what the paper studies in isolation, keeping the encoder fixed across all compared methods.
The design of the regression head is not a trivial choice. It determines:
- What class of probability distributions the model can express (a single point estimate? a Gaussian? an arbitrary shape?)
- How the model's uncertainty is represented and whether that uncertainty can be multi-modal or asymmetric
- Whether the model handles unbounded outputs gracefully or requires careful normalization
- How the model learns about numeric quantities—through continuous error gradients or through discrete token-level feedback
The paper's position is that decoding-based heads—which treat numbers as token sequences that can be generated autoregressively—offer a principled and practical alternative to the dominant approaches, but that this alternative has been underexplored in controlled settings where the encoder and decoder are evaluated independently.
Why This Problem Matters Now: The Convergence of Two Trends
The practical importance of this work stems from two converging trends in machine learning that make numeric output representation suddenly relevant beyond traditional regression tasks.
Trend 1: Language Models as General-Purpose Regressors. Large language models are increasingly being used for tasks that require numeric predictions, not just text generation. The paper points to several recent developments in Section 2:
"The idea of text-to-text regression is especially relevant as LLMs are currently being fine-tuned as 'Generative Reward Models' (Mahan et al., 2024; Zhang et al., 2024), i.e. end-to-end scoring methods for reinforcement learning feedback"
In reinforcement learning from human feedback (RLHF), models need to output scalar reward values that assess the quality of generated text. These reward models are conventionally trained with Bradley-Terry preference losses, but this requires custom architectures (additional prediction heads appended to the language model) and custom training objectives. A text-to-text approach—where the model simply outputs a score as a string like "7.3"—would be architecturally simpler and could reuse the same cross-entropy training pipeline used for everything else the model does. But this raises the question: does treating a reward score as a string of tokens actually work well in terms of the quality of the resulting scalar predictions? The paper's controlled analysis of decoding-based heads provides evidence that it does, anchoring this emerging practice in rigorous benchmarking.
Trend 2: The Implicit Assumption That Numbers-as-Text Is "Unprincipled." There is a natural skepticism about using token-based representations for continuous quantities. The paper captures this skepticism directly:
"Understandably, one could argue that regular supervised fine-tuning over numbers represented as strings is unprincipled, considering that there is no notion of numeric distance when using cross-entropy loss."
Standard cross-entropy loss treats all token-level mistakes equally: predicting "7" when the ground truth is "8" incurs the same loss as predicting "0" when the ground truth is "8," even though the first mistake is numerically closer. There is no built-in inductive bias that neighboring numbers should be treated similarly. This stands in sharp contrast to mean squared error (MSE) loss on scalar outputs, where the gradient naturally reflects numeric distance. The paper's theoretical analysis in Section 3.3 directly addresses this concern by showing that under a tree-based tokenization scheme, the cross-entropy-trained decoder can still achieve consistent density estimation with known bias-variance tradeoffs—numeric distance is not lost; it emerges from the hierarchical structure of the representation.
The Practical Stakes. The choice of regression head has concrete downstream consequences that the paper highlights through several lenses:
-
For self-improvement pipelines and RLHF: If decoding-based reward models work well, the entire RLHF pipeline could be simplified—no need for separate reward model architectures, no need for pairwise preference data construction, no need for custom loss functions. Everything becomes next-token prediction, which is what language models are optimized to do.
-
For density estimation: Many applications require not just a point estimate of given , but a full conditional distribution —think uncertainty quantification, risk assessment, or generative modeling of continuous outcomes. Traditional parametric heads (Gaussians) force unimodal, symmetric assumptions. The paper shows that decoding-based heads can capture multi-modal, asymmetric, or oddly-shaped distributions that Gaussian mixtures struggle with or require many components to approximate.
-
For multi-task settings: As noted in Song et al. (2024) and referenced in Section 3.1, when a single model must handle regression tasks with very different -scales (some outputs in , others in , still others unbounded), normalizing all outputs into a common range becomes tedious and can introduce numeric instability. The unnormalized tokenization scheme (Section 3.1) handles this by representing numbers in a floating-point-like format that naturally covers many orders of magnitude.
Prior Approaches and Their Limitations
The paper organizes existing regression heads into three families, each with distinct strengths and failure modes that motivate the exploration of decoding-based alternatives (Section 2).
Pointwise Heads: Simple but Rigid
Pointwise heads are the workhorse of deep learning regression. They consist of a deterministic function—typically a linear layer or small feed-forward network—that maps directly to a scalar , trained by minimizing a pointwise loss like mean squared error (MSE).
What they do well: They are simple, computationally cheap, and well-understood. For problems where a single best guess is sufficient and the output range is well-behaved, they work excellently.
Where they fall short—and why it matters for this paper:
-
-normalization is essential for stability. The paper notes that "the -values must be normalized in space, e.g. within " (Section 2). Without this, training can become unstable because MSE gradients scale with prediction error magnitude. But normalization introduces its own problems: if the true -range is large or unknown, min-max scaling can compress values into a tiny sub-interval of , leading to floating-point precision issues. The paper observes this directly in Figure 4, where the pointwise head struggles with functions that have "very high or unbounded -ranges" (Section 4.1). Additionally, in low-data regimes, the pointwise head can produce predictions outside due to undertraining, requiring the authors to "append a sigmoid activation to enforce the normalized output to be within to avoid extremely high MSE errors" (Section 4.2).
-
Cannot represent multi-modal or asymmetric uncertainty. A pointwise head produces a single number. It cannot express "the output is either around 2.3 or around 8.7, but definitely not in between." For applications requiring uncertainty quantification, this is a fundamental limitation.
-
Struggles with abrupt changes. Functions with large Lipschitz constants—steep slopes, sharp transitions—are challenging for pointwise heads because the continuous MSE loss smooths over discontinuities. Figure 4 demonstrates this on stepwise and hyperbolic functions where the pointwise head's fit is visually poor while the decoder head captures the shape.
Parametric Distribution Heads: Probabilistic but Constrained in Shape
To address the uncertainty-representation limitation of pointwise heads, parametric distribution heads output the parameters of a probability distribution rather than a single scalar. The most common example is a Gaussian head: , where both and are learned functions of . Training proceeds by maximizing log-likelihood rather than minimizing MSE.
What they do well: They provide principled uncertainty estimates (predictive variance) alongside point predictions. The Gaussian assumption is reasonable for many natural phenomena.
Where they fall short:
The paper focuses on Gaussian Mixture Models (Mixture Density Networks, or MDNs) as the more flexible variant of this family. An MDN with components models:
where the mixture weights , means , and standard deviations are all functions of . In principle, with enough components, an MDN can approximate any continuous density arbitrarily well (Bishop, 1994).
However, the paper's empirical results reveal significant practical limitations:
-
High variance across tasks. Table 2 (Section 4.3) shows MDN negative log-likelihood ranging from excellent (0.05 ± 0.12 on Wine) to catastrophic (7.49 ± 0.73 on Kin 40K). This variance—standard deviations that sometimes exceed means—indicates that MDNs can fail badly on certain datasets, requiring careful tuning of component count and training procedure.
-
Component count is a critical hyperparameter. Too few components and the model can't capture the true distribution shape; too many and optimization becomes difficult (mixture weights collapse, components diverge). The paper sweeps over [1, 2, 5, 10, 20, 50, 1000] (Appendix C), but even with this sweep, MDN performance deteriorates on certain UCI datasets (Table 3 in Appendix A.5 shows negative NLL values on Challenger, Fertility, Solar, and Stock, which can occur when the fitted density assigns probability mass where test points fall but the distribution is extremely peaked or the model overfits).
-
Normalization still required. Like pointwise heads, MDNs typically benefit from -normalization to keep means and variances in well-behaved ranges.
Riemann (Histogram) Distribution Heads: Flexible but Data-Hungry
Histogram-based heads, which the paper terms Riemann heads following Hollmann et al. (2025), discretize the output space into equally-spaced bins over a finite support set . The head outputs a categorical distribution over these bins via softmax:
where contains a learned embedding for each bin. This approach has proven effective in distributional reinforcement learning (Bellemare et al., 2017) and tabular data settings (Chen et al., 2022; Hollmann et al., 2025).
What they do well: They can represent arbitrary distribution shapes without parametric assumptions—multi-modal, skewed, heavy-tailed distributions are all possible. The discrete bin structure also makes them natural for value-based RL where actions have discrete returns.
Where they fall short—and why this directly motivates decoding-based heads:
- The curse of dimensionality in bin count. To achieve fine-grained resolution, you need many bins. But bins means separate learnable embedding vectors in , and the model must learn to distinguish between all pairs of bins. The paper observes:
"a drawback is that learning numeric distances between all of the bins requires more data as the size of the vocabulary increases" (Section 2)
This is the core empirical finding: Riemann heads are data-inefficient when bin counts are high. Figure 7 (Section 4.2) shows Riemann heads plateauing on several AMLB tasks while the decoder head continues to improve with more data. On task 233212, for instance, the Riemann head with 1024 bins essentially flatlines below the decoder's performance, never catching up even with training points.
- The exponential reduction insight. The paper's key observation about Riemann heads is that they represent a special case of decoding with sequence length —a single token selects one bin among choices. The natural extension is to use longer sequences:
"By extending the sequence length instead, there can be an exponential reduction in bin count – e.g. 1000 (= 10^3) bins can be expressed instead using 10 bins and 3 decoding steps."
This is the central architectural argument for decoding-based heads: instead of separate embedding vectors to represent bins, you represent the same resolution using only vocabulary items that are combined sequentially. The model learns to represent numbers hierarchically—first the most significant digit, then the next, refining the estimate step by step—rather than learning each level of precision independently.
Connection to extreme classification. The paper notes that this sequential decomposition idea has been studied in extreme multi-label classification (Wydmuch et al., 2018) where the output space is combinatorially large, but:
"it has not been thoroughly examined for numeric regression, which is the focus of our work" (Section 2)
This situates the paper at the intersection of two literatures that rarely speak to each other: the classification community's work on hierarchical output spaces, and the regression community's work on continuous density estimation.
The Gap: No Systematic Analysis of Decoding-Based Regression in Controlled Settings
The paper identifies a clear gap in the literature. Prior work on text-to-text regression with language models (Akhauri et al., 2025; Song et al., 2024; Vacareanu et al., 2024) treats both input and output as text, processing everything through a single language model pipeline. This makes it impossible to disentangle whether performance comes from the language model's understanding of the input (via text representations and pretraining) or from the decoding-based representation of the output.
Other work keeps the traditional regression head but varies the input representation: Tang et al. (2024) use LLM embeddings of text inputs attached to standard feed-forward regression heads, while Nguyen et al. (2024) attach Gaussian heads to these same embeddings. These studies investigate whether LLM representations improve regression but preserve the conventional output head.
The paper's framing of this gap is precise:
"However, there has not been work investigating the inverse situation, i.e. is represented as text or structured tokens." (Section 1)
The "inverse situation" is: keep the encoder generic (a standard MLP, no language model pretraining), and only vary the regression head. This isolates the contribution of the decoding-based representation itself, separate from any benefits of LLM-based input processing. The paper's encoder is intentionally simple—a ReLU MLP with 2–5 layers and up to 2048 hidden units—and the decoder head uses only 1 layer and 32 units (less than 10% of total parameters), making it clear that any differences in performance come from the output representation strategy rather than from additional model capacity.
How This Paper Positions Itself
The paper explicitly positions itself as a controlled, principled study of output representation for regression, distinct from two neighboring research programs:
Distinct from LLM-as-regressor work. While acknowledging work like Vacareanu et al. (2024) that shows GPT-4 and Gemini can perform regression, the paper distances itself from the "prompt engineering" approach to numeric prediction. Instead, it asks: if we design a small autoregressive Transformer from scratch with a numeric tokenization scheme, how does it compare to standard regression heads when given identical features? This is a question about representation, not about scale.
Distinct from input-representation work. Studies like Tang et al. (2024) and Nguyen et al. (2024) ask whether better input representations (from LLMs) improve regression. This paper asks the complementary question: whether better output representations improve regression, even with simple input encoders. The two lines of work are orthogonal and could be combined—one could use an LLM encoder with a decoding-based head—but the paper keeps them separate for clarity of analysis.
Building on but generalizing Riemann/histogram approaches. The paper treats Riemann heads not as competitors to be beaten but as a special case ( sequence length) of a more general framework. The theoretical analysis in Section 3.3 and Appendix B explicitly derives risk bounds that apply to both Riemann and decoding heads, unifying them under a common bias-variance decomposition. The practical advantage of decoding heads over Riemann heads then emerges naturally from the theory: when is small relative to (the low-data, high-resolution regime), the decoder head's implicit regularization allows it to achieve lower risk by learning smoother distributions than the histogram's bin-conditional empirical frequencies.
A note on terminology. The paper includes a careful distinction in Section 2 to prevent conflation with Variational Autoencoders, where "decoder" refers to the network that reconstructs inputs from latent codes. The paper's "decoder head" is instead an autoregressive sequence model that generates the output token-by-token given . This clarification matters because both VAEs and decoding-based regression heads map from a feature vector to a distribution, but they serve fundamentally different purposes (reconstruction vs. prediction) and are evaluated differently.
The Theoretical Motivation: Why Tokenization Might Work Despite Lacking Numeric Distance Inductive Bias
The paper anticipates and addresses the most natural objection to decoding-based regression head-on:
"one could argue that regular supervised fine-tuning over numbers represented as strings is unprincipled, considering that there is no notion of numeric distance when using cross-entropy loss"
The resolution comes from the tree structure implicit in positional numeric tokenization. Section 3.3 develops this idea formally, but the intuition is worth understanding now because it underpins the entire approach:
Consider binary (base-2) tokenization of numbers in . The first bit determines whether is in or . The second bit determines which quarter-interval within that half. The third bit determines which eighth-interval. And so on. This means:
- Mistakes at early bits are numerically large: predicting the wrong first bit puts the estimate in the wrong half of the interval, incurring a large numeric error.
- Mistakes at later bits are numerically small: predicting the wrong fifth bit shifts the estimate by at most .
- Cross-entropy loss reflects this hierarchy implicitly: the model can get the first few bits right (coarse placement) with high confidence while being more uncertain about the final bits (fine details). The loss decomposes across bit positions, and the model learns to allocate its representational capacity such that the most significant digits are predicted most reliably.
This is visualized in Figure 3 (Section 3.3), which shows how a truncated Gaussian distribution is iteratively refined through a binary tree: each additional bit subdivides each existing bin into two, doubling the resolution. The decoding process traverses this tree from root to leaf, making increasingly fine-grained decisions.
The theoretical contribution formalizes this intuition. Theorem 1 (Section 3.3) proves that under a "-bit universality" assumption—the model class is flexible enough to fit any discrete distribution over categories—the risk decomposes into:
R(f, f_k^*_N) = \underbrace{\frac{2^{-2k}}{12} \int_0^1 f'(y)^2 dy}_{\text{Bias}} + \underbrace{\frac{2^k}{N}}_{\text{Variance}} + O(2^{-4k} + 1/N)
The bias term decays exponentially with (more bits = better resolution = less smoothing error), while the variance term grows exponentially with (more bins = fewer samples per bin = noisier estimates). The optimal balances these two terms. This is exactly the classical histogram density estimation tradeoff, but the crucial difference is that the decoder head can operate at different effective resolutions for different parts of the distribution—it's not committed to a single fixed because the autoregressive decoding can stop early (in principle) or the model can learn to be more certain about some branches than others.
The paper also observes an intriguing phenomenon in Figure 2: in the low-data regime (, large ), the decoder head significantly outperforms the theoretical risk curve that the Riemann head closely follows. The authors hypothesize:
"a combination of the inductive bias of our model class and the implicit bias of our SGD training procedure makes the decoder less likely to fit noise; a concrete example would be that the model is biased to learn smooth distributions"
This is a crucial empirical finding that goes beyond the formal theory: the decoder head appears to possess implicit regularization that the simpler histogram estimator lacks. When there are too few samples to reliably estimate all bin probabilities from empirical frequencies, the decoder head smooths across bins rather than fitting the noisy empirical distribution exactly. This smoothness bias emerges from the combination of the Transformer architecture and SGD training dynamics, making the decoder more data-efficient than the theory (which assumes perfect optimization) would predict.
Summary of the Case for Investigating Decoding-Based Regression
The paper's motivation can be distilled into a set of linked claims about why this approach merits serious study:
-
Architectural simplicity: Decoding-based heads use the same cross-entropy training objective as classification and language modeling—no custom losses, no distributional assumptions, no normalization constraints. This makes them drop-in replacements for standard heads with minimal engineering overhead.
-
Distributional flexibility: Unlike Gaussian or mixture heads, decoding-based heads can represent arbitrary distribution shapes without committing to a parametric family. Unlike Riemann heads, they achieve this without an explosion in learnable parameters.
-
Unbounded output support: The unnormalized tokenization scheme (Section 3.1) naturally handles outputs across many orders of magnitude, avoiding the normalization sensitivity that plagues pointwise and parametric heads.
-
Theoretical grounding: The tree-based tokenization structure provides a clean bias-variance decomposition that connects decoding-based regression to classical nonparametric density estimation, giving principled guidance on choosing the representation granularity.
-
Empirical competitiveness: The experiments aim to show that these benefits do not come at the cost of pointwise prediction accuracy—decoding heads are competitive with or better than pointwise heads on standard benchmarks—making them a practical choice, not just a theoretically interesting one.
The next sections of the paper (Sections 3 and 4) develop these claims through formal analysis and extensive benchmarking, testing whether the promises of decoding-based regression hold up under controlled experimental conditions.
3. Technical Approach
3.1 Reader Orientation
This paper designs and evaluates decoding-based regression heads: small autoregressive Transformer decoders that represent numeric outputs as sequences of discrete tokens (like "1", ".", "2", "3" for the number 1.23) and are trained with standard next-token prediction cross-entropy loss. The core insight is that attaching such a head to any feature encoder solves a fundamental tension in regression architecture design—pointwise heads are simple but cannot represent uncertainty, while flexible distributional heads (like histograms or Gaussian mixtures) either scale poorly with output resolution or impose restrictive parametric assumptions on the shape of the conditional density $p(y|x)$.
3.2 Big-Picture Architecture (Diagram in Words)
The system has three major components arranged in a feed-forward pipeline:
-
Feature encoder
$\phi(x)$: A standard multi-layer perceptron (MLP) with ReLU activations that processes the input vector$x$and produces a fixed-dimensional representation$\phi(x) \in \mathbb{R}^d$. This component is shared identically across all regression head types compared in the paper, ensuring that any performance differences come from the output representation strategy rather than from additional model capacity or better input processing. -
Decoder head (autoregressive Transformer): A small Transformer decoder that takes
$\phi(x)$as its initial embedding and generates a sequence of tokens$(t_1, t_2, \ldots, t_K)$from a finite vocabulary$\mathcal{V}$, one token at a time, left-to-right. Each token is sampled from the conditional distribution$p_\theta(t_k | \phi(x), t_1, \ldots, t_{k-1})$produced by the Transformer's output softmax at position$k$. The sequence of tokens is a string representation of a number—either a base-$B$expansion of a normalized value in$[0,1]$or a floating-point-like representation with separate sign, exponent, and mantissa components for unnormalized values. Because the tokens are generated autoregressively, the head implicitly defines a probability distribution over all representable numbers through the chain rule:$p_\theta(t_1, \ldots, t_K | \phi(x)) = \prod_{k=1}^K p_\theta(t_k | \phi(x), t_1, \ldots, t_{k-1})$. -
Token-to-number mapper: A deterministic post-processing step that converts the decoded token sequence back into a real number. For normalized tokenization (base-
$B$expansion in$[0,1]$), this is a simple base-conversion operation. For unnormalized tokenization (IEEE-754-style floating-point), this reassembles the sign, exponent, and mantissa tokens using the formula$\text{sign} \times B^{\text{exponent}} \times \text{mantissa}$. The mapper also handles edge cases: if the decoded sequence is invalid (e.g., violates the constrained token grammar), the output might beNaNor an out-of-range sentinel value.
Information flow during inference: Input $x$ passes through the encoder to produce $\phi(x)$ → the decoder head takes $\phi(x)$ as its initial context and autoregressively samples tokens $(t_1, \ldots, t_K)$ (potentially with constrained decoding to enforce valid number representations) → the token-to-number mapper converts the token sequence to a real-valued prediction $\hat{y}$. If the application requires a point estimate, an aggregation function (mean, median, or mode) is applied over multiple sampled token sequences. If the application requires a density estimate, the token-level probabilities from the decoder directly define $p_\theta(y|\phi(x))$ over the discrete grid of representable numbers.
3.3 Roadmap for the Deep Dive
-
First, the numeric tokenization schemes (Section 3.1): These are the foundation of the entire approach—how real numbers get encoded as token sequences and decoded back. I explain both the normalized (base-
$B$expansion within$[0,1]$) and unnormalized (floating-point with sign, exponent, mantissa) schemes, because the choice between them determines whether the head can handle unbounded outputs and multi-task settings with varying scales. -
Second, pointwise estimation from the decoder (Section 3.2): Once the decoder defines
$p_\theta(y|\phi(x))$, how do we extract a single scalar prediction (mean, median, mode)? This is non-trivial because the decoder can assign non-zero probability to arbitrarily large outliers (especially in unnormalized mode), and naive sample-mean aggregation can be catastrophically sensitive to these outliers. I then explain error-correction tokenization and alternative decoding techniques that mitigate this problem. -
Third, the density estimation theory (Section 3.3): This provides formal guarantees that decoding-based regression works despite using cross-entropy loss (which has no built-in notion of numeric distance). Theorem 1 decomposes the risk into bias and variance terms that depend on tokenization depth
$k$, and I walk through why the decoder's hierarchical tree structure naturally captures numeric proximity even though the loss treats tokens as discrete categories. I also explain the implicit regularization phenomenon in Figure 2 where the decoder outperforms the theoretical risk bound in low-data regimes. -
Fourth, training procedure and loss: How the decoder head is trained end-to-end with the encoder using standard cross-entropy loss across all token positions simultaneously, and how constrained decoding is enforced during inference (but not training) to guarantee valid number representations.
3.4 Detailed, Sentence-Based Technical Breakdown
This is primarily an empirical analysis paper with theoretical grounding whose core idea is that an autoregressive decoder trained with cross-entropy loss over tokenized numeric representations can serve as a flexible, data-efficient regression head that is competitive with pointwise heads for scalar prediction and superior to parametric distribution heads for density estimation.
Numeric Tokenization Schemes
The fundamental design choice in decoding-based regression is how to map between real numbers (a continuum) and discrete token sequences (a finite combinatoric space). The paper proposes two schemes for different use cases, both based on the intuition that numbers can be represented hierarchically—coarse structure first, then increasingly fine detail—and that an autoregressive decoder can naturally exploit this hierarchy.
Normalized Tokenization (for outputs with known, bounded range):
When the output $y$ can be normalized to $[0, 1]$ (e.g., via min-max scaling using training set statistics), the tokenization is simply the base-$B$ expansion of $y$ truncated to $K$ digits. The vocabulary $\mathcal{V}$ contains tokens <0>, <1>, ..., <B-1>, and $y$ is represented as a length-$K$ sequence $t_1 t_2 \ldots t_K$ where $t_1$ is the most significant digit and $t_K$ is the least significant:
where $t_k \in \{0, 1, \ldots, B-1\}$ is the token at position $k$ and $B$ is the chosen base (swept over 2, 4, or 8 in experiments; see Appendix C).
What it computes: given a real number $y \in [0, 1]$, the tokenization procedure converts it to its first $K$ base-$B$ digits by iteratively multiplying by $B$, taking the integer part as the next token, and continuing with the fractional remainder. The inverse operation (token-to-number mapping) reconstructs the value as a $K$-digit fraction in $[0, 1]$. The representable values form a uniform grid of $B^K$ points spaced $B^{-K}$ apart.
Why this form: this representation has a tree structure where each additional digit subdivides the remaining interval into $B$ equal pieces. The first digit $t_1$ determines which of $B$ coarse intervals $[j/B, (j+1)/B)$ the number falls into; the second digit refines within that interval; and so on. This means the decoder learns to make coarse predictions first (high confidence on $t_1$), then progressively finer adjustments (potentially higher uncertainty on later digits). A flat histogram head would need $B^K$ separate parameters to achieve the same resolution, while the decoding head uses only $B \times K$ vocabulary embeddings shared across positions, with the autoregressive Transformer learning to compose them. The hyperparameter choices $B \in \{2, 4, 8\}$ and $K \in \{4, 6, 8\}$ (Appendix C) provide resolutions from $2^4 = 16$ bins up to $8^8 = 16,777,216$ bins, spanning the range from very coarse to extremely fine-grained.
Unnormalized Tokenization (for outputs with unknown or varying scale):
When $y$ can span multiple orders of magnitude or when multiple tasks with different scales share the same model, normalizing to $[0,1]$ becomes impractical. The unnormalized scheme generalizes the base-$B$ expansion to a floating-point representation inspired by the IEEE-754 standard. The token sequence has the form:
<sign_s> <sign_e> <e_1> ... <e_E> <m_1> ... <m_M>
where:
<sign_s>is the sign of the significand (mantissa), encoded as either a dedicated<+>or<->token or optionally reusing<0>and<1>tokens from the vocabulary.<sign_e>is the sign of the exponent, similarly encoded.<e_1> ... <e_E>is the base-$B$representation of the exponent magnitude.<m_1> ... <m_M>is the base-$B$representation of the mantissa's most significant digits.
The decoded numerical value is:
where $\text{sign}_s \in \{-1, +1\}$ is the decoded sign of the output, $\text{sign}_e \in \{-1, +1\}$ is the decoded sign of the exponent, and $\text{value}(e_1 \ldots e_E)$ is the base-$B$ integer represented by the exponent digits.
The paper gives a concrete example (Section 3.1): if $B=10$, $E=3$, $M=4$, then the number $10^{-222} \times 1.23456789$ is represented as <+> <-> <2> <2> <2> <1> <2> <3> <4>. The trailing digits 56789 are truncated—this is the rounding error inherent in any finite-length representation of real numbers.
What it computes: the unnormalized tokenization decomposes a real number into three parts: a sign (positive or negative), an exponent that determines the order of magnitude, and a mantissa that determines the precise value within that order of magnitude. The range of representable values is enormous: with $E$ exponent digits in base $B$, the exponent can range from $0$ to $B^E - 1$ (times the sign), so the output can span from $B^{-(B^E-1)}$ to $B^{B^E-1}$. For $B=10, E=2$, this is approximately $10^{-99}$ to $10^{99}$—more than sufficient for practically any regression task. The resolution within any decade is determined by $M$: with $M$ mantissa digits, there are $B^M$ representable values per decade.
Why this form: the key advantage over normalized tokenization is scale insensitivity. The model can learn to output numbers at very different magnitudes without the encoder needing to know the output range in advance. This matters for multi-task regression (Song et al., 2024) where a single model might predict both completion time in seconds (range $10^0$ to $10^3$) and memory usage in bytes (range $10^6$ to $10^{10}$). Normalizing these to a common $[0,1]$ range would require task-specific scaling factors and careful balancing. The unnormalized scheme handles this naturally: the exponent tokens adapt to the scale, and the mantissa tokens handle the precision. A secondary advantage is stability with unbounded functions: in the curve fitting experiments (Figure 4, Section 4.1), functions with vertical asymptotes (hyperbolic, tangent) have $y$ values that diverge to $\pm \infty$—these would require aggressive clipping under normalized tokenization, but the floating-point representation can express very large (though finite) values without special handling.
Design choices and their justifications:
-
Base
$B$is swept over small values ($B \in \{2, 4, 8, 10\}$): larger bases mean fewer tokens per number (shorter sequences) but more vocabulary items to learn. Base 2 gives maximal sequence length for a given resolution, which means more autoregressive steps but simpler per-step decisions (binary choice). Base 10 aligns with human-readable decimal representation. The sweeps in Appendix C let the paper determine empirically whether the tradeoff matters significantly. -
Sign tokens can be dedicated or reused: the paper notes that
<+>and<->can be separate tokens or reuse<0>and<1>(Section 3.1) and "this made little difference in results." This is a robustness check suggesting the decoder is not sensitive to the exact vocabulary design as long as the positional structure is preserved. -
Special
<NaN>token for invalid inputs: not used in the paper's experiments (which assume$y \in \mathbb{R}$), but mentioned as a practical extension for downstream applications where$x$might be an invalid input that should produce a "not a number" output. This is a capability that standard regression heads cannot gracefully express without ad-hoc sentinel values. -
Constrained decoding ensures valid sequences: during inference, the decoder's sampling is restricted to only those token sequences that correspond to valid numbers. For example, in normalized tokenization with
$B=10$and$K=4$, any sequence of four digit tokens is valid. In unnormalized tokenization, the grammar is more complex—the sign token must come first, followed by exponent-sign, then$E$exponent digits, then$M$mantissa digits—and constrained decoding prevents the model from producing syntactically malformed outputs. The paper does not enforce constrained decoding during training (the model learns the grammar from data), only during inference.
Vocabulary and architecture details:
The vocabulary $\mathcal{V}$ contains all tokens needed for the chosen tokenization scheme (digits 0 through $B-1$, signs if dedicated, plus any special tokens like <NaN>). The autoregressive decoder is a Transformer decoder (Vaswani et al., 2017) with:
-
1 layer and 32 hidden units by default, or an alternative configuration of 3 layers, 128 units, 4 attention heads (Appendix C). The paper emphasizes that the small configuration "makes up for less than 10% of the total network parameter count" (Section 4, preamble), meaning the decoder contributes negligible capacity relative to the encoder. This is a deliberate design choice to isolate the effect of the output representation: if the decoder were large enough to significantly increase total model capacity, performance improvements could be attributed to having more parameters rather than to the decoding-based approach itself.
-
The initial token embedding is set to the encoder output
$\phi(x) \in \mathbb{R}^d$—the decoder does not have a learned start-of-sequence token. Instead,$\phi(x)$is projected (if necessary) to match the decoder's hidden dimension and used as the query context for generating the first output token$t_1$. This is the standard way to condition an autoregressive decoder on an external representation: the encoder output serves as the "prompt" from which the decoder generates. -
Causal (autoregressive) attention masking ensures that each token
$t_k$can only attend to$\phi(x)$and previously generated tokens$t_1, \ldots, t_{k-1}$, not to future tokens. This is the standard left-to-right generation constraint that enables efficient training via teacher forcing and consistent sampling at inference.
Pointwise Estimation from the Decoder
Once the decoder defines a distribution $p_\theta(y|\phi(x))$ over the discrete grid of representable numbers, many applications require a single scalar prediction $\hat{y}$ rather than a full distribution. The paper frames this as selecting a pointwise functional $M(p_\theta)$ of the model's predictive distribution that minimizes expected loss under some error function $\ell : \mathbb{R}^2 \to \mathbb{R}$:
where $\ell$ is the pointwise loss function (L2, L1, L0, etc.), $p(\cdot|x)$ is the true conditional distribution, and $M(p)$ is the optimal point estimate under that loss.
What it computes: the mapping from loss function to optimal point estimate. For standard choices:
- L2 loss (squared error):
$M(p)$is the mean of$p(\cdot|x)$. - L1 loss (absolute error):
$M(p)$is the median of$p(\cdot|x)$. - L0 loss (0-1 error):
$M(p)$is the mode of$p(\cdot|x)$.
Why this form: this framing decouples the modeling problem (learning $p_\theta$ that approximates the true conditional distribution) from the decision problem (choosing a point estimate from that distribution). The same decoder head can serve different downstream needs—mean for MSE-sensitive applications, median for outlier-robust applications, mode for applications requiring the most likely value—without retraining. Standard pointwise heads trained with MSE directly learn to output the conditional mean, but they cannot trivially provide the conditional median or mode; decoding heads learn the full distribution and can extract any of these functionals post-hoc.
Practical estimation from samples:
The decoder head does not provide these functionals in closed form—$p_\theta$ is implicitly defined through the token-level probabilities. To estimate $M(p_\theta)$, the paper discusses several approaches:
-
Mode estimation via beam search: The mode (most likely complete sequence) can be approximated by running beam search during decoding rather than sampling—at each step, keep the top-
$W$highest-probability partial sequences and expand them. The final highest-probability sequence is the approximate mode. This is computationally efficient but only provides the mode, not the mean or median. -
Sample-based estimation of mean/median: Draw
$S$independent samples$y^{(1)}, \ldots, y^{(S)}$from$p_\theta(\cdot|\phi(x))$using vanilla temperature sampling, decode each to a real number via the token-to-number mapper, and compute the sample mean or sample median. The sample mean is:
The sample median can be estimated efficiently using the Harrell-Davis estimator (Harrell and Davis, 1982), which the paper cites as a method that "can efficiently approximate the true median from pure temperature samples" (Section 3.2).
The outlier problem with unnormalized tokenization:
A critical practical issue arises with sample-based mean estimation under unnormalized tokenization. The paper explains:
"Especially for unnormalized tokenization, additional care needs to be taken, since in practice, the model can have a miniscule but non-zero probability of decoding an arbitrarily large outlier, even if the underlying true distribution is bounded. Such outliers can easily sway non-robust estimators such as the sample mean."
The floating-point representation can express numbers up to $B^{B^E}$—for $B=10, E=2$, this is $10^{99}$. Even if the model assigns extremely low probability to such values (say, $10^{-9}$ per token sequence), the variance contribution to the sample mean is proportional to the outlier magnitude squared, which can overwhelm the signal from typical samples. A single $10^{99}$ sample among 100 typical samples around $10^0$ would make the sample mean approximately $10^{97}$—catastrophically wrong.
Solutions discussed:
-
Robust point estimators: Use the sample median instead of the sample mean, since the median is insensitive to outliers regardless of their magnitude (as long as fewer than half the samples are outliers). More generally, the paper cites Lehmann (1983) on statistical point estimators for robust alternatives.
-
Error-correction tokenization (Section 4.5): Inspired by coding theory, the paper introduces token repetition as a form of error correction. During training, the target sequence is repeated multiple times:
$(t_1, \ldots, t_K, t'_1, \ldots, t'_K, t''_1, \ldots, t''_K, \ldots)$, where each repetition encodes the same number. At inference, majority voting is performed independently at each token position$k \in \{1, \ldots, K\}$: for position$k$, take the most frequently sampled token across all repetitions, then assemble the consensus sequence. Formally, if$R$is the number of repetitions and$t_k^{(r)}$is the token sampled at position$k$in repetition$r \in \{1, \ldots, R\}$, the consensus token is:
The decoded number is then computed from $(t_1^*, \ldots, t_K^*)$.
Why this form: token-level majority voting reduces the effective probability of outlier sequences. For an outlier to affect the final prediction, the model must sample an outlier token at the same position in the majority of repetitions. If outlier tokens are individually low-probability (as is typical—the model assigns high probability to correct or near-correct tokens), the probability that they win a majority vote decays exponentially with the number of repetitions. Figure 10 (Section 4.5) demonstrates this: with $R=1$ (no repetition), the relative MSE increases with more samples because outliers become more likely to appear at least once; with $R=3$ or $R=5$, the MSE decreases with sample count as the majority vote filters out rare errors.
- Decoding techniques from LLM literature: The paper mentions top-
$k$sampling (Fan et al., 2018), top-$p$(nucleus) sampling (Holtzman et al., 2020), and temperature reduction as ways to "filter out possible outliers" by restricting the decoder to high-probability tokens. Lowering the temperature sharpens the token-level softmax:
where $T < 1$ makes the distribution more peaked around the mode. Top-$k$ sampling restricts the sampling to the $k$ highest-probability tokens, clipping the tail entirely. These techniques reduce the chance of sampling outlier tokens but may also reduce the diversity needed for accurate density estimation—the paper reports that "vanilla temperature sampling with temperature ≈ 1.0 is the best way to match $p(y|x)$" (Section 4.3) for density estimation, while for pointwise estimation, more aggressive filtering may be appropriate.
- RAFT (Regression-Aware Fine-Tuning) approach (mentioned but not used): An alternative that avoids sampling altogether by computing expectations in closed form over a fixed finite evaluation set
$\mathcal{Y}$:
where $p_\theta(y')$ is the full-sequence probability from the decoder. The paper notes two limitations: "the choice of $\mathcal{Y}$ may be non-trivial to obtain an unbiased estimate, especially over unnormalized tokenizations" (since you would need to cover an enormous range), and "this may also defeat the purpose of using a decoding head, which offers several density estimation benefits"—since RAFT requires enumerating the support explicitly, you lose the decoder's ability to implicitly represent distributions without exhaustive bin enumeration.
The paper's stance: The paper treats "the choice of method for computing pointwise representations" as "a hyperparameter to be tuned depending on the application" (Section 3.2). There is no single best method—the tradeoff depends on whether outliers are a concern, whether density estimation or pointwise prediction is the primary goal, and computational constraints (beam search vs. sampling).
Density Estimation and Theory
The density estimation capability is what distinguishes decoding-based heads from pointwise heads and (arguably) from Riemann histogram heads. This section develops the formal framework and theoretical guarantees.
Training objective: sequence-level cross-entropy loss.
The decoder head is trained with the standard next-token prediction cross-entropy loss, summed over all token positions in the sequence. For a single training example with target token sequence $y = (t_1, \ldots, t_K)$ and model output probabilities $p_\theta(t_k | \phi(x), t_1, \ldots, t_{k-1})$:
where $K$ is the sequence length (number of tokens per number), $\mathcal{V}$ is the vocabulary, $\mathbb{1}(\cdot)$ is the indicator function that is 1 when $\hat{t}_k$ equals the true token $t_k$ and 0 otherwise, and $p_\theta(\hat{t}_k | \cdot)$ is the model's predicted probability for token $\hat{t}_k$ at position $k$.
What it computes: for each position $k$ in the target sequence, the model produces a categorical distribution over the vocabulary $\mathcal{V}$. The loss sums the negative log-probability assigned to the correct token at each position. Because the model is autoregressive, the loss at position $k$ is computed after feeding the ground-truth tokens $t_1, \ldots, t_{k-1}$ as context (teacher forcing). The total loss for the example is the sum of per-position losses.
Why this form: this is the standard maximum-likelihood objective for autoregressive sequence models. It has two important properties: (1) it decomposes the joint sequence probability into conditionals, making optimization tractable (the alternative—modeling all $B^K$ possible sequences directly—would require a flat softmax over $B^K$ categories, which is exponential in $K$); (2) it provides a per-position training signal, so the model gets explicit feedback on every digit—if it predicts $t_3$ incorrectly, it receives a loss gradient at position 3 even if $t_1$ and $t_2$ were correct. This is in contrast to a scalar MSE loss, which provides only a single number summarizing the total error without indicating which part of the prediction was wrong.
The expected loss over the true data distribution is:
Minimizing this expected loss is equivalent to minimizing the KL divergence between the true conditional distribution $p(\cdot|x)$ and the model distribution $p_\theta(\cdot|\phi(x))$ over the discretized output space.
Formal density estimation guarantees (Theorem 1).
The paper provides a rigorous risk bound for decoding-based density estimation under a binary (base-2) tokenization scheme. The analysis uses the following definitions setup:
Definition 1 (K-bit universality): A parametric model $p_\theta$ is K-bit universal if it can perfectly fit any discrete distribution over $K$-bit strings (equivalently, $2^K$ categories). Formally:
where $H(p, q) = \mathbb{E}_{y \sim p}[-\log q(y)]$ is the cross-entropy between two discrete distributions, and $H(p, p)$ is the Shannon entropy of $p$. In words: the model class is flexible enough that there exists some parameter setting $\theta^*$ that achieves the theoretical minimum loss (the entropy of the true distribution) for any target distribution $p$.
What this assumption means practically: it requires that (1) the Transformer decoder with its learned parameters can represent any distribution over $2^K$ categories, and (2) the training procedure (SGD on cross-entropy loss) can find such a $\theta^*$. This is a strong assumption that won't hold exactly in practice, but it provides a useful idealized baseline against which real performance can be compared—if the actual decoder outperforms the theory under this assumption (as observed in Figure 2 at low $N$), that indicates the presence of beneficial inductive biases or implicit regularization beyond what the idealized analysis captures.
Definition 2 (Marginal distribution over first $k$ bits): For a model trained on $K$-bit sequences, $p_\theta^k$ is the distribution over the first $k$ bits obtained by marginalizing (summing) over all possible values of the remaining $K-k$ bits:
What this computes: if we run the autoregressive decoder for only $k$ steps instead of the full $K$, what distribution do we get over $k$-bit prefixes? This matters because we may choose to use a coarser resolution (smaller $k$) at inference time than what we trained with—the decoder supports this naturally since each step's prediction conditions only on previous steps.
Definition 3 (Mapping bits to histogram bins): Define $\lambda_k : [0, 1) \to \{0,1\}^k$ as the operation that returns the first $k$ bits after the binary point in the binary expansion of $y$. For example, if $y = 0.10110...$, then $\lambda_3(y) = (1, 0, 1)$. The bit sequence $\lambda_k(y)$ can be interpreted either as a sequence or as the real number it represents: $\sum_{i=1}^k b_i 2^{-i}$. The key geometric interpretation: a $k$-bit sequence $(b_1, \ldots, b_k)$ identifies a specific bin (interval) $[y, y + 2^{-k})$ of width $2^{-k}$ in $[0, 1]$, where $y = \sum b_i 2^{-i}$.
The maximum likelihood estimator: Given $N$ i.i.d. samples $Y_1, \ldots, Y_N$ from true density $f$ on $[0,1]$, define $\theta^*$ as the parameters that minimize cross-entropy loss on the $K$-bit truncated representations:
The resulting density estimator at resolution $k \leq K$ is:
f_k^*_N(y) = 2^k \cdot p_\theta^k(\lambda_k(y))
where $p_\theta^k$ is the marginal distribution over the first $k$ bits (Definition 2), and the factor $2^k$ converts from probability mass (which sums to 1 over $2^k$ bins) to probability density (which integrates to 1 over $[0,1]$).
Risk definition: The risk is the expected mean integrated squared error between the true density $f$ and the estimator $f_k^*_N$:
R(f, f_k^*_N) = \mathbb{E}_{Y_1, \ldots, Y_N \sim f} \left[ \int_0^1 (f(y) - f_k^*_N(y))^2 dy \right]
The expectation is over the randomness in the training data—different draws of $N$ samples will produce different $\theta^*$ and hence different density estimates.
Theorem 1 (Main theoretical result): Under the K-bit universality assumption and assuming $f$ is twice continuously differentiable on $[0,1]$, the risk for the estimator $f_k^*_N$ at resolution $k \leq K$ decomposes as:
R(f, f_k^*_N) = \underbrace{\frac{2^{-2k}}{12} \int_0^1 f'(y)^2 dy}_{\text{Bias}^2} + \underbrace{\frac{2^k}{N}}_{\text{Variance}} + O(2^{-4k} + 1/N)
where $k$ is the number of bits used at inference, $N$ is the number of training samples, $f'(y)$ is the first derivative of the true density (measuring how quickly $f$ changes), and the integral of $(f')^2$ is a measure of the density's roughness (a smooth density has small squared derivative).
What it computes: the expected integrated squared error when using a $2^k$-bin histogram estimator derived from the decoding head's marginal distribution over $k$ bits. The bias term captures systematic error from approximating a smooth density with a piecewise-constant histogram—it decreases as $2^{-2k}$ (doubling the number of bins reduces bias by factor 4). The variance term captures random error from having finite data to estimate each bin's probability—it increases as $2^k/N$ (doubling the number of bins doubles the variance for fixed $N$). The optimal $k$ balances these terms: $k^* \approx \frac{1}{3} \log_2\left(\frac{N \int (f')^2}{6}\right)$.
Why this form matters: this theorem establishes that decoding-based regression is statistically consistent—as $N \to \infty$ with appropriately chosen $k \to \infty$, the risk goes to zero. More importantly, it quantifies the bias-variance tradeoff explicitly in terms of the tokenization depth $k$. This gives practitioners a principled way to choose the number of tokens: if they have $N = 10^4$ training points and expect a relatively smooth conditional density, they should use $k \approx 5$ bits ($2^5 = 32$ bins); if they have $N = 10^6$ points, $k \approx 7$ bits ($2^7 = 128$ bins) becomes optimal.
The tree-structured interpretation (Figure 3): The paper visualizes this process as a binary tree traversal. The root node represents the entire interval $[0, 1)$. The first bit $b_1$ splits this into left child $[0, 0.5)$ (if $b_1 = 0$) and right child $[0.5, 1)$ (if $b_1 = 1)$. The second bit $b_2$ splits whichever half was selected into quarters. After $k$ bits, we arrive at a leaf representing an interval of width $2^{-k}$. The probability mass in that leaf, multiplied by $2^k$ to convert to density, gives the histogram estimate at that point. The autoregressive decoding process is literally walking down this tree: at each step, the Transformer outputs probabilities for going left (0) or right (1), and these probabilities are conditioned on the path taken so far.
The implicit regularization phenomenon (Figure 2): The theorem assumes perfect optimization—under K-bit universality, the model finds $\theta^*$ that exactly matches the empirical bin frequencies. In practice, this doesn't always hold, and Figure 2 reveals something striking: when $N$ is small (1024) and $K$ is large (≥9), the decoder's actual risk is substantially lower than the theoretical risk, while the Riemann head's risk closely tracks the theory. The decoder head is somehow avoiding the high-variance regime that the theory predicts.
The paper's explanation (Section 3.3):
"a combination of the inductive bias of our model class and the implicit bias of our SGD training procedure makes the decoder less likely to fit noise; a concrete example would be that the model is biased to learn smooth distributions, and so when asked to fit the highly discontinuous empirical distribution arising from dropping few samples into a large number of bins, it refuses to, instead opting to learn a smooth approximation, and thereby driving down the variance term and hence the overall risk."
What this means concretely: with 1024 samples and $K=10$ bits (1024 bins), the empirical histogram has on average 1 sample per bin, with many bins empty and a few bins having 2–3 samples. This empirical distribution is extremely noisy and discontinuous. The Riemann head, which essentially memorizes bin frequencies via its learned bin embeddings, faithfully reproduces this noisy distribution—high variance, high risk. The decoder head, constrained by its architecture (shared vocabulary embeddings across positions, attention-based composition) and trained with SGD (which has an implicit bias toward smoother solutions), learns a smoothed version that doesn't fully fit the bin-to-bin noise. This smoothing reduces variance at the cost of some additional bias, but the net effect is lower total risk in the low-data regime.
This is a practically crucial finding because it means the decoder head has built-in regularization against overfitting that the simpler Riemann head lacks. It explains why Figure 7 (Section 4.2) shows the Riemann head plateauing on several tasks while the decoder continues to improve—the decoder's inductive biases prevent it from wasting capacity on noise, leaving more effective capacity for genuine signal.
The $k$ vs. $K$ distinction: An important subtlety is that the theorem allows inference at a different resolution $k$ than the training resolution $K$. The model is trained on $K$-bit sequences, but at inference we can compute $p_\theta^k$ for any $k \leq K$ by marginalizing out the remaining bits. This means we can train with a high $K$ (to give the model the option of fine granularity) but evaluate at a lower $k$ when data is scarce, adapting the effective resolution to the available training budget without retraining. The paper doesn't extensively exploit this capability in experiments (it appears to use $k=K$ by default), but the theoretical framework supports it.
Extension to unnormalized tokenization: The formal theory is presented for binary expansions in $[0, 1]$ (normalized case), but the intuition extends to the unnormalized floating-point scheme. The exponent bits determine the "scale" (which tree level to start at), while the mantissa bits refine within that scale. The bias-variance tradeoff in the unnormalized case would depend on both $E$ and $M$, with a more complex interaction: small $E$ limits the expressible range (bias toward limited dynamic range), while large $E$ with small $N$ might struggle to learn the exponent distribution accurately (variance in scale estimation).
Training Procedure and Architecture Details
Training algorithm: standard supervised learning with teacher forcing. For each training example $(x, y)$:
- The encoder
$\phi$(MLP) processes$x$to produce feature vector$\phi(x) \in \mathbb{R}^d$. - The true output
$y$is tokenized using the chosen scheme (normalized or unnormalized) to produce target sequence$(t_1, \ldots, t_K)$. - The decoder Transformer is fed
$\phi(x)$as its initial context and the ground-truth prefix$t_1, \ldots, t_{k-1}$as input when predicting token$t_k$. - Cross-entropy loss is computed at each position
$k$and summed. - Gradients flow through both the decoder and the encoder, updating all parameters jointly. This means the encoder learns to produce representations
$\phi(x)$that are useful specifically for the decoding-based head, not for some generic regression task.
Optimization hyperparameters (Appendix C):
- Optimizer: Adam
- Learning rate: swept over
[1e-4, 5e-4] - Training epochs: maximum 300
- Early stopping: patience=5 on validation loss, with validation split = 0.1 of training set
- Batch size: not explicitly stated, but implied to be full-batch or large-batch given "at most 20 minutes on a single Nvidia P100 GPU" (Section 4 preamble)
Input normalization: $x$ values are standardized using training set statistics: $x \leftarrow (x - x_{mean}) / x_{std}$ where $x_{mean}$ and $x_{std}$ are computed coordinate-wise over all training data. This is standard practice for MLP encoders.
Output normalization (for applicable heads):
- Normalized decoder, Riemann:
$y \leftarrow (y - y_{min}) / (y_{max} - y_{min})$using training set min and max. This maps all$y$values to$[0, 1]$. - Pointwise, MDN: additionally shift by
$y \leftarrow y - 0.5$to center values in$[-0.5, 0.5]$. - Unnormalized decoder: no output normalization—the raw
$y$values are used directly, tokenized via the floating-point scheme.
Decoder architecture specifications:
- Default configuration: 1 Transformer layer, 32 hidden units, 1 attention head. This is deliberately tiny.
- Alternative configuration (for size ablation): 3 layers, 128 hidden units, 4 attention heads.
- Parameter count: less than 10% of total network parameters (including the MLP encoder with up to 5 layers and 2048 hidden units).
- Attention mechanism: standard scaled dot-product attention with causal masking.
- Positional encoding: not explicitly mentioned, but standard Transformer decoders typically use learned or sinusoidal positional embeddings. Given the short sequence lengths (
$K$≤ 8 in most experiments), positional encoding may not be critical.
Constrained decoding at inference: During training, the model sees only valid token sequences (since training targets are always valid numbers). At inference, to prevent the model from sampling invalid sequences (e.g., sign token in the wrong position, incorrect number of exponent digits), the sampling is constrained to tokens that produce valid prefixes according to the tokenization grammar. The paper doesn't detail the constraint mechanism, but standard approaches include:
- Maintaining a finite-state automaton of valid token sequences and masking out tokens that would lead to dead-end states.
- Simpler approach: for normalized tokenization, any sequence of
$K$digit tokens is valid, so no constraints are needed beyond "sample exactly$K$tokens." For unnormalized tokenization, enforce that<sign_s>and<sign_e>are from{<+>, <->}, exponent digits appear in positions 3 through$E+2$, and mantissa digits appear in positions$E+3$through$E+M+2$.
Loss aggregation: The total loss for a batch is the mean (or sum) of per-example cross-entropies $H(y^{(i)}, p_\theta^{(i)})$ across all examples in the batch. No additional regularization terms (like weight decay) are mentioned for the decoder, though the encoder hyperparameter sweep includes weight decay of [0.0, 0.1, 1.0] (Appendix C, Pointwise section—presumably also applied to encoder when used with other heads).
Why train with cross-entropy rather than a numeric-aware loss: this is the central philosophical choice of the paper. Alternative losses could incorporate numeric distance—for example, a loss that penalizes predicting "7" more heavily when the truth is "9" than when it's "1," reflecting the numeric proximity. The paper argues that such losses are unnecessary: the tree structure of the tokenization already embeds numeric distance into the representation (nearby numbers share long common prefixes), and the model's inductive biases (smoothness preference) cause it to learn the numeric structure from the token-level supervision alone. Avoiding a custom loss keeps the approach simple and compatible with standard deep learning infrastructure (no need to implement a new loss function with custom gradients).
Why a Transformer decoder rather than a simpler autoregressive model: the Transformer's self-attention mechanism allows each token position to directly attend to all previous token positions and to the encoder output $\phi(x)$. This global receptive field is important for learning dependencies between digits at different positions—for example, the model might learn that $t_3$ (the third digit) should be influenced by $t_1$ (the most significant digit) because numbers with a large first digit tend to have small later digits (this is a distribution-specific correlation). A simple RNN would propagate information sequentially, which can work but may be less effective at capturing long-range digit interactions.
Hardware and runtime: the paper reports that "for the vast majority of tabular regression problems, we found that the process of training and tuning only requires at most 20 minutes on a single Nvidia P100 GPU" (Section 4 preamble). This establishes that decoding-based heads are practical in terms of computational cost—they don't require specialized hardware or days of training—and can be used as drop-in replacements without significant overhead. The P100 is a 2016-era GPU with 16GB memory, so the training regime is accessible on commodity hardware.
Summary of Design Choices and Their Justifications
-
Tree-based tokenization (base-
$B$digit-by-digit): encodes numeric hierarchy explicitly—most significant digits first, progressive refinement. Avoids the flat$B^K$-way classification that Riemann heads require, reducing parameters from$O(B^K)$to$O(B \times K)$. -
Unnormalized floating-point scheme: handles multi-scale outputs without task-specific normalization. Enables a single model to predict across many orders of magnitude. Extends the tree-based idea with an additional "exponent tree" that selects the scale before the "mantissa tree" refines within that scale.
-
Cross-entropy loss on all token positions: standard autoregressive training, no custom numeric loss needed. The tree structure and model inductive biases allow the model to recover numeric smoothness from token-level supervision.
-
Tiny decoder (1 layer, 32 units): isolates the effect of the output representation from model capacity. Any performance differences must come from the decoding-based approach, not from having more parameters.
-
Shared encoder across all head types: ensures fair comparison. The encoder features
$\phi(x)$are the same quality regardless of which head is attached, so head performance differences reflect the head's capabilities, not the encoder's. -
Constrained decoding at inference only: during training, the model learns the valid token grammar from data. At inference, constraints prevent amortized sampling errors from producing nonsense outputs. This decouples representation learning (which handles the token structure implicitly) from deployment robustness.
-
Hyperparameter sweeps over base, length, decoder size: the paper doesn't commit to a single "best" configuration but explores the tradeoff space, letting practitioners choose based on their data volume and precision needs.
4. Key Insights and Innovations
Innovation 1: Decoding-Based Heads as a Unifying Framework That Reveals Histogram Heads Are a Degenerate Special Case
The paper's most intellectually distinctive contribution is not proposing a new head architecture, but rather reframing the relationship between existing approaches in a way that exposes an unexplored frontier. Prior work treated Riemann (histogram) distribution heads and autoregressive decoders as fundamentally different model classes — one discretizes the output into flat categorical bins via softmax, the other generates token sequences for text. The paper collapses this distinction by showing that a Riemann head with $n$ bins is exactly a decoding head with sequence length $K=1$ and vocabulary size $B=n$.
This is a genuine conceptual reframing, not just an architectural trick. Once you see the Riemann head as $K=1$, the natural question becomes: what happens when $K > 1$? The answer, which the paper verifies empirically and theoretically, is that you get an exponential reduction in parameter count for the same output resolution. To represent $B^K$ bins, a Riemann head needs $O(B^K)$ embedding vectors — one per bin — that must each be learned independently. A decoding head needs $O(B \times K)$ vocabulary embeddings shared across sequence positions, with the autoregressive Transformer learning to compose them hierarchically. For $B=10, K=4$, that's 10,000 bins vs. 40 vocabulary items — a 250× reduction in learned bin representations.
Why this matters beyond parameter counting: the parameter reduction fundamentally changes how the model generalizes across bins. A Riemann head treats neighboring bins as independent categories unless the encoder's features happen to place them close in embedding space. The decoding head, by contrast, forces structural generalization: the digit "3" in the third position means the same fractional contribution regardless of what the first two digits were, because the vocabulary embedding for "3" is shared across all positions. This is a strong inductive bias that the paper shows leads to better sample efficiency in Figure 7 — the decoding head continues improving with data while the Riemann head plateaus — and in Figure 2, where the decoder outperforms the theoretical risk curve in low-data regimes through implicit smoothness regularization.
The paper also notes that this framing connects to extreme multi-label classification (Wydmuch et al., 2018) but argues the regression application has been "not thoroughly examined" (Section 2). This is a fair claim: the extreme classification literature focuses on discrete labels with no metric structure, while the regression setting introduces a continuous target with meaningful numeric distances, making the hierarchical decomposition both more natural (since numbers inherently have positional significance) and more demanding (since the model must learn that "7" and "8" are close while "0" and "9" are far, without being told this explicitly).
Significance: This is a fundamental rather than incremental contribution. It doesn't just propose a new method — it shows that an existing class of methods (histogram heads) is a boundary case of a broader design space parameterized by sequence length $K$. This opens a research direction: what other "flat" output representations in machine learning can be productively factorized into autoregressive sequences? The paper's own suggestions include multi-objective regression (Section 5) where $y^{(1)}, \ldots, y^{(M)}$ could be decoded as a concatenated sequence, which cannot be done with a single flat head.
Innovation 2: Implicit Regularization as an Empirically Demonstrated Property, Not Just an Assumption
The paper's theoretical framework (Theorem 1, Appendix B) relies on a $K$-bit universality assumption — that the model can perfectly fit any discrete distribution over $2^K$ categories. Under this assumption, the decoding head achieves exactly the same risk as a classical histogram estimator with the standard bias-variance decomposition: bias $\propto 2^{-2k}$, variance $\propto 2^k/N$. This is a useful baseline but not novel in itself — histogram density estimation risk is textbook material (the proof in Appendix B is a careful Taylor-expansion-and-integration exercise that any statistics graduate student could reproduce).
What is novel, and what makes the paper's contribution more than an application of known theory, is what happens when the universality assumption fails — specifically, in the low-$N$, high-$K$ regime of Figure 2. With $N=1024$ samples and $K \geq 9$ bits (512–1024 bins), the empirical distribution is extremely noisy — most bins contain 0 or 1 sample, a few contain 2–3. The Riemann head, which is explicitly parameterized as an independent softmax over bins, fits this noise faithfully and suffers high risk as the theory predicts. The decoder head does not. Its actual risk is substantially lower than the theoretical curve, implying it learns something smoother than the empirical bin frequencies.
The paper's interpretation is that this reflects "a combination of the inductive bias of our model class and the implicit bias of our SGD training procedure" that "makes the decoder less likely to fit noise" (Section 3.3). This is not proven — the paper offers no formal characterization of what the decoder's implicit bias actually is or why it favors smoothness — but it is empirically demonstrated with clear evidence. Figure 2 shows the effect robustly across 10 runs, and Figure 7 shows the downstream consequence: on multiple AMLB tasks, the Riemann head's performance plateaus at a level below the decoder's best performance, even with increasing data — suggesting the Riemann head is fitting noise in the bin structure while the decoder is learning a more generalizable representation.
Comparison to prior work: Implicit regularization in neural networks is a well-studied phenomenon (gradient descent is known to favor low-complexity solutions, flat minima, etc.), but its application to density estimation via autoregressive decoders has not been previously characterized. Prior work on histogram distribution heads in RL (Bellemare et al., 2017) and tabular data (Hollmann et al., 2025) did not analyze or exploit this regularization effect — they treated the histogram head as simply a flexible non-parametric estimator and accepted its data requirements as a cost of that flexibility. The paper shows that the decoder architecture provides a "free lunch" of sorts: better regularization without explicit smoothing penalties, simply as a consequence of the autoregressive factorization and Transformer inductive biases.
Significance: This is a diagnostic insight more than a theoretical breakthrough — it identifies why decoding heads work better than histograms in practice, not just that they do. The finding has practical implications for practitioners: if you're in a low-data regime, prefer the decoder over the Riemann head not just because it has fewer parameters, but because its training dynamics actively resist overfitting to bin-level noise. It also suggests a research direction: can we characterize this implicit bias more precisely? Is it the shared vocabulary embeddings? The attention mechanism? The teacher-forcing training? Understanding the source would let us design even better regularized output representations.
Innovation 3: The Floating-Point Tokenization Scheme as a Practical Solution to Multi-Scale Regression Without Task-Specific Normalization
The paper's unnormalized tokenization scheme (Section 3.1) — a base-$B$ generalization of IEEE-754 floating-point with separate sign, exponent, and mantissa tokens — addresses a problem that prior regression head designs handle inelegantly: how to predict outputs across many orders of magnitude without knowing the range in advance. Standard pointwise and parametric heads require $y$-normalization (typically min-max scaling to $[0,1]$ or standardization to zero mean and unit variance), which demands knowing the output distribution of the training data — and, critically, assumes the test distribution has the same range. In multi-task settings where different tasks have wildly different $y$-scales (Song et al., 2024), this requires either task-specific normalization (breaking the unified model architecture) or aggressive clipping (losing information about extreme values).
The floating-point tokenization solves this by design. The exponent tokens $e_1 \ldots e_E$ determine the order of magnitude — with $E=2$ and $B=10$, the exponent ranges from 0 to 99, so the representable values span $10^{-99}$ to $10^{99}$ (times the mantissa). The mantissa tokens $m_1 \ldots m_M$ provide $B^M$ bins of resolution within each decade. The model learns to first output the appropriate scale (exponent), then refine within that scale (mantissa), without any preprocessing step needing to know that scale in advance.
Why this is distinct from the normalized case: The normalized decoder with $K$ digits in $[0,1]$ is essentially a fixed-precision representation — it allocates $K$ digits of precision uniformly across the output range. This means it spends the same representational capacity distinguishing between 0.0001 and 0.0002 as it does between 0.9000 and 0.9001, even if the data mostly lies in $[0.5, 0.7]$ with high precision needed there and coarse values elsewhere. The floating-point scheme allocates precision adaptively: the mantissa resolution is constant relative to the current decade, so the absolute precision scales with the magnitude. This is precisely the property that makes IEEE-754 floating-point useful for scientific computing — it provides roughly constant relative error across the entire representable range.
Empirical validation: Figure 4 (Section 4.1) provides the clearest evidence. The "Vertical Asymptote (Tangent)" and "Vertical Asymptote (Hyperbolic)" functions have $y$ values that diverge to $\pm \infty$ near asymptotes. The pointwise head (with min-max normalization) fails catastrophically on the tangent case — the paper notes "Riemann prediction for 'Vertical Asymptote (Tangent)' went out of range" — because the normalization compresses an unbounded range into $[0,1]$, causing floating-point precision issues. The normalized decoder performs "decently" (Section 4.1) but is still constrained by its fixed range. The unnormalized decoder captures the shape well because its floating-point representation naturally handles the large dynamic range. Table 1 further shows that across 20 BBOB synthetic functions with varying scales and shapes, the unnormalized decoder achieves Kendall-Tau correlations competitive with pointwise heads (e.g., 89.56 vs. 89.08 at input dimension 5) without needing $y$-normalization.
Significance: This is an incremental rather than fundamental contribution — the IEEE-754 floating-point format is 40-year-old engineering — but its application as a regression head output representation for neural networks is novel and practically useful. It removes a tedious and error-prone preprocessing step ($y$-normalization) that practitioners currently accept as necessary. More interestingly, it points toward a general principle: representing numbers as structured token sequences with explicit scale/position decomposition may be more robust than representing them as continuous scalars when the model architecture (Transformers) is designed to process discrete tokens. This principle could extend beyond regression to other tasks where numeric outputs are required but their range is uncertain — e.g., reward modeling in RLHF, where reward values might drift significantly during training.
Innovation 4: Decoupling the Modeling Problem from the Decision Problem in Regression via Pointwise Functionals
Standard regression with a pointwise MSE-trained head conflates two distinct objectives: learning the conditional distribution $p(y|x)$ and producing a point estimate that minimizes expected squared error. The head never learns the distribution; it directly outputs the conditional mean, which is optimal under L2 loss but useless if the downstream task requires a different loss function (e.g., L1 for robustness to outliers, or L0 for the most likely value). You cannot extract the conditional median or mode from a model trained only to output the mean — you'd need to retrain with a different loss.
Parametric distribution heads (Gaussians, MDNs) partially decouple these objectives by explicitly modeling $p_\theta(y|x)$, but they impose a parametric form that limits what distributions can be represented. A Gaussian head can output mean and variance, but cannot represent a bimodal conditional distribution (e.g., "the output is either around 2 or around 8, depending on an unobserved variable"). An MDN can in principle approximate arbitrary distributions with enough components, but Table 2 shows this capability is unreliable in practice — NLL can range from 0.05 (excellent) to 7.49 (catastrophic) across datasets.
The decoding head fully separates the two problems by design. Training with cross-entropy loss learns $p_\theta(y|\phi(x))$ — the full conditional distribution over the discretized output space, with no parametric restrictions beyond the tokenization granularity. Once this distribution is learned, the paper's Section 3.2 framework lets the practitioner extract any pointwise functional $M(p_\theta)$ post-hoc:
- Mean (L2-optimal): average sampled decoded values
- Median (L1-optimal): Harrell-Davis estimator or sample median
- Mode (L0-optimal): beam search
The same trained model serves all three, and the choice can even be made at inference time per-example without retraining. This is a conceptual shift from the prevailing approach where the loss function is baked into the training objective and cannot be changed later.
Why this matters practically: In many real-world scenarios, the appropriate loss function depends on the deployment context, not the training context. A model predicting house prices might be deployed in a setting where overestimates are 10× more costly than underestimates (asymmetric loss), or where the most likely value is needed rather than the expected value. A pointwise MSE-trained model cannot adapt to these changes; a decoding-based model can, by switching from mean to a different functional or even computing the full conditional distribution for risk analysis.
Connection to the error-correction results (Section 4.5): The paper's token-repetition scheme can be understood as an instance of this decoupling. The raw decoder distribution $p_\theta$ may have high variance (outlier risk) that makes the sample mean unreliable. Rather than changing the training procedure to reduce variance, the paper introduces a separate inference-time procedure (majority voting across repeated tokens) that produces a more robust point estimate from the same distribution. This modularity — improve the distribution, or improve the estimator, or both independently — is only possible because modeling and decision-making are architecturally separated.
Significance: This is a reframing contribution more than a novel technical mechanism. The idea that $M(p)$ is the optimal point estimate under loss $\ell$ is standard in statistical decision theory (Lehmann, 1983, cited in Section 3.2), but it has not been the dominant paradigm in deep learning regression, where end-to-end training with a fixed loss has been the default. The paper shows that when the model architecture makes full distribution learning practical (which decoding heads do), this classical principle becomes actionable. The limitation, which the paper acknowledges implicitly, is that the model only learns distributions over the discretized output grid — if the true $p(y|x)$ has features finer than the tokenization resolution, the learned distribution will be an approximation, and the extracted point estimates will inherit this approximation error.
Innovation 5: Error-Correction Tokenization as an Inference-Time Robustness Mechanism Inspired by Coding Theory
The paper's error-correction tokenization scheme (Section 3.2, Section 4.5) is a distinctive contribution because it applies a concept from communication over noisy channels (coding theory) to a problem in machine learning inference (outlier robustness in autoregressive decoding). The connection is non-obvious and the mechanism is simple enough to be practically useful.
The setup: Under unnormalized tokenization, the decoder can assign a tiny but non-zero probability to tokens that produce extreme outliers — e.g., an exponent token of "9" when "1" is correct, shifting the prediction by 8 orders of magnitude. A single such outlier in $S$ samples can dominate the sample mean, making mean aggregation unreliable. Standard solutions from the LLM literature (top-$k$, top-$p$, low temperature) reduce the probability of drawing outlier tokens but also reduce distributional fidelity — the paper notes that "vanilla temperature sampling with temperature ≈ 1.0 is the best way to match $p(y|x)$" for density estimation (Section 4.3), so these techniques solve the outlier problem at the cost of degrading the density estimate.
The coding theory connection: In digital communication, you transmit bits over a noisy channel that may flip some bits. To make the transmission robust, you add redundancy — encode the message with extra bits so the receiver can detect and correct errors (e.g., majority voting across repeated transmissions of the same message). The paper applies this exact idea to token sequences: instead of generating $(t_1, \ldots, t_K)$ once and decoding, generate the same sequence $R$ times $(t_1^{(1)}, \ldots, t_K^{(1)}, \ldots, t_1^{(R)}, \ldots, t_K^{(R)})$ (training the decoder to repeat its output), and at inference, perform per-position majority voting: $t_k^* = \arg\max_v \sum_{r=1}^R \mathbb{1}[t_k^{(r)} = v]$. If outlier tokens are individually rare (low per-token probability), the chance that they win a majority vote at any position decays exponentially with $R$.
What makes this novel in the regression context: The standard approach to reducing outlier sensitivity in autoregressive models is to modify the sampling distribution (via temperature, top-$k$, etc.). The error-correction approach instead modifies the representation — the model learns to embed redundant information in its output, and the inference procedure exploits this redundancy. This is a more principled separation of concerns: the model's probabilistic knowledge (which includes uncertainty about whether a given token might be wrong) is preserved in the token-level distributions, and the inference procedure handles robustness independently.
Empirical evidence: Figure 10 (Section 4.5) shows that with $R=1$ (no repetition), relative MSE increases with more samples (more samples → higher chance of drawing at least one outlier), while with $R=3$ or $R=5$, the error monotonically decreases — the majority vote filters outliers effectively. The paper also reports a negative result in Appendix A.4: Hamming-distance-based binary representations (Qin, 2018), which are designed to bound numeric distortion as a function of bitwise edit distance, do not improve regression performance. The authors hypothesize this representation is "more difficult to learn" (Appendix A.4), which is an interesting finding in itself — not all coding-theoretic representations that are theoretically robust are learnable by gradient descent.
Significance: This is an incremental contribution measured by theoretical depth — majority voting is the simplest possible error-correcting code — but it demonstrates a productive cross-pollination between coding theory and autoregressive model inference that has not been widely explored in the regression literature. The negative result with Hamming-distance representations suggests the space of "good numeric tokenizations" is not simply those that minimize worst-case numeric distortion, but those that are learnable by Transformer decoders trained with cross-entropy loss. This opens a research question at the intersection of representation theory and optimization.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. Experiments span three categories: (1) synthetic 1D functions for curve-fitting visualization, (2) the Black-Box Optimization Benchmarking (BBOB) suite (Elhara et al., 2019) with 24 continuous objective functions evaluated at input dimensions 1–20, (3) real-world tabular regression benchmarks: AMLB (Gijsbers et al., 2024) with 28 tasks and OpenML-CTR23 (Fischer et al., 2023) with 31 tasks, both using up to 10K training points, and (4) the UCI regression repository (Dua & Graff, 2017) with 25 datasets for density estimation, using preprocessed versions from a public repository. Training sets are used as-is from the benchmark sources; for AMLB data scaling experiments, random subsets of the training data are drawn at specified sizes.
-
Base model(s). All methods share an identical feature encoder: a multi-layer perceptron (MLP) with ReLU activations, swept over 2–5 layers and 256–2048 hidden units (Appendix C). This encoder is deliberately generic—no pretraining, no domain-specific architecture—to isolate the effect of the regression head. The decoder head itself is a small Transformer with either 1 layer and 32 hidden units (default, "less than 10% of the total network parameter count" per Section 4 preamble) or a larger variant of 3 layers, 128 hidden units, 4 attention heads for ablation studies.
-
Metrics. For tabular regression, the paper uses relative mean squared error (MSE) within individual tasks (lower is better) and Kendall-Tau rank correlation for aggregate comparisons across tasks (higher is better, scale-invariant). For density estimation, the metric is negative log-likelihood (NLL) on held-out test data (lower is better), computed as the negative log-probability the model assigns to the true
$y$values. For curve fitting and density visualization, qualitative visual fit is the primary evaluation. For the BBOB benchmark, Kendall-Tau correlation is reported at each input dimension. -
Baselines. Four regression head types are compared: (1) Pointwise head: a deterministic feed-forward network (often a single linear layer) mapping
$\phi(x)$to a scalar, trained with MSE loss. (2) Riemann (histogram) distribution head: a softmax-parameterized categorical distribution over uniformly spaced bins in$[0,1]$, following Hollmann et al. (2025) and the distributional RL literature (Bellemare et al., 2017). Bin counts are swept over 16–16384. (3) Mixture Density Network (MDN): a Gaussian mixture model where mixture weights, means, and standard deviations are learned functions of$\phi(x)$, following Bishop (1994). Mixture counts are swept over 1–1000. (4) Normalized decoder head: the decoding-based approach with min-max$y$-normalization to$[0,1]$, using base-$B$digit expansion with$K$tokens. (5) Unnormalized decoder head: the decoding-based approach with floating-point tokenization, no$y$-normalization. -
Generation budget / compute accounting. For pointwise estimation from the decoder head, the paper draws
$S$independent samples from$p_\theta(\cdot|\phi(x))$and aggregates them (mean, median, or mode). The sampling size$S$is varied in the error-correction experiments (Figure 10, Section 4.5) and in data scaling experiments (Figure 11, Appendix A.1). For density estimation, the full token-level probability distribution is used directly without sampling. Training compute is measured in wall-clock time: "at most 20 minutes on a single Nvidia P100 GPU" per task (Section 4 preamble). All models are trained for a maximum of 300 epochs with early stopping (patience=5) using a 0.1 validation split. -
Cross-validation / statistical protocol. For tabular regression tasks, each bar in Figures 5, 6, and 13 is averaged over 10 training runs with different random initializations (Section 4.2). For the UCI density estimation experiments (Table 2, Table 3), results are averaged over 10 train-test splits. For the BBOB experiments (Table 1, Figure 12), each point is averaged over 10 runs with 100K training points where
$x$is sampled uniformly from$[-5,5]$coordinate-wise. For the risk experiments (Figure 2), results are averaged across 10 runs each. Error bars or standard deviations are reported where applicable (e.g., ± StdDev in Table 2). The risk calculations in Figure 2 use the truncated Gaussian distribution$\mathcal{N}_{[0,1]}(0.5, 0.25^2)$as the ground truth.
Main Quantitative Results
Curve Fitting and Synthetic Function Approximation (Section 4.1, Table 1)
The paper first establishes the fundamental representational capacity of decoding-based heads through qualitative curve fitting and quantitative BBOB benchmarks.
Figure 4 (Section 4.1) demonstrates visual fit quality across six 1D functions spanning qualitatively different shapes: smooth (Weierstrass variants), discontinuous (stepwise alternating and increasing), and unbounded (vertical asymptotes via hyperbolic and tangent functions). The paper reports:
"the unnormalized decoder head is able to successfully capture the shapes of various functions with which both the Riemann and pointwise head struggle."
The pointwise head fails on functions with "very high or unbounded $y$-ranges" (Section 4.1), suffering from the $y$-normalization required for training stability. The Riemann head produces an out-of-range prediction on the tangent function. The normalized decoder performs "decently" but is constrained by its fixed $[0,1]$ range. Only the unnormalized decoder successfully captures all six function shapes without range-related artifacts. The paper notes these results "occur regardless of xy-scales, which are omitted for brevity," suggesting robustness to the specific input/output magnitudes.
Table 1 quantifies BBOB performance across 24 synthetic functions at input dimensions 5–20, using mean Kendall-Tau correlation (higher is better) aggregated over all functions:
| Input Dimension | Unnormalized Decoder | Normalized Decoder | Pointwise | Riemann |
|---|---|---|---|---|
| 5 | 89.56 | 89.40 | 89.08 | 88.94 |
| 10 | 88.71 | 88.54 | 88.25 | 88.30 |
| 15 | 87.49 | 86.90 | 88.06 | 87.42 |
| 20 | 86.11 | 86.02 | 86.78 | 86.78 |
The unnormalized decoder achieves the highest or near-highest correlation at every dimension, with differences typically within 1 percentage point of the best method. This establishes that decoding heads are not sacrificing approximation quality for their distributional flexibility — they are competitive with pointwise heads even on pure point-prediction tasks. The full per-function breakdown in Figure 12 (Appendix A.2) reveals function-specific patterns: the decoder excels on some landscapes (Sphere, Ellipsoidal, Discus) while the pointwise head leads on others (Katsuura, Lunacek), suggesting the optimal head depends on function structure but no method dominates universally.
The key interpretive point: these synthetic experiments validate that the tokenization schemes do not fundamentally limit representational capacity. The decoder can fit smooth, discontinuous, and asymptotically unbounded functions — a range broader than pointwise or Riemann heads handle gracefully. The remaining experiments test whether this flexibility translates to real-world regression tasks where data is limited and noise is present.
Real-World Tabular Regression: Pointwise Comparison (Section 4.2, Figures 5, 6, 7)
The central practical claim is that decoding heads are competitive with or superior to standard pointwise heads on tabular regression benchmarks given the same training data and encoder architecture.
Figure 5 (Section 4.2) compares unnormalized decoder vs. pointwise head across all AMLB and OpenML-CTR23 tasks. Each bar represents the average Kendall-Tau correlation over 10 runs, with bars from the same task stacked and sorted by the performance gap. The paper reports:
"in the majority of tasks, the decoder outperforms the pointwise head, and in a few cases, the gap can be quite significant (>0.3)."
This is the paper's strongest claim for pointwise regression. The full breakdown in Figure 13 (Appendix A.3) extends this to all four heads (unnormalized decoder, normalized decoder, Riemann, pointwise), showing both decoder variants remain competitive throughout — they are not simply outperforming pointwise heads on a few tasks while failing on others. Tasks are sorted by pointwise head performance, and the decoder bars consistently match or exceed the pointwise bars across the sorted index.
Figure 6 (Section 4.2) directly compares decoder heads against the Riemann head via paired scatter plots. Each point is a task; the x-coordinate is the decoder's Kendall-Tau score and the y-coordinate is the Riemann head's score. Points above the diagonal indicate tasks where the decoder outperforms Riemann:
"both decoding heads outperform the Riemann head in the vast majority of tasks as well, suggesting improved sample efficiency from minimizing vocabulary / bin sizes."
This is the paper's primary empirical evidence for the exponential-reduction-in-parameters argument from Section 2. The decoder uses $B \times K$ vocabulary embeddings while the Riemann head uses $B^K$ — for the normalized decoder with $B=10, K=4$, that's 40 vs. 10,000 parameters — and the decoder's superior performance supports the claim that this factorization is not just parameter-efficient but also statistically efficient.
Figure 7 (Section 4.2) provides the most detailed evidence for the sample efficiency claim, plotting relative MSE vs. training data size for the normalized decoder, Riemann, and pointwise heads on selected AMLB tasks. The key patterns are:
-
Riemann head plateaus: On tasks 233212, 359938, 359939, and 360945, the Riemann line is essentially flat or near-flat across the data range, unable to achieve the same error level as the decoder even with
$10^4$training points. This is direct evidence for the "data inefficiency of using the histogram head" (Section 4.2) — too many bins spread too few samples, and the softmax parameterization cannot learn to interpolate between bins effectively. -
Decoder consistently improves with data: On tasks 233215, 359930, 359938, and others, the decoder's error decreases monotonically with training data, showing it can exploit additional samples effectively.
-
Non-obvious pointwise weakness: In low-data regimes (around
$10^1$to$10^2$points), the pointwise head can perform worse than the decoder head. The paper notes this is due to "numeric instabilities of its own" — "the pointwise head required appending a sigmoid activation to enforce the normalized output to be within [0,1] to avoid extremely high MSE errors" (Section 4.2). This is a practical finding: the decoder head, despite needing to learn numeric token representations, is more stable with very small datasets because its output is constrained to the representable grid by construction — it cannot output values outside$[0,1]$in the normalized case or extreme outliers in the unnormalized case when using constrained decoding. -
Decoder can outperform pointwise in high-data regimes: On task 359948, the pointwise head plateaus while the decoder continues improving, crossing below the pointwise error at around
$10^3$points and maintaining lower error at$10^4$. The paper attributes this to the decoder's ability to model "more complex conditional distributions that cannot be captured by a single conditional mean estimate" (implicit in the density estimation motivation).
The full data scaling results over all 28 AMLB tasks appear in Figure 11 (Appendix A.1), confirming these patterns are not limited to the selected subset shown in Figure 7. The paper states:
"We confirm the data-efficiency of the decoder head against the Riemann distribution head on nearly every regression task. Furthermore, we observe numerous cases where both distributional methods outperform the pointwise head, especially in low data regimes."
Quantifying the gaps: The paper does not provide aggregate summary statistics (mean/median gap across tasks) for the decoder-vs-pointwise or decoder-vs-Riemann comparisons in Figures 5–7, which is a limitation. The reader must visually assess the per-task bar plots and scatter points. However, the statement that the decoder outperforms pointwise "in the majority of tasks" is visually supported by Figure 5, where the decoder bars are taller than the pointwise bars for most stacked pairs.
Density Estimation (Section 4.3, Figure 8, Table 2)
The density estimation experiments test the claim that decoding heads can capture flexible conditional distributions $p(y|x)$ without parametric shape constraints.
Figure 8 (Section 4.3) provides qualitative density estimation visualization on four 1D conditional density shapes: Half Moons (bimodal), Zig-Zag (variance that varies with $x$), Spiral (heteroscedastic with $x$-dependent mean shifts), and Hollow Square (bimodal with a low-density region between modes). The unnormalized decoder head with vanilla temperature sampling (temperature ≈ 1.0) successfully captures all four shapes:
"Given unbounded training data it is able to capture the overall distribution
$p(y|x)$well, although there can be slight outlier noise as shown by lighter points."
The "lighter points" refer to occasional samples far from the main density mass — the outlier problem discussed in Section 3.2. The paper notes that lowering temperature reduces this noise but can "reduce expressivity" (Section 4.3), and that "vanilla temperature sampling with temperature ≈ 1.0 is the best way to match $p(y|x)$." This is an important practical guidance: if density estimation fidelity is the goal, accept some outlier noise rather than distorting the distribution with aggressive temperature scaling.
Appendix A.6, Figure 15 extends this visualization to compare MDN (100 mixtures), decoder with various sampling strategies (vanilla, temp=0.1, top-p=0.9, top-k=5), and Riemann (1000 bins). The MDN captures the overall shape but can produce artifacts (spurious modes, misplaced mass). The decoder with vanilla sampling best matches the ground truth. Top-k and top-p sampling produce biased estimates that miss parts of the distribution. The Riemann head with 1000 bins produces noisy, discontinuous density estimates due to the curse of dimensionality in bin count.
Table 2 (Section 4.3) provides quantitative NLL comparisons on 11 representative UCI datasets, with the MDN, unnormalized decoder (UD), normalized decoder (ND), and Riemann (R) heads. The full results over all 25 datasets are in Appendix A.5, Table 3. Key patterns:
-
MDN variance is extreme: The MDN achieves the best NLL on some datasets (Airfoil: 0.12 ± 0.11, Wine: 0.05 ± 0.12) but catastrophic NLL on others (Kin 40K: 7.49 ± 0.73, Slice: 7.09 ± 0.09 from Table 3). The standard deviations are often comparable to or larger than the means, indicating high run-to-run instability. The paper presents this as evidence that MDNs are "at times able to perform the best but also extremely poorly depending on the task."
-
Decoder heads are reliable: Both UD and ND achieve NLL below 0.7 on every dataset in Table 2, and below 0.69 on every dataset in Table 3 (the worst is ND on TamiElectric at 0.69 ± 0.00). This is a striking result: the decoder never catastrophically fails in the way the MDN can. The paper summarizes: "both decoding heads remain reliable overall (NLL<0.7 always)."
-
Riemann head consistently underperforms: The Riemann NLL is always worse than both decoder variants, often by a factor of 2–5× (e.g., Housing: UD=0.41, R=1.56; Wine: UD=0.24, R=1.67). This is consistent with the Figure 7 finding that Riemann heads are data-inefficient — the UCI datasets are relatively small (typically a few hundred to a few thousand points), and the high bin counts needed for reasonable density resolution produce noisy estimates.
-
Normalized vs. unnormalized decoder tradeoffs: On most datasets, ND slightly outperforms UD (e.g., Bike: ND=0.10 vs. UD=0.12; Kin 40K: ND=0.12 vs. UD=0.19). The paper doesn't extensively analyze this, but it's consistent with the normalized decoder having a simpler task (no scale learning) when
$y$-normalization is appropriate. However, the unnormalized decoder is competitive and doesn't require knowing the$y$-range in advance, which is the practical tradeoff.
A methodological note on negative NLL values: Table 3 (Appendix A.5) reports negative NLL for the MDN on several datasets (Challenger: -0.29, Fertility: -0.06, Solar: -1.40, Stock: -0.15). Negative NLL can occur when the model assigns probability density greater than 1 to regions where test points fall — this indicates the density is extremely peaked (low variance) and may reflect overfitting or model misspecification. It's a subtle indicator that the MDN's density estimates are not well-calibrated on these datasets, while decoder and Riemann heads (which are discrete distributions over a bounded grid) always produce non-negative NLL by construction.
Ablation: Decoder Head Size (Section 4.4, Figure 9)
The paper ablates the effect of decoder capacity on density estimation quality using the normalized decoder with fixed tokenization ($B=10, K=4$).
Figure 9 varies three architectural parameters from a default of (3 layers, 4 heads, 128 units): number of layers (1–4), number of attention heads (2–8), and number of hidden units (500–1000). The metric is NLL on three UCI datasets (Bike, Energy, Housing).
The paper reports:
"larger decoding heads do sometimes help, but only up to a certain point, at which overfitting can occur."
On Bike, NLL decreases from ~1.2 (1 layer) to ~0.8 (2 layers), then increases at 3–4 layers (overfitting). On Housing, the optimal is 2 layers. For attention heads, the behavior is dataset-dependent: Bike benefits from more heads (NLL drops from ~1.0 at 2 heads to ~0.92 at 8 heads), while Energy and Housing show slight U-shaped curves. Hidden units show the clearest overfitting pattern — on Housing, NLL increases from ~1.45 at 500 units to ~1.48 at 1000 units.
The paper states this "was also observed over regression over BBOB functions and with the unnormalized decoding head, but we omitted these results for brevity." This ablation primarily serves as a sanity check — it confirms that the default small configuration (1 layer, 32 units) is not limiting performance on these datasets, and that increasing capacity provides negligible or negative benefit. This supports the paper's methodological choice to keep the decoder tiny: the benefits of decoding-based regression come from the representation, not from additional model capacity.
Ablation: Error Correction via Token Repetition (Section 4.5, Figure 10)
This ablation tests whether training the decoder to repeat its output sequence multiple times and applying per-position majority voting at inference improves pointwise estimation robustness.
Figure 10 shows relative MSE on selected AMLB tasks (unnormalized decoder, mean aggregation) as a function of sampling size $S$, for repetition counts $R \in \{1, 2, 3, 4, 5\}$.
The paper reports two key findings:
"when using regular tokenization (repeat count=1), as more samples are drawn, the likelihood of drawing outliers increases the error."
This is visible in the $R=1$ curves: on tasks 233213, 359934, and 359944, MSE increases or stays flat as sampling size grows, because more samples → higher probability of drawing at least one extreme outlier that corrupts the sample mean.
"the error can be substantially decreased by training the decoding head to decode the same tokens repeatedly and allow better scaling with samples."
With $R=3$ or $R=5$, the MSE curves become monotonically decreasing with sampling size — the majority vote filters outlier tokens while the increased sample count reduces estimation variance for inlier tokens. The paper notes that "repeating too many times may make learning more difficult" — visible in the $R=5$ curve on task 359945, which underperforms $R=3$ until larger sampling sizes.
The paper also reports a negative result in Appendix A.4: Hamming-distance-based binary representations (Qin, 2018), which are theoretically designed to bound numeric distortion under bitwise errors, "may not lead to better regression results" (Appendix A.4). Figure 14 shows that the tree-based tokenization outperforms or matches the Hamming representation on all tested tasks. The authors hypothesize "this representation being more difficult to learn."
Critical Assessment
Claim 1: "Decoding-based heads are competitive with pointwise heads on tabular regression."
What the experiments demonstrate: Across AMLB and OpenML-CTR23 (28 + 31 = 59 tasks), the unnormalized decoder achieves higher Kendall-Tau correlation than the pointwise head "in the majority of tasks" (Section 4.2, Figure 5), with some gaps "quite significant (>0.3)." The BBOB results (Table 1) show the unnormalized decoder achieving the highest mean Kendall-Tau at input dimensions 5, 10, and 15, and within 0.67 points of the best method at dimension 20.
What the experiments do NOT demonstrate: The paper provides no formal statistical significance test for the per-task comparisons. The "majority" claim is based on visual inspection of Figure 5 — counting bars — and the 10-run averaging is stated but confidence intervals are not shown. This matters because with 59 tasks and relatively small per-task datasets (up to 10K points), some of the "decoder beats pointwise" results could be within noise. Additionally, the pointwise head's hyperparameter sweep includes weight decay [0.0, 0.1, 1.0] (Appendix C), but the decoder's regularization comes primarily from its architectural inductive biases — it's not clear whether the pointwise head was given equal tuning opportunity. The pointwise head requiring a sigmoid activation to avoid extreme MSE errors in low-data regimes (Section 4.2) suggests it may not have been optimally regularized.
What would strengthen the claim: Pairwise statistical tests (e.g., Wilcoxon signed-rank across tasks and runs), reporting the fraction of tasks where each method wins with confidence intervals, and ensuring the pointwise head's regularization (weight decay, early stopping, architecture) was tuned as extensively as the decoder's tokenization hyperparameters.
Conditional boundary: The claim holds for tabular regression with MLP encoders and dataset sizes typical of OpenML benchmarks. The paper does not test whether this competitiveness extends to very large datasets (>10^6 points) where pointwise heads with sufficient capacity might close any gap, or to domains where the encoder architecture differs substantially (CNNs for images, Transformers for text). The BBOB results (Figure 12) show function-specific variation — the pointwise head wins on Katsuura and Lunacek — so the competitiveness is not universal even within the tested domain.
Claim 2: "Decoding heads outperform Riemann (histogram) heads in the vast majority of tasks due to sample efficiency."
What the experiments demonstrate: This is the most strongly supported claim in the paper. Figure 6 shows paired scatter plots where decoder heads dominate Riemann heads on the large majority of tasks for both AMLB and OpenML-CTR23. Figure 7 (and the extended Figure 11 in Appendix A.1) shows Riemann heads plateauing on many tasks while decoder heads continue to improve with data — the sample efficiency argument has direct empirical support. Figure 2 provides complementary evidence from the controlled risk experiment: the decoder's actual risk is lower than the theoretical histogram risk in the low-$N$, high-$K$ regime.
What the experiments do NOT demonstrate: The paper does not systematically control for the "effective bin count" between Riemann and decoder heads to ensure a fair comparison. A Riemann head with 1024 bins and a decoder head with $B=2, K=10$ both have 1024 representable values, but the Riemann head uses 1024 embedding vectors while the decoder uses 20 vocabulary items. The performance gap could be attributable to the parameter count difference (Riemann overfitting with too many parameters) rather than to the hierarchical decomposition per se. A properly regularized Riemann head — with, say, weight decay tuned as carefully as the decoder's architecture — might close part of the gap.
What would strengthen the claim: A controlled experiment where Riemann and decoder heads are matched on both output resolution AND total parameter count (e.g., by limiting the Riemann head's embedding dimension or adding a bottleneck), to isolate the effect of the autoregressive decomposition from the effect of raw parameter count. Also, testing whether a Riemann head with a lower bin count (sacrificing resolution for statistical efficiency) can match the decoder's performance — the decoder is effectively doing adaptive resolution via the tree structure, which the flat Riemann head cannot do.
Conditional boundary: The claim is restricted to the tabular regression datasets tested (AMLB, OpenML-CTR23, UCI). For very large datasets where the Riemann head's variance is no longer the bottleneck (the high-$N$ regime of Figure 2, where both methods converge to the same risk), the advantage may disappear. The paper does not test this regime extensively.
Claim 3: "Decoding heads provide reliable density estimation with competitive NLL compared to MDNs, without catastrophic failure modes."
What the experiments demonstrate: Table 2 and Table 3 show that both decoder variants achieve NLL below 0.7 on all 25 UCI datasets, while MDN NLL ranges from -1.40 (overfitting) to 7.49 (extremely poor fit), with large standard deviations. This is clear evidence for the "reliability" claim. Figure 8 and Figure 15 visually demonstrate the decoder capturing complex conditional densities (bimodal, heteroscedastic, hollow) that would require many Gaussian components to approximate.
What the experiments do NOT demonstrate: The paper does not compare the decoder against the MDN with equal hyperparameter tuning effort. The MDN mixture count is swept over [1, 2, 5, 10, 20, 50, 1000] (Appendix C), but MDNs are notoriously sensitive to initialization, optimization, and component collapse. The paper does not report whether techniques like component-conditional learning rates, specialized initialization, or EM-pretraining were attempted — any of which might stabilize the MDN and reduce the catastrophic failure rate. The NLL comparison may be unfair to the MDN if the MDN training procedure was less optimized.
Additionally, the NLL metric has a subtle interpretation issue: the decoder and Riemann heads define distributions over a discrete grid of representable values, while the MDN defines a continuous density. Computing NLL for continuous densities on discrete data points involves evaluating the density function, which can produce arbitrarily negative values (overfitting). A fairer comparison might use a discretized version of the MDN's density evaluated on the same grid as the decoder, or report a metric like the continuous ranked probability score (CRPS) that doesn't have this discretization mismatch.
What would strengthen the claim: A more extensive MDN tuning protocol (learning rate warmup, gradient clipping, component dropout, multiple restarts), reporting MDN failure rates (how many runs diverged or produced degenerate components), and using metrics that treat discrete and continuous densities symmetrically.
Conditional boundary: The "reliability" claim is valid for the UCI dataset regime (small to medium datasets, relatively smooth conditional densities). For datasets with genuinely multi-scale or heavy-tailed conditional distributions, the decoder's fixed grid (normalized) or fixed-precision floating-point (unnormalized) might miss fine tail structure that a well-tuned MDN could capture. The paper does not test extreme distribution shapes (e.g., power-law tails) where the floating-point exponent range might need to be very large.
Claim 4: "The unnormalized decoder handles multi-scale and unbounded outputs without $y$-normalization."
What the experiments demonstrate: Figure 4 provides visual evidence: the unnormalized decoder captures vertical asymptote functions (tangent, hyperbolic) while the pointwise head (requiring normalization) fails and the Riemann head goes out of range. Table 1 shows the unnormalized decoder achieving the highest Kendall-Tau on BBOB functions without $y$-normalization — it handles the output range natively.
What the experiments do NOT demonstrate: The experiments do not test the unnormalized decoder on genuinely multi-task regression where different tasks have dramatically different $y$-scales. The paper motivates this capability in Section 3.1 by citing Song et al. (2024), but the experiments are all single-task. The BBOB functions have different output ranges per function, but they are trained and evaluated per-function, not in a multi-task setting. The claim that the unnormalized decoder solves the multi-task normalization problem is therefore not experimentally validated — it's a capability argument supported by the tokenization design, not by multi-task benchmarks.
What would strengthen the claim: A multi-task regression experiment where a single model is trained on multiple datasets (e.g., multiple UCI or AMLB tasks simultaneously) with unnormalized decoder vs. task-normalized pointwise heads. This would directly test whether the unnormalized decoder avoids the "tedious" normalization balancing that the paper claims is a motivation.
Conditional boundary: The floating-point tokenization's range is finite — with $E$ exponent digits in base $B$, the maximum representable value is approximately $B \times B^{-1} \times B^{B^E-1} = B^{B^E-1}$. For $B=10, E=2$, this is $10^{99}$, which is enormous but not infinite. Functions that diverge faster than exponential (e.g., double-exponential) could exceed this range, though this is unlikely in practice.
Claim 5: "Error-correction tokenization improves pointwise estimation robustness against outliers."
What the experiments demonstrate: Figure 10 clearly shows that with $R=1$ (no repetition), MSE increases with sampling size on several tasks (233213, 359934, 359944), while $R \geq 3$ makes MSE decrease monotonically. This is clean, causally interpretable evidence: the majority voting mechanism is responsible for the improvement.
What the experiments do NOT demonstrate: The paper tests only mean aggregation with error correction. It doesn't show whether error correction helps median-based estimation (where outliers are already naturally handled) or mode-based estimation (where beam search doesn't sample outliers at all). The benefit may be specific to mean aggregation, which is the most outlier-sensitive estimator — a practitioner who uses median estimation may not need error correction at all. The paper also doesn't compare error correction against simpler alternatives like top-k sampling or lowered temperature, which would address outliers without requiring the model to learn repetitive outputs.
What would strengthen the claim: A comparison of error correction vs. top-k/top-p sampling vs. median estimation for pointwise prediction, showing whether error correction provides benefits beyond these simpler alternatives. Also, testing whether the repetition training harms density estimation quality — the paper only reports MSE, not NLL, for the error-corrected models.
Conditional boundary: The error correction benefit is demonstrated only for unnormalized tokenization where outlier risk is highest. For normalized tokenization (bounded to $[0,1]$), the outlier problem is less severe, and error correction may be unnecessary.
Overall experimental design assessment:
Strengths:
- The shared encoder design rigorously isolates the regression head's contribution — any performance difference is attributable to the output representation, not to better input processing.
- The head size is deliberately minimal (<10% of parameters), preventing the interpretation that decoder heads work simply because they add capacity.
- The multi-benchmark approach (synthetic functions, BBOB, AMLB, OpenML-CTR23, UCI) provides robustness across qualitatively different problem types.
- The 10-run averaging with standard deviations provides some protection against run-to-run variance, though confidence intervals would be more informative.
- The theoretical analysis (Theorem 1, Figure 2) complements the empirical results by providing a bias-variance framework that makes specific predictions testable in the controlled risk experiment.
Weaknesses:
- No statistical significance testing: The paper makes comparative claims ("majority of tasks," "significantly outperform") without formal tests. With 59 tasks, visual inspection of bar charts is insufficient to establish reliable superiority.
- Limited encoder architecture: All experiments use an MLP encoder. The claims about decoding heads as "swap-in replacements" for pointwise heads would be stronger with at least one non-MLP encoder (e.g., a simple CNN on image regression, or a Transformer on sequence regression) to show the representation isn't MLP-specific.
- No pre-registration or held-out test protocol: Hyperparameters are swept on the same benchmarks used for evaluation, with early stopping on validation splits. The 10-fold averaging provides some protection, but the risk of overfitting to benchmark-specific characteristics remains.
- Unclear tuning equity: The pointwise head gets weight decay
[0.0, 0.1, 1.0]; the decoder gets tokenization sweeps over base, length, and Transformer size. It's not obvious that the pointwise head received equal hyperparameter optimization effort — the paper's central claim would be undermined if the pointwise head was simply less tuned. - Missing comparison to simple alternatives: The paper doesn't compare against quantile regression (which can produce non-parametric conditional distributions), Gaussian Processes (which provide principled uncertainty), or ensemble methods (which can capture multi-modality). The baseline set is reasonable but not comprehensive for the density estimation claims specifically.
6. Limitations and Trade-offs
The Difficulty Estimation Cost Is Unaccounted for in the Headline Efficiency Gains
The assumption or constraint. The paper's motivating comparison between decoding heads and Riemann heads hinges on an exponential reduction in parameter count—representing $B^K$ bins using $B \times K$ vocabulary items rather than $B^K$ independent embedding vectors. This is presented as a source of sample efficiency: fewer parameters mean less overfitting when data is scarce. The paper states:
"a drawback is that learning numeric distances between all of the bins requires more data as the size of the vocabulary increases" (Section 2)
and the experiments in Figures 6 and 7 are designed to demonstrate that this parameter efficiency translates to statistical efficiency.
The consequence. A practitioner evaluating whether to use a decoding head over a Riemann head needs to know which regime each is appropriate for—not just that one dominates the other on the tested benchmarks. The paper shows that Riemann heads plateau on many AMLB tasks (Figure 7) and perform worse than decoders in the paired scatter plots (Figure 6), but it does not systematically characterize the boundary conditions: at what dataset size $N$ does the Riemann head catch up? For a given resolution $B^K$, what is the minimum $N$ needed for the Riemann head to outperform the decoder? Without this boundary characterization, the paper's practical guidance reduces to "use the decoder head," which may over-generalize from the tested benchmarks to regimes where Riemann heads with sufficient data are perfectly adequate.
An additional, unaccounted practical cost arises from the hyperparameter tuning burden of decoding heads. While the paper notes that training requires "at most 20 minutes on a single Nvidia P100 GPU" (Section 4 preamble), this is per configuration. The decoder head requires sweeping over base $B$, sequence length $K$ (for normalized) or exponent count $E$ and mantissa count $M$ (for unnormalized), Transformer size (layers, heads, units), and for unnormalized decoders, the sampling strategy for pointwise estimation (error-correction repetition count $R$, top-$k$, top-$p$, or temperature). The Riemann head requires only a bin count sweep. A practitioner with a fixed tuning budget might get better results from an extensively-tuned Riemann head than from a coarsely-tuned decoder head, but the paper's comparison does not control for tuning budget.
What evidence exists in the paper. Figure 7 (and extended Figure 11 in Appendix A.1) shows the Riemann head plateauing below the decoder's performance on specific AMLB tasks (233212, 359938, 359939, 360945). However, these plots show performance up to $10^4$ training points—a relatively small data regime. The paper does not test at larger $N$ (e.g., $10^5$ or $10^6$) to determine whether the Riemann head eventually closes the gap. Figure 2 provides the only controlled comparison of decoder vs. Riemann risk as a function of $N$ and $K$, but this is on a simplified synthetic density (truncated Gaussian), not on the tabular benchmarks where the practical claims are made. The crossover point where Riemann becomes competitive with the decoder is therefore unknown for real-world regression tasks.
Mitigation status. The paper does not attempt to characterize the $N$-vs-resolution boundary where the decoder advantage disappears. The data scaling experiments (Figures 7 and 11) show decoder superiority in the tested range but do not extrapolate. The paper's suggestion that future work could explore "how these changes affect the performance of numeric decoding" (Section 5, Modern LLM Architectures) acknowledges architectural sensitivity but not the data regime question specifically.
The Single Encoder Architecture Limits Generality of the "Swap-In Replacement" Claim
The assumption or constraint. Every experiment in the paper uses the same encoder: a multi-layer perceptron (MLP) with ReLU activations, swept over 2–5 layers and 256–2048 hidden units. The encoder is deliberately generic to isolate the regression head's contribution. The paper frames decoding heads as broadly applicable:
"decoder heads can be effective swap-in replacements to common pointwise regression heads" (Section 4, experimental goals)
This framing implies that the findings generalize across encoder architectures—CNNs for image regression, Transformers for sequence regression, graph neural networks for structured data.
The consequence. There are at least two reasons the MLP-only evaluation threatens generalizability:
-
Interaction between encoder architecture and output representation. MLPs produce features
$\phi(x)$that have no inherent sequential or spatial structure. Transformers, CNNs, and GNNs produce features with built-in inductive biases (positional, translational, relational). It is unknown whether a decoding head conditioned on a CNN feature map behaves differently from one conditioned on an MLP feature vector. For instance, if the CNN features are localized (corresponding to image patches), the decoder's attention over these features might learn to attend to different regions when predicting different digit positions (coarse digits from global features, fine digits from local features). This interaction could either enhance or degrade regression performance relative to a pointwise head that simply pools all features—the paper provides no evidence either way. -
The decoder head's implicit regularization depends on optimization dynamics. The paper's explanation for the decoder's data efficiency relies on "a combination of the inductive bias of our model class and the implicit bias of our SGD training procedure" (Section 3.3). The MLP encoder is trained jointly with the decoder via SGD; the interplay between encoder gradients and decoder gradients determines what features
$\phi(x)$learns to produce. With a different encoder architecture (e.g., a pretrained vision Transformer that is fine-tuned rather than trained from scratch), the optimization trajectory would differ, and the implicit regularization that prevents the decoder from fitting bin-level noise might weaken or disappear.
What evidence exists in the paper. None. The paper does not test any encoder architecture other than an MLP. The synthetic curve fitting (Figure 4, Section 4.1), BBOB (Table 1, Section 4.1), tabular regression (Section 4.2), and UCI density estimation (Section 4.3) all use the MLP encoder. The paper acknowledges this implicitly by not claiming otherwise—the "swap-in replacement" language in Section 4 is aspirational, describing the experimental goal rather than a tested conclusion.
Mitigation status. Not addressed. The paper's discussion section (Section 5) focuses on modern LLM architectures as a future direction for the decoder itself, not on encoder diversity. Testing with at least one non-MLP encoder—even a simple convolutional network on an image regression task—would substantially strengthen the generality claim. The current evidence supports only that decoding heads work as MLP-attached regression heads on tabular and synthetic data; the "swap-in replacement" framing overstates what the experiments demonstrate.
Density Estimation Quality Is Not Benchmarked Against Properly-Regularized or Modern Alternatives
The assumption or constraint. The paper's density estimation experiments compare the decoder against a Mixture Density Network (MDN) and a Riemann (histogram) head. The paper reports that the MDN "has high variability, at times able to perform the best but also extremely poorly depending on the task" while "both decoding heads remain reliable overall (NLL<0.7 always)" (Section 4.3).
The consequence. The MDN baseline may be substantially weaker than what a well-tuned practitioner would deploy for density estimation. The paper sweeps mixture counts $M \in \{1, 2, 5, 10, 20, 50, 1000\}$ (Appendix C) but does not report whether standard MDN stabilization techniques were used:
-
Component collapse prevention: MDNs are notorious for "component collapse" where all but one mixture component become degenerate (zero weight or infinite variance), reducing the model to a single Gaussian. Standard mitigations include component-specific learning rates, KL regularization toward a prior, or variance flooring. The paper does not mention any such technique. The catastrophic NLL values on some datasets (e.g., Kin 40K: 7.49 ± 0.73) are consistent with component collapse or numerical instability.
-
Initialization sensitivity: MDN optimization is highly sensitive to the initialization of mixture parameters, particularly variance parameters. A poor initialization can trap the model in a high-loss basin. The 10-run averaging (Table 2) partially accounts for this, but standard practice is to run multiple restarts and select the best, not to average over runs—averaging will inflate the reported NLL if some runs are catastrophic failures.
The paper's negative NLL values for MDN on several datasets (Challenger: -0.29, Fertility: -0.06, Solar: -1.40, Stock: -0.15; Table 3 in Appendix A.5) are a red flag. Negative NLL for a continuous density indicates the model assigns probability density >1 to regions where test points fall—this can happen when the fitted variance is extremely small (the density becomes a near-delta spike at the training points, overfitting). A properly regularized MDN should not produce negative NLL on held-out test data; this suggests the MDN training procedure lacks adequate regularization, making it an artificially weak baseline for the "reliability" claim.
Furthermore, the paper does not compare against modern alternatives that have been developed specifically to address MDN instability:
- Normalizing flows: Transform a simple base distribution through a series of invertible mappings to produce flexible densities with stable training (no mixture components, just deterministic transformations).
- Deep ensemble or Monte Carlo dropout for uncertainty: If the goal is reliable predictive distributions, ensembling multiple pointwise networks or using dropout at inference provides well-calibrated uncertainty without parametric density assumptions.
- Gaussian Processes with deep kernels: Provide principled uncertainty with strong theoretical guarantees, though at higher computational cost.
What evidence exists in the paper. Table 2 and Table 3 provide quantitative NLL comparisons. Figure 15 (Appendix A.6) provides qualitative visualization of density estimates, but only compares decoder variants against MDN with 100 mixtures and Riemann with 1000 bins—no flow-based or ensembling alternative is shown. The paper does not report training stability metrics for the MDN (component collapse rate, variance distribution, convergence diagnostics). The sweep over mixture count $M$ (Appendix C) is the only MDN-specific tuning mentioned.
Mitigation status. The paper acknowledges the limitation implicitly in Section 5 under "Other Regression Architectures," where it mentions stochastic networks, Bayesian neural networks, and energy-based models but states "we did not compare to significantly more complex methods" and notes these methods "have not seen wide adoption due to their complex architectures." This framing—that simpler is better—is a reasonable design choice for the paper's scope, but it means the density estimation results establish only that decoding heads beat an under-regularized MDN baseline, not that they are the best available approach for flexible density estimation. A practitioner choosing between a decoding head and a modern normalizing flow or deep ensemble cannot use this paper's evidence to decide.
Hard Extrapolation Problems and Distribution Shift Receive No Evaluation
The assumption or constraint. The paper evaluates all methods using random train-test splits from the same underlying dataset. For AMLB and OpenML-CTR23, the test data comes from the same distribution as the training data (i.i.d. splits). For the UCI density estimation experiments, 10 train-test splits are used (Appendix A.5), but again these are random partitions of a fixed dataset, not distribution-shift evaluations. The paper does not test any setting where the test distribution differs from the training distribution—covariate shift, temporal drift, or out-of-distribution inputs.
The consequence. This is a critical omission because decoding heads and pointwise heads handle extrapolation fundamentally differently. A pointwise MLP head, being a continuous function, will extrapolate smoothly beyond its training range—linearly if it is a linear layer, or according to its learned basis functions if deeper. In contrast, a decoding head is a discrete distribution over a finite representable grid. When the test $y$ values fall outside the range of representable numbers (which can happen under distribution shift even for normalized outputs if the normalization range was set from training data that doesn't cover the test range), the decoding head cannot represent them. For the normalized decoder, the output is strictly bounded to $[0,1]$; for the unnormalized decoder, the floating-point range is enormous but finite—$B^{B^E-1}$—and the model may assign negligible probability to extreme values, making density estimates unreliable in the tails even if the value is technically representable.
In an i.i.d. evaluation, this limitation is invisible because train and test share the same $y$-range. The practical consequence emerges in deployment scenarios where $y$-normalization parameters are estimated from training data and applied at test time. If a production system encounters inputs that produce larger $y$ values than seen during training, a normalized decoder will saturate at its maximum representable value (or produce undefined behavior if constrained decoding tries to enforce $[0,1]$ bounds), while a pointwise head will extrapolate (potentially unreliably, but at least continuously). Which failure mode is preferable depends on the application, but the paper provides no evidence to inform this choice.
A related extrapolation concern arises from the decoder's token-level uncertainty structure. In the tree-based representation, numbers that are close numerically share long common prefixes (e.g., 0.1234 and 0.1235 differ only in the last digit). This means the decoder can interpolate smoothly within the training distribution by learning that nearby numbers have similar prefix probabilities. However, when extrapolating to $y$ values that require the decoder to produce high-confidence tokens for digits it has rarely or never seen in certain positional contexts, the behavior is unpredictable—the model might assign uniform probability across all digits (high uncertainty, flat density) or might latch onto specific tokens based on spurious training correlations (producing a sharp but incorrect density). The paper's theoretical analysis (Theorem 1) assumes $f$ is supported on $[0,1]$ and twice differentiable—it provides no guidance for the boundary or extrapolation behavior.
What evidence exists in the paper. None. The curve fitting experiments (Figure 4) show the decoder capturing asymptotically unbounded functions (tangent, hyperbolic), but these are trained on the full function range—the model sees $y$ values across the entire range during training. This is in-distribution prediction, not extrapolation. The BBOB experiments (Table 1) train on uniformly sampled $x \in [-5,5]$ and test on the same distribution. The tabular benchmarks use standard random splits. No experiment tests the decoder on inputs whose corresponding $y$ values lie substantially outside the training $y$ distribution.
Mitigation status. Not addressed. The paper's discussion (Section 5) mentions multi-task regression and modern LLM architectures as future directions but does not flag extrapolation or distribution shift as a concern. The <NaN> token mentioned for invalid inputs (Section 3.1) is a partial mitigation—it allows the model to signal "out of supported range"—but (1) it is not used in any experiment, and (2) it requires the model to learn to detect out-of-range conditions from the input $x$ alone, which is a harder problem than continuous extrapolation (a pointwise head can output $y > y_{\max}$ without needing to detect anything).
The Tradeoff Between Density Fidelity and Pointwise Robustness Is Observed but Not Resolved
The assumption or constraint. The paper identifies a tension between two modes of using the decoder head:
-
Density estimation: The goal is to faithfully represent
$p(y|x)$. The paper reports that "vanilla temperature sampling with temperature ≈ 1.0 is the best way to match$p(y|x)$" (Section 4.3). -
Pointwise estimation: The goal is to extract a reliable scalar prediction
$\hat{y}$from$p_\theta(y|\phi(x))$. The paper identifies a specific failure mode: under unnormalized tokenization, "the model can have a miniscule but non-zero probability of decoding an arbitrarily large outlier, even if the underlying true distribution is bounded" (Section 3.2), which can corrupt the sample mean.
The consequence. The paper demonstrates two approaches to mitigate the outlier problem for pointwise estimation—error-correction tokenization via repetition and majority voting (Section 4.5), and broader sampling modifications like top-$k$, top-$p$, and reduced temperature (Section 3.2)—but neither approach is evaluated for its impact on density estimation quality. It is highly plausible that the same mechanisms that suppress outliers also distort the learned distribution:
- Training with repeated tokens forces the model to commit to a single deterministic sequence multiple times, which may encourage sharper (lower-entropy) token distributions that suppress genuine uncertainty in the density estimate.
- Lowering temperature at inference directly sharpens the token-level softmax, reducing distributional entropy.
- Top-
$k$and top-$p$sampling truncate the token distribution, removing probability mass from the tails that may represent genuine (if unlikely) output values.
The paper observes this tradeoff in passing but does not quantify it:
"While one can enforce the sampling to be tighter (e.g. lowering temperature) to remove noise, this tighter sampling can unfortunately also reduce expressivity" (Section 4.3)
The implication is that a practitioner must choose between faithful density estimates (vanilla temperature, potentially unreliable mean estimates) and reliable point estimates (error correction or aggressive sampling, potentially distorted density estimates). The paper offers no guidance on how to navigate this tradeoff—no Pareto frontier, no metric that jointly evaluates both objectives, and no method that achieves both simultaneously (e.g., density estimation with trimmed means, or adaptive temperature that varies by token position).
What evidence exists in the paper. The evidence is scattered across sections:
- Figure 8 (Section 4.3) shows density estimation quality with vanilla temperature sampling.
- Figure 10 (Section 4.5) shows pointwise MSE improvement with error-correction tokenization and mean aggregation.
- Figure 15 (Appendix A.6) compares density estimation with decoder (vanilla, temp=0.1, top-p=0.9, top-k=5) alongside MDN and Riemann baselines, showing visually that the non-vanilla sampling strategies produce biased density estimates.
However, no experiment jointly evaluates density fidelity and pointwise robustness on the same model and same task. The density estimation experiments (Figure 8, Table 2) use temperature ≈ 1.0 vanilla sampling. The error-correction experiments (Figure 10) use mean aggregation but do not report NLL for the same models. A reader cannot determine from the paper whether a model trained with 5× token repetition for robust mean estimation still produces reasonable density estimates, or whether the repetition training degrades NLL to the point where the decoder loses its advantage over simpler baselines.
The paper also does not explore position-dependent strategies that could partially resolve the tradeoff. For instance, the first few tokens (most significant digits, or exponent tokens in unnormalized mode) determine the coarse structure of the prediction—outliers in these positions cause large errors. Aggressive top-$k$ or error correction on early tokens while allowing higher entropy on later tokens might reduce outlier risk with minimal density distortion. The paper's error-correction scheme applies repetition uniformly to all token positions, making no distinction between significant and insignificant digits.
Mitigation status. The paper acknowledges the tradeoff implicitly by noting that vanilla temperature is best for density matching while also describing outlier-robust alternatives for pointwise estimation, but it does not frame this as a limitation to be resolved. The discussion of RAFT (Section 3.2), which computes expectations in closed form over a fixed evaluation set, is presented as an alternative that avoids sampling noise but "may also defeat the purpose of using a decoding head, which offers several density estimation benefits." This suggests the authors are aware of the tension but consider it outside the current paper's scope. No future work direction is proposed for resolving the density-fidelity vs. pointwise-robustness tradeoff.
The Computational Overhead of Autoregressive Decoding at Inference Is Not Quantified
The assumption or constraint. The paper evaluates decoding heads primarily in terms of predictive accuracy and distributional fidelity, with training cost discussed qualitatively ("at most 20 minutes on a single Nvidia P100 GPU," Section 4 preamble). Inference cost receives almost no analysis. The paper mentions that the decoder uses "only 1 layer and 32 units" (Section 4 preamble) and constitutes "less than 10% of the total network parameter count," implying it is cheap.
The consequence. This parameter-count argument is misleading for inference latency. A pointwise head is a single forward pass through a small feed-forward network—constant time regardless of the output resolution. A decoding head with sequence length $K$ requires $K$ sequential forward passes through the Transformer decoder, one per token. Each step conditions on all previously generated tokens, so they cannot be parallelized. For $K=4$ (a typical normalized setting), this is 4× the forward passes of a pointwise head. For $K=8$ with a larger base, it is 8×. For density estimation where $S$ samples are drawn (e.g., $S=100$ for Monte Carlo estimation of the mean), the total cost is $K \times S$ forward passes.
Comparing against a Riemann head makes the latency difference even starker. A Riemann head with $B^K$ bins performs a single matrix multiplication to produce logits for all bins—the computation is $O(B^K \times d)$ where $d$ is the feature dimension, but it is a single parallel operation. A decoding head with the same resolution performs $K$ sequential matrix multiplications, each $O(B \times d_{\text{decoder}})$. The total FLOPs may be lower for the decoder (fewer parameters, smaller matrices), but the wall-clock time on modern hardware will be dominated by the sequential dependency chain. A GPU with thousands of cores can compute all $B^K$ bin logits in the Riemann head simultaneously; it cannot accelerate the $K$ sequential steps of the decoder.
This latency overhead is particularly consequential for the paper's recommended practices:
- Error correction via repetition: If the decoder generates
$R$repetitions of the output sequence for majority voting (Figure 10), the sequence length becomes$R \times K$, multiplying latency by$R$. For$R=5, K=4$, this is 20 sequential decoding steps. - Sample-based pointwise estimation: Computing the sample mean from
$S$sampled sequences requires$S \times K$sequential steps (assuming samples are drawn independently and could be batched, but batched autoregressive generation still requires$K$steps for the longest sequence in the batch). - Vanilla temperature sampling for density estimation: Producing a faithful density estimate requires many samples (Figure 8 shows qualitative results but does not specify sample counts; typically hundreds to thousands for reliable density visualization), each requiring
$K$sequential steps.
What evidence exists in the paper. The paper provides no inference latency measurements. The 20-minute training time on a P100 (Section 4) is the only computational cost mentioned. There are no throughput comparisons, no latency-vs-accuracy plots, and no discussion of how the sequential decoding cost scales with $K$, $B$, or the number of samples needed for reliable pointwise estimation. The "less than 10% of total parameters" framing implicitly equates parameter count with computational cost, which is appropriate for training FLOPs but not for inference latency on parallel hardware.
Mitigation status. Not addressed. The paper's discussion (Section 5) mentions modern LLM architectures (sparse attention, mixture of experts) as future directions for the decoder itself, which could affect inference cost, but does not frame latency as a current limitation or propose mitigation strategies. A practitioner deploying decoding heads in a latency-sensitive application (real-time prediction, high-throughput batch inference) would need to independently measure and optimize the inference cost, as the paper provides no guidance.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper does not propose a new model architecture or a state-of-the-art result on a competitive benchmark. Its contribution is more structural: it reframes the relationship between classification and regression in deep learning by showing that the boundary between them is an artifact of output representation, not a fundamental modeling constraint. Prior to this work, the prevailing design pattern treated regression and classification as distinct problems requiring different heads (linear/MSE vs. softmax/cross-entropy), different loss functions, and different normalization procedures. This paper collapses that distinction by demonstrating that a cross-entropy-trained autoregressive sequence model can serve as a drop-in regression head with competitive pointwise accuracy and superior distributional flexibility—without custom losses, without $y$-normalization (in the unnormalized case), and without parametric assumptions on the output distribution.
The magnitude of this shift is a reframing with practical consequences, not a paradigm shift. It does not overturn any foundational beliefs about regression; it instead reveals that an architecture already widely deployed for language modeling (Transformer decoders) can productively replace components that the field has treated as categorically separate for decades. The paper's most significant conceptual move is to show that Riemann (histogram) distribution heads—which have a substantial track record in distributional RL (Bellemare et al., 2017) and tabular regression (Hollmann et al., 2025)—are a degenerate special case of decoding heads with sequence length $K=1$. This reframing does two things simultaneously: it retroactively explains why histogram heads work well (they are exploiting a restricted form of autoregressive decomposition), and it immediately implies that $K>1$ should be better (exponential reduction in parameters for the same resolution, hierarchical structure matching the natural tree decomposition of numeric quantities).
The paper resolves a latent contradiction that has been simmering in the LLM-as-regressor literature. On one side, Vacareanu et al. (2024) and Akhauri et al. (2025) have shown that language models prompted or fine-tuned to output numbers as text can perform regression competitively with traditional methods. On the other side, there has been widespread skepticism—articulated in the paper's own introduction—that "regular supervised fine-tuning over numbers represented as strings is unprincipled, considering that there is no notion of numeric distance when using cross-entropy loss." The paper's contribution is to show that numeric distance does not need to be baked into the loss function. The tree structure of positional tokenization embeds it into the representation: nearby numbers share long common prefixes, and the autoregressive factorization forces the model to learn coarse structure before fine detail. The theoretical analysis (Theorem 1) and the empirical finding that the decoder smooths rather than memorizes bin-level noise (Figure 2) provide complementary evidence that the unprincipled-seeming approach is, in fact, well-behaved. This should reduce the barrier to adopting decoding-based approaches in contexts where the alternative is architecting custom regression pipelines around language model outputs—most notably, generative reward modeling for RLHF, where Zhang et al. (2024) and Mahan et al. (2024) are already moving toward text-to-text reward prediction.
Several research directions become more attractive in light of this reframing:
-
Joint input-output autoregressive modeling becomes more natural. If the regression head can be a small Transformer decoder conditioned on encoder features, there is no architectural reason to keep the encoder and decoder as separate model families. A unified sequence model could process input tokens (text, tabular features serialized as tokens, image patches) and then autoregressively decode numeric targets—all within the same Transformer, trained end-to-end with cross-entropy. This is already how LLMs perform text-to-text regression, but the paper's controlled experiments suggest the performance gains come from the output representation, not from LLM-specific scale or pretraining. A small, purpose-built unified model might be sufficient for many regression tasks.
-
Verifier-guided search over numeric outputs becomes thinkable. The paper's error-correction experiments (Section 4.5) hint at a broader capability: once the output is a token sequence, the entire toolkit of LLM inference-time strategies—beam search, best-of-
$N$verification, iterative refinement—can be applied to regression problems. A process reward model (PRM) could score partial numeric sequences (e.g., "the first three digits look reasonable given the input") and guide search toward high-likelihood answers, analogous to how PRMs guide mathematical reasoning in LLMs. This would be qualitatively new: traditional regression has no notion of "step-by-step verification" because the output is a single scalar. -
Numeric output representations as a design space open up. The base-
$B$tree representation and the floating-point scheme are the simplest possible tokenizations. The paper's unsuccessful experiment with Hamming-distance-based representations (Appendix A.4) suggests the space of "good" tokenizations is constrained by learnability under gradient descent, not just by information-theoretic properties. Exploring alternative representations—Gray code (where adjacent numbers differ by exactly one bit), balanced ternary, or learned tokenizations where the vocabulary embeddings themselves are optimized—could yield representations that are both numerically robust and easy for Transformers to learn. This is a genuinely new research question that did not exist before decoding-based regression was formalized.
Conversely, some research directions become less attractive:
-
Custom numeric-aware loss functions for text-to-number tasks may be unnecessary. If cross-entropy with tree-based tokenization already recovers numeric smoothness via model inductive biases, the engineering effort of designing and implementing losses that incorporate
$|y - \hat{y}|$into token-level supervision may not be worth the complexity. The paper's results suggest that practitioners should first try standard cross-entropy with an appropriate tokenization before reaching for custom losses. -
Flat histogram heads for high-resolution density estimation are now clearly suboptimal except in the
$K=1$large-data limit. Anyone currently using Riemann/histogram distribution heads with more than a few hundred bins should consider replacing them with a$K>1$decoding head of equivalent resolution, since the paper shows the decoder dominates in low-to-moderate data regimes (Figures 6, 7, 11) and the implementation complexity is comparable (a small Transformer vs. a large softmax layer).
Follow-Up Research This Work Enables
Characterizing the implicit regularization that makes decoding heads data-efficient. The paper's most intriguing empirical finding—the decoder's actual risk falling substantially below the theoretical histogram risk in the low-$N$, high-$K$ regime (Figure 2)—is attributed to "a combination of the inductive bias of our model class and the implicit bias of our SGD training procedure" (Section 3.3), but the mechanism is not identified. A strong follow-up would systematically vary the decoder architecture and training procedure to isolate which component causes the smoothing effect: (1) Do shared vocabulary embeddings across positions force smoothness? Replace shared embeddings with position-specific embeddings and measure whether the risk returns to the theoretical curve. (2) Is it the autoregressive factorization? Compare against a non-autoregressive decoder that predicts all tokens simultaneously (via a single softmax over $B^K$ categories) but with a low-rank factorization of the output matrix, to determine whether the sequential dependency or the parameter sharing drives regularization. (3) Is it SGD specifically? Train with full-batch gradient descent, Adam, and sharpness-aware minimization (SAM) to see if the implicit bias changes. The controlled experimental setup from Figure 2 (truncated Gaussian, varying $N$ and $K$, known ground truth) makes this analysis tractable and reproducible. The outcome would tell us whether the decoder's data efficiency is a robust architectural property or a contingent effect of specific training choices.
Testing the "swap-in replacement" claim across diverse encoder architectures. The paper restricts all experiments to an MLP encoder, but frames decoding heads as "effective swap-in replacements to common pointwise regression heads" (Section 4). A natural extension would evaluate the normalized and unnormalized decoder heads on image regression tasks (e.g., predicting object coordinates or crowd counts from images) using a standard CNN encoder (ResNet-18 or similar), and on sequence regression tasks (e.g., predicting scalar properties of molecules from SMILES strings) using a Transformer encoder. For each domain, train pointwise, Riemann, and decoder heads with identical encoders and compare both pointwise accuracy (MSE, Kendall-Tau) and distributional quality (NLL, calibration) on standard benchmarks. The key question is whether the decoder's inductive biases interact with encoder-specific feature structures—do CNNs produce features that the decoder head can exploit differently than MLPs? Are there domains where the decoder systematically underperforms pointwise heads? This would establish the boundary conditions for the "swap-in" claim and give practitioners domain-specific guidance.
Combining decoding heads with process reward model (PRM) search for numeric prediction. The paper establishes that decoding heads produce token-level probability distributions over numeric digits. This is exactly the interface that PRMs require: given a partial numeric sequence (e.g., tokens for "1.2" generated so far), a PRM could score whether this prefix is on track toward the correct answer, enabling beam search over the space of representable numbers. The experiment would: (1) train a small verifier model that takes $\phi(x)$ and a partial token sequence $(t_1, \ldots, t_k)$ and predicts whether the final decoded number will be within some tolerance of the ground truth; (2) at inference, use this verifier to guide beam search over the $B$-ary tree of possible numeric completions, keeping the top-$W$ beams at each step; (3) compare the resulting point predictions (mean/mode of the beam outputs) against standard sample-based aggregation from the paper. The hypothesis is that verifier-guided search could reduce the outlier problem (since improbable branches with high numeric error would be pruned early) without the density distortion caused by temperature reduction or top-$k$ truncation. This would connect the decoding regression framework to the rapidly growing literature on test-time compute scaling.
Stress-testing the unnormalized decoder in genuine multi-task and distribution-shift settings. The paper motivates the unnormalized tokenization by citing multi-task regression (Song et al., 2024) where different tasks have vastly different $y$-scales, but the experiments are all single-task. A direct follow-up would: (1) construct a multi-task regression benchmark from existing UCI and OpenML datasets with heterogeneous output ranges—e.g., one task with $y \in [0, 1]$, another with $y \in [10^2, 10^5]$, a third with $y$ spanning both positive and negative values; (2) train a single model with shared encoder and unnormalized decoder head on all tasks jointly (with task-identifying tokens prepended to the decoder input sequence, or separate task embeddings added to $\phi(x)$); (3) compare against per-task pointwise heads requiring task-specific $y$-normalization, and against a single pointwise head with careful global normalization. The paper's claim that normalization is "tedious" and the unnormalized decoder handles multi-scale outputs natively would be validated or refuted. Furthermore, a distribution-shift evaluation would test the decoder on held-out tasks whose $y$-ranges differ from any training task, probing whether the floating-point exponent structure generalizes to unseen scales. Negative results here—e.g., the decoder failing to learn appropriate exponent distributions for novel scales—would refine our understanding of what the decoder actually learns about numeric magnitude.
Evaluating decoding heads on probabilistic forecasting and time series. The paper evaluates density estimation on static UCI datasets with i.i.d. train-test splits. Time series forecasting—where the goal is to predict $p(y_{t+1} | y_{1:t})$ and where distributions are often heavy-tailed, multi-modal, or heteroscedastic—is a natural stress test for the decoder's distributional flexibility. An experiment would: (1) replace the MLP encoder with a causal Transformer or RNN encoder that processes historical sequences; (2) train normalized and unnormalized decoder heads alongside MDN, Riemann, and quantile regression baselines on standard probabilistic forecasting benchmarks (e.g., M4 competition, electricity, traffic datasets); (3) evaluate using both pointwise metrics (sMAPE) and distributional metrics (CRPS, pinball loss at multiple quantiles). Time series forecasting is a setting where the "outlier problem" in unnormalized decoders is particularly salient—extreme events like demand spikes or sensor failures produce $y$ values far from the central mass—so the error-correction mechanisms from Section 4.5 would be tested under realistic conditions. The hypothesis is that the decoder's ability to represent multi-modal conditional distributions (e.g., "tomorrow's demand is either around 100 or around 1000, depending on an unobserved promotion") would give it an edge over unimodal or limited-mixture alternatives.
Learning adaptive tokenization and resolution. The paper treats tokenization hyperparameters ($B$, $K$, $E$, $M$) as fixed architectural choices that must be swept manually. An ambitious extension would learn the tokenization itself: (1) replace the fixed base-$B$ vocabulary with a continuous embedding space where each "digit" is a vector in $\mathbb{R}^d$, and the decoder outputs a soft assignment over these embeddings at each position; (2) the token-to-number mapping becomes a learned function (a small network) that maps the concatenated embedding sequence to a scalar; (3) train end-to-end with a combination of cross-entropy (for learning the discrete-like structure) and MSE (for numeric accuracy), potentially with an entropy regularizer that encourages the soft assignments to be nearly one-hot. This would combine the benefits of decoding heads (distributional flexibility, hierarchical decomposition) with the benefits of continuous representations (no discretization error, gradient-based optimization of the representation itself). The risk experiment from Figure 2 could be replicated with this learned tokenization to see if it achieves lower risk than the fixed-base decoder at the same effective resolution. This direction also connects to VQ-VAE-style discrete representation learning, where codebook vectors are learned jointly with the model.
Practical Applications and Downstream Use Cases
Generative reward modeling for RLHF without custom architectures. Current reward modeling in RLHF typically requires appending a scalar regression head to a language model and training with a Bradley-Terry preference loss—a custom pipeline that diverges from the standard next-token prediction framework. The paper's results suggest a simpler alternative: fine-tune the language model to output reward scores as tokenized numbers (e.g., "4.7" or "8.2") using standard cross-entropy loss, treating reward prediction as a text-to-text task. The controlled experiments provide evidence that this approach does not sacrifice accuracy relative to pointwise heads—the decoder matched or exceeded pointwise performance on the majority of AMLB and OpenML-CTR23 tasks (Figure 5)—while also providing distributional information (predictive variance over reward estimates) that could be useful for active learning or uncertainty-weighted policy updates. A team implementing this would need only to format reward values as strings during data preparation; no architectural changes or custom loss functions are required. The reliability finding from the UCI experiments (decoder NLL < 0.7 on all 25 datasets, Table 3) suggests this approach is unlikely to catastrophically fail in the way MDN-based reward models might when faced with out-of-distribution responses.
Scientific and engineering surrogate modeling with built-in uncertainty quantification. Many scientific applications require not just a point prediction of some quantity of interest (fluid drag coefficient, material stress, molecular binding energy) but a reliable estimate of predictive uncertainty. The paper's density estimation results on UCI datasets (Table 2) show that the decoder head produces well-behaved predictive distributions without the component-collapse instability of MDNs (which returned negative NLL on several datasets, indicating pathological overconfidence). An engineer deploying a surrogate model for design optimization could use a normalized decoder head (if the output range is known from physical constraints) or an unnormalized decoder head (if simulating phenomena across multiple scales) attached to whatever encoder processes the design parameters. At inference, the full conditional density $p(y|x)$ is available for computing credible intervals, expected improvement for Bayesian optimization, or risk metrics. The computational cost is modest—the paper reports "at most 20 minutes on a single Nvidia P100 GPU" per task (Section 4)—meaning the approach is practical for iterative design workflows where models must be retrained as simulation data accumulates.
Edge deployment of regression models with adaptive output precision. The hierarchical nature of decoding-based regression—coarse prediction from early tokens, progressive refinement from later tokens—enables a capability that pointwise and parametric heads lack: anytime prediction. On an edge device with strict latency constraints, the decoder can be stopped after generating only $k < K$ tokens, producing a coarser but still meaningful prediction. The theoretical framework (Theorem 1) shows that this corresponds to evaluating the density estimator at a lower resolution $k$, with known bias-variance properties. A deployment system could dynamically choose $k$ based on available compute budget, current battery level, or the criticality of the prediction (a rough estimate may suffice for non-safety-critical decisions). This capability is built into the autoregressive decoding process by design; no architectural changes or retraining are needed. The paper's BBOB results (Table 1) showing that the decoder achieves Kendall-Tau correlations of 89.56 at dimension 5 and 86.11 at dimension 20—competitive with pointwise heads—suggest that even coarser predictions (from fewer tokens) would retain substantial rank correlation with ground truth.
When to Prefer This Method
The paper articulates a clear set of tradeoffs between decoding heads and the alternatives it benchmarks. The decision rule can be extracted from the experimental results and the design motivations:
-
Prefer an unnormalized decoding head when: (1) the output
$y$spans multiple orders of magnitude or the range is unknown at training time (multi-task regression, surrogate modeling across physical scales); (2) you need the full conditional density$p(y|x)$without parametric assumptions; (3) you want a single unified architecture for pointwise and distributional prediction; and (4) you can tolerate$K$sequential decoding steps per output sample at inference. The unnormalized decoder's floating-point tokenization handles the curve-fitting cases where pointwise and Riemann heads fail (Figure 4) and achieves the highest Kendall-Tau on BBOB functions without$y$-normalization (Table 1). Use error-correction token repetition with$R \geq 3$if the sample mean is the desired point estimate and outliers are a concern (Figure 10); use the median or beam search for the mode if you need alternative point estimates without the repetition overhead. -
Prefer a normalized decoding head when: (1) the output
$y$can be reliably normalized to a bounded range (e.g.,$[0,1]$via min-max scaling) and the test distribution is expected to stay within that range; (2) you need distributional flexibility without the scale-learning burden of the unnormalized scheme; and (3) the data volume is small to moderate, where the decoder's implicit regularization provides an advantage over Riemann heads (Figures 7, 11). The normalized decoder achieves NLL below 0.7 on all UCI datasets (Table 2) and slightly outperforms the unnormalized variant on several tasks, suggesting the simpler representation is beneficial when the output range is known. The normalized decoder is also simpler to configure (sweep over base$B$and length$K$rather than exponent/mantissa decomposition). -
Prefer a pointwise head when: (1) only a single best-guess prediction is needed, with no uncertainty quantification; (2) inference latency is critical—a pointwise head requires one forward pass versus
$K$sequential passes for the decoder; (3) the output range is well-behaved and normalization is straightforward; and (4) you are operating in the very large data regime where the decoder's data efficiency advantage is diminished and the pointwise head's architectural simplicity dominates. The pointwise head remains competitive on BBOB and several AMLB tasks (Figures 5, 12), and its inference cost is strictly lower. It is the safe default when distributional modeling is not required. -
Prefer a Riemann head when: the paper provides no experimental setting where the Riemann head outperforms a decoder head of equivalent resolution. The one theoretical regime where Riemann could be preferable is when
$N$is very large relative to$B^K$—the high-data, low-resolution corner of Figure 2 where both methods converge to the same risk and the Riemann head's architectural simplicity (single matrix multiply, no sequential decoding) gives it a latency advantage. But the paper does not test this regime on real-world tasks, and the default recommendation from the experimental evidence is to replace Riemann heads with decoder heads. -
Be cautious about MDNs for density estimation without extensive stabilization effort. The paper's MDN results (Table 2, Table 3) show high run-to-run variance and catastrophic failures on several datasets. If you have the expertise to tune MDN training (component initialization, variance flooring, KL regularization, multiple restarts), they can achieve excellent NLL on some tasks (0.05 on Wine, 0.12 on Airfoil). But the decoder head provides reliable performance without this tuning burden, making it the pragmatic choice for non-experts or for applications where robust density estimates matter more than squeezing out the last bit of likelihood on well-behaved datasets.