Hi-fi functional priors by learning activations

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-05 13:40:46

Abstract: Function-space priors in Bayesian Neural Networks (BNNs) provide a more intuitive approach to embedding beliefs directly into the model’s output, thereby enhancing regularization, uncertainty quantification, and risk-aware decision-making. However, imposing function-space priors on BNNs is challenging. We address this task through optimization techniques that explore how trainable activations can accommodate higher-complexity priors and match intricate target function distributions. We investigate flexible activation models, including Pade functions and piecewise linear functions, and discuss the learning challenges related to identifiability, loss construction, and symmetries. Our empirical findings indicate that even BNNs with a single wide hidden layer when equipped with flexible trainable activation, can effectively achieve desired function-space priors.

Main paper: 
 introduction: 
 Models trained in a function-space rather than in the space of weights and biases (parameters space) exhibit flatter minima, better generalization, and improved robustness to overfitting (Qiu et al., 2024) . Better properties achieved by focusing directly on the output space are particularly advantageous in scenarios where the relationship between parameters and function behavior is not straightforward, especially for Bayesian Neural Networks (BNNs). Moreover, function-space priors offer a direct method of specifying beliefs about the functions being modeled by BNNs, rather than just the parameters, leading to intuitive and often more meaningful representation of prior knowledge (Tran et al., 2022) .Finding accurate posteriors for deep and complex models such as BNNs is notoriously challenging due to their high-dimensional parameter spaces and complex likelihood surfaces. On the other hand, for single-hidden-layer wide BNNs, it has been recently shown that posterior sampling via Markov Chain Monte Carlo (MCMC) can be performed efficiently (Hron et al., 2022) . Moreover, research has demonstrated that the exact posterior of a wide BNN weakly converges to the posterior corresponding to the Gaussian Process (GP) that matches the BNN's prior (Hron et al., 2020) . This allows BNNs to inherit GP-like properties while preserving their advantages. For example, BNNs scale better to large datasets as they can also leverage deep learning methods, reducing computational burden compared to GPs. The relation between NNs and BNNs has been a topic of significant research interest, which we briefly discuss in Section A.A classic result by Neal (1996) ; Williams (1996) , later extended to deep NNs by Lee et al. (2017) ; Matthews et al. (2018) , shows that infinitely wide layers in NNs behave a priori like GPs by identifying the kernel of a GP that matches covariance of a (B)NN. The past studies were conducted for better understanding of NNs. We are interested in a similar setting, but our objective is the opposite: we aim to implement the function-space priors (in particular, GP-like behavior) in BNNs, by matching both their priors on parameters and activations. GP prior specification generally offers greater interpretability compared to the prior applied to the weights of a BNN. This is because the kernel clearly governs key characteristics of the prior functions, such as shape, variability, and smoothness.We provide an extended reasoning on BNNs mitigating GPs behaviour in Appendix B. Finding BNNs that exhibit the behavior of GPs is a notoriously difficult problem, with analytical solutions existing for only a few GP kernels. For example, Meronen et al. (2020) found a solution (a BNN's activation function) for the popular Matern kernel. To address the challenge of imposing function-space a priori behavior to BNNs, Flam-Shepherd et al. (2017 ), Flam-Shepherd et al. (2018) , and especially Tran et al. (2022) , attempted to find both parameters and priors on parameters by using gradient-based optimization. Their approach, which requires deep networks equipped with complex priors, presents an additional undisclosed challenge in learning posteriors. Instead, we demonstrate that by learning also activations, we can achieve faithful function-space priors using just a shallow, wide BNN. In this work, we establish a practical and straightforward alternative to finding closed-form solutions (activations) to impose function-space priors on BNNs.Results in the paper are supplemented in the Appendix, which covers Related Work (Section A), a discussion on the relationship between BNNs and GPs (Section B), challenges in optimization (e.g., identifiability; Section C), and details of the Experimental Setting followed by Additional Results. 
finding weights and activations to match function-space priors: 
 The covariance between two inputs x and x ′ of a (single-hidden layer wide) BNN isEQUATIONwhere l reads here as last (or output) and l -1 as input layer. The BNN matches, a priori, a GP with an appropriate kernel κ realised by the covariance cov. In this work, we focus on zero-centered i.i.d.priors for weights and biases, e.g., E[w ij ] = 0 and E[b] = 0. To ensure that the asymptotic variance neither vanishes nor explodes, marginal variances are assumed to be inversely proportional to layerwidth, i.e., V[w l ij ] = σ l w 2 H l .For the full discussion, please see Section B.We consider the inverse problem of imposing behavior akin to a GP(0, κ) on a BNN. The task is to identify hyperparameters in Eq. 1, i.e., priors on w, b (where f l-1 is also expressed via w l-1 and b l-1 ) and activation ϕ, to align them with the desired GP, such that on an input subspace (=index set)X , the covariance would match the kernel. Ideally,f l i ∼ GP(0, κ) [X ] , i.e., p nn (f l i ) [X ] = p gp (f l i ) [X ] .However, in practice we strive for the approximate similarity:p nn (f l i ) [X ] ≈ p gp (f l i ) [X ]. The matching is performed on an input subspace denoted here by the functional-space measurement set [X ](which can generally be defined in multiple ways (Sun et al., 2019) ) and is posed as an optimization task:p * nn = argmin pnn D(p nn (f l i )[X ], p gp (f l i )[X ]), where D is an arbitrary (but in our case differentiable) divergence measure. While p gp can be both sampled and evaluated (for input X , the GP reduces to a Gaussian distribution), p nn is known only implicitly, which means that we can sample from it but its density cannot be evaluated directly.We proceed by reparameterizing the prior distributions on model weights and biases. In particular, we use the basic zero-centered factorized Gaussians with learned variances, e.g., p(w|σ 2 w ), p(b|σ 2 b ).However, we assume additionally, parametric and differentiable activations ϕ(•|η). Thus, p nn is fully specified by λ = {σ, η}, leading to the final optimization objective:EQUATIONwhich we address through gradient-based optimization with respect to λ. This formulation involves deciding on the loss function D and the model for ϕ.The problem of learning activations for NNs involves designing functions that enable networks to capture complex and non-linear relationships effectively. In principle, any function ϕ : R → R can serve as an activation, but the choice significantly impacts the network's expressiveness and training dynamics. We explore several models for ϕ, including Rational (Pade) (Molina et al., 2019) , Piece-wise linear (PWL) 1 , and activations realized by a NN with a single narrow hidden layer and ReLU/SiLU own activations. These functions are efficient to compute and introduce desirable nonlinear properties, striking a balance between efficiency and the capacity to model intricate patterns.The loss D measures the divergence between two distributions over functions. Due to the implicit nature of p nn , it must be specified between sets of samples {f i } (we used 512 samples from each distribution) and approximated using Monte Carlo. We discovered that the standard losses, including KL-divergence, fail for this task. For example, in KL, estimating the empirical entropy term ( p nn (f ; ϕ) log p nn (f ; ϕ)df ) presents a challenge. Consequently, we followed Tran et al. (2022) and relied on the Wasserstein distance:W p (π, ν) = inf γ∈Γ(π,ν) X ×Y D(x, y) p γ(x, y)dxdy 1/p = sup |ψ| L ≤1 E π ψ(x) -E ν ψ(x) (3)Unlike Tran et al. (2022) , we used the 2-Wasserstein metric, which for Multivariate Gaussians (applicable in the case of GPs and wide BNNs, at least approximately) has a closed-form solution (Mallasto and Feragen, 2017) and can be efficiently computed based on results by Buzuti and Thomaz (2023) . If the assumptions cannot be fulfilled, we can always revert back to the nested optimization for |ψ| L in Eq. 3.3 Learning activations enables more accurate function-space priors.To map GP function-space priors to a BNN, we can: train its priors on parameters, train both priors and activation, or learn just the activation function. We empirically show that learning also activations provides better solutions to the inverse problem of imposing function-space priors on BNNs. Figure 1 illustrates the results of an extensive study where we compare the quality of the GP prior fit for various models of parameters' priors and activations. The target was GP (0, Matern(ν = 5/2, l = 1)). For each configuration, we conducted several training iterations and measured the final (converged) loss value multiple times. Generally, learning activations alone is not sufficient for good fits, but when combined with learning priors, it significantly improves the solutions. We evaluate models with trained parameter priors (denoted by w), activations (denoted by a), and both (denoted by a+w). Each label specifies whether a fixed activation (e.g., ReLU) or a specific activation model (e.g., Rational) was employed. Gaussian parameter priors were used by default. If not trained, we set variances to 1., and for the hidden layer, we normalized the variance by its width. The label w: Matern refers to a BNN with the closed-form (fixed) activation as derived by Meronen et al. (2020) .4 Do expressive function-space priors require networks to be deep? Tran et al. (2022) takes a practical approach to learning function-space priors. Instead of considering BNNs with wide layers, they postulate that a deep network as a whole can model a functional prior.They then tune its parameters to reflect this functional prior. This involves modeling complex priors on all weights considered jointly, for example, using Normalizing Flows (Rezende and Mohamed, 2015) to ensure sufficient fidelity. This approach reveals a challenge in finding posteriors for deep networks. While MCMC (Chen et al., 2014; Del Moral et al., 2006) algorithms are computationally expensive, approximate posteriors (Hoffman et al., 2013; Ritter et al., 2018) do not possess sufficient expressiveness to adapt to such complex priors. We demonstrate that a BNN with a single wide hidden layer, basic Gaussian priors, and learned activation can match or even outperform the results of Tran et al. (2022) , both in terms of priors and posteriors quality. For the description of the experimental setting, see Appendix D. We present the results in Fig. 2 and in Tab. 1 (in Appendix).5 Can learned activations match performance of the closed-form ones?Complexity of deriving suitable neural activations makes matching BNNs to GPs a formidable challenge. We instead propose a gradient-based learning method that adapts both activations and parameters priors. This approach departs from the traditional methods that rely on fixed, analytically derived activations, offering instead a flexible alternative that empirically can match or even surpass the performance of the more conventional approaches. The relation between NNs/BNNs has been a topic of significant research interest. For example, Meronen et al. (2021) explored periodic activation functions in BNNs to connect network weight priors with translation-invariant GP priors. Pearce et al. (2020) derived BNN architectures to mirror GP kernel combinations, showcasing how BNNs can produce periodic kernels. Furthermore, Karaletsos and Bui (2020) introduced a hierarchical model using GP for weights to encode correlated weight structures and input-dependent weight priors, aimed at regularizing the function space. Matsubara et al. (2021) proposed using ridgelet transforms to approximate GP function-space distributions with BNN weight-space distributions, providing a non-asymptotic analysis with finite sample-size error bounds. Finally, Tsuchida et al. (2019) extended the convergence of NN function distributions to GPs under broader conditions, including partially exchangeable priors. For a more detailed discussion of related topics see for example, Sections 2.3 and 4.2 in Fortuin ( 2021). Here, we included just a brief review of the selected works.Behavior akin to GPs has been observed or postulated within neural networks. In particular, let's consider a Multilayer Perceptron:f 0 i (x) = I j w 0 ij x j + b 0 i , i = 1 . . . H 0 ; f l i (x) = H l j w l ij ϕ(f l-1 j (x)) + b l i , i = 1 . . . H l (4)where x denotes I-dimensional input, andf L denotes O-dimensional output. In BNNs, z ij (x) = w l ij ϕ(f l-1 j (x)) is a r.v. and assuming common prior distributions for w and b, f l i (x) becomes a sum of i.i.d r.v.-s. Then, from Central Llimit Theorem follows p(f l i (x)) H l - →∞ ------→ N (µ(x), σ 2 (x)). For two different inputs x, x ′ we can write p(f l i (x), f l i (x ′ )) H l - →∞ ------→ N (µ l f [x, x ′ ], κ l f [x, x ′ ]) where we use [•]to denote that we limit the mean and covariance potentially specified on a bigger space to these two inputs. In particular, by Kolmogorov extension theorem, the prior distribution of functions induced in the l-th layer, weakly converges to a GP as p(f l i )H l - →∞ -----→ GP(µ l f , κ l f ).We focus on zero-centered i.i.d. priors for weights and biases, e.g., E[w l ij ] = 0 and E[b l ] = 0.Furthermore, to ensure that asymptotic variance neither vanishes nor explodes, marginal variances are assumed to be inversely proportional to layer width, i.e., V[w l ij ] = σ l w 2 H l . Then, following from the definition, the covariance for two inputs x and x ′ can be obtained as (we repeat here Eq. 1)cov(f l i (x), f l i (x ′ )) = κ l f (x, x ′ ) = σ l b 2 + σ l w 2 • E f l-1 j [ϕ(f l-1 j (x))ϕ(f l-1 j (x ′ ))]The expectation is taken over realisations of the r.v. f l-1 i . We writeE f l-1 j meaning E p(f l-1 j (x)),p(f l-1 j (x ′ )) to signify r.v.-s which are integrated out.MLPs consisting of several (wide) hidden layers were considered by Lee et al. (2017) and Matthews et al. (2018) . Such the network can be considered a compound GP, where outputs of each layer are modeled by a GP. Then, f l-1 i ∼ GP(0, κ l-1 ) and subsequent (l-th) layer pre-activations(f l j , f l k ) T ∼ N (0, Σ l •δ jk ) where Σ l = κ l f (x, x ′ ) = σ l b 2 +σ l w 2 •E f l-1 i ∼GP(0,κ l-1 ) [ϕ(f l-1 i (x))ϕ(f l-1 i (x ′ ))].For an even more basic network with only one wide hidden layer, as considered by Neal (1996); Williams (1996) , the behavior converges to that of a single individual GP, inheriting its capacity for arbitrary function approximation (Rasmussen and Williams, 2005) . The approximation capability will depend on the chosen activation function, the nature and scale of the prior distributions on the weights and biases, and the width of the layer. For this case, the r.v.-s f 0 i are expressed with w 0 and b 0 and the expectation in Eq.( 1) is taken over these variables asE f 0 i [•] = E w 0 ∼p 0 w ,b 0 i ∼p 0 b [•].Models with a larger number of hyperparameters and a higher degree of freedom allow for achieving more complex distributions and potentially better solutions to Eq.( 2). On the other hand, overparam-eterization and complication of shape of the optimization manifold, may create identifiability problems or increase the risk of getting stuck in a local optimum for the steepest-gradient optimization.To alleviate this, the optimization task in Eq.2 can be slightly simplified by appropriately fixing parameters in Eq.(1). In particular, bias b l i can be removed (set to 0) and variance of weights between layer l -1 and l can be set to V[w l ij ] = 1 H l effectively simplifying Eq.(1) toEQUATIONWe cannot fix the weights as we propose to do for biases as this would mean thatf l i (x) = f l j (x)effectively reducing such the layer to a single output.On the other hand, the problem of finding BNNs behaving like GPs (e.g. inverting covariance equation) in unidenfiable. There exist multiple solutions assuring similar quality of the final match. In particular, activations ϕ solving Eq.( 2) are not unique, regardless if p nn is given by Eq.( 1) or Eq.( 5) (below). For example:• The solutions are symmetric w.r.t activity values (y-axis), i.e., for ϕ ′ (f ) = -ϕ(f ), values of covariance given by Eq.( 1) or Eq.( 5) are not changed, simply because (-1) 2 = 1.• For p(f l-1 i (x)) symmetric around 0 (for example, Gaussians), activations with flipped arguments ϕ ′ (f ) = ϕ(-f ) result in the same covariances.• Scaling activations ϕ ′ (f ) = αϕ(f ) leads to the same covariances as scaling variances of the output weights as(σ l w ) ′ = α • σ l wGiven a sufficiently flexible model (like a neural network itself), one can learn to approximate any target activation function to an arbitrary degree of accuracy on a compact domain. This is in line with the universal approximation theorem. However, in practice there are multiple limitations and challenges. A model may require an impractically large number of parameters to approximate certain complex functions to a desired level of accuracy. More complex models may be harder to fit and may require more training data. Some functions might require high numerical precision to be approximated effectively, and even if a model can fit a target activation function on a compact set, it might not generalize well outside the training domain. Overall, one needs to consider factors like the gradient behavior (for backpropagation), computational efficiency, and numerical stability. These factors can limit the practicality of using certain models of activations.1-dimensional regression -comparison with Tran et al. (2022) In order to show the abilities of our approach on a standard regression task and for a fair comparison with another method optimizing Wasserstein distance, we follow the exact experimental setting presented in Tran et al. (2022) . We used a GP with RBF kernel (length scale=0.6, amplitude=1.0 and noise variance=0.1) as the ground truth.For a baseline, we used a method from Tran et al. (2022) consisting of a BNN (3 layers with 50 neurons each) with TanH activation function. We consider all three priors realizations presented in that work: simple Gaussian prior, hierarchical prior and prior given by a Normalizing Flow. On the other hand, our model configuration consists of a BNN with Gaussian priors centered in 0 with trained variances, where the activation function was modeled by a NN with a single hidden layer consisting of 5 neurons using SiLU activation. Posterior distributions were obtained using a HMC sampler.We find that our approach allows for achieving the same or better results than Tran et al. (2022) baseline both by visual and numerical comparisons. In general, we obtain better results in terms of distributional metrics as presented in Tab. 1 and getting better (visually) posterior distributions (see Fig. 2 ) without utilizing any specific activation function or computationally heavy prior realization (like, e.g., Normalizing Flow).For the comparison against Meronen et al. (2020) , we followed their experimental setting and used a GP with Matern kernel (ν = 5/2, length scale=1) as the ground truth. For a baseline we utilized the method by Meronen et al. (2020) , where the authors derived analytical activation to match the Matern kernel. Our model configuration consists of a BNN with Gaussian priors with trained variances, where the activation function was modeled by a NN with a single hidden layer consisting of 5 neurons using SiLU activation. Posterior distributions were generally obtained using a HMC sampler, except in the case of (Meronen et al., 2020) , where the original implementation employed MC Dropout to approximate the posterior. However, for completeness and fairness in comparison, we also generated results using a HMC-derived posterior for the model by Meronen et al. (2020) . Additional tests (see Tab. 2) were conducted on BNNs with fixed parameters priors (Default=Gaussian priors with variance of 1, and Normal=Gaussian priors normalized by the hidden layer width) and with activation functions including ReLU and TanH.In this section, we present the comprehensive numerical evaluation for the experiments from the main paper. The results for regression experiments are presented in Tab. 1, whereas the results for classification are presented in Tab. 2. 2022)). We compared both prior and posterior predictive distributions of functions over a range of X covering regions with and without data (to account also for overconfidence far from data). The lower, the better. The posteriors were obtained using the HMC sampler in all cases, except for (Meronen et al., 2020) , where the original code was used (however, results for this model with a HMC-derived posterior are also included below). The lower, the better. For our method, we present results using weights and priors pretrained for 1D inputs, as well as those trained on functional priors for 1D/2D inputs matched to the task on X (see Fig. 3 ). https://pypi.org/project/torchpwl/ 
 
 Appendix:  



Reviews: review: title: title: Good paper.
main_review: main_review: Strengths:

1. The paper is well organized, and the presentation is good.

2. The theoretical contribution of this paper is sufficient, which is a bright spot.

Weaknesses:

1. If more experimental results in the appendix can be added to the main text, the experimental part will be more complete.

2. At the end of this paper, the summary of the full text seems to be missing.

3. In the relevant work, the investigation of previous studies is insufficient.
title: title: Interesting ideas, acceptable presentation
main_review: main_review: The paper proposes a novel approach to learning activation functions by minimizing a loss function between the GP probability and the NN probability. Below are my major comments:

The paper lacks a discussion on how this approach compares to prior works that integrate domain knowledge, such as Dylan Sam et al.'s Bayesian Neural Networks with Domain Knowledge Priors (ArXiv:2402.13410).

Typically, researchers transition from GP to BNN when the input space becomes large. How does this approach scale with respect to the size of the input space?