Bayesian Forecasting of LLM Outputs from Pre-Generation Hidden States

The stochasticity of next-token prediction in large language models (LLMs) implies that LLM outputs should be represented probabilistically. Yet for problem settings in which outputs take on a structured format, obtaining such probabilities can be costly, as this requires repeated generation in response to a user's query. In this paper we propose a computationally cheap alternative to approximate LLM-derived probabilities. Using pre-generation hidden state representations of an LLM, our method forms predictive distributions that govern predefined concepts of the output space, aiming to forecast empirically-derived LLM probabilities. The approximating nature of the approach, a consequence of limited information and lack of foresight in hidden states, implies that a proper treatment of uncertainty is essential if we wish to place confidence in the resulting distributions. To this end, we utilize Bayesian classifiers in order to decompose a predictive distribution into aleatoric uncertainty, namely the LLM stochasticity that we intend to model, and epistemic uncertainty, or the model uncertainty arising from a lack of evidence in the LLM hidden states. We ground our work in the domain of visualization specification generation, representative of a task for which outputs adhere to a predefined grammar, and output concepts, such as graphical marks and visual channels, are well-defined. Experimental results illustrate that hidden states encode substantial information required for forecasting visualization design concepts, where forecasting difficulty increases towards composite tasks that require both context understanding and response planning. Hidden states also provide meaningful evidence for estimating uncertainty, correlated with inherent difficulty of the given query, suggesting that uncertainty can act as a meaningful criteria for evaluating forecasting quality.