Evaluating High-Order Predictive Distributions in Deep Learning
Abstract
Most work on supervised learning research has focused on marginal predictions.
In decision problems, joint predictive distributions are essential for good performance.
Previous work has developed methods for assessing low-order predictive distributions with inputs sampled i.i.d.
from the testing distribution.
With low-dimensional inputs, these methods distinguish agents that effectively estimate uncertainty from those that do not.
We establish that the predictive distribution order required for such differentiation increases greatly with input dimension, rendering these methods impractical.
To accommodate high-dimensional inputs, we introduce dyadic sampling , which focuses on predictive distributions associated with random pairs of inputs.
We demonstrate that this approach efﬁciently distinguishes agents in high-dimensional examples involving simple logistic regression as well as complex synthetic and empirical data.
INTRODUCTION
We consider learning agents that are trained on data pairs (( X t ,Y t +1 ): t =0 , 1 ,...,T − 1) .
At a new input X T , such an agent can generate a predictive distribution of the outcome
Y T +1 that is yet to be observed.
This distribution characterizes the agent’s uncertainty about Y T +1 .
We refer to such a prediction as marginal to distinguish it from a joint predictive distribution over a sequence of prospective outcomes ( Y T +1 ,...,Y T + τ ) with inputs ( X T ,...,X T + τ − 1 ) .
Predictive distributions express uncertainty about future observations.
The importance of such uncertainty estimation has motivated a great deal of research over recent years, much of which in the Bayesian deep learning community [Neal, 2012].
With the proliferation of agents that generate predictive distributions, it is increasingly important to systematically study and improve their performance.
Recent theoretical work has highlighted the importance of joint predictive distributions in driving effective decisions [Wen et al., 2022].
This theory is supported by experiments that assess and compare agents using synthetic data generated by a random neural network and 2D inputs [Osband et al., 2022].
That work evaluates the quality of joint predictive distributions over ten inputs sampled i.i.d.
from the training distribution.
The results clearly distinguish agents that effectively estimate uncertainty from those that do not.
This evaluation predicts agent performance when used to guide decisions in high-dimensional ‘neural bandits’.
However, as the input dimension increases, the aforementioned approach to evaluating agents becomes uninformative.
As we will later discuss, the reason lies in the order of the predictive distributions being evaluated.
With a twodimensional input, the tenth order distribution sufﬁces, but the predictive distribution order required to produce meaningful assessments increases rapidly with the input dimension.
We could consider scaling the predictive distribution order as needed, but the evaluation algorithms of Osband et al.
[2022] become computationally intractable.
To accommodate high-dimensional inputs, we introduce dyadic sampling , which focuses on predictive distributions associated with random pairs of inputs rather than those scattered according to the training input distribution.
We demonstrate that this approach efﬁciently distinguishes agents in high-dimensional examples involving simple logistic regression as well as complex synthetic and empirical data.
For example, agent assessments based on dyadic sampling are predictive of performance in high-dimensional neural bandits presented in results of Osband et al.
1.1 RELATED WORK
We are motivated by the importance of joint predictions in driving effective decisions [Wen et al., 2022].
Empirical analysis of joint predictions for deep learning in 2D supports this theory [Osband et al., 2022].
We provide a practical heuristic to scale these insights to high dimensions.
Our research is closely related to topics in Bayesian deep learning [MacKay, 1992, Wilson and Izmailov, 2020], and robustness [Hendrycks and Dietterich, 2019].
For the most part, these communities have focused on the problem of marginal prediction [Nado et al., 2021, Wilson et al., 2021].
Recent work has also highlighted a notion of crosscorrelation in regression and related decision problems [Wang et al., 2021].
Our paper provides a related perspective that scales to classiﬁcation and high dimensional data.
1.2 KEY CONTRIBUTIONS
We propose dyadic sampling , which evaluates high-order joint predictions at random pairs of inputs.
Section 2 motivates the approach, and shows that it can mitigate some challenges in evaluating high-order predictive distributions.
Section 3 shows that dyadic sampling provides useful assessments in logistic regression.
As input dimension scales, i.i.d.
sampling from the training distribution does not offer a feasible approach.
Dyadic sampling offers a viable path where the evaluation of Osband et al.
[2022] is inadequate.
Section 4 extends these insights to The Neural Testbed – an opensource package for the evaluation of joint predictions in deep learning.
As in logistic regression, the neural network generative process is not amenable to evaluation via i.i.d.
sampling when the input dimension exceeds three.
In contrast, dyadic sampling scales gracefully as the input dimension grows large.
As part of this project, we submit all agent and evaluation code to github.com/deepmind/neural_testbed .
Section 5 shows that our methodology can extend beyond synthetic data.
Dyadic sampling can feasibly evaluate joint predictions on high-dimensional real datasets.
We evaluate benchmark approaches to Bayesian deep learning and show that the insights from the Testbed carry over to real data.
We see that, after tuning, all agents perform similarly in terms of marginal predictions.
However, there are signiﬁcant differences in the quality of joint predictions per agent, evaluated via dyadic sampling.
Further, Testbed performance is highly predictive of performance on empirical data.
EVALUATING PREDICTIVES
This section introduces notation for the standard supervised learning framework we will consider as well as our evaluation metric: KL-loss.
We show that estimating KL divergence in high dimensional distributions can be challenging, and present dyadic sampling as an effective heuristic.
2.1 ENVIRONMENT AND PREDICTIONS
Consider a sequence of pairs (( X t , Y t +1 ) :
t = 0 , 1 , 2 , . . . ) , where each X t is a feature vector and each Y t +1 is its target label.
Each target label Y t +1 is produced by an environment E , which we formally take to be a conditional distribution E ( ·| X t ) .
The environment E is a random variable; this reﬂects the agent’s uncertainty about how labels are generated.
Note that P ( Y t +1 ∈ ·|E , X t ) =
E ( ·| X t ) and P ( Y t +1 ∈ ·| X t ) =
E [ E ( ·| X t ) | X t ] .
We consider an agent that learns about the environment from training data D T ≡ (( X t , Y t +1 ) :
t = 0 , 1 , . .
, T − 1) .
After training, the agent predicts testing class labels
Y T +1: T + τ ≡ ( Y T +1 , . . .
, Y T + τ ) from unlabeled feature vectors X T : T + τ − 1 ≡ ( X T , . . .
, X T + τ − 1 ) .
We describe the agent’s predictions in terms of a generative model, parameterized by a vector θ T that the agent learns from the training data D T .
Speciﬁcally, θ T parameterizes a distribution P ( ˆ E ∈ ·| θ T ) over imagined environment ˆ E , which is also a conditional distribution.
For any inputs X T : T + τ − 1 , to generate the imagined labels ˆ Y T
+1: T + τ , the agent ﬁrst samples an imagined environment ˆ E from
P ( ˆ E ∈ ·| θ T ) , then generates ˆ Y t +1 ∼ ˆ E ( ·| X t ) conditionally i.i.d.
for each t = T, . . .
, T + τ − 1 .
The agents τ th -order predictive distribution is given by
[Equation]
P
[Equation]
) , which represents an approximation to what would be obtained by conditioning on the environment:
P
[Equation]
If τ = 1 , this represents a marginal prediction of a single label for a single feature vector.
For τ > 1 , this is a joint prediction over τ labels for τ different feature vectors.
2.2 EVALUATING JOINT PREDICTIONS
A learning agent can be assessed through the quality of its predictive distribution
ˆ P T +1: T
+ τ .
A canonical approach is to evaluate the KL-divergence [Wen et al., 2022],
[Equation]
Recall that the expectation represents an integral over all random variables.
The minimum of d τ KL over all agents that depend on the environment only through D T is attained by the posterior agent, whose predictive distribution is
[Equation]
[Equation]
Let d τ KL denote the minimum achievable KL-divergence.
Algorithm 1 provides a simple Monte-Carlo approach to evaluate d τ KL .
As the order of the predictive distribution τ grows, this provides a more nuanced evaluation of agent beliefs than just marginals.
However, even for simple problems, the magnitude of τ required to provide additional insight beyond marginals can become intractably large.
To anchor our thinking on this matter we consider a simple coin tossing example.
Example 1 (Bag of coins) .
Let each X t be a sample from coins { 1 , . . .
, M } .
Let the probability of heads p x ∼
Unif(0 , 1) i.i.d.
for each coin x .
Each observation Y t +1 is the outcome from tossing coin X t , so that E (1 | X t ) =
p X t .
Let us consider a predictive distribution for which P ( ˆ Y 1: τ | X 0: τ − 1 )=
(cid:81)
τ − 1 t =
P ( ˆ Y t +1 | X t )
=1 / 2 τ .
Suppose an agent uses this to select coins sequentially with the aim of maximizing the expected number of heads.
While ˆ Y 1: τ accurately minimizes d 1KL , the agent assumes that toss outcomes are independent and therefore does not learn from history to improve successive choices.
Accounting for dependencies arising in the joint distribution, as would be captured by d τ KL , is essential to maximizing performance.
Proposition 1 (Small τ approximately marginal) .
If the agent deﬁned above is applied to Example 1 with τ (cid:28) M , d τ KL =
d τ KL + O (cid:0) τ 3 /M (cid:1) .
Proof.
Note that under the event that there are no repeated inputs in X 0: τ − 1 , the posterior agent is equivalent to the agent deﬁned above.
For τ (cid:28) M , this event occurs with high probability.
The detailed proof is in Appendix A.1.
Proposition 1 shows that if τ (cid:28) M , then d τ KL is unable to distinguish agents that only match marginals from those that are useful for decision making.
When the cardinality of the input space M is much larger than the test distribution order τ it is unlikely that any correlated inputs will be sampled.
The metric d τ KL punishes agents that impose an erroneous correlation, but is unlikely to reward agents that correctly capture this dependence until τ is sufﬁciently large.
In Example 1, it may sufﬁce to use a value of τ that grows cubicly in M .
However, due to the curse of dimensionality, the required magnitude of τ can grow exponentially in problem dimension.
To handle such cases, we introduce a practical evaluation metric that correctly identiﬁes high quality predictive distributions with modest values of τ .
For this result, we introduce notation for assignment: for random variables
A, B, C and a function f ( c )
≡ E [ A | B = c ] , let E [ A | B ← C ] = f ( C ) .
Note that, in general, if C is a random variable then E
[ A | B ← C ] (cid:54) =
E [ A | B = C ] .
Deﬁnition 1 (Polyadic test sampling (of order κ )) .
For any κ ∈ N , let ‘anchor points’ X 1: κ be drawn i.i.d.
from P ( X t ∈ · ) , and let ˜ X κT : T + τ − 1 ∼ Unif { X 1: κ } .
We deﬁne,
[Equation] [Equation]
Polyadic sampling is motivated by a desire to investigate an agent’s predictions in situations where correlation between predictions is more likely to play an important role.
Under reasonable regularity assumptions lim κ →∞ d τ,κ KL = d τ KL for all τ .
In the case of κ < τ , we can ensure that at least one input will be sampled multiple times.
In the special case of κ = 1 , we call this monadic sampling.
Proposition 2 (Monadic sampling cannot spot bad agents) .
Consider an agent that ignores the inputs and predicts
[Equation]
where ˆ Y 1 , . .
, ˆ Y τ are sampled independently from Ber(ˆ p ) with a shared parameter ˆ p ∼ Unif(0 , 1) .
Then, for any τ ∈ N in Example 1 this agent achieves the minimum d τ, 1KL over all agents.
Proof.
This agent is constructed so that for any repeated inputs X 0: τ − 1 = ˜
X 10: τ − 1 , this agent’s predictive distribution matches that of the posterior agent.
Monadic sampling can examine whether an agent understands the correlation structure at a single input.
However,
Proposition 2 shows that it does not punish agents that erroneously ascribe correlation to independent input-output pairs.
This agent achieves the best possible score in d τ, 1KL but is useless for driving decisions.
In order to weed out these agents it is crucial to sample more than one input point.
2.3 DYADIC SAMPLING ( κ = 2)
This paper introduces dyadic test input sampling ( κ = 2) as a practical heuristic for assessing the quality of joint predictions in high dimensions.
This sampling scheme samples two random anchor points from the input space, and then randomly resamples the τ inputs from these anchor points.
Even with moderate τ = 10 , we can be sure that most batches will contain a mix of points that are highly correlated to each other, as well as some others which may be quite different.
Dyadic sampling is a heuristic approach designed to work well in practical problems.
The choice of κ = 2 addresses the extreme shortcomings of d τ,κ KL by Propositions 1 andin the settings κ → ∞ and κ = 1 respectively.
However, it is certainly not a perfect substitute for evaluating d τ KL with very large τ .
Depending on the setting, it is certainly possible to design agents that fare very well according to d τ, 2KL , but very poorly according to d τ KL .
One might ask, ‘Does some other 1 < κ < ∞ provide a better candidate for practical evaluation of posterior predictives?’.
Could there be an analagous result to Propositionwhen considering κ = 2 , but evaluating posterior predictions at three anchor points?
Note that, since d τ, 2KL already evaluates the quality of the joint predictions at any pair of inputs, then for most problems the distribution over any three inputs will also be estimated well.
In particular, for any Gaussian process, the ﬁrst two moments are enough to determine the entire distribution of E .
We push details to
Appendix A.2.
2.4 JOINT PREDICTIONS AND INFORMATION
So far, we have motivated dyadic sampling mostly through appeal to Example 1, together with some heuristic arguments.
In this subsection we expand on this intuition through the lens of information theory.
To illustrate this, let’s consider the posterior agent, which is optimal for generating predictive distributions.
Note that under the posterior agent,
[Equation]
where I ( · ) denote the (conditional) mutual information [Cover, 1999] and D T : t ≡ ( X T : t − 1 , Y T +1: t ) .
Note that the second equality follows from the chain rule of mutual information.
On the other hand,
[Equation]
where the second equality follows from X T : T + τ − 1 are i.i.d.
For d τ KL to be signiﬁcantly different from τ d 1KL , we need for at least one t , the dataset D T : t is informative about the target label Y t +1 at X t .
For practical problems with τ small relative to the input space, the D T : t is not informative about
Y t +1 .
In such cases, we have d τ KL ≈ τ d 1KL .
One way to think about dyadic sampling is a heuristic approach to sample X T : T + τ − 1 so that D T : t is particularly informative about
Y t +1 and so evaluate the quality of the posterior approximation.
Depending on the problem settings, other input sampling schemes may also be appropriate to accomplish this goal.
LOGISTIC REGRESSION
The results of Section 2 provide a motivation for dyadic sampling where it can sidestep the curse of dimensionality in higher-order predictive distributions.
In this section, we show that this effect can occur in practical settings, not just obtuse problems cooked up for theory.
In fact, even for the canonical problem of logistic regression, the beneﬁts of dyadic sampling can be signiﬁcant.
3.1 PROBLEM FORMULATION
We consider the familiar problem of D -dimensional logistic regression.
Inputs are sampled i.i.d.
X t ∼ N (0 , I D ) and the environment E is determined by parameter φ ∼ N (0 , I D ) .
Outputs Y t +1 ∈ { 0 , 1 } are then sampled according to
P ( Y t +1 =
[Equation]
|
[Equation] [Equation]
Here ρ > 0 is the temperature controlling signal to noise ratio (SNR).
We set ρ =0 .
01 for a high SNR setting.
In this simple setting, we can compare three agents that make predictions ˆ Y 1: τ given inputs X 0: τ − 1 .
uniform : P ( ˆ Y t +1 =1 | X t )
= 12 for t =0 , 1 ,..
marginal :
Samples λ ∼ N (0 , 1) , and then predicts
[Equation]
prior :
Samples ˆ φ ∼ N (0 ,I D ) , and then predicts
[Equation]
The agents are chosen to highlight speciﬁc properties of the logistic regression problem.
The uniform agent makes the correct marginal predictions at any input, but does not capture any correlation among Y 1: τ .
The marginal agent makes the correct marginal predictions, and it also gets the correct joint distribution if inputs X 0: τ − 1 are all sampled at a single point (monadic sampling).
However, it introduces spurious correlation among the predicted outputs if the inputs are not all equal.
The prior agent samples from the true prior, and so is optimal for all d τ,κ KL .
We would like to have a practical evaluation metric that can separate this optimal agent from these sub-optimal approximations.
We consider metrics d τ KL and d τ,κ KL for κ = 1 , 2 , all of which are estimated through Monte Carlo sampling according to
Algorithm 1.
Despite the simplicity of this problem, only dyadic sampling ( d τ,κ =2 KL ) can correctly separate the agents once the input dimension grows.
3.2 RESULTS
As the prior agent makes optimal predictions in this problem, in principle, this agent will outperform all others according to d τ KL as the order of the predictive distribution τ grows.
However, to separate the agents, the required τ can quickly become intractable in the input dimension.
Figure 1 shows that, in logistic regression, for dimension
D ≥ 5 , even τ = 10 , 000 is insufﬁcient to give a factor of 2 separation between the optimal prior agent and the uninformed uniform agent.
The computational cost of evaluating d τ KL grows with τ , so that this can quickly becomes impractical even for relatively small-scale problems.
By contrast, evaluation with d τ,κ =2 KL is able to identify this separation with only τ = 10 even as the input dimension grows.
[Figure]
1e 1e m i n i m u m s u c h t h a t d K L ( p r i o r ) <
12 d K L ( un i f
o r m )
[Figure] [Figure]
[Figure]
dyadic can distinguish these agents with a value
[Figure]
Figure 2 shows that this scaling carries over to high dimensions, ﬁxing τ = 10 .
Sampling test inputs i.i d. cannot distinguish the uniform agent from the prior agent in dimensions greater than 100 .
At a high level, this result matches the spirit of Proposition 1.
Figure 2 also shows that monadic
[Figure]
cannot distinguish the prior agent from the marginal agent.
This mirrors Proposition 2, but in a setting with generalization.
Dyadic sampling
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
correctly identiﬁes that prior agent is a superior agent across all input dimensions.
These results clearly demonstrate that the theoretical concerns raised in Section 2 actually occur in practical problems.
Further, these concerns can occur even in the most simple settings of logistic regression, rather than contrived scenarios.
We push the details on the robustness/sensitivity of these results to Appendix B.
THE NEURAL TESTBED
In this section we show that the insights observed in the linear setting of Section 3 extend to nonlinear function approximation and neural networks.
Osband et al.
[2022] introduce the Neural Testbed as a simple synthetic 2D problem to evaluate posterior predictives in deep learning.
We show that, using the exisitng d τ KL evaluation, this approach does not scale to higher dimensions.
However, using dyadic sampling we are able to extend these insights to practical scales.
As part of our work we contribute these changes to github.com/deepmind/neural_testbed . 4.1 PROBLEM FORMULATION
The Neural Testbed works with a synthetic data generating process around random 2-layer MLPs [Osband et al., 2022].
For each random seed, a random neural network is sampled according to standard Xavier initialization [Glorot and
Bengio, 2010].
Then, random train/test inputs are sampled
X 1: T + T (cid:48) ∼ N (0 , I ) and labels assigned randomly according to the probabilities of the generative MLP.
We follow the exact settings in the existing opensource package except for two key changes.
First, we supplement the existing evaluation by d 1KL , d 10KL to also evaluate according to d 10 ,κ =2 KL .
Then, we vary the input dimension of the problem (which is ﬁxed at D = 2 in the original Neural Testbed release).
To account for the different data requirements in higher dimensions we similarly increase the number of training pairs in low, medium, high data regimes to scale with the input dimension.
The full testbed sweep is deﬁned over input dimensions
D ∈ { 2 , 10 , 100 } , number of training pairs T = λD for λ ∈ { 1 , 10 , 100 , 1000 } , temperature ρ ∈ { 0 . 01 , 0 . 1 , 0 . 5 } with 5 random seeds in each setting.
We push full details, together with opensource implementation, to Appendix C.
[Table]
[Figure]
[Table]
4.2 BENCHMARK AGENTS
To compare the performance of benchmark agents we make use of the opensource agents developed by Osband et al.
Table 1 lists agents that we study and compare as well as hyperparameters that we tune.
In our experiments, we optimize these hyperparameters via grid search.
The choices start from the defaults released in github.com/deepmind/neural_testbed , but extend and tweak some hyperparmeter choices for high dimensional problems.
Further detail on these agents is provided in Appendix C.2.
4.3 OVERALL RESULTS
Figure 4 shows the KL estimates for these agents, normalized so that the baseline MLP has a score of 1.
In each case, these agents are tuned for performance on the Neural
Testbed for input dimension 100.
We can see that in this setting evaluation in d 10KL is statistically indistinguishable from that of marginal predictions.
We also see that, for the most part, the quality of these marginal predictions is not massively improved versus the MLP.
However, unlike the 2D testbed results, we do see that some of these more advanced approaches can improve marginal predictions.
However, we see that evaluating agents according to dyadic sampling leads to massive distinctions in their evaluations.
Interestingly, these qualitative results match the d 10KL with i.i.d.
test sampling ordering in the 2D setting.
Osband et al.
[2022] showed that this order was highly correlated with performance in sequential decision problems, even for high input dimension.
Our results provide a signiﬁcant new ﬁnding; in high dimensional problems dyadic sampling sampling can provide a more targeted signal for the suitability in downstream tasks.
4.4 PRIORS IN HIGH DIMENSIONS
One of the most clear and interesting pairs of agents to compare is ensemble and ensemble+ .
These agents are identical except for the addition of randomized, ﬁxed prior networks.
Prior work has shown that this difference can be crucial in high-dimensional decision problems [Osband et al., 2018, Burda et al., 2019].
Comparison of joint predictions d 10KL in 2D problems also showed a signﬁcant difference, but only for very small training sets T ≤ 30 .
The question remained, do these randomized priors provide value in large scale supervised learning?
Figure 4 shows that, according to d 10KL the beneﬁts of ensemble+ appear to evaporate for input dimensions ≥ 2 .
However, using dyadic sampling and κ = 2 we can see there are huge differences in the quality of their posterior approximation that extend to high dimensional problems.
Figure 5 shows that, as we increase the dimensionality of the problem, so too we increase the size of the largest training sets where prior functions afford signﬁcant advantages.
Rather than becoming irrelevant in large problems, the importance of good inductive bias actually increases with input dimension.
[Figure] [Figure]
REAL DATA
In this section we show that the key insights gained from the synthetic neural testbed can carry over to real datasets.
We replace the neural network generative process of Sectionwith small challenge datasets drawn from the deep learning literature.
We then tune the agents of Table 1 for each of these settings and analyse the results.
We ﬁnd that all agents can be tuned to perform roughly equivalently in terms of marginal predictions.
However, their performance difference greatly in terms of their joint performance as measured by dyadic sampling.
Further, agent performance on the testbed is highly correlated with performance on real datasets.
5.1 PROBLEM FORMULATION
Progress in the ﬁeld of deep learning has been driven in large part through evaluation on shared, ﬁxed datasets [Krizhevsky et al., 2012].
We repeat the analysis of Section 4 but replace the synthetic data generating process with a collection of datasets drawn from the literature [TFD].
Table 2 outlines the ten datasets we include in our analysis.
We wanted to choose datasets that might provide an analagous challenge to the Neural Testbed and so selected them based on their popularity in the literature, and suitability for training with a 2-layer MLP.
For this reason, large scale challenges such as ImageNet or language modelling, which typically require different classes of models were not included in our selection [Deng et al., 2009].
To mirror our evaluation in the Neural Testbed we begin with datasets D nT n = (( X t , Y t +1 ) :
t = 0 , .., T n − 1) for n = 1 , .., 10 .
To evaluate different data regimes we create subsampled datasets ˜ D
nT for T = 10 , 100 , 1000 , T n to evaluate different data regimes.
We then evaluate d τ,κ KL in the ‘low temperature’ limit, taking the labels in the supplied test set as probability 1 or, equivalently, the negative log-likelihood [Wen et al., 2022].
As in Section 4.2, we evaluate the agents outlined in Tableacross each of these datasets in each data regime.
We then tune the hyperparameters per dataset, per data regime and aggregate the performance by taking the average over all evaluations.
This mirrors the procedure that we applied in
Section 4.
We push full details to Appendix D.
5.2 RESULTS
We begin by assessing the quality of the agents’ performance in marginal predictions, when averaged over all datasets, for all data regimes.
Figure 6 shows that, once agents are optimized for each setting, the differences between agents is not statistically signiﬁcant.
This ﬁnding mirrors our observation in the case of synthetic data and Figure 3.
These agents perform similarly at marginal prediction in the testbed, and overall they perform similarly in the real datasets as well.
Once you consider the quality of joint predictions however, there is a signiﬁcant difference in the quality of predictive distributions evaluated on real data.
Further, Figure 7 shows
[Table]
[Figure]
l p e n
[Figure]
e m
[Figure]
l
[Figure]
d
[Figure]
o p
[Figure] [Figure]
N LL o n r e a l d a t a ( = 1 )
[Figure] [Figure]
that this difference is highly correlated with performance on the Neural Testbed.
Agents that perform better in the setting with synthetic data also tend to perform better when evaluated on real data.
This ﬁnding is particularly signiﬁcant since the differences in d 10 , 2 KL are quite large even for these state of the art agents.
These results provide strong indications that the issues observed in sequential decision problems [Osband and Van Roy, 2015] and synthetic data [Osband et al., 2022] can extend to real data.
Now, in some sense the results we have presented are ‘nonstandard’ in that our evaluation includes averages over restricted-data versions of the canonical datasets in Table 2.
We believe that this is a sensible approach if you are interested in designing learning agents that work in online decision making and are robust to different data regimes.
However, in some supervised learning settings it is more common from practitioners to care only about the ‘full’ datasets with T = T n .
In fact, the ﬁndings of Figure 6 and
Figure 7 are essentially unchanged when restricting only to the ‘full data’ regime.
That is, the differences in marginal predictions τ = 1 are quite minor, but the differences in τ = 10 , κ = 2 are extreme.
Further, that these differences in joint performance are highly correlated with agent performance in the testbed.
We push full details to Appendix D.
CONCLUSION
Good predictions are essential for good decisions.
Crucially, the quality of these decisions depends on the quality of joint predictions and not just the marginals [Wen et al., 2022].
In this paper, we highlight the difﬁculties in evaluating high-order predictive distributions that are essential for decision making.
We introduce dyadic sampling as an practical heuristic to sidestep the curse of dimensionality.
We motivate dyadic sampling through a simple discrete example, and show that the key insights extend to linear and then nonlinear systems.
We show that the Neural Testbed cannot effectively scale to high dimensions with i.i.d.
sampling, but that it can with dyadic sampling.
Importantly, this approach also scales to challenge datasets, and we show that testbed performance is highly correlated with real data.
A major contribution of our work is the opensource effort at github.com/deepmind/neural_testbed .
This includes all the code used to generate the paper, and helps to provide clear and reproducible benchmarks for the community.
We believe that this paper can provide an stimulating base for future research into agents that make predictions in high-dimensional problems, and drive effective AI systems.
A EVALUATING PREDICTIVE DISTRIBUTIONS
This section contains supplementary material for Section 2.
Importantly, we provide the proof for Proposition 1and discuss why dyadic sampling is sufﬁcient for Gaussian process.
A.1 PROOF FOR PROPOSITION
Proposition 1 (Small τ approximately marginal) .
If the agent deﬁned above is applied to Example 1 with τ (cid:28) M ,
[Equation]
Proof.
Note that by deﬁnition, d τ KL ≥ d τ KL .
We now prove that d τ KL ≤ d τ KL + O (cid:0) τ 3 /M (cid:1) .
Note that
[Equation]
where τ log (cid:0) 12 (cid:1) is the log-likelihood under the uniform agent, and
[Equation]
Consequently, we have
[Equation]
We deﬁne the event G as
G =
{ there are no repeated inputs in X 0:
[Equation]
− 1 } .
One key observation is that conditioning on G , the posterior predictive distribution is i.i.d.
across inputs, and
[Equation]
conditioning on G .
Hence
[Equation]
where ¯ G is the complement of G , and the inequality follows from log ( P ( Y 1: τ | X 0: τ − 1 )) ≤ 0 .
Hence we have
[Equation]
Finally, note that
[Equation]
Hence P ( ¯ G ) =
O (cid:0) τ 2 /M (cid:1) and we have
[Equation]
The conclusion follows from
[Equation]
A.2 DYADIC SAMPLING AND GAUSSIAN PROCESSES
In this section, we discuss why dyadic sampling is sufﬁcient for Gaussian processes (GPs).
In particular, we show that when both the environment E and the imagined environment ˆ E of an agent follow GP, then with sufﬁciently large τ and under suitable regularity conditions, performing well under d τ,κ =2 KL is sufﬁcient to ensure that the posterior distribution of E and the agent’s belief over ˆ E are close.
Assume that both E and ˆ E are GPs with the same ﬁnite domain X and that the training input distribution is uniform over X .
Speciﬁcally, under the environment E ,
[Equation]
and under the imagined environment ˆ E ,
[Equation]
where W t +1 ’s and ˆ W t +1 ’s are i.i.d.
observation noises according to N (0 , σ 2 ) , and f and ˆ f are functions over X .
We assume that P ( f ∈ ·|D T ) =
N ( µ, Σ) and P ( ˆ f ∈ ·| θ T ) =
N (ˆ µ, ˆΣ) .
Note that by deﬁnition d τ,
[Equation]
Note that the ﬁrst term in the above equation is irreducible and independent of the agent, hence, performing well under d τ,κ =2 KL is equivalent to performing well under ˜ d τ,κ =2 KL .
Under suitable regularity conditions, for sufﬁciently large τ , we have
[Equation]
where ˜ X 1:2 = (cid:0) ˜
X 1 , ˜ X 2 (cid:1) and ˜ X 1 and ˜ X 2 are i.i.d.
sampled from P X .
Thus, if the RHS of the above equation is small, then it implies that
[Equation]
) is small for all ˜ X 1:2 .
Let µ ( ˜ X 1:2 ) ∈ (cid:60) 2 and Σ( ˜ X 1:2 ) ∈
(cid:60) 2 × 2 respectively denote µ and Σ restricted to ˜ X 1:2 , and ˆ µ ( ˜ X 1:2 ) and ˆΣ( ˜ X 1:2 ) are deﬁned similarly, then we have
[Equation]
Consequently, if equation 7 is small, then µ ( ˜ X 1:2 ) is close to ˆ µ ( ˜ X 1:2 ) and Σ( ˜ X 1:2 ) is close to ˆΣ( ˜ X 1:2 ) .
Since this holds for all ˜ X 1:2 , this further implies that µ is close to ˆ µ and Σ is close to ˆΣ .
In other words, the posterior distribution of E and the agent’s belief over ˆ E are close.
B LOGISTIC REGRESSION
This appendix provides supplementary details for Section 3.
We include all of the code necessary to generate Figures 1 and 2 as part of our opensource submission github.com/deepmind/neural_testbed .
Results are averaged overrandom seeds per problem setting.
Figure 8 provides another kind of insight to the scaling observed in Figure 1.
In these plots we show the KL ratio of a perfect prior agent when compared to uniform .
We can see that, for any input dimension, the empirical KL ratio decreases with τ .
However, as the input dimension grows reasonably large ( D = 10) , that even large τ = 10 , 000 are not enough to observe this ratio under 0.5.
We know that, as τ → ∞ this ratio will tend to zero for these two agents.
By contrast, dyadic sampling is able to clearly distinguish these agents even for moderate values of τ .
Figure 9 provides some insight to the robustness of Algorithm 1 under varying number of agent samples.
We make use of the epistemic neural network notation introduced by Osband et al.
We can see that these monte carlo estimates converge empirically as we increase the number of samples.
Therefore, for the purposes of our experiments in this
1e 1e 1e 1e  0.
[Figure]
( p r i [Figure]
d K L ( un i f o r m ) test: i.i.d.
input_dim         101e 1e 1e
[Figure]
test: dyadic ( =2)
[Figure]
K L e s t i m a t e ( = 10 , = 2 )
[Figure]
input_dim
[Figure] [Figure]
C NEURAL TESTBED
This appendix provides supplementary details for Section
C.1 PROBLEM FORMULATION
We build on the opensource code of the Neural Testbed github.com/deepmind/neural_testbed .
Our testbed sweep is deﬁned over input dimensions D ∈ { 2 , 10 , 100 } , number of training pairs T = λD for λ ∈ { 1 , 10 , 100 , 1000 } , temperature ρ ∈ { 0 . 01 , 0 . 1 , 0 . 5 } with 5 random seeds in each setting.
We replace the d 10KL evaluation with dyadic sampling d 10 ,κ =2 KL .
We release all of our code and implementation at github.com/deepmind/neural_testbed .
C.2 BENCHMARK AGENTS
We make use of the benchmark agents introduced in Osband et al.
[2022] and opensourced at github.com/deepmind/neural_testbed .
Since our testbed includes settings with number of training pairs as small as 2 (when D = 2 , λ = 1 ) and as large as 100,000 (when D = 100 , λ = 1000 ), in order to improve agent performance over all settings, we allow agents to adjust their number of training steps based on the problem setting.
Agents implementation can be found in our open source code under the path /agents
/factories .
We make small alterations to the tuning sweeps proposed in Osband et al.
[2022] in an effort to improve agent performance in high dimension problems.
This change strictly improved the agent performance as we only added hyperparameter choices and did not restrict them.
Our sweeps can be found in our open source code under the path /agents
/factories/ sweeps/testbed , but we highlight the differences that helped to improve agent performance.
For mlp , ensemble , dropout , bbb , hypermodel , ensemble+ agents, we found out that their performance improves by allowing them to adjust their default number of training steps based on the problem setting: increase it by 5x when λ = 1000 and decrease it by 5x when λ = 1 .
For sgmcmc agent, we found out that we can improve the performance of this agent by allowing it to increase prior variance parameter by 2x when D = 100 .
C.3 OVERALL RESULTS
Figure 3 provides an overview of the agent performance on the testbed in terms of d KL .
These numbers are normalized so that the baseline MLP has a value of 1.
In classiﬁcation problems it is common to also consider the classiﬁcation accuracy, or the percentage of inputs for which the agent correctly labels the input.
Figure 10 conﬁrms that, after tuning, none of the agents perform signiﬁcantly differently from baseline MLP.
[Figure]
D REAL DATA
This section provides supplementary details regarding the experiments in Section 5.
As before, we include full implementation and source code in our open source code under the path /real_data .
D.1 PROBLEM FORMULATION
Table 2 outlines the datasets included in our experiments.
For each dataset, we perform a standard preprocessing on inputs to be mean zero and unit variance.
Full details are available in our open source code under the path /real_data
/utils.py .
In the testbed we are able to evaluate a wide range of SNR regimes by varying temperature.
This means that we can query a given input X t multiple times and potentially obtain different class labels Y t .
For these ﬁxed dataset there is only one testing dataset, with deterministic labels given for each input.
We map this setting to the low temperature limit (and high SNR) setting of our testbed.
As such, we evaluate the negative log-likelihood in place of d τ KL .
This is equivalent to assuming the underlying world model was deterministic at these testing points, and is standard practice in deep learning.
We note that this ‘high SNR’ assumption appears to be reasonable in practice, since for all of the datasets considered in
Table 2 the benchmark mlp agent is able to obtain high classiﬁcation accuracy on held out data.
This would not be possible if the underlying system was fundamentally stochastic, due to the irreducible error due to chance.
D.2 RESULTS
In this section we provide some supplementary results that analyze the performance of our benchmark agents on real data.
To allow for hyperparameter tuning separately on the testbed and real datasets, we included different sweeps for the testbed and real datasets.
Our sweeps for real data can be found in our open source code under the path /agents
/factories/ sweeps/real_data .
One of the headline results in our paper is Figure 7, which shows that the quality of joint predictions on the testbed is highly correlated with performance in real data.
Figure 11 shows that this result is still true when you restrict the evaluation to the ‘full training data’ setting in each dataset.
Further, this aggregate correlation is not driven by just one outlier dataset
[Figure]
but actually occurs in each dataset individually.
In fact, after bootstrapping only the results on Iris were not signiﬁcant at the 95% conﬁdence levels.
This gives some additional reassurance that the relationship between joint performance on testbed and real data is robust
[Figure]
Our results in this paper allow for hyperparameter tuning separately on the testbed and real datasets.
We believe that this is reasonable practice, and reﬂects the way machine learning algorithms are usually used in practice.
However, one natural question might be if tuning an agent’s performance on the testbed leads to good hyperparameter settings on real data.
Figure 12 shows the results of this analysis across a wide range of agent-hyperparameter pairs.
Agent-hyperparameter pairs that perform better on the testbed generally also perform better on real data.
This result is statistically signiﬁcant in both τ = 1 and τ = 10 dyadic sampling.
However, we do see a stronger correlation in joint predictions rather than marginals.
So while we do not necessarily recommend tuning your agent for real datasets using the Neural Testbed, these results say that it will provide a better answer on average than random chance
[Figure]