Decoding-based Regression (2501.19383v1)

Abstract: LLMs have recently been shown capable of performing regression tasks wherein numeric predictions are represented as decoded strings. In this work, we provide theoretical grounds for this capability and furthermore investigate the utility of causal auto-regressive sequence models when they are applied to any feature representation. We find that, despite being trained in the usual way - for next-token prediction via cross-entropy loss - decoding-based regression is as performant as traditional approaches for tabular regression tasks, while being flexible enough to capture arbitrary distributions, such as in the task of density estimation.

• The paper formalizes decoding-based regression by converting numeric prediction into a tokenized, autoregressive sequence modeling task.
• It demonstrates data-efficient performance in tabular regression and competitive density estimation, matching Gaussian mixture models.
• The work provides theoretical bounds and error correction strategies to mitigate numeric instabilities inherent in token quantization.

Here is a detailed summary of the paper.

The paper introduces decoding-based regression, a technique that leverages causal auto-regressive sequence models for regression tasks. Instead of directly predicting numerical values, the method represents them as decoded strings, offering a flexible approach to modeling numeric outcomes. The authors provide theoretical underpinnings for this capability and evaluate its performance on tabular regression and density estimation tasks.

The contributions of the paper are:

• Formalization of decoding-based regression including tokenization schemes.
• Theoretical bounds on fitting smooth one-dimensional densities.
• Empirical demonstration of data-efficient performance competitive with traditional pointwise heads in tabular regression.
• Demonstration of expressive power for matching Gaussian mixture models in density estimation tasks.

The paper begins by highlighting the recent application of LLMs to regression tasks, where numeric predictions are represented as decoded strings. The performance of a regression model depends on how it processes the input features $x$ and how it models the output $y$, including the representation of $y$ and distributional assumptions on $p(y|x)$. The authors note that while previous work has explored text-to-text regression and text-to-anything regression (where LLM embeddings are attached to traditional numeric tensor heads or in-context neural processes), the inverse case of "anything-to-text" regression using decoding-based regression heads alone has been relatively unexplored.

The authors formalize decoding-based regression by specifying a tokenization scheme and providing bounds on its ability to fit arbitrary smooth one-dimensional densities. They explore methods for pointwise estimation and demonstrate empirically that decoding-based regression heads, with appropriate settings, are data-efficient and competitive with regular pointwise heads on tabular regression tasks. Furthermore, they show that these heads are expressive enough to match Gaussian mixture models for density estimation tasks.

The paper argues that using tokens to represent floating-point numbers is a natural outcome, despite the lack of a direct notion of numeric distance when using cross-entropy loss. They provide an overview of relevant work on regression heads, including tensor-based representations, parametric distribution heads (e.g., Gaussians), and histogram (Riemann) distributions. The authors propose that decoding a sequence can simplify learning numeric distances, offering an exponential reduction in bin count compared to histogram distributions.

In decoding-based regression, a token representation maps a real number $y \in \mathbb{R}$ to a fixed-length sequence of tokens $(t_1, \dots, t_K) \in \mathcal{V}$, where $\mathcal{V}$ is the vocabulary of possible tokens. This mapping is lossy and introduces rounding error. Given a feature representation $\phi(x) \in \mathbb{R}^d$, a decoding-based regression head is represented as an auto-regressive prediction model $$p_\theta(t_{k} \mid \phi(x), t_{1}, \ldots, t_{k-1})$$, from which an end-to-end model $$p_\theta(y|x) = p_\theta(t_1, \ldots, t_K \mid \phi(x)) = \prod_{k=1}^{K} p_\theta(t_{k} \mid \phi(x), t_{1}\ldots t_{k-1})$$ is obtained.

The paper discusses normalized and unnormalized tokenization schemes. In normalized tokenization, $y$ values are restricted to $[0,1]$, and any smooth density $p(y|x)$ can be represented with increasing granularity as more tokens are used. This aligns with standard data-science practices of normalizing $y$-values. In unnormalized tokenization, the authors generalize the IEEE (Institute of Electrical and Electronics Engineers)-754 floating-point representation for base-2 to any base $B$. Each number is represented as $s \cdot B^{e} \cdot m$, where $s \in \{-1, +1\}$ is the sign, $e \in \mathbb{Z}$ is the exponent, and $m \in [0, B)$ is the mantissa. The tokenization is {$s$}{$s_e$}{$e_1$}\ldots{$e_{E}$}{$m_1$}\ldots{$m_M$}, where $s_{e}, e_1, e_2, \ldots, e_E$ are the sign and base-B representation of the exponent $e$, and $m_1, m_2, \ldots, m_M$ are the most significant base-B digits of the mantissa $m$.

Any autoregressive model can be used, provided it supports constrained token decoding to enforce valid numeric sequences. The authors use a small Transformer due to its autoregression capabilities, with the initial token embedding as $\phi(x)$. The Transformer size is kept minimal compared to the encoder.

For pointwise estimation, the paper discusses estimating a scalar quantity of interest $M(p_\theta)$ from the model's distribution. For common error functions like mean squared error and mean absolute error, the optimal values are the mean and median of $p(\cdot | x)$, respectively. The mode can be approximated using beam search. The authors note that for unnormalized tokenization, outliers can significantly impact non-robust estimators like the sample mean, and suggest using alternative tokenizations based on coding theory or decoding techniques from the LLM literature (e.g., top-$k$, top-$p$) to mitigate this issue.

For density estimation, the authors apply the standard cross-entropy loss over all sequence tokens during training. They provide formal guarantees for estimating one-dimensional densities on $[0,1]$ using a tree-based tokenization and training loss.

The authors present a theorem that provides a risk bound, assuming the decoding-based regression model is $K$-bit universal. The risk $R$ is defined as the mean integrated squared error between the true density $f$ and its estimator $\widehat{f}_N(Y_1,\dots,Y_N)$: $$R(f, \widehat{f}_N) = \mathop{E}_{Y_1,\dots, Y_N \sim f} \int_0^1 \left(f(y)-\widehat{f}_N(y)\right)^2 dy$$. Given density $$\widehat{f}_N^k(y) = 2^k p_{\widehat\theta(Y_1,\dots,Y_N)}^k(\lambda_k(y)) \;\text{for}\;y \in [0,1]$$, the risk is: $R\left(f,\widehat{f}_{N}^k\right) \approx \frac{2^{-2k}{12} \int_0^1 f'(y)^2 dy + \frac{2^k}{N}, \;\;\; \forall k \leq K$. Here, $k$ is the number of bits, $K$ is the maximum number of bits, $N$ is the number of data points, $f'(y)$ is the derivative of the density function, and $\lambda_k(y)$ is the operation that returns the first $k$ bits after the radix point in the binary representation of $y$. The theorem implies that a higher resolution $K$ is needed to capture the curvature of $f$, but as the number of bins increases, more data points $N$ are required to learn to separate these $2^K$ bins.

The paper details several experiments designed to demonstrate the effectiveness of decoders as replacements for pointwise regression heads, establish the density estimation capabilities of decoding-based heads, and ablate the effect of decoder size and sequence-specific methods like error correction on performance.

The authors find that the unnormalized decoder can effectively capture the shapes of various functions, even with unbounded training data, while pointwise heads struggle. For synthetic continuous objectives from the Black-box Optimization Benchmarking (BBOB) suite, both unnormalized and normalized decoder variants can sufficiently fit functions over various input dimensions.

Over real-world OpenML regression tasks from OpenML-CTR23 and AMLB, the unnormalized decoder is competitive with regular pointwise heads, often outperforming them. A normalized decoder, Riemann histogram distribution head, and pointwise head are compared by varying the amount of training data. The authors observe the data inefficiency of using the histogram distribution and note that while the decoder may struggle more in low data regimes, the pointwise head can perform worse due to numeric instabilities.

The paper also visualizes the decoder's ability to perform density estimation over various shapes, capturing the overall distribution $p(y|x)$ well. They also display the negative log-likelihood (NLL) on real-world datasets from the UCI regression repository, finding that Mixture Density Network (MDN) performance has high variability, while decoding methods remain reliable.

The effect of the decoder size on performance is ablated, with results showing that larger decoder models can help up to a certain point, beyond which overfitting can occur. Finally, the paper explores improving regression behavior using error correction techniques, such as having the decoder repeat its output multiple times and performing majority voting at inference.

The paper concludes by summarizing the benefits and drawbacks of decoding-based regression. It is also capable of density estimation over a variety of conditional distributions $p(y|x)$, and can outperform common baselines such as Gaussian mixtures and Riemann distributions. The authors suggest that future work could explore improved tokenization schemes, alternative basis distributions, and applications in computer vision and multi-target regression.