On the Importance of Uncertainty in Decision-Making with Large Language Models
Authors: Nicolò Felicioni
Lucas Maystre
Sina Ghiassian
Kamil Ciosek
Venue: Accepted by TMLR ()

Abstract: We investigate the role of uncertainty in decision-making problems with natural language as input. For such tasks, using Large Language Models as agents has become the norm. However, none of the recent approaches employ any additional phase for estimating the uncertainty the agent has about the world during the decision-making task. We focus on a fundamental decision-making framework with natural language as input, which is the one of contextual bandits, where the context information consists of text. As a representative of the approaches with no uncertainty estimation, we consider an LLM agent with a greedy policy, which picks the action corresponding to the largest predicted reward. We compare this baseline to LLM agents that make active use of uncertainty estimation by integrating the uncertainty in a Thompson Sampling policy. We employ different techniques for uncertainty estimation, such as Laplace Approximation, Dropout, and Epinets. We empirically show on real-world data that the greedy policy performs worse than the Thompson Sampling policies. These findings suggest that, while overlooked in the LLM literature, uncertainty improves performance on bandit tasks with LLM agents.

Published in Transactions on Machine Learning Research (07/2024)
On the Importance of Uncertainty in Decision-Making with Large Language Models
Nicolò Felicioni *
nicolo..elicionn@@pppoiii.
Politecnico di Milano
Lucas Maystre Spotify
lucasm@@spolify.coom com
Sina Ghiassian Spotify
sinag@Aspotify.com com
Kamil Ciosek
kamelle @spolify.....coommmmm com
Spolify
Reviewed OpenReview:
https: openreview. net/ forum? id= x/PxU
Abstract
We investigate the role of uncertainty in decision-making problems with natural language as input. For such tasks, using Large Language Models as agents has become the norm. However, none of the recent approaches employ any additional phase for estimating the uncertainty the agent has about the world during the prcision-making task. We focus on fundamental decision-making framework with natural language input, which is the one of contextual bandits, where the context information consists of text. As representative of the approaches with no uncertainty estimation, we consider an L.M agent with greedy policy, which picks the action corresponding to the largest predicted reward. We compare this bascline to LLM agents that make active use uncertainty estimation by integrating the uncertainty a Thompson Sampling policy. We employ different techniques for uncertainty estimation, such as Laplace Approximation, Dropout, and Epincts. We empirically show on real-world data that the greedy policy performs worse than the Thompson Sampling policies. These findings suggest that, while overlooked in the LLM literature, uncertainty improves performance on bandit tasks with LLM agents.
1 Introduction
0.350
Greedy
Large language models (LLMs) have emerged domi- nant paradigm in natural language processing (Ouyang
0.325
Last LA T5
Dropout
et al. performance across a wide range of tasks (Rae et al. 2021). To reach this progress, LLMs have pushed model scale and dataset. size unprecedented levels. By opti- mizing such immense models exclusively predict text, perhaps surprisingly, LLMs have achieved strong perfor- mance broad range of datasets and tasks (Bubeck et al. 2023), including translation, question--nsswerngg, and dialogue.
2022 OpenAl 2023). achieving state-of-the-art
0.300
Epinet TS
Average regret.
0.275
0.200
0.175
0.50
Observed data points
Figure 1: Average regret obtained on toxic con- tent detection bandit task.
Research completed while internship Spotify.

---

Published in Transactions on Machine Learning Research (07/2024)
In parallel, many critical real-world systems are increasingly relying on these models to make decisions Yang et al. 2023b), where the consequences of a particular action are typically reflected by a reward signal. This usually works by fine-tuning an LLM to serve as a reward model, where we try to teach the network to prodict the mcan reward for cach action. A common way of using this model to make decisions is to use the predictions greedily (Riquelme et al. 2018), laking the action with the greatest estimated reward. However, this approach ignores the fact that reward estimates may be inaccurate, which can lead to never-ending pathological behavior (Russo et al.) 2018). On the other hand, the literature on contextual bandils (Wang) et al. 2005) provides a principled approach to deal with inaccurate reward estimates in decision-makinn task. A popular modelling tool for doing this in the Bayesian framework is to maintain probabilistic reward model that relies on two separate notions of uncertainty: epistemic and aleatoric.
Epistemic uncertainty is the one that reflects the fact that we have not yet seen enough fine-tuning data to estimate the mean reward well, while alcatoric uncertainty reflects the irreducible noise associated with observing a reward. When such a reward model coupled with an action-seloction algorithm like Thompson Sampling (Thompson 1933). high-quality estimates of epistemic uncertainty can be leveraged to produce better balance between exploration and exploitation in the decision- making task, evidenced by lower regret.
In this paper, we want to investigate the importance of employing epistemic uncertainty models in decision- making lasks with natural language and LLMs. In particular, we focus on one of the most fundamental satisionn-mainn tasks with natural language as input: conteztual bandits. In a contextual bandit problem, the agent must continually take actions based on observed context, which in our case is text, and observes the reward only for the chosen action. This framework is applicable to a myriad of real-world scenarios.
As concrete example application, let consider the problem of automated content moderation on online platforms (Gorwa et al. 2020 Vaccaro et al., 2020: Ma & Kou, 2021 Avadhanula et al., 2022). In this scenario, users post comments and the agent decides whether to publish each comment. The user comments serve as the context, in text form. The agent takes an action (e.g. "publish" or "not publish") ") based on the context. After acting, the environment provides the reward ouly for the selected action, which the agent uses to update its policy. For example, if the agent publishes a toxic comment, it receives a large negative reward since other users will report the comment. Instead, if the agent does not publish a comment, it receives a neutral reward regardless of the comment's 's content, since the comment is not seen. To learn effectively in this setting, the agent must balance exploration and exploitation.
In this paper, we study the contextual bandit scenario empirically, combining the insights from the LLM and bandit communities. The greedy approach outlined above is our main baseline and was chosen for its simplicity and ease of implernentation. We benchmark it against approximate Bayesian models coupled with the Thompson Sampling algorithm for action selection. Since Thompson Sampling requires epistemic uncertainty to be of high quality to work well, we investigate several different approaches of doing so. Specifically, we compare techniques for epistemic uncertainty estimation that can be scaled up be used with current LLMs. These are the Laplace Approximation (Daxberger et al. 2021), Dropout (Gal & Ghahramani 2016) and Epincts (Osband et al.) 2023a We show that all our TS policies significantly and consistently oulperform the greedy baseline on a real-world bandit. benchmark in terms of final regret. A preview of the results is shown in Figure[] [ These findings suggest that, while often overlooked in the LLM literature, epistemic uncertainty improves performance on bandit tasks with LLM agents.
Contributions
We provide an extensive empirical study of LLM-based contextual bandits;
We adapt existing neural epistemic uncertainty estimation techniques to LLMs. In particular, we identify the best variant of the Laplace approximation and dropout to use, and we show one possible epinet architecture that can also tackle this task;
We empirically show on real-world data that the greedy policy is sub-optimal in terms of final regret, and that epistemic uncertainty is helpful in contextual bandit, problems with text data.
2

---

Published in Transactions on Machine Learning Research (07/2024)
To the best of our knowledge, this is the first time that such empirical analysis on bandit learning with LLMs has been provided. This empirical evaluation sheds light on critical aspects of LLMss, such as their performance on decision-making tasks and the importance of epistemic uncertainty for those models.
2
Preliminaries
In this section, we describe the problem statement and we introduce the topic of uncertainty Machine Learning models.
2.1 Batch Contextual Bandit problem
The contextual bandit problem is sequential decision-making framework where an agent interacts with an environment over number of time steps. In particular, we focus on the more general case where, for each time step, we observe a batch of contexts and the agent has to select batch of actions (Kandasamy et al. 2018).
Let us call the total number of time steps T. At each time step t t= 1,2 the agent observes a batch of contexts, which we denote with ,,....,, Each context x contains useful information that should be used to solve the task. In this paper, we focus on dccision-making problems driven by text. Hence, this case, each context = consists of text information. The agent. then uses the observed data until now, which we call D, to select action at from set A of K possible actions, conditioned on the given context zt with b E {1...., B}. For cach selected action of, the agent receives a reward ry that depends on the context z and action a². At the end of the time step, we have observed a set of tuples that we denote as D. = {( r(),)..)... (o B rr])- This set of tuples added to the observed dataset +- DUD, The aim of the agent is to learn a policy that minimizes the regret over the T time steps. We define the average regret as the difference between the cumulative reward from the learned policy and the cumulative reward from an optimal policy, all divided by the number of time steps. We define the optimal policy x² the policy that, within the considered policy space II, maximizes the expected reward: x* arg max Ea [r(z, a)]. Let us call rī the expected rewards obtained by the optimal policy at any time step t. We can define the average regret as follows:
T Rr =
B Σ
T
B
Since we focus on bandits with text. information as contexts, we would like use the slate-of-the-art natural language processing capabilities of pre-trained LLMs and build bandit policies based on these models. We will explain in more detail how to design such deep bandits in Section 3
2.2 Aleatoric and Epistemic Uncertainty in Machine Learning
In the field of machine learning, the concept of uncertainty plays crucial role in both model interpretation and reliability. A common framework to characterize uncertainty, given particular model class, is to catego- rize it into two types: aleatoric and epistemic. Aleatoric uncertainty arises from the inherent randomness in the data or the environment. It represents the variability in the outcome that cannot be reduced even if more data is provided. Hence, this type of uncertainty is irreducible and intrinsic to the current task or dataset. On the other hand, epistemic uncertainty stems from the lack of knowledge. This type of uncertainty is roducible by collecting more data. It reflects the uncertainty in the model parameters due to limited data limited model capacity.
As concrete example, consider contextual bandit setting with observed data D {(=1 01,)),,.)) and parametric model fo(x, a) that we wish to train to predict expected reward given context and action, with parameters 0. In this regression setting, a common approach to model the aleatoric uncertainty is to assume that there is additive noise in the rewards: r = E[r] + e, where the noise is distributed according to a generic distribution NO, 2, where 2he the observation variance. This kind of uncertainty independent of the model, it is irreducible even with infinite data.
P(c). A simple but effective choice set the noise distribution be Gaussian:
3

---

Published in Transactions on Machine Learning Research (07 (07/2024) /024)
Algorithm 1 Greedy
Require: Bandit model fa.
1: Initialize D マ
2: for time 1 .,..., T do
3: Observe contexts =t,-. Zt
4: for b 00=1, ,,...,, B do
5: Select d; = arg maxa fo Observe reward rg.
end for
Create D. a
a a,f,tt)))
Add to the observed dataset D t- (D,,...,Di?
10: Update the paramcters f training on the observed data D minimizing Eq. 3
11: end for
Epistemic uncertainty, instead, stems from uncortainty in the model parameters 0. A principled Bayesian approach is to maintain posterior distribution over the parameters, P(O)D), that is updated new dala arrives. The posterior covariance quantifies epistemic uncertainty. As the dataset grows, the posterior distribution becomes more peaked, reducing the epistemic uncertainty.
Maintaining a posterior over allows the model to "know what it does not know", critical for balancing explo- ration and exploitation in bandits. We will expand on using Bayesian techniques for principled exploration in Section
3 Large Language Model Bandits
In this section, we describe how to build bandit agents with pre-trained Large Language Models.
We approach the contextual bandit problem with a regression model. Specifically, we train a model to predict the expected reward for each action, given the context. To leverage recent advances large pre- trained language models, we initialize our model with a pre-trained LLM. As these models are pre-trained for next-token prediction, which is classification task, discard the final classification layer and append new linear regression output layer to predict expected rewards.
Specifically, let Tüpr denote the pre-trained LLM with parameters 0pr. Such a model is a function that maps any sequence of tokens into a probability distribution over the vocabulary: Tep (x) E AVV), where A(V) denotes the simplex over the vocabulary of tokens V. For our task, we do not need the final classification head; hence, we consider only the features before the final layer, which we denote as #0pm (x) E Rd. Now, for each tuple (z,a,r) in our dataset D, we feed æ into the pre-trained feature extractor म and we pass those features through a final linear layer produce the expected reward for each a E A. We call our model fo, which is defined as follows:
fo(x: linear(@opp (x)) R.K where K -AA. Also, we denote the a-th output of our model as fo(x, a) e R.
(1)
3.1 Greedy policy
As representative approach with no additional uncertainty estimation phase, consider the greedy policy, The bchavior of this policy is illustrated in Algorithm I In particular, in our case, we use a pro-trained LLM as a feature extractor with a linear layer on top to create a regression model fa, as described in Section 3 Given a context, the greedy policy selects the action for which the model predicts the highest reward. Let us consider a certain time step f. After the action selection phase, we observe the real reward for the batch of actions the policy has selected. The dataset will be composed of sets of tuples: = (DI...., Dz). We exploit the observed data to update the parameters of the regression model before the beginning of the
4

---

Published in Transactions on Machine Learning Research (07/2024)
next time step. A typical way to do so is to compute maximun-a-possteriar (MAP) estimate. This is done by minimizing loss corresponding to the negative log posterior:
[(O:
Σ 1(0:2,0,7) +
(2)
(z,,r))D
negative log-prior
negative log. rng-likelhhood
Usually, for regression task, it is assumed that there is Gaussian aleatoric noise with zero mean and fixed variance o2ts This implies that the negative log-likelihood part of the loss is mean-squared error (MSE); 1(0:z, a,r) 26 -r fo(x,a) a))². Regarding the prior, another typical assumption is to set the prior belief on the parameters as indepentent Gaussians with O mean and of variance. However, in our case, it is inconsistent to assume that all the weights have 0 prior since we initialize every weight (except for the last layer) to have the same weight as a pre-trained LLM. Hence, we use a different prior mean: 8p = [θPT. O], where with the notation [V1, U2] we denote the concatenation of vectors i1 and 122 Therefore, this means that every weight (except for the last layer) has prior centered in the pre-trained weight value, while the last layer is regularized towards zero. The negative log- prior term hence becomes r(0) 203 10 8p112.
Minimizing the loss shown in Equation Ltt the end of each time step, however, can have some drawbacks in a bandit task. Indeed, it can be highly expensive for large neural models, especially because the complete datasct is constantly increasing at each time step. An alternative and more scalable loss to minimize the following, which considers only the data Dr observed at the current time step d:
B
C(0 Di)
(ou (ll0(t) g(t 0(t-1)112
(3)
where A = 2o and we set 0()== and we update the prior at each time step with the weights obtained at the previous time step. This will be the loss we will use to optimize the greedy bandit our experimental analysis in Section 5 Using this loss, we are doing MAP inference many times, applying the Bayes rule anew for each new batch of data, In this way, perform MAP, which is well-known established principle in deep learning, while avoiding looping throughout the whole dataset. For more detailed discussion, we refer to Appendix B This loss is also present in the literature. For instance, (Daxberger et al. 2021 use this same loss in their continual learning experiment (Daxberger et al. 2021), Appendix C.4.1).
= Op- This loss now only assumes that the data points in Dr are i.i.d.,
3.2 Thompson Sampling
Thompson Sampling (TS) (Thompson 1933 Russo el al. 2018 is a probabilistic algorithm for the contextual bandit problem. The key idea of Thompson Sampling is to maintain an epistemic uncertainty estimate in the form of a posterior distribution over the parameters of the model P(O[D). Initially, this distribution is set prior distribution P(0) that represents the agent's initial uncertainty. On each time step, for each observed context, the agent samples a set of parameters i ~ P(0 from the posterior distribution, and selects the action with the highest reward according to the sampled model. At the end of the time step, the agent updates the posterior distribution of the parameters with the observed data. A high-level description of this procedure is provided in Algorithm 2
This approach balances exploration and exploitation by sampling from a posterior that quantifies the epis- temic uncertainty. Actions with more uncertain estimates will be explored, while actions that are currently believed to have higher rewards will be exploited. TS can be thought as taking actions according to the (cpis- temic) probability that they are optimal, which leads to a good balance between exploration and exploita- tion. The posterior distribution concentrates over time, automatically adjusting the exploration-explotaa. trade-off.
Thompson Sampling has been shown to be effective both in theory and in practice (Chapelle & Li 2011 When the Bayesian updates are exact, there are theoretical guarantees on the performance of TS agents
This technique was also proposed (Xuhong al. 2018 and has been shown reduce the fine-tuning loss when using pre-trained neural models.
5

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Algorithm 2 Thompson Sampling
Require: Bandit model fg.
Require: Prior distribution the parameters P(0) NVpp, Ep).
1: Initialize D + 0
2: for time t = 1 l..... T do
3: Observe context s) .
4: for b == 1....... do
Sample parameters 1 ~~ P() Select d arg max, fe a)) Observe reward rg. end for
Create D, {(x) ( rt),, (xt a, ,r 10: Add to the observed dataset D - DDD
11:
Update the posterior distribution P))
12: end for
(Agrawal & Goyal 2017). For instance, Russo & Van Roy (2016) show that Thompson Sampling has a Bayesian regret bound. However, in our case, we are dealing with a complex bandit model that uses a pre-trained LLM. Therefore, an exact Bayesian update is computationally infeasible, and we have to resort to approximations. Section 4 we will show different tochniques to estimate postcrior distribution of the parameters of a neural nelwork and how adapt such techniques our case of LLM bandits.
4 Epistemic uncertainty estimation for pre-trained LLMs
In order to use a Thompson Sampling policy, need to estimate the epistemic uncertainty of our model. This additional step is not required for greedy policy. Also, it is not trivial step for deep neural networks, for which there is no (tractable) closed-form solution to update the posterior distribution of the parameters.
Crucially, for our particular task, the uncertainty estimation technique is required to scale to very large models, which are typically pre-trained at great expense, such as LLMs. This means that we cannot employ techniques that require us to change the training process |Graves 2011), or ensembles (Lakshminarayanan ct al., 2017), which require us to run training multiple times. Instead, we rely on more scalable techniques, such Dropout, Laplace Approximation, and Epinets. In the following, we discuss these epistemic uncer- tainty estimation techniques in further detail, and show how to adapt them to LLM bandits.
4.1 Dropout
The dropout technique (Srivastava et al. 2014) consists of randomly setting proportion of neuron outputs to zero at each forward pass through the network during training. This random "dropping out" of neurons allows the network to sample and train on different. (but overlapping) architectures.
In standard supervised learning, dropout is then deactivated during inference, and dropout neurons are re-scaled account for the fact that all dropout neurons are active. However, in our case, we still apply dropout during the action selection (Gal & Ghahramani 2016 Riquelme al.) 2018). Using dropout in this phaso, we randomly seleet (according the dropout probability) a set of paramcters Ô use. This procedure can be seen as sampling parameters from an approximate posterior distribution: Ô P() and it has connections with variational inference (Gal & Ghahramani] 20166). With such approximate postcrior distribution, we can apply Thompson Sampling. This technique valuable for large models, such as our LLM bandit, because it does not add any overhead the procedure and does not require additional memory.

---

Published in Transactions on Machine Learning Research (07 (024))
4.2 Laplace Approximation
The Laplace Approcimation (LA) is technique that can be used to approximate the posterior distribution of the parameters of a neural model and does not require changing the training process or training multiple
models.
LA exploits the fact that, from a Bayesian point of view, minimizing a regularized loss can be seen as finding a marimum-a-Paa p-sttrrioo (MAP) estimate: 0MAP = arg ming L(D; 0).
Then, LA consists of replacing the loss with its second-order Taylor expansion around 0MAP:
1 [ 0MAP) H 0MAP), = (DD: (():0) x L(D; 2
(4)
Notice that the first-order term vanishes because we are expanding around 0MAP, which is point of minimum of the loss. After some algebraic manipulations, it can be shown that the posterior distribution can be expressed a
P(O = (OMMP, H
which means that, after training, the posterior distribution of the parameters is Gaussian distribution, centered on the parameters obtained with training (i.c., dMAP), and with the inverse Hcssian as covariance matrix. Therefore, to derive the approximate posterior in practice, we need first to identify the weights OMAP by training our LLM bandit regression model. Subsequently, the only additional step is the calculation of the Hessian matrix Н. This posterior distribution can be used to apply the Thompson Sampling policy by sampling set of parameters ~ P(O[D) at each time step. Notice, however, that there still significant drawbacks in using LA with LLM bandits: (1) computing the Hessian requires looping through the whole fine-tuning dataset, which is always increasing; (2) storing the Hessian is infeasible due to its size, which is quadratic with the number of parameters; (3) computing the Hessian may be computationally infeasible, and the Hessian may be indefinite. In the following, we show that there are many ways to cope with these limitations.
Recursive computation of the Hessian By applying Bayesian reasoning, we notice that the Hessian at given time step t can be computed recursively, exploiting the Hessian computed at t- Let us assume that we are at time step l. Hence, the dataset will look like this: composed of i.i.d. data points. The posterior distribution of the parameters can be re-written follows:
[Di...., Dr), where each D, is
P(OID) x P(D.)0) P(D2(0) P(D.10)-PU)
(6)
<P(()).)
x <((((D).Da)
where P(D.(0) = P(rtlat, o). This implies that: P(O D.) x P(D..)) Dz-1).
(7)
Hence, the Hessian of the negative pog-posttrior time which we call (1:1), can be re-written as:
(1:)) V2 - log P(De 1)
+
Σ 22 log P(Dr 0) ee
log P(O)
(8)
MAP
MAP
neg. ligeliool Hessian ((())
previous neg. log-likellthod Hessians H(0 neg. log-prior Hessian H
This means that, at the end of any time step f, we just need to compute the log likelihood Hessian H0 with respect the current data D, and sum the previous Hessian: H(1:)) H(()) ((a--1),
For more detailed derivation, Appendix

---

Published in Transactions on Machine Learning Research (07/2024)
Furthermore, this formulation gives rise to loss that uses only the current data when training at time t:
[(0(0); 20 (æ,a,r)eD;
Σ (r- 10 (z, (x)
(0 - MAP) (1()) ((t-1)
(9)
neg. log prior / updated posterior: -log ('00DD...........
neg. log likelibood: -log P8D:(8)
If we are time & == 1, assume that @ are the initial weights, and H(.:0) 2 log P(O)
MAP
Diagonal approximation Even if computed recursively, storing the full Hessian may be infeasible even for small neural worwsrks. Hence, it is not viable option for the case of LLM bandits. A practical approximation is to consider only the diagonal of the matrix, which is equivalent to maintain posterior distribution for each parameter independently. This reduces the memory complexity from quadratic to linear in the number of parameters.
Fisher Hessian approximation To further reduce the computational complexity, we can replace the true Hessian by the expected Fisher matrix. If we assume a model where the aleatoric noise is Gaussian with mean zero and fixed variance obs: we can compute this approximation of the Hcssian as follows:
diag(fi Jbss (z,a)ED,
Σ (Vio(z,a))
(10)
This approximation has the advantage of being positive semi-definite and (Kunstner al. 2019 showed that it is accurate when the regression residuals are small. For more details this approximation, see Appendix C
4.3 Last-Layer Laplace Approximation
Alternatively, instcad of using Hessian approximations, another way to scale LA for large models compute the full Hessian but only for subset of the parameters of the network, which typically is the last layer of the model. In our case, this means that we will compute the Hessian for the randomly initialized final layer, while the pre-trained LLM layers remain fixed during the sampling phase of TS (every parameter is still fine-tuned during training). This approximation can reduce the exploration due to the fixed weights, but at the same time it allows a full quadratic Hessian computation, leading to a better quality of exploration. Last layer LA is equivalent to Bayesian linear regression (Box Tiao 2011] on the features before the last layer within a single batch. Also, with Gaussian likelihood, last-layer LA is exact within a single batch.
4.4 Epinets
A different approach to estimating epistemic uncertainty is the epinel (Osband el al. 2023a is a heuristic approach that tries to estimate the epistemic uncertainty with a separate neural network. It consists of neural network added to a base network, which, in our case, is the LLM bandit model. The epinet takes as inputs both features (x)) derived from the base network and an epistemic inder 2, which is a random vector sampled from reference distribution Pz. In the case of our L.M bandit model, #)(x) denotes the feature vector extracted by the model before the final regression layer. The prediction then obtained by adding the predictions of the base network and the epinet:
An epinet
(x; fo(x) + epi, (sg[
where with sg denote the stop-gradient operator, and with 7 denote the additional parameters of the epinet.
3 We the expected Fisher the fact that different from the empirical Fisher matrix Kunstner
2019

---

Published in Transactions on Machine Learning Research (07/2024)
The epinet epi, comprises two parts: learnable network epity and a prior network epif that represents prior uncertainty. The prior network has no trainable parameters. This allows the epinet to adapt uncertainty estimates to observed data. In principle, every network that takes features and epistemic index as inputs could be an epinct. However, for our use case, we need to use a small epinet so that we do not add excessive computational overhead the LLM bandit model. In Section we describe in further detail the epinet architecture we selected for our experiments.
During training, the epinet model 90.n is trained as normal deep learning model, with the addition of sampling phase of an epistemic index Pz for each data point. Also during the action selection phase, for each action to select, an epistemic index is sampled. While this is not Bayesian approach, the sampling phase can be seen as an approximate Thompson Sampling, as shown by Osband et al. (202bb).
5 Experiments
In this section, we provide an extensive empirical study of L.LM-based contextual bandits on real-world data. We show how to adapt the epistemic uncertainty estimation techniques we described to bandits with pre-trained LLMs In particular, we identify the best variant of the Laplace approximation and dropout to use, in addition to a proposed Epinet architecture that we show works well for TS with LLM bandits. Finally, we empirically show that the greedy policy sub-optimal (in terms of final regret), shedding light on the importance of epistemic uncertainty for LLM bandits.
5.1 Experimental Methodology
Tasks We evaluate the bandit policies on various weal-world tasks. These tasks consist of:
An open-source dataset, called "measuring hate speech", which is openly available HuggingFace This dataset consists of around 136,000 comments. Each of them is associated with a
Dalasel c continuous score, called the "hate speech score", provided for cach comment. Comments with score 0.5 are considered "toxie" (around 36% of the comments are labeled as toxic, while the others are not toxic). We will refer to this task as toxic.
>
An open-source dataset, called 9 "MDE"" (Maas al. 2011), which is openly available on Hugging Dataset This dataset consists of 50,000 movie reviews. Each of them is associated with class, which can be "positive" or "negative", according to the sentiment expressed in the review. The dataset consists of exactly 50% "positive" reviews and 50% "negative" reviews. We will refer to this task as imdb.
An open-source dataset, called " Offfensive Language Identification" (Zampieri al. 2019). which is openly available on HuggingFace Datasets This dataset consists of over 14,000 English tweets. Each of them is classified "non-oflensive" or "offensive". The datasel consists of 66.9% "non-offensive" tweets and 33.1% "offensive" tweets. We will refer to this task as offensive.
An soun-s dataset, called "HatEval" Basile et al. 2019 which is openly available on Hug- gingFace Datasets This dataset consists of 13,000 English and Spanish tweets, where some of them contain hate speech against immigrants and women. Each of them is classified "non-hateful" "hateful". The dataset consists of 58% "non-hateful" tweets and 42% "hateful" tweets. We will refer to this task as hate.
These tasks are framed as a contextual bandit problem as follows. The context the text input, and the actions available to the agent are "not publish" or "publish". For each time step, the agent will observe batch of B 32 comments. The reward function is the following: if the agent decides not to publish

---

Published in Transactions on Machine Learning Research (07 /022)
0.350
Greedy
0.375
Greedy
Diag. LA TS
Diag. LA TS
Drapout TS
Dropout TS
0.300
Epinct 15
0.325
Epinet TS
Average regret
Average regret
0.300
0.250
0.275
0.200
0.225
0.175
0.200
0.150
1500
2500
3000
500
1000
2500
Observed data points
Observed data points
Figure 2: Average regret (+ std. err.) obtained on the toxic bandit task.
Figure 3: Average regret (#: std. err.) obtained on the imdb bandit task.
comment, reward of 0.5 is observed regardless of the actual toxicity of the comment. If the agent publishes non-toxic comment, reward of 1 is observed. If the agent publishes toxic comment, reward of -0.5 is observed. This asymmetric reward function represents a possible example suitable for this real-world scenario. The goal of the agent is to minimize the regret (compared to a clairvoyant optimal publishing policy) over time by learning to make optimal publish/not publish decisions based on the text content. For the imdb task, we consider "negative" reviews as toxic content; for the offensive task, we consider "offensive" tweets as toxic content; for the hate task, we consider "hateful" tweets toxic content.
Bandit models To investigate the role of uncertainty in LLM bandits, we compare different TS variants with the greedy baseline. Every bandit model is initialized as in Eq. [ with pre-trained GPT2 model with 124M parameters (Radford et al. 2019). We also include additional experiments with GPT2-XL, (1.5B) in order to investigate if our findings generalize to larger models.
As TS variants, we include dropout, Diagonal Fisher LA (which we will call Diag. LA), LA with full Hessian on the last layer (which we will call Last I.A), and Epinet TS. Regarding the epinet architecture, we follow prior work on epinets (Osband al.] 2023a b) and select an architecture with a multi-layer perceptron h which is multiplied with dot product with the epistemic index: epi, ()) h, (1))()(())) =)) For more details the architecture, except for Diag. LA, which updates the prior with the new posterior at cach time step, as shown in Eq. M All the parameters are updated during training. We train each model at the end of each time step for 50 epochs with the Adam optimizer (Kingma & Ba 2014), with learning rate set to 10-5, For each model, hyperparameters are tuned on 10 random runs (with different seeds than the testing ones) with = 50 on the toxic task. We did not tune the dropout probability because we wanted exploit the fact that GPT2 was pre-trained with dropout p 0.1 and the same p. We investigate the effectiveness of this choice in the results section. We describe the training phase and the hyperparameter tuning procedure in more detail
Appendix D Every model is trained with regularized MSE loss as in Eq. B.
in Appendix D
5.2 Results
We show the experimental results (20 random runs, T- 100) for the toxic task (Figure 2), imdb task (Figure B) offensive task (Figure 4) and hate task (Figure 5) with GPT2 (124M) LLM bandits. These results empirically exhibit the importance of actively using epistemic uncertainty in contextual bandit problems with Large Language Models. Overall, we notice how the approaches leveraging epistemic uncertainty tend to achieve lower regret compared to the greedy policy in all the considered tasks.
Analyzing the results in more detail, we observe that Dropout TS cxhibits strong performance on the toxic and hate tasks, achieving lower regret compared the greedy policy. However, its performance is comparable to the greedy approach the imdb and offensive tasks. This behavior is expected, Dropout TS and the
10

---

Published in Transactions on Machine Learning Research (07 2222))
Greedy
Greedy
0.34
Diag. LA T5
0.36
Diag. LA TS
0.32
Drapout TS
0.34
Dropout TS
Epinct TS
Epinet TS
0.32
Average regret
0.30
Average regret
0.30
0.26
0.24
0.24
0.22
0.22
500
1000
1500
2000
2500
3000
1000
1500
2000
2500
3000
Observed data points
Observed data points
Figure 4: Average regret (+ std. err.) obtained on the offensive bandit task.
Figure 5: Average regret (# std. err.) obtained on the hate bandit task.
greedy policy are very similar, with the only distinction being the activation of dropout during inference in Dropout TS. It is important to notice that we did not tune the dropout probability, using the same one used during the pre-training of the base LLM. The performance of Dropout TS may possibly improve after luning the dropout probability. In Appendix D we investigated the effect of changing the dropout probability, Also, Dropout TS incurs no additional computational cost, requiring only a single forward pass, akin to the greedy method.
Diag. LA TS demonstrates excellent results the toxic task and attains lower regret than the greedy policy on the offensive and hate tasks. However, its performance on the imdb task is comparable to the greedy approach. This may be attributed to the fact that hyperparameter selection was performed on the toxic task, and may not generalize well to the imdb task.
Both Last LA TS and Epinet TS consistently achieve lower regret compared to the greedy policy across all four tasks, albeit with some distinctions. Last LA TS emerges as the best-performing method in all cases, but it requires computing the full Hcssian the last layer, which can be computationally expensive. On the other hand, Epinet TS, while slightly inferior to Last LA TS, offers more computationally efficient alternative by employing an additional small neural network. These results confirm previous findings in the literature (Osband et al. 2022 2023b). Also, we conjecture that the choice of the epinet architecture may inffuence results. Indeed, epincts are a very broad class of ncural networks, and exploring epinct design was not the focus of this work. Hence, we did not perform an extensive architecture search, but we believe that our work can shed light on the potential of epinets, and we will leave further investigation future work.
Moreover, notice that the confidence interval of the greedy policy is generally large. This suggests that there is a high variance in the results obtained by the greedy policy. To further investigate the cause of this finding, let us define the action selection ratio as 8I (a) 7 E I(a, = a). We selected two sample runs for the toxic task and plotted the action selection ratio for the action "publish" during these two sample runs Figure g Notice that the action "publish" is a risky action, meaning that it can lead to a strongly negative reward toxic content published, while the action "not publish" safer and gives a constant reward. Figure [ clearly shows that there are cases where the greedy policy suffers from under-exploration and persists in choosing a suboptimal arm, while Thompson Sampling (in this case, Diag. LA) exhibits more balanced behavior. Due to this issue, the greedy algorithm will incur constant regret in the random runs which require more balanced exploration. This issue makes the greedy policy perform worse than TS policies in our experiments.
Experiments with 1.5B LLM To further investigate if our findings generalize to larger models, we performed additional experiments using GPT2-XL, larger language model with 1.5B parameters. For this set of experiments, we train our models for 5 epochs for each batch of data. We perform this new set of experiments the hate task. The results of these experiments (20 random runs, T 100) are shown
11

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Action selection ratio for action "publish"
Diag. LA TS
0.6
0.4
0.4
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Figure 6: Action selection ratio for the action "publish" for two particular sample runs (toxic task).
in Figure 7 These results confirm our findings, showing how TS policies generally achieve a lower regret compared to the greedy policy. In particular, we again see how Last LA TS achieved the lowest regret among the TS variants. These results further highlight the importance of leveraging epistemic uncertainty in contextual bandit problems, even when scaling up to larger language models.
Additional ablation study on model dimen- sion and pre-training We also performed an additional ablation study investigate the effect of using large pro-trained models to initialize the bandit policies in our experiments. We conduct
0.38
Greedy
Last LA TS
0.36
Dropaut TS
Epinet TS
Average regret
those studies by rumning bandit algorithm using a smaller transformer model with weights initialized from scratch. We report those experiments in Ap- pendix D.4
0.32
0.30
0.28
6 Conclusion
0
500
1000
1500
2000
2500
3000
In this paper, we investigated the role of epis- temic uncertainty estimation in decision-making tasks that use natural language input. For such tasks, using Large Language Models agents has become the norm. However, none of the recent ap- proaches estimates the epistemic uncertainty of the agent. We focused on a fundamental prcision-makinng task: the contextual bandit problem, where context consists of text. We approached the bandit task with deep regression model initialized with a pre-trained I.L.M As representative of the approaches with no uncertainty estimation, we considered an LLM bandit with greedy policy, which picks the action correspond- ing to the largest predicted reward. We compared the greedy baseline with various approaches integrating uncertainty estimates into the decision process via Thompson Sampling. We adapted several epistemic un- certainty estimation techniques to LLMs, such as dropout, Laplace Approximation, and epinets. Finally, we provided an empirical analysis of bandit. learning with LLMs on real-world data. Our experiments showed that using uncertainty information leads to greatly improved performance over the greedy approach. These improvements highlight the benefits of modeling uncertainty for exploration bandit problems with textt and Large Language Models. Our work suggests that uncortainty should play more central role in developing LLM--based agents for decision-making. As future work, natural extension would be to study
Observed data points
Figure 7: Average regret (+t std. err.) obtained on the hate bandit task with 1.5B LLM bandits.
12

---

Published in Transactions on Machine Learning Research (07/2024)
LLM bandits stepping stone to more complex settings like text- reinforcement learning. Another interesting direction would be to study the scaling behavior of LLM bandits in systematic way.
References
Shipra Agrawal and Navin Goyal. Near-optimal regret bounds for thompson sampling. Journal of the ACM (JACM), 64(5):1 24, 2017.
Vashist Avadhanula, Omar Abdul Baki, Hamsa Bastani, Osbert Bastani, Caner Goemen, Daniel Haimovich, Darren Hwang, Dima Karamshuk, Thomas Leeper, Jiayuan Ma, et al. Bandits for online calibration: An application to content moderation on social media platforms. arXiv preprint wwXi:2221.0661111 2022.
Valerio Basile, Cristina Bosco, Elisabetta Fersini, Debora Nozza, Viviana Patti, Francisco Manuel Rangel Pardo, Paolo Rosso, and Manuela Sanguinetti. SemEval-2019 task 5: Multilingual detection of hate speech against immigrants and women in Twitter. In Proceedings of the 13th International Workshop on Semantic Evaluation, pp. 54 63, Minneapolis, Minnesota, USA, 2019. Association for Computational Linguistics. doi: 10.18653/vv /v155//1//// /////112200 2007. URL https: http:::/wwww.cccc.c http://wwww......aaoooooooooomee /ttp://///oooooog///9-0000000
George EP Box and George C Tiao. Bayesian inference in statistical analysis. John Wiley & Sons, 2011.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877 1901, 2020.
Sébastien Bubeck, Varun Chandrasckaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, al. Sparks of artificial general intelligence: Early experiments with gpt-4, arXiv preprint arXi::2303. 12712, 2023.
Thomas Carta, Clément Romac, Thomas Wolf, Sylvain Lamprier, Olivier Sigaud, and Pierre- Oudeyer. Grounding large language models in interactive environments with online reinforcement. learning. Tn Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (cds.), International Conference Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawati, USA, volume 202 of Proceedings of Machine Learning Research, pp. 3676 3713. PMLR, 2023. URL htttp:://proceedingg mlr. press/v202/ carta23a. html
Olivier Chapelle and Lihong Li. An empirical evaluation of thompson sampling. Advances in newral infor- mation processing systems, 24, 2011.
Liting Chen, Lu Wang, Hang Dong, Yali Du, Jie Yan, Fangkai Yang, Shuang Li, Pu Zhao, Si Qin, Saravan Rajmohan, et al. Introspective tips: Large language model for in-context decision making, arXie preprint arXiv:2305115508 2023.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortless bayesian deep learning. Advances in Newral Information Processing Systems, 34:20089-20 20103,, 2021.
Vikranth Dwaracherla, Seyed Mohammad Asghari, Botao Hao, and Benjamin Van Roy. Elficient exploration for lims. arXiv preprint arXic:2402.0033 00396, 2024.
Yarin Gal. Uncertainty in deep learning. 2016.
Yarin Gal and Zoubin Ghahramani. Bayesian convolutional neural networks with bernoulli approximate variational inference. arXiv preprint arXit:1506.0210 02158, 2015.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pp. 1050-1055 1059. PMLR, 2016.
13

---

Published in Transactions on Machine Learning Research (07/2024)
Luyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon, Pengfei Liu, Yiming Yang, Jamie Callan, and Graham Neubig. Pal: Program-aided language models. In International Conference on Machine Learning, pp. 10764 10799. PMLR, 2023.
Gemini Team, Rohan Anil, Scbastian Borgeaud, Yonghui Wu, Yean-Baptiste Alayrac, Jiahui Yu, Radu Sori- cut, Johan Sch,lkwyk, Andrew M Dai, Auja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXtv:2812.1180: 2023.
Robert Gorwa, Reuben Binns, and Christian Katzenbach. Algorithmic content moderation: Technical and political challenges in the automation of platform governance. Big Data & Society, 7(1)20053551779997994, 2020.
Alex Graves. Practical variational inference for neural networks. In J. Shawe-Taylor, R. Zemel, P. Bartlett, F. Pereira, and K.Q. Weinberger (eds.), Advances in Neural Information Processing Systems, volume 24. Curran Associates, Inc., 2011. URL httts://ppooedagggs..r..c..ppppp http:///apap..elee//p//////////0
file//28bcbeadd111e0000000044423 Paper. pdf
Shibo Hao, Tianyang Liu, Zhen Wang, and Zhiting Hu. Toolkengpt: Augmenting frozen language models with massive tools via tool embeddings. arXiv preprint rrXi:22305.1155. 2023.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint (rXi::1606.08455, 2016.
Eyke Hüllermeier and Willem Waegeman. Aleatoric and epistemic uncertainty in machine learning: An introduction to concepts and methods. Machine Learning, 110:457-506, 2021.
Kirthevasan Kandasamy, Akshay Krishnamurthy, Jeff Schneider, and Barnabás P6czos. Parallelised bayesian optimisation via thompson sampling. In International Conference on Artificial Intelligence and Statistics, pp. 133-142. PMLR, 2018.
Alex Kendall and Yarin Gal. What uncertainties do need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
Diederik P Kingma and Jimmy Ba.
Adam: A method for stochastic optimization. ar preprinl
arXiv:1412 6980, 2014.
Martin Klissarov, Pierluca D'Oro, Shagun Sodhani, Roberta Raileanu, Pierre-Luc Bacon, Pascal Vincent, Amy Zhang, and Mikael Henaff. Motif: Intrinsic motivation from artificial intelligence feedback. arXiv preprint (rXi::2310. 001600 2023.
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Linitations of the empirical fisher approximation for natural gradient descent. Advances in neurol information processing systems, 32, 2019.
Balaji lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncer- tainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Shuang Li, Xavier Puig, Chris Paxton, Yilun Du, Clinton Wang, Linxi Fan, Tao Chen, De-An Huang, Ekin Akyürek, Anima Anandkumar, et al. Pre-trained language models for interactive decision-making, Advances in Neural Information Processing Systems, 35:31199 199-31212, 2022.
Renkai Ma and Yubo Kou. "how adverliser-friendll is my video?": Youtuber's socioeconomic interac- tions with algorithmic content moderation. Proceedings of the ACM on Coman--ommpttr Interaction, 5 (CSCW2):1 25, 2021.
Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang. Andrew Y Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annvall meeting of the association for computational linguistics: Humon longuage technologies, pp. 142-150, 2011.
David JC MacKay. Bayesian interpolation. Neural compulation, 4(3):415 447, 1992.
14

---

Published in Transactions on Machine Learning Research (07/2024)
Grégoire Mialon, Roberto Dessi, Maria Lomeli, Christoforos Nalmpantis, Ram Pasumuru, Roberta Raileanu, Baptiste Rozière, Timo Schick, Jane Dwivedi-Yu, Asli Celikyilmaz, et al. Augmented language models: a survey. arxiv preprint arXiv:2302.0784.. 2023.
OpenAL. Gpt-4 technical report, 2023.
Tan Osband, Seyed Mohammad Asghari, Benjamin Van Roy, Nat. McAleese, John Aslanides, and GeofTrey Irving. Fine-tuning language models via epistemic neural networks. arXiv preprint arXiv:2211 1..211111166 2022.
Ian Osband, Zheng Wen, Scyed Mohammad Asghari, Vikranth Dwaracherla, Morteza Ibrahimi, Xiuyuan Lu, and Benjamin Van Roy. Epistemic neural networks. Advances in Neurul Information Processing Systems, 36, 2023a.
Ian Osband, Zheng Wen, Seyed Mohammad Asghari, Vikranth Dwaracherla, Morteza Ibrahimi, Xiuyuan Lu, and Benjamin Van Roy. Approximate thompson sampling via epistemic neural networks. In Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intellignnce, UAI '23. JMLR.org, 2023b.
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkinn Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Frascr Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback, 2022.
Alec Radford, Karthik Narasimhan, Tim Salimans, Πya Sulskever, et al. Improving language understanding by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Tlya Sutskever, et al. Language models are unsupervised multitask learners. OpenA/ blog, 1(8):9, 2019.
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint ar arxiv:2112 12. 122.21144 11446, 2021.
Carlos Riquelme, George Tucker, and Jasper Snoek. Deep bayesian bandits showdown: An empirical com- parison of bayesian deep networks for thompson sampling. In International Conference Learning Representations, 2018.
Daniel Russo and Benjamin Van Roy. An iennaatiionthoortti analysis of thompson sampling. Journal of Machine Learning Research, 17(68):1 30, 2016.
Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, Zheng Wen, et al. A tutorial on thompson sampling. Foundations and Trendso in Machine Leurning, 11(1):1 96, 2018
Timo Schick, Jane Dwivedi- Yu, Roberto Dessi, Roberta Raileanu, Maria Lomeli, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools. arXiv preprint 2rxi::2302.0 04 761, 2023.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimiza- tion algorithms. arXiv preprint "xXix:107.003// 2017.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1): 1929-1958, 2014.
William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, (3(3-4))285-294, 1933.
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazoer, Apoorv Kulshreshtha, Heng- Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint or arXii22201.08223 2022.
15

---

Published in Transactions on Machine Learning Research (07/2024)
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Maric-Anee Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Fric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. orXiv preprint arXim:2302.1399 13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bash- lykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and finc-tuned chat models. Xiv preprinl arXiv:2307.09288 2023b.
Kristen Vaccaro, Christian Sandvig, and Karrie Karahalios. "at the end of the day facebook does what it wants" how users experience contesting algorithmic content moderation. Proceedings of the ACM human-computer interaction, 4(CSCW2))1-22, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances neural information processing systems, 30, 2017.
Chil-Chun Wang, Sanjeev R Kulkarni, and H Vincent Poor. Bandit problems with side observations. 1EEE Transactions on Automatic Control, 50(3):338 355, 2005.
LI Xuhong, Yves Grandvalet, and Franck Davoine. Explicit inductive bias for transfer learning with convo- lutional networks. In International Conference on Machine Learning, pp. 2825-2833 PMLR, 2018.
Rui Yang, Lin Song, Yanwei Li, Sijie Zhao, Yixiao Ge, Xiu Li, and Ying Shan. Gpt4tools: Teaching large language model tools via self-nsstruction. arXiv preprint 1187:22505.18752, 2023a.
Sherry Yang, Ofir Nachum, Yilun Du, Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foundation models for decision making: Problems, methods, and opportunities. orXiv preprint arXiv:2303.04 04129, 2023b.
Marcos Zampieri, Shervin Malmasi, Preslav Nakov, Sara Rosenthal, Noura Farra, and Ritesh Kumar. Semeval-2019 task 6: Identifying and categorizing offensive language in social media (offenseval). In Proceedings of the 13th International Workshop on Semantic Evaluation, pp. 75-86, 2019.
16

---

Published in Transactions on Machine Learning Research (07/2024)
A Prior Work
Decision-making with Large Language Models Large language models (LLMs) recently emerged as dominant paradigm in natural language processing (Ouyang et al.
2022
OpenAl 2023), achieving state-
of-the-art performance across a wide range of tasks (Rae et al. 2021
pushing model scale and dataset size
to unprecedented levels. Models such as the OpenAl's GPT LLM series (Radford et al., 2018 2019 Brown
2020 OpenAl 2023). Google's PaLM (Chowdhery et al. 2023) and Gemini (Gemini Team et al. 2023). , or Meta's LLaMA (Touvron et al.
et al.
2023a have leveraged the transformer architecture (Vaswani iet al. 2017) with model sizes ranging from hundreds of millions to hundreds of billions of parameters and are trained on up to hundreds of billions of text examples.
Due to their remarkable capabilities in text processing; LLMs have also been applied to decision-making tasks Yaan et al.) 2023b), and there is a plethora of papers in the research literature that investigated this idea (Li et al. 2022 Carla et al. 2023 Chen et al. 2023: Klissarov et al. 2023; inter alia). One of the most prominent examples is that of dialogue agents, Many recent papers model the dialogue between the LLM and a user as a sequential decision-making problem, where the action is the answer that the LLM should provide to the user after receiving the user's message. In particular, those works typically use Reinforcement Learning (RL) techniques (Ouyang et al., 2022) to fine-tune language models for dialogue applications. Hence, they use the LLM as a policy to solve the RL problem with state-of- the-r-rttthe-tttto RL algorithms, such as Prozimal Policy Optimization (Schulman et al.
2017). Another example is enhancing LLMs by allowing them to use Yang et al. 2023a. Hao et al. 2023: Gao et al., 2023: Mialon et al. 2023 Schick et al. 2023). In this case, the LLM-based agent has an action space which is the set of external tools at disposal and the various interactions possible with those tools.
external lools (Thoppilan et al. 2022
The aforementioned approaches typically assume that a sufficient amount of data has been collected priori to train an effective policy. They do not explicitly address the exploration-exploitaaioo trade-off and they do not syslematically explore issues around the role of epistemic uncertainty in decision-making. In contrast, in our work, we explicitly consider the problem of learning from interaction, and we focus one of the most fundamental and natural decision-making tasks, which is the one of contertual bandits (Chapelle & Li 2011). We comprehensively investigate the role of epistemic uncertainty for bandit models with traned LLMs. The importance of exploiting epistemic uncertainty is emphasized also in (Dwaracherla, et al.) 2024). where they propose an approach to approximate epistemic uncertainty for the alignment problem of LLMs.
Uncertainty in Deep Learning Deep learning models provide statc-of-tho-art performance in several different tasks, ranging from image recognition to natural language processing. However, these models usually provide poor uncertainty estimates (Kendall & Gal, (017). For this reason, several techniques have been proposed in the litcrature that allow the estimation of the uncertainty of deop learning algorithms. A useful mathernttica model for characterizing uncertainty in deep learning is to calegorize it into two types: epistemic and aleatoric (Gall 2016 Hüllermeier & Waegeman 2202). p. Epistemic uncertainty stems from our lack of knowledge about the best model to describe a process. It is reducible as more data or knowledge is gathered. Alcatoric uncertainty, in contrast, is due to the inherent randomness in the data environment and is irreducible even with more data. In particular, epistemic uncertainly has been proven to be particularly helpful when it comes to decision-makinn problems.
A possible way to capture epistemic uncertainty in deep learning is to equip neural networks with a distri- bution over the parameters, which is then updated upon seeing new data. This kind of neural network is usually called Boyesian Neural Network (BNN). From a thcoretical perspective for decision making under uncertainty, Bayesian neural networks (BNNs) are appealing since they provide a full posterior distribu- tion over models, allowing the derivation of formal regret bounds to guide exploration (Agrawal & Goyal, 2017). However, exact BNN inference is intractable for large models like LLMs. Thus, we must rely on approximations.
One common approach to approximate the posterior distribution is Dropout (also called Monte Carlo Dropout) (Gal & Ghahramani 2016). Dropout is a technique initially proposed for standard supervisod learning (Srivastava et al.]|211, which consists of randomly setting a proportion p of neuron outputs to zero at each forward pass through the network during training. In standard supervised learning, dropout
17

---

Published in Transactions on Machine Learning Research (07/2024)
is then deactivated during inference, and all dropout neurons re-scaled to account for the fact that all dropout neurons are active. However, we can still apply dropout during inference. In this way, we randomly select (aceording to the dropout probability) a set of parameters to use. This procedure can be seen approximately as sampling paramcters from a posterior Bayesian distribution, and it is proven to have links with variational inference techniques (Gal & Ghahramani 2015).
Another popular technique is Variational Inference (VI) (Graves 2011). VI a technique in machine learning used for approximating complex posterior distributions in Bayesian inference. The goal is to find simpler, parameterized distribution (the variational distribution) that is close to the true posterior distribution of interest by solving an optimization problem. However, it requires changes in the training procedure, which is not possible if we want to use pre-trained models.
Laplace Approximation (LA) (MacKay 1992) is another technique that is used to approximate the posterior distribution of the parameters of a neural network. It consists of assuming that the training consists of a maximun-a-posterion estimate of the parameters and replacing the loss with its second-order Taylor expan- sion around the MAP estimate. With these two steps, an analytical solution for the posterior distribution of the weights can be found. In particular, the weights are distributed as multivariate Gaussian with the MAP estimate as the mean and the inverse of the Hessian as the covariance matrix. The key advantage of LA is that the MAP estimate of the weights is usually available after standard deep learning training. A more problematic issue is the computation of the Hessian. However, a recent paper (Daxberger et al. 2021 investigates different techniques to approximate the Hessian, making the computation feasible even for modern neural networks.
There are also non-Baycsian approaches to estimating epistemic uncertainty. Deep Ensembles (Lkkshmi- parayanan et al. 2017) provide conceptually simple way to capture uncertainty, but are expensive and do not yield a well-defined posterior. The authors show that the degree of disagreement among the NNs within the ensemble is indicative of the epistemic uncertainty of the ensemble. While promising, this approach is not suitable for LLMs, for which an ensemble would be prohibitively expensive for both training and inference.
A more affordable approach is the use of epinets (Osband et al., 2023a). An epinet estimates the epistemic uncertainty with separate neural network that takes as inputs both features derived from the base network (usually before the last layer) and an epistemic indez, which is random vector sampled from fixed reference distribution. The prediction is then obtained by adding the predictions of the base network and the epinet. Though not Bayesian, epinets has been shown to provide useful uncertainty estimates for guiding exploration combined with a Thompson Sampling policy, with much lower overhead than ensembles (Osband et al. 2023b
In our paper, we focus on scalable approaches estimating epistemic uncertainty, and we adapt them to the case of pre-trained LLM: used the pre-trained weights as prior, exploited the dropout probability used in the pre-training phase, and used different Hessian approximations. Finally, we empirically show the importance of using epistemic uncertainty by embedding it into Thompson Sampling policies (Thompson) 1933). This significantly outperforms a greedy policy that does not account for uncertainty in decision- making.
B Additional details on the used loss
In this section, we elaborate in more detail the loss described in Eq. 3
Using the loss shown in Eq. we are doing MAP inference many times, applying the Bayes rule anew for each new batch of data. In this way, we perform MAP, which is a well-known established principle in deep learning, while avoiding looping throughout the whole dataset. Indeed, at a given time step 6, the dataset will look like this: D = ( Dr), where each D; is composed of i... data points. The posterior distribution of the parameters can be re-written as follows:
P(ODD P(D.)0).... .... P(D210) P(D,)0) (P()).
(11)
xx P((())
18

---

Published in Transactions on Machine Learning Research (07 2024)
where P(D.(0) II P(rilal, 0). This implics that:
Pi. D,) x P(D.)))((0DD............ De-1))
(12)
This justifies the loss in Eq. 8 as a MAP method: is equivalent to minimizing the negative logarithm of Eq.
12
C
Additional details on Laplace Approximation
In this section, describe in more detail the Laplace Approximation technique.
Laplace Approximation The Laplace Approximation (LA) technique can be used to approximate the posterior distribution of weights of a neural network. As we will see, Laplace Approximation does not require any change in the training process nor training multiple models; thus, is feasible for modern deep learning neural networks (such as LLMs) that are typically pre-trained (hence, their training procedure can not modified) and expensive to train (hence, training multiple models unfeasible).
Typically, deep neural networks trained via the minimization of loss. Let say that dataset D = (2,, 2:1 We call the neural network want to train fo. which is parametric function.
given
The parameters of the neural network are E RR, Standard deep learning losses can usually be decomposed in regularization term and sum of empirical loss terms on single data points:
((D:6)=rlee + Σ [x 9::0)
(13)
From a Bayesian point of view, it is straightforward to interpret the regularization term the negative log
prior: r(0) log P(O), and the sum of empirical losses the negative log-likelihoods: Σ (Xx; y4i 0) -
loog P(D - -b log P(y))f((xz))- This means that minimizing such a loss is actually leading a Muacimumn-a-posterioon (MAP) estimate: OMAP arg ming L(D; From these considerations, can also re-write the posterior distribution as follows:
P(O = P(DIO)P(0) = exp(- z P(D(0)P(0)d0)
(14)
Z
Now, Laplace Approximation consists in replacing the loss with its second-order Taylor expansion around OMAP:
1 _ 0MAR) H(O OMAP), H = ()
C(D;0) ~ L(D;OMAP) +
(15)
Notice that the first-order term vanishes because are expanding around 0MAP, which is point of minimum of the loss.
Starting from this approximation, it can be shown that the posterior distribution is a multivariate Gaussian. First, we obtain a flosed-form solution for the normalizing constant Z:
= P(D\0) = ((D)(c(D::))
/
200 -0MAP) H - (MAr))dd
1 _ @MAP) TH() BMAR))
-xp(--(DD
exp(
(2TT) 7
exp(- (()(-L(D;AAA) (P)(-A()))MAA))
(det H)
19

---

Published in Transactions on Machine Learning Research (07/2024)
where the last equality derives from the multivariate normal density.
Now, if we come back to the posterior distribution, we obtain:
1 P(O =
Z P(D)O))(() =
Z exp( exp(-I(D:0)
(det H)
1 expl- (0 HMAP) T 0MAP))
(16)
(2T) $
= -N(00;MAR -
Therefore, to derive the approximate posterior in practice, we need first identify the weights OMAP that maximize the log- oog-oosterior function. This, in deep learning terms, corresponds to the training phase, where regularized loss is minimized. Subsequently, the only additional step is the calculation of the Hessian matrix at the point IAPP. This means that LA can also be applied to pre-trained neural networks with no need to change the training procedure. This is a fundamental requirement. we want to apply this procedure to LLMs.
Expected Fisher Matrix The Laplace approximation as outlined above requires the knowledge of the Hessian of the loss function. Typically, a further approximation is used. First of all, let us rewrite the Hessian as the sum the Hessian of the likelihood and the Hessian of the prior: Hi + Hp. Now, the true Hessian of the likelihood replaced by the expected Fishe matrix.
H1 R D Eg w) [-V log Po(y]=)]
(17)
Despite the fact that the Hessian on the left hand side depends on the true regression targets, while the right hand side does not, the approximation is accurate when the regression residuals are small (Kunstner et al. 2019). In order to make the formula in equation 17 practical, replace the expectation with respect to the datapoints using Monte- Carlo estimate.
[D] E, PywPool))))) [- 2 log Pa(y)=))
~ Σ
Polulo)) [-V log Po(ulæ:)
(18)
FED
We can further use the identity (Kunstner et al. 2019
Σ Ey-ro(ule) [-v log Po(y)=) = RDD
22
E (g)P)( V log log
(19)
ZED
to obtain a formula that does not require us to compute second-order derivatives. This highlights one benefit of the expected Fisher matrix: it is posilive-semii-deliiiee by construction.
Diagonal Approximation While equation 19 is possible to evaluate in principle, the resulting size of the Hessian matrix still way too large store for even small-scale LLMs. We therefore make another approximation Hi, computing only the diagonal entries in equation 19
diag(fii) Σ Eyy-Pollle (V log (oyl))))
(20)
FED
The diagonal approximation corresponds to taking the Taylor expansion behind the Laplace approximation for each coordinate separately, which leads to posterior being normal distribution with diagonal co- variance. However, this distribution is not necessarily the best KT-projection of the Gaussian arising from equation 19 on the space of diagonal Gaussians
We the term expected Fisher to stress the fact that is different from the empirical Fisher matrix
Kunstner et al.
2019
Such a projection could be obtained by inverting the formula in equation and taking diagonal, which intractable.
20

---

Published in Transactions on Machine Learning Research (07/2024)
Analytic Solution for Gaussian Alcatoric Noise If we assume model where the aleatoric noise is Gaussian with mean zero and fixed variance Laber can compute the expectation in equation 20 analytically. For the sake of clarity, let us consider the one-dimensional case. The parametric likelihood (o(y)) is defined as
Po(ylæ) V 2TT
cxp (- 20
(u - fo(x))²
Hence, the gradient of the log-likelihood is:
log Po(yl=)
(v
(oo()))
obs
The expected value of the squared gradient of the log- log-likelihood (g-l)( log Po(ol))))
[V]x(z)))
1 -v- fo(x)))
(- 1 (v (V- dy
exp
(-
o0bs
obs
v /2T0 abs
20 obs
(21)
Now let us focus on the quantity 1:
1
(y fo(=))²) dy
(y 80-fal fa(x))
(y fa(x)) (xxx)))xxp exp
V 2TO
2002
obs
ohs
(y - [o(x)) exp
(y So(x)))
V 22TO2 obs
20
(22)
obs
(y - fo(x))²) dy 20
1
+
exp
2TTo obs
obs
=0 + +1=
Hence, we obtained an analytical solution for the quantity of interest:
E Po [VVlog Po(ul=))
(Vio(x))²
(23)
D Additional experimental details
D.1 Experimental setting
D.1.1 Computational resources
Our experiments were conducted on one NVIDIA A100 GPU with 80GBs of VRAM. One random run with GPT2 (124M) bandit, T 100, took about 30 minutes for the fastest bandits (i.c., greedy and dropout) and up 45 minutes for the bandits that had to compute the Hessian (i.e., Diag. LA and Last LA). We ran different bandit models for 20 random runs during the testing phase, and repeated those runs for all four tasks. In total, the experiments with GPT2 bandits took around 220 hours of computation. The experiments done with GPT2-XL (1.5B) bandits, T = 100, 20 random runs, 5 different bandit models, hate task, took around 50 hours of computation.
D.1.2 Tasks and Data processing
We evaluate the bandit policies on various real-world tasks. These tasks consist of:
An open-source datasct, called "measuring hate speech", which is openly available on HuggingFace
Dalasel TO This datasel consists of around 136,000 comments. Each of them is associated with
htttp:://hugginggfccccooo/ http::///...atttee....///////tta hate-seeehhteeeeeehehhh
21

---

Published in Transactions on Machine Learning Research (07/2024)
continuous score, called the "hate speech score", provided for each comment. Comments with score > 0.5 are considered "toxie" (around 36% of the comments are labeled as toxic, while the others are not toxic). We will refer to this task as toxic.
An open-source dataset, called "IMDE" (Maas et al. 2011), which is openly available on HuggingFace Datasel This dataset consists of 50,000 movie reviews. Each of them is associated with a class, which can be "positive" "negative", according to the sentiment expressed in the review. The dataset consists of exactly 50% "positive" reviews and 50% "negative" reviews. We will refer to this task as imdb.
An open-source dataset, called "Offensive Language Identification" Zampieri et al. 2019 which openly available on HuggingFuce Dulasels 12 This dataset consists of over 14,000 English tweets. Each of them is classified as "non-offensive" " or "offensive". The dataset consists of 66.9% "non- offensive" tweets and 33.1% "offensive" tweets. We will refer to this task as offensive.
An open-source dataset, called " "HatEval" (Basile et al. 2019), which is openly available on Hugging- Face Datasets | This dataset consists of 13,000 English and Spanish tweets, where some of them contain hate speech against immigrants and women. Each them is classified as "non-hateful" or "hatcful". The dataset consists of 58% "non-hateful" tweets and 42% "hateful" tweets. We will refer to this task as hate.
For each task, the text is Lokenized with the GPT2 Tokenizer. In particular, we rely on the implementation provided by HaggingFace For each task, we selected a maximum number of tokens, and we applied left padding for the comments with lengths less than the selected number of tokens. For the imdb task, we selected 256 as the maximum number of tokens; for the toxic task, we selected 128 as the maximum number of tokens; for the hate and of offensive tasks, we selected 64 the maximum number of tokens. The difference is due to the difference in the length of each data point of those datasets.
These tasks are framed as a contextual bandit problem as follows. The context is the text input, and the actions available to the agent are "not publish" or "publish". For each time step, the agent will observe a batch of B = 32 comments. The reward function the following: if the agent decides not to publish a commentt, a reward of 0.5 is observed regardless of the actual loxicity the comment. If the agent publishes a non-toxic comment, a reward of 1 is observed. If the agent publishes toxic comment, a reward of -0.5 is observed. This asymmetric reward function represents a possible example suitable for this real-world scenario. The goal of the agent is to mininize the regret (compared to clairvoyant optimal publishing policy) over time by learning to make optimal publish/no /notishh/ottti publish decisions based on the text content. For the imdb task, we consider "negative" reviews as toxic content; for the offensive task, we consider "offensive" tweets as toxic content; for the hate task, we consider "hateful" tweets as toxic content.
D.1.3 Bandit models
To investigate the role of uncertainty in LLM bandits, compare different TS variants with the greedy baselinn. Every bandit model is initialized as in Eq. [ with a pre-trained GPT2 model (Radford et al. 2019) (either the base one with 124M parameters, or the XL version with 1.5B parameters). We use the implementation provided by the HuggingFace library To GPT2 is a Causal Language Model, which means that the attention of cach token can only look to the current token and the previous oncs. For our purposes, we remove the final classification head, and we take the embedded features of the last loken. This is because the last token can look to all the tokens in the sentence with the attention mechanism. In the GPT2 model, this final feature vector has dimension of 768. We then add final lincar layer with 2 output neurons, which will be trained to solve our regression task. All the parameters are updated during training. We train each model at the end of each time step (for 50 epochs when we use GPT2, for 5 epochs when we
https://hugringtten httpp///aw....ee/aaaaaaaasa/taaa
12 http:///wwaaaaaaeeco////atttttts
GPT2Tokenizer
22

---

Published in Transactions on Machine Learning Research (07 (27/))
Algorithm 3 Dropout TS
Require: Bandit model fe, regularization scale A.
Og t- 10rr. 0|
2: Initialize D 0
3: for time t = .,...,
T do
4: Observe context 2/,....;22
5: for b = b do
Apply dropout to parameters: 4- dropout(0). Select d arg max, fg(z); Observe reward rg. end for
10:
Create D, [(x) (( rt) (xt a r
11:
Add to the observed dataset D + DUD,
12:
Train the model fo with Adam optimizer for 50 epochs minimizing:
B 1(000;;)) = Soo + (l10(0 - 0(1-1)112
b=1
13: end for
use GPT2-XL) with the Adam optimizer [Knnmaa & Ba, 2014), with learning rate set to 3 10-5. As TS variants, we include dropout, Diagonal Fisher LA (which we will call Diag. LA), LA with full Hessian on the last layer (which we will call Last LA), and Epinet TS. In the following paragraphs, we describe in further detail the different TS bandit models.
Dropout TS The dropout TS differs from standard supervised learning because we still apply dropout during the action selection phase (Gal & Ghahramani 2016 Riquelme et al. 2018). Using dropout in this phase, randomly select (according to the dropout probability) set of parameters 0 to use. This procedure can be seen approximately as sampling parameters from posterior distribution: Ó P(ODD). With such approximate posterior distribution, observe the rewards for the selected action, and we update all the parameters of the model by minimizing a MSE loss. A detailed explanation of the Dropout TS is provided in Algorithm
apply Thompson Sampling. After the action selection phase, we
Diag. LA TS With the Diag. LA TS, we have to maintain a diagonal Hessian for all the weights, At the beginning, the Hessian initialized as a diagonal matrix with all the entries equal to the inverse of the prior variance of Then, action selection is performed by sampling every time different of weights. Finally, a maximun-a-posterrort loss minimized. Once we finished the training phase, we oblain
of weights which call MAP. We update Hessian using the recursive formula and the expected Fisher approximation. With these quantities, we update the posterior distribution P() as P(0 DD = ~( H-1). Then, we can re-start observing context and selecting actions, repeating the loop. An algorithmic description of this bandit procedure is provided in Algorithm
Last LA TS With the Last LA TS, we have to maintain full Hessian only for the last layer parameters. At the beginning, the Hessian initialized a diagonal matrix with all the entries equal to the inverse of the prior variance of and prior weights equal to zero. Then, action selection is performed by sampling every time different set of weights for the last layer, while all the other parameters stay fixed. Finally, regularized MSE loss is minimized. Once we finished the training phasc, we obtain set of weights which we call 0MA for the last layer. We update our Hessian using the recursive formula and computing the full Hessian on the likelihood part the loss. With these quantities, we update the posterior distribution on the last layer parameters P(O)D) as P(@ = ((0MAP, H-1). Then, we can re-start observing context and selecting actions, repeating the loop. An algorithmic description of this bandit procedure is provided in Al Algorithm 5
23

---

Published in Transactions on Machine Learning Research (07/2024)
Algorithm 4 Diag. LA TS
Require: Bandit model fo, prior variance of observation variance dobs
1: Op 6- 0PTT 0]
H H(1:0) diag (1/o2)
P(O) (0 () (0,,,, 11-1)
Initialize D 4
5: for time 1 = = , do
6:
Observe context ....... Zt
7:
for b= b......, B do Sample parameters: 0~ P)) Select a arg max, fg(2),a).
8:
10:
Observe reward rg. 11: end for
12:
Create D. [(12, at, (xl ag rr
13:
Add to the observed dataset D + -DDD.
14:
Train the model so with Adain optimizer for 50 epochs minimizing:
0MAP= + arg
arg min C(0()
E (r- (x,) + 1 2000) (
(1:11) -((AAP)
20
(x,a,r)ED,
Compute current Hessian of the likelihood: diag diag(ff()) (=(1)== o² (tE,))D, (Vfo(z,a))
15:
Update Hessian ((1:t) +- (1(t) (1=-11)
16:
Update posterior distribution P(O)D) + N(MMAP, H-)), where - H((1)
17:
18:
end for
Model
Hyperparameter
Range
Selected
Greedy
Regularization factor X
0.1, 0.5,
Prior variance o Obs. variance dobs
0.0001, 0.0005, 0.001, 0.005, 0.01
0.0001
Diag. LA TS
0.0001, 0.0005, 0.001, 0.005, 0.01
0.01
Prior variance o
0.0001, 0.0005, 0.001, 0.005, 0.01
0.01
Last LA TS
Obs. variance o
0.0001, 0.0005, 0.001, 0.005, 0.01
0.01
Dropout
Regularization factor X
0.1, 0.5,
Epinet
Regularization factor X
0.1, 0.5,
Table 1: List of the tuned hyperparameters
Epinet TS Regarding the epinet archilecture, we follow prior work on epinets (Osband el al. 2023a b and select architecture with multi-layer perceptron h which is multiplied with dot product with the =))"2. 2. In particular, with GELU activation function (Hendrycks & Gimpel 2016) (the same as GPT2). The last layer of the epinet. is a linear layer with a two-dimensional output of dimension 32 x 2. While we conjecture that the choice the epinet architecture may influence results, exploring epinet design was not the focus of this work. Therefore, we used this architocture, which is inspired by prior work (Osband al. (023b)) Epincts are a very broad class of neural networks. We believe that our work can shed light on the potential of epinets and leave further investigation to future work.
epistemic index: epi,(x; z)
insert hidden layer with 256 neurons
D.2 Hyperparameter tuning
For each GPT2 (124M) bandit model, hyperparameters are tuned on 10 random runs (with different seeds than the testing ones) with T 50 on the toxic task. The set of hyperparanotors are shown Table In particular, for all the hyperparameter configurations, average, measured by cumulative regret, except for the TAA TS models. For the LA TS models, selected
selected the ones that best performed on
24

---

Published in Transactions on Machine Learning Research (07/2024)
Algorithm 5 Last LA TS
Require:
Bandit model fo, prior variance of observation variance dobs'
1: Op 6 1 0]
Initialize the Hessian for the last layer H(1:0) diag(1/o?)
Initialize the prior distribution for the last layer P(0) P(010) N(0,111-
Initialize D + /
T do 5: for time 1 .,...,
6: Observe context z) r/,.....z?
7:
for b 1 1,....,, B do
8: Sample last layer parameters: 0 P())D) (the remaining parameters stay fixed).
Select at arg maxa fa(s); a).
10:
Observe reward r/. 11: end for
12:
Create D. [(x2 at, (xl ag r1
13: Add to the observed dataset D + D-DUL DD.
14:
Train the model fo with Adain optimizer for 50 epochs minimizing:
B
OMAP +- arg min ; Dj) - Σ (2 di)) + (l11( 000 10
b=1
where À /o. Compute the full Hessian of the likelihood (only for the last layer parameters):
15:
A 2
1
Σ
(r- (o(
2012
obs (2,a,r)ED,
Update last layer Hessian (1::t) +- H1(0) + H(I:t-1)
16:
Update posterior distribution P( + N(OMAR, H-1,, where = ((1a)
17:
18: end for
the best configuration and the ones that are not statistically different from the best (measured with t- lest, p=value=0.05). Among those configurations, we selected the ones with the highest, values for prior and observation variances, in order to induce exploration.
D.3 Additional ablation experiment dropout
Among the Thompson Sampling techniques, the Dropout method delivers strong results even without tuning the dropout probability. By relying on the same dropout rate used during pre-training, same uncertainty that the original model had in learning to generate natural language. Therefore, it is not necessarily the best dropout probability to use in a bandit task. To investigate the importance of this hyperparameter in our experiments, we show dropout TS policics across various dropout probabilities in Figure S on the toxic bandit lask. From these results, the rate originally used for pre-training appears optimal for the bandit task. Also, we notice that, for a smaller value of p (p 0.05), the reduced exploration induces larger variance, as expected.
using the
D.4 Additional experiment on model dimension and pre-training
In this experiment, we want to investigate the effect of using a large pre-trained model to initialize the bandit policies in contextual bandits with text as input. To this end, we evaluate different greedy policies on the same four tasks presented in Section 5) To investigate the role of the dimension of the model, we include smaller transformer with a GPT2-like architecture, with attention blocks, embedding dimension of 128, and 4 heads, for total of around 7M parameters. The training procedure analogous to the one of the
25

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Dropout TS p= 0.05
Dropout -5.1-0.1
Average regret
0.275
0.250
0.225
0200
0.175
200
600
800
1000
1400
1600
Observed data points
Figure 8: Average regret obtained with different dropout values.
0.350
Greedy
Greedy Small Transformer
0.32
rearet
Average regrit
0.275
0.30
0.250
0.28
0.225
0.26
0.200
0.24
Greedy
Greedy Small Transtormer
0.22
0.175
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Observed data points
Observed data points
Figure 9: Average regret (+ std. err.) obtained the toxic bandit task.
Figure 10: Average regret (+ std. err.) obtained the imdb bandit task.
Greedy Greedy Small Transformer
0.32
0.34
0.30
Average rigret
0.30
0.26
0.28
0.26
Greedy
Greedy
0.24
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Observed data points
Observed data points
Figure 11: Average regret (#: std. err.) obtained the offensive bandit lask.
Figure 12: Average regret (+ std. err.) obtained the hate bandit task.
26

---

Published in Transactions on Machine Learning Research (07/2024)
Algorithm 6 Epinet TS
Require: Epinet bandit model ge.n, regularization scale X, reference distribution Pz.
0g < 10rr.00 0|
2: Initialize D 0
3: for time t = ,,..., do T
4: Observe context .......
5: for b = ,, do
Sample epistemic index 2 Pz Select di arg maxa I0.n (2t a) Observe reward r.. end for
10:
Create D, (t)),, (xt a r Add to the observed dataset D + DD.
11:
12:
Train the model 90.n with Adam optimizer for 50 epochs minimizing:
()) = > (r) b=1
gor no (0
a =²)) + 2110 - 0/1-11113
where 2 and Pz. 13: end for
These results show how using a pre-trained LLM is essential for solving contextual bandit tasks with text as input. Indeed, the policies that use small transformer always get stuck with suboptimal actions without learning the real bandit task. By using pre-trained LLMs instead, we can leverage model that can understand human language (since it was pre-trained on natural language generation) and focus on learning the task-specific reward function using expensive bandit data. This transfer learning approach improves the sample-effciency of our bandit.
pro-trained greedy policy. The results are shown in Figure 9 10
27


Reviews:
[{"review_id": "rZnFOJPTLk", "reviewer": "Reviewer_aHHM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "S: The overall idea is sensible\n\nW: The presentation is confusing, and I struggled to follow the paper. The research claims, setting and goals could be clarified. \n\nW: The experiments are very limited.\n\nW: The paper seems to just do standard bandit approaches on text tasks. The novelty and significance is limited.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes doing contextualised bandits on top of LLM features.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "My main issues are:\ni) I'm not sure this is a bandit problem in the first place. The paper needs to include a classification baseline, and demonstrate why it is useful to tackle this problem via bandits. \nii) The presentation needs dramatic improvement in terms of the research goals and setting. \niii) The research claims do need more empirical evidence. Another real-world task should be added. The paper needs to demonstrate why the epistemic/aleatoric stuff is \"fundamental\". \n\n\n- In sec 3.0 the connection between LLM and bandits is not concretised sufficiently. I’m not sure what is x (current token or all previous, or something else), a (??), r (??) in this context. I’m not sure how the LLM’s are used, or if they are (Are we still doing token-by-token generation? Since we discarded the softmax, I guess not..? But what are we then doing?) I’m not sure what problem are we solving; what is happening; or why stuff that happens happens. The paper would benefit a lot by having a motivational example to concretise the notation. The paper claims content moderation as example, but does not show how it connects with the methods. One would expect moderation to be a simple binary classification problem, and I don’t see why we need the bandits for that. I don’t understand what is “greedy baseline” for that either. If you want uncertainty, then surely simplest is to do BNNs to predict from the LLM features to the moderation outcome. Can you help me understand?\n- In Sec 3.1. it is implied that the LLM weights are also being optimised during the bandit system. Can you comment on the relationship between updating the LLM and updating the bandit reward estimator f: it seems that you are doing two very different things at once, with massive difference in complexity or size. Are you really finetuning the LLM model as you go?\n- In sec 3.2. it is implied that you want to do posterior sampling. Again, do you really sample the LLM parameter posterior? This seems unrealistic: can you comment on what kind of estimate are you expecting to get? How can you sample a posterior of billions of parameters?\n- In sec 3.2. there is a discussion of exploration and exploitation. I don’t see how this applies to the example of content moderation. How can you “explore” or “exploit” in this setting? What do you “explore”? I think you are just given a piece of text, and you need to approve or disapprove it. This is a simple f>0.5 classification problem and there is nothing to explore or exploit. Surely in bandits you need a settting where you have many actions to take, or you need to take many actions repeatedly to achieve some goal. I also don’t see how the uncertainty plays into this: I don’t think content moderation has much meaningful aleatoric uncertainty. I’m again quite confused of what is the research setting, and what is the research goal, and how do you use LLMs and how do you use bandits. Can you help me understand?\n- Sec 4: what is the epistemic uncertainty of the “model”? Uncertainty of what?\n- I don’t understand why you can’t do ensembles. Surely you could freeze the LLM and learn an ensemble of single-layer reward predictors on top of this, and train on whatever task data you have. Can you explain what are the bottlenecks assumed to be in this paper?\n- What does theta mean in sec 4? Do you apply dropout to both LLM and reward regressor, or only to regressor? The paper says that its both, but I find this very hard to believe. Surely if you have some billion-parameter network and you remove 10% of neurons without the network expecting it, things will get screwed up. Paper also claims it does not add any overhead. Surely you have to now run the dropout many times to get the distribution estimates, which can be very expensive for LLMs.\n- Do you really do the Laplace for the LLM parameters?\n- Can you overall discuss the finetuning and how expensive it is. Why do you want to finetune the LLM at all? How do you prevent the finetuning from degrading the LLM?\n- In experiment I don’t see why the content moderation problem is a bandit problem. It seems that you just have to classify text comments, and that’s it. Where is the bandit angle? In bandits you select from multiple choices iteratively to uncover which choices are good. In the text moderation problem there is no iteration: each comment is decided once and never returned to. Each comment is independent of other comments. There are no multiple choices. It’s not an active learning problem either, since you don’t control what is the next batch. I would then argue this is just a standard single-class classification problem under streaming data.\n- What’s the difference between greedy and non-greedy policy. Since this is a single-class problem, the greedy policy just goes for f>0.5 action. So does the other model then choose against the predictor..? This would make no sense. I don’t understand what are you demonstrating in the experiments, or what are you comparing. This is just a classification problem, so how/why does “last LA TS” classify better than “greed”?\n- It is implied that greedy policy is doing some kind of exploration in the experiments. What is this exploration? What do you explore? I think it’s just the yes/no action, but why do you need to explore this? Surely your loss will tell you that always deciding “reject” is a bad thing, and would increase the probability of “accept” without any exploration.", "broader_impact_concerns": "No issues.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "FPBTe9onHv", "reviewer": "Reviewer_5dit", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is interesting in that it applies uncertainty based decision making strategies to LLM applications, that of content moderation. This combination feels interesting and might be relevant to application driven audience.\n\nThe main weakness of the paper seems to be that the combination seems quite \"simplistic\" and does not highlight any specific nature of LLMs. There is no technical evidence in the paper that suggests why combining uncertainty based methods with LLM is a sensible choice in the first place, why is better than combining with just regular non-pretrained models. The evaluation is also quite narrow.\n\nSo far the results read like \"we combine approximate thompson sampling with pretrained LLMs on language bandit tasks and it just works better\". This does not shed light on if there is anything special about the fact that the underlying models are pretrained LLMs, nor do we need pretrained models in the first place. Details below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposed to combine various uncertainty based decision making algorithms with a large language model pretrained reward function (or value function). They test out a range of algorithms such as greedy, dropout, laplace approximation and epinet, and benchmark the performance on content moderation task (a contextual bandit task). They show that uncertainty based methods outperform greedy in terms of average regret.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "=== **why pretrained LLMs** ===\n\nMaybe a central question is: why uncertainty based methods with pretrained LLMs? What specific properties of pretrained LLMs are being leveraged in designing the combination? \n\nThe paper starts with the premise that it \"makes sense\" to just combine uncertainty methods with pretrained LLMs. What kind of difference does the fact that the LLMs are \"pretrained\" make? What if we start with a non-pretrained LLMs, or just a regular neural net, what will the difference be like? Running LLMs does not come without its downside, that is the computational burden. This comparison will illustrate the importance of combining such methods with pretrained LLMs in the first place\n\n=== **the combination does not leverage anything specific about LLM** ===\n\nThroughout the design, there is no specificity to the fact that the underlying models are pretrained LLMs. It is as if the combination with uncertainty based method does not account for the fact that the underlying method is pretrained at all. Maybe a concrete question is, what would you do differently if the underlying model is not pretrained, or not a LLM at all?\n\n=== **narrow evaluation** ===\n\nThe eval is only on content moderation task, which is a bit too narrow to illustrate the technical solidity of the approach, ideally there should be more language related benchmarking.\n\n=== **results** ===\n\nIt actually comes at little surprise that greedy is outperformed by alternative methods, given evidence in prior work. I think an interesting experiment to run is: if we use non-pretrained LLMs or just a regular transformer based model, how would the performance gap be between different methods? Do we expect pretrained LLMs to enlarge the performance gap or make the gap smaller - there is \"argument\" that can make the results go either way, and it is interesting to characterize this.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "fV8ERBrLfp", "reviewer": "Reviewer_XoFs", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths:\n1. The paper is well organized and well written. The main research question and the approaches to answer it are clear.\n\n## Weaknesses:\nI will organize this section in accordance with the TMLR acceptance criteria.\n\n### **Are the claims made in the submission supported by accurate, convincing and clear evidence?**\n1. **Empirical evidence is too narrow in coverage to be convincing.** The experiments are run with only 2 datasets (1 relegated to the appendix) and 1 model. The claim that \"uncertainty plays a fundamental role in bandit tasks with LLMs\" is a large one that is not sufficiently supported by the given experiments. The study needs to include several datasets and several LLMs in order for the results to be plausibly generalizable. \n\n2. **Non-standard loss function.** Equation 3 is a non-standard loss function used to train the greedy bandit, justified by the claim that the standard loss function Equation 2 is not scalable with LLMs. It is not clear whether this loss function is a novel contribution. If so, there is no evidence to believe that this new loss function results in a properly trained greedy bandit (one that is close to that trained by the standard loss function). If this is not a novel contribution and exists in prior continual learning literature, this needs to be clearly cited so that its use is justified.\n\n### **Would some individuals in TMLR's audience be interested in the findings of this paper?**\n1. **Expected result.** The study answers the claim: \"Equipping a pre-trained language model with uncertainty estimates and using Thompson sampling is better than using the greedy algorithm on 2 language-based contextual bandit tasks.\" This result is not surprising to the bandits/sequential decision making/probabilistic models audience who would reasonably expect that Thompson sampling is superior to the greedy algorithm no matter the domain.\n2. **Uninteresting use of LLMs.** The LLM used here is a GPT-2 (124M parameters) model that is very small by modern standards. Furthermore, the LLM's generation capabilities are not used at all in this paper, as the LLM's classification head is replaced with a regression head (compare this to \"Efficient Exploration for LLMs\" (Dwaracherla et al., 2024) in which the language generation plays a central role in the decision making loop). For these reasons, the study may not be interesting to the NLP/LLM audience.\n3. **Narrow empirical coverage**. Practitioners may not find this study useful in its current form as it only presents results with 1 model and 2 datasets. There is little evidence that the results are generalizable. A practitioner would not be adequately convinced that any particular model and/or uncertainty method is useful for general contextual bandit tasks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper investigates the use of large language models (LLMs) as decision making tools, and how equipping LLMs with uncertainty estimates affects their performance. The central claim is that uncertainty estimates play a \"fundamental role in bandit tasks with LLMs\". This claim is empirically supported by contextual bandit experiments using the \"measuring hate speech\" dataset with a GPT-2 (124M parameters) LLM with its classification head replaced by a regression head in order to predict rewards. The methods used to equip the LLM with uncertainty estimates are the Laplace approximation, dropout, and epinet. The action selection algorithms tested are the greedy algorithm which does not require uncertainty estimates, and Thompson sampling which requires uncertainty estimates. The experiments show that Thompson sampling outperforms the greedy algorithm across all uncertainty methods.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "## Critical to recommendation\n1. **Run experiments with a large variety of datasets and LLMs**. Studies whose claims are supported purely by empirical evidence need to have comprehensive experiments in order to be convincing.\n2. **Either justify or remove non-standard loss function**. Examples of justification include: citing prior work that use it, or run experiments to show that the greedy bandit trained with it is similar to that trained with the standard loss function. I highly recommend removing it and using other methods for scalability. For instance, it is common to freeze the \"torso\" of the LLM during training and to only update the parameters of the linear regression head, as in Dwaracherla et al. (2024). This effectively turns the LLM torso into a feature extractor, and the extracted features of previous contexts can be cached so that only a small number of LLM forward passes need to be done at each iteration.\n\n\n## Would strengthen work\n1. **Move experiment in appendix to main body and do a similar analysis**.\n2. **Downgrade claim.** The word \"fundamental\" is very strong: the authors should consider downgrading their claim to \"uncertainty improves bandit tasks with LLMs\".\n\n## Additional thoughts\nThe authors may notice that the above requested changes only pertain to the question \"Are the claims made in the submission supported by accurate, convincing and clear evidence?\" The question \"Would at least some individuals in TMLR's audience be interested in knowing the findings of this paper?\" is subjective, and my answers to that question may not be addressable without changing the research question completely. I advise the authors to consider more interesting problems at the intersection of LLMs and sequential decision making. For example, how does the number of parameters of the LLM affect performance at contextual bandit tasks? Is a 124M parameter model all you need, or does performance keep improving as the models get larger? How are these scaling laws different with Thompson sampling or UCB or the greedy algorithm? Are there contextual bandit tasks in which the generation capabilities of the LLM are useful?", "broader_impact_concerns": "There are no broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper combines uncertainty estimation techniques with LLMs bandits methods. The author compare variants of dropout and Laplace approximation, evaluated on decision-making textual tasks.

This is a useful empirical analysis, and all the reviewers agree that this combination of methods is relevant for the TMLR audience. The authors have addressed the main concerns of the reviewers. The last concern that stand out is the lack of apparent novelty, but I think this should not be an important factor in the decision. I thus recommend acceptance.
Decision: Accept (poster)