Text-to-Distribution Prediction with Quantile Tokens and Neighbor Context

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Zhu, Yilun, Zhuang, Yuan, Vedula, Nikhita, Dhyani, Dushyanta, Xu, Shaoyuan, Li, Moyan, Bayati, Mohsen, Wang, Bryan, Malmasi, Shervin
Format: Preprint
Publié: 2026
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866911614083530752
author Zhu, Yilun
Zhuang, Yuan
Vedula, Nikhita
Dhyani, Dushyanta
Xu, Shaoyuan
Li, Moyan
Bayati, Mohsen
Wang, Bryan
Malmasi, Shervin
author_facet Zhu, Yilun
Zhuang, Yuan
Vedula, Nikhita
Dhyani, Dushyanta
Xu, Shaoyuan
Li, Moyan
Bayati, Mohsen
Wang, Bryan
Malmasi, Shervin
contents Many applications of LLM-based text regression require predicting a full conditional distribution rather than a single point value. We study distributional regression under empirical-quantile supervision, where each input is paired with multiple observed quantile outcomes, and the target distribution is represented by a dense grid of quantiles. We address two key limitations of current approaches: the lack of local grounding for distribution estimates, and the reliance on shared representations that create an indirect bottleneck between inputs and quantile outputs. In this paper, we introduce Quantile Token Regression, which, to our knowledge, is the first work to insert dedicated quantile tokens into the input sequence, enabling direct input-output pathways for each quantile through self-attention. We further augment these quantile tokens with retrieval, incorporating semantically similar neighbor instances and their empirical distributions to ground predictions with local evidence from similar instances. We also provide the first theoretical analysis of loss functions for quantile regression, clarifying which distributional objectives each optimizes. Experiments on the Inside Airbnb and StackSample benchmark datasets with LLMs ranging from 1.7B to 14B parameters show that quantile tokens with neighbors consistently outperform baselines (~4 points lower MAPE and 2x narrower prediction intervals), with especially large gains on smaller and more challenging datasets where quantile tokens produce substantially sharper and more accurate distributions.
format Preprint
id arxiv_https___arxiv_org_abs_2604_20216
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Text-to-Distribution Prediction with Quantile Tokens and Neighbor Context
Zhu, Yilun
Zhuang, Yuan
Vedula, Nikhita
Dhyani, Dushyanta
Xu, Shaoyuan
Li, Moyan
Bayati, Mohsen
Wang, Bryan
Malmasi, Shervin
Computation and Language
Many applications of LLM-based text regression require predicting a full conditional distribution rather than a single point value. We study distributional regression under empirical-quantile supervision, where each input is paired with multiple observed quantile outcomes, and the target distribution is represented by a dense grid of quantiles. We address two key limitations of current approaches: the lack of local grounding for distribution estimates, and the reliance on shared representations that create an indirect bottleneck between inputs and quantile outputs. In this paper, we introduce Quantile Token Regression, which, to our knowledge, is the first work to insert dedicated quantile tokens into the input sequence, enabling direct input-output pathways for each quantile through self-attention. We further augment these quantile tokens with retrieval, incorporating semantically similar neighbor instances and their empirical distributions to ground predictions with local evidence from similar instances. We also provide the first theoretical analysis of loss functions for quantile regression, clarifying which distributional objectives each optimizes. Experiments on the Inside Airbnb and StackSample benchmark datasets with LLMs ranging from 1.7B to 14B parameters show that quantile tokens with neighbors consistently outperform baselines (~4 points lower MAPE and 2x narrower prediction intervals), with especially large gains on smaller and more challenging datasets where quantile tokens produce substantially sharper and more accurate distributions.
title Text-to-Distribution Prediction with Quantile Tokens and Neighbor Context
topic Computation and Language
url https://arxiv.org/abs/2604.20216