Understanding the Intuition Behind Formulas in "Epistemic Neural Networks" [closed]

Question

I've been reading the paper "Epistemic Neural Networks" by Ian Osband et al. (2023), and I'm having trouble grasping the main intuition behind some of the formulas presented. The paper states the usefulness of joint predictions in estimating uncertainty. Figure 1 of the paper is an intuition of the benefit of using joint predictions.

They introduce a reference distribution $P_{z}$ from which the epistemic index $z$ is sampled and used to express epistemic uncertainty.

An ENN architecture, on the other hand, is specified by a pair: a parameterized function class $f$ and a reference distribution $P_{z}$ . The vector-valued output $f_{\theta} (x, z)$ of an ENN depends additionally on an epistemic index $z$, which takes values in the support of $P_{z}$ . Typical choices of the reference distribution $P_{z}$ include a uniform distribution over a finite set or a standard Gaussian over a vector space. The index $z$ is used to express epistemic uncertainty. In particular, variation of the network output with $z$ indicates uncertainty that might be resolved by future data.

In the context of classification, Osband takes $\tau$ inputs $x_1, ... , x_\tau$ and considers the joint prediction $$\hat{P}_{1:\tau}(y_{1:\tau})$$ which assigns a probability to each class combination $y_1, ... , y_\tau$.

The benefit, with respect to conventional neural networks, should be the following:

• Conventional neural networks are not designed to provide joint predictions. They can only be obtained by multiplying predictions of single data points together (assuming each data point as independent) $$\hat{P}_{1:\tau}^{NN}(y_{1:\tau}) = \prod_{t=1}^\tau softmax(f_\theta (x_t))_{y_t}$$
• ENNs create more expressive joint predictions integrating over the epistemic indices $$\hat{P}_{1:\tau}^{ENN} (y_{1:\tau}) = \int_z P_Z(dz) \prod_{t=1}^\tau softmax(f_\theta (x_t, z))_{y_t}$$

What I don't understand

• How can the choice of the reference distribution be so arbitrary?
• What is the intuition behind integrating over the epistemic indices in the formula for $\hat{P}_{1:\tau}^{ENN} (y_{1:\tau})$ ?

The paper continues and describes the architecture more in detail, but any help in understanding the intuition behind the above formulas would be very appreciated. Thank you.

Answer 1

Regarding the typical choices of the reference distribution $P_z$ as a uniform distribution over a finite set or a standard Gaussian over a vector space, they're similar idea as Bayesian neural networks (BNNs) to account for epistemic uncertainties in the model's inferred parameters but their ENNs implementation is supposed to be much more efficient than BNN ensemble particles.

When $P_z$ is a uniform distribution over a finite set, it essentially means that each element in the set corresponds to a distinct neural network configuration or weight initialization or some other hyperparameter values. A finite set allows for a manageable number of models making it computationally feasible to work with the ensemble.

When $P_z$ is a standard Gaussian over a vector space, it means that the parameters of the neural network are sampled from a mathematically convenient Gaussian distribution which has maximum entropy given a fixed variance and thus represents a more uninformative prior for the usual continuous parameters.

Finally for the intuition behind integrating out the epistemic index $z$ of ENNs is just to get marginal probability distribution from its conditional distribution by integrating out the conditioning random variable. Note the LHS marginal joint prediction $\hat{P}_{1:\tau}^{ENN} (y_{1:\tau})$ is essentially a (joint) "marginal" class distribution $p(y_{1:\tau}|x_{1:\tau})$, and the RHS term $\text{softmax}(f_\theta (x_t, z))_{y_t}$ is nothing but its corresponding "conditional" class distribution $p(y_{1:\tau}|x_{1:\tau},z)$ depending on any possible value of the epistemic index $z$ as defined in the support of $P_z$. By interpreting this way you can see clearly the whole equation is just marginalizing out the conditioning variable $z$.