Elamur Research⌕
☰
home / papers / 2026 / Text-to-Distribution Prediction with Quantile Tokens and Neighbor Context

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

Paper · 2026 · Yilun Zhu, Yuan Zhuang, Nikhita Vedula et al.

Summary

Most LLM-based text regression outputs a single number. Tasks like pricing or demand forecasting need the whole conditional distribution. The authors treat this as text-to-distribution prediction. Each text input has several observed outcomes, which are converted into an empirical quantile function and interpolated onto 99 levels (0.01–0.99).

They argue that the multi-head design of Vedula et al. (2025) creates a bottleneck, because every quantile is read from one shared hidden state. Their alternative is Quantile Token Regression: learnable quantile tokens are appended to the input, so each quantile collects its own evidence through self-attention. They also retrieve semantically similar neighbors and supply each one with its full empirical distribution, not a single label. A theoretical analysis of loss functions favors Wasserstein matching over pinball variants.

Two benchmarks are built from Inside Airbnb and StackSample and tested with Qwen3 models (1.7B–14B). Tokens plus neighbors cut MAPE by roughly 4 points and halve interval width. The largest gains appear on the smaller, harder StackSample data. This matters wherever spread and tails drive decisions.

Key points

  • Retrieval with neighbor distributions lowers average MAPE by about 8% relative on Airbnb and 63% on StackSample. Quantile tokens beat the shared-head baseline on StackSample by 14% relative MAPE, with intervals about 6× narrower.
  • ℓ1 Wasserstein loss gives the best accuracy–sharpness tradeoff. Pinball-based objectives score higher on CRPSS but produce RCIW@95 above 150, versus 3–4 for Wasserstein.
  • Neighbor ablation (Qwen3-4B, Airbnb dev):
    • Going from K=0 to K=16 reduces average MAPE from 29.54 to 25.85.
    • K=16 needs about 2× memory and 1.4× training time compared with K=8.
  • Predicted quantiles are not guaranteed to be ordered. Sorting them at inference time (post-processing) gives the lowest MAPE (26.35) without widening intervals, unlike a cumulative-sum parameterization.
  • Code is released at github.com/yilunzhu/text2distribution. The stated limitations are interpolation error when an instance has few labels and testing on only two backbone families (Qwen3, Phi-3).

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Stanford UniversityStanford · United States506 works · 330,053 citations🏢AmazonSeattle · United States128 works · 41,700 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 works

Cite

Cite as Yilun Zhu, Yuan Zhuang, Nikhita Vedula, Dushyanta Dhyani, Shaoyuan Xu, Mohsen Bayati, Bryan Wang and Shervin Malmasi. 2026. Text-to-Distribution Prediction with Quantile Tokens and Neighbor Context. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.acl-long.758

@inproceedings{zhu2026text,
  title = {{Text-to-Distribution Prediction with Quantile Tokens and Neighbor Context}},
  author = {Yilun Zhu, Yuan Zhuang, Nikhita Vedula, Dushyanta Dhyani, Shaoyuan Xu, Mohsen Bayati, Bryan Wang and Shervin Malmasi},
  booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  year = {2026},
  doi = {10.18653/v1/2026.acl-long.758},
  url = {https://doi.org/10.18653/v1/2026.acl-long.758}
}

References

36 references, 9 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Marah Abdin, Jyoti Aneja, Hany Awadalla, Ahmed Awadallah, Ammar Ahmad Awan, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Jianmin Bao, Harkirat Behl, Alon Benhaim, Misha Bilenko, Johan Bjorck, Sébastien Bubeck, Martin Cai, Qin Cai, Vishrav Chaudhary, Dong Chen, Dongdong Chen, and 110 others. 2024. Phi-… — Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone
  2. Anthropic. 2025. Claude sonnet 4. https://www.anthropic.com. Accessed: 2026. ↗
  3. Siddharth Arora, James W. Taylor, and Ho-Yin Mak. 2023. Probabilistic forecasting of patient waiting times in an emergency department. find
  4. Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. 2023. Self-rag: Learning to retrieve, generate, and critique through self-reflection. arXiv preprint arXiv:2310.11511.
  5. R. R. Bahadur. 1966. A note on quantiles in large samples. The Annals of Mathematical Statistics, 37(3):577–580. find
  6. Zsolt Bitvai and Trevor Cohn. 2015. Non-linear text regression with a deep convolutional neural network. In Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 2: Short Papers), page… find
  7. Tingting Chen and Shijing Si. 2024. Predicting rental price of lane houses in Shanghai with machine learning methods and large language models. arXiv preprint arXiv:2405.17505. — [arXiv:2405.17505](https://arxiv.org/abs/2405.17505) find
  8. Cheng-Han Chiang, Hung-yi Lee, and Michal Lukasik. 2025. TRACT: Regression-aware fine-tuning meets chain-of-thought reasoning for LLM-as-a-judge. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 2934–2952, Vienna, Austria. Asso… find
  9. Nicolai Dorka. 2024. Quantile regression for distributional reward models in rlhf. arXiv preprint arXiv:2409.10164. — [arXiv:2409.10164](https://arxiv.org/abs/2409.10164) find
  10. Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. 2022. What can transformers learn in-context? a case study of simple function classes. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS ’22, Red Hook, NY, USA. Curran Associates Inc.
  11. Nate Gillman, Daksh Aggarwal, Michael Freeman, Saurabh Singh, and Chen Sun. 2025. Fourier head: Helping large language models learn complex probability distributions. In Proceedings of ICLR. find
  12. Tilmann Gneiting and Adrian E. Raftery. 2007. Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association, 102(477):359–378. find
  13. Nate Gruver, Marc Finzi, Shikai Qiu, and Andrew Gordon Wilson. 2023. Large language models are zero-shot time series forecasters. In Advances in Neural Information Processing Systems, volume 36.
  14. Wenjun Gu, Yihao Zhong, Shizun Li, Changsong Wei, Liting Dong, Zhuoyue Wang, and Chao Yan. 2024. Predicting stock prices with FinBERT-LSTM: Integrating news sentiment analysis. In 2024 8th International Conference on Cloud and Big Data Computing. find
  15. Ragıp Gürlek, Francis de Véricourt, and Donald K.K. Lee. 2024. Boosted generalized normal distributions: Integrating machine learning with operations knowledge. Working paper, SSRN. Available at SSRN, Posted: August 1, 2024. find
  16. Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2022. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations.
  17. Joseph Marvin Imperial. 2021. BERT embeddings for automatic readability assessment. In Proceedings of the International Conference on Recent Advances in Natural Language Processing (RANLP 2021), pages 611–618, Held Online. INCOMA Ltd. find
  18. Inside Airbnb. 2025. Get the data. https://insideairbnb.com/get-the-data/. Dataset downloads (city snapshots). Licensed under CC BY 4.0. Accessed 2025-12-17. ↗
  19. Ryan Jacobs, Maciej P. Polak, Lane E. Schultz, Hamed Mahdavi, Vasant Honavar, and Dane Morgan. 2024. Regression with large language models for materials and molecular property prediction. arXiv preprint arXiv:2409.06080. — [arXiv:2409.06080](https://arxiv.org/abs/2409.06080) find
  20. Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. Preprint, arXiv:2001.08361.
  21. Thomas Kneib, Alexander Silbersdorff, and Benjamin Säfken. 2023. Rage against the mean–a review of distributional regression approaches. Econometrics and Statistics, 26:99–123. find
  22. Roger W Koenker and Gilbert Bassett. 1978. Regression quantiles. Econometrica, 46(1):33–50. find
  23. Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. 2020. Retrieval-augmented generation for knowledge-intensive nlp tasks. In Proceedings of the 34th Internation…
  24. Michal Lukasik, Zhao Meng, Harikrishna Narasimhan, Yin-Wen Chang, Aditya Krishna Menon, Felix Yu, and Sanjiv Kumar. 2025. Better autoregressive regression with llms via regression-aware fine-tuning. In Proceedings of ICLR. find
  25. Xingyou Song, Oscar Li, Chansoo Lee, Bangding Yang, Daiyi Peng, Sagi Perel, and Yutian Chen. 2024. Om-nipred: Language models as universal regressors. find
  26. Trans. Mach. Learn. Res., 2024. models for price prediction. In Findings of the Association for Computational Linguistics: ACL 2025, pages 12396–12415, Vienna, Austria. Association for Computational Linguistics. find
  27. Cédric Villani. 2009. Optimal Transport: Old and New, volume 338 of Grundlehren der mathematischen Wissenschaften. Springer. find
  28. Hairu Wang, Sheng You, Qiheng Zhang, Xike Xie, Shuguang Han, Yuchen Wu, Fei Huang, and Jufeng Chen. 2025. Llp: Llm-based product pricing in ecommerce. Preprint, arXiv:2510.09347. — [arXiv:2510.09347](https://arxiv.org/abs/2510.09347) find
  29. Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pier-ric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, and 3 others. 2020. Huggin…
  30. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, and 41 others. 2025. Qwen3 technical report. arXiv preprint arXiv:2505.09388.
  31. Ruosong Yang, Jiannong Cao, Zhiyuan Wen, Youzheng Wu, and Xiaodong He. 2020. Enhancing automated essay scoring performance via fine-tuning pre-trained language models with combination of regression and ranking. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 1560–1569… find
  32. Stack Overflow. 2019. Stacksample: 10% of stack overflow q&a. Kaggle dataset. Kaggle-hosted dataset. The page shows "Updated 6 years ago" when accessed on 2025-12-17, implying a last-updated year of approximately 2019. find
  33. Stack Overflow. 2025. What is the license for the content i post? Stack Overflow Help Center. States Stack Overflow user contributions are licensed under CC BY-SA, with version depending on contribution date (2.5/3.0/4.0). find
  34. Eric Tang, Bangding Yang, and Xingyou Song. 2025. Understanding LLM embeddings for regression. Transactions on Machine Learning Research. find
  35. Robert Vacareanu, Vlad-Andrei Negru, Vasile Suciu, and Mihai Surdeanu. 2024. From words to numbers: Your large language model is secretly a capable regressor when given in-context examples. Preprint, arXiv:2404.07544. — [arXiv:2404.07544](https://arxiv.org/abs/2404.07544) find
  36. Nikhita Vedula, Dushyanta Dhyani, Laleh Jalali, Boris N. Oreshkin, Mohsen Bayati, and Shervin Malmasi. 2025. Quantile regression with large language find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zhu-2026-text-distribution-prediction-quantile-tokens/