Text-to-Distribution Prediction with Quantile Tokens and Neighbor Context
Yilun Zhu, Yuan Zhuang, Nikhita Vedula, Dushyanta Dhyani, Shaoyuan Xu, Mohsen Bayati, Bryan Wang, Shervin Malmasi
Abstract
Many applications of LLM-based text regression require predicting a full conditional distribution rather than a single point value. We study distributional regression under empiricalquantile 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 2× narrower prediction intervals), with especially large gains on smaller and more challenging datasets where quantile tokens produce substantially sharper and more accurate distributions. 1 * Work done while at Amazon. Currently at Apple. 1 Our code is publicly available at https://github.com/ yilunzhu/text2distribution/ .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext dece2374-23c8-4b45-901b-eeebadcdce43Builds on8
- Retrieval-Augmented Generation for Knowledge-Intensive NLP TasksPatrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni et al.NeurIPS 2020 · 19,162 citations
- LoRA: Low-Rank Adaptation of Large Language ModelsEdward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu et al.ICLR 2022 · 18,833 citations
- Self-RAG: Learning to Retrieve, Generate, and Critique through Self-ReflectionAkari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil et al.ICLR 2024 · 1,798 citations
- Large Language Models Are Zero-Shot Time Series ForecastersNate Gruver, Marc Finzi, Shikai Qiu, Andrew Gordon WilsonNeurIPS 2023 · 898 citations
- What Can Transformers Learn In-Context? A Case Study of Simple Function ClassesShivam Garg, Dimitris Tsipras, Percy Liang, Gregory ValiantNeurIPS 2022 · 883 citations
Related papers
- Teaching Metric Distance to Discrete Autoregressive Language ModelsJiwan Chung, Saejin Kim, Yongrae Jo, Jaewoo Park et al.ICLR 2026 · 5 citations
- Multi-modal Auto-regressive Modeling via Visual TokensTianshuo Peng, Zuchao Li, Lefei Zhang, Hai Zhao et al.ACM MM 2024 · 1 citation
- DisPPO: Quantile-Based Distributional Reinforcement Learning for Large Language ModelsZhijian Zhou, Long Li, Xuan Zhang, Zongkai Liu et al.ICML 2026
- Eliciting Numerical Predictive Distributions of LLMs Without Auto-RegressionJulianna Piskorz, Kasia Kobalczyk, Mihaela van der SchaarICLR 2026 · 2 citations
- Enhancing Numerical Prediction of MLLMS With Soft LabelingPei Wang, Zhaowei Cai, Hao Yang, Davide Modolo et al.ICCV 2025 · 1 citation
