summaries: The paper proposed a new class of neural process algorithm called neural bootstrapping attention for neural processes (NeuBANP). This method utilizes efficient Neural Bootstrapping (NeuBoots) to improve Bootstrapped Attentive Neural Processes (BANP) in capturing functional uncertainty. Authors show that NeuBANP achieves state-of-the-art performance in benchmark experiments including Bayesian optimization and contextual multi-armed bandits.

summaries: The authors introduce NeuBANP, an extension of B(A)NP that replaces the iterative prediction method of B(A)NP with a single prediction step following the idea behind the neural bootstrapper. This new model is computationally more efficient and able to produce more accurate uncertainty estimates as well as correctly model heteroscedastic uncertainties. The authors provide experimental results on nonparametric regression tasks, Bayesian optimisation, contextual multi-armed bandit tasks and image inpainting (in the appendix).

summaries: The paper proposes Bootstrapping Attentive Neural Processes (NeuBANP), which utilize random sum-to-one weight in the encoder following from Neural Bootstrapper (Shin et al., 2021). The authors argue that this utilization resolves the overfitting problem in attentive NPs, and the proposed NeuBANP is more efficient in computation and memory perspective. Various experiments, such as synthetic examples and contextual multi-armed bandit, are conducted to compare against previous works of NPs and GP.

summaries: The authors propose a new bootstrapping method in NP family called NeuBANP. The authors acknowledge the limitations of BANP and used the neural bootstrapping method to make development from BANP. While BANP does better functional uncertainty estimation than ANP, it carries a higher computational burden compared to ANP since it requires multiple computations of the encoder network, the adaptation layer, and additional heuristics. By incorporating neural bootstrapping to the bootstrapping procedure, NeuBANP is a more computationally efficient model, estimates functional uncertainty better than BANP, and also alleviates the overfitting problem that ANP and BANP carries. Experimental results show that it achieves state-of-the-art performance on stochastic optimization problems, including multidimensional Bayesian optimization and contextual multi-armed bandit.

abstractText: Neural Processes learn to fit a broad class of stochastic processes with neural networks. Modeling functional uncertainty is an important aspect of learning stochastic processes. Recently, Bootstrapping Neural Processes (B(A)NP) propose a bootstrap method to capture the functional uncertainty which can replace the latent variable in (Attentive) Neural Processes ((A)NP), thus overcoming the limitations of Gaussian assumption on the latent variable. However, B(A)NP conduct bootstrapping in a non-parallelizable and memory-inefficient way and fail to capture diverse patterns in the stochastic processes. Furthermore, we found that ANP and BANP both tend to overfit in some cases. To resolve these problems, we propose an efficient and easy-to-implement approach, Neural Bootstrapping Attentive Neural Processes (NeuBANP). NeuBANP learns to generate the bootstrap distribution of random functions by injecting multiple random weights into the encoder and the loss function. We evaluate our models in benchmark experiments including Bayesian optimization and contextual multi-armed bandit. NeuBANP achieves the best performance in the sequential decision-making tasks among NP methods, and this empirically shows that our method greatly improves the quality of functional uncertainty modeling.

authors:
Maximilian Balandat
Brian Karrer
Daniel R. Jiang
Samuel Daulton
Benjamin Letham
Andrew Gordon Wilson
Eytan Bakshy

title: Botorch: A framework for efficient Monte-Carlo bayesian optimization

venue: In Advances in Neural Information Processing Systems,

authors:
Yutian Chen
Matthew W. Hoffman
Sergio Gómez Colmenarejo
Misha Denil
Timothy P. Lillicrap
Matt Botvinick
Nando de Freitas

title: Learning to learn without gradient descent by gradient descent

venue: In International Conference on Machine Learning,

authors:
Gregory Cohen
Saeed Afshar
Jonathan Tapson
André van Schaik

title: Emnist: Extending MNIST to handwritten letters

venue: In International Joint Conference on Neural Networks,

title: Better bootstrap confidence intervals

venue: Journal of the American statistical Association,

authors:
Chelsea Finn
Pieter Abbeel
Sergey Levine

title: Model-agnostic meta-learning for fast adaptation of deep networks

authors:
Alexandre Galashov
Jonathan Schwarz
Hyunjik Kim
Marta Garnelo
David Saxton
Pushmeet Kohli
S.M. Ali Eslami
Yee Whye Teh

title: Meta-learning surrogate models for sequential decision making

authors:
Jacob Gardner
Geoff Pleiss
Kilian Q. Weinberger
David Bindel
Andrew Gordon Wilson

title: Gpytorch: Blackbox matrix-matrix gaussian process inference with gpu acceleration

authors:
Marta Garnelo
Dan Rosenbaum
Christopher Maddison
Tiago Ramalho
David Saxton
Murray Shanahan
Yee Whye Teh
Danilo Rezende
S.M. Ali Eslami

title: Conditional neural processes

authors:
Jonathan Gordon
Wessel P. Bruinsma
Andrew Y.K. Foong
James Requeima
Yann Dubois
Richard E. Turner

title: Convolutional conditional neural processes, 2020

authors:
Hyunjik Kim
Andriy Mnih
Jonathan Schwarz
Marta Garnelo
Ali Eslami
Dan Rosenbaum
Oriol Vinyals
Yee Whye Teh

title: Attentive neural processes

venue: In International Conference on Learning Representations,

title: Benchmark functions for bayesian optimization

authors:
Jungtaek Kim
Seungjin Choi

title: BayesO: A Bayesian optimization framework in Python

venue: https: //bayeso.org,

authors:
Diederik P. Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: In International Conference on Learning Representations (Poster),

authors:
Balaji Lakshminarayanan
Alexander Pritzel
Charles Blundell

title: Simple and scalable predictive uncertainty estimation using deep ensembles

authors:
Tuan Anh Le
Hyunjik Kim
Marta Garnelo
Dan Rosenbaum
Jonathan Schwarz
Yee Whye Teh

title: Empirical evaluation of neural process objectives

venue: In Advances in Neural Information Processing Systems Workshop on Bayesian Deep Learning,

authors:
Juho Lee
Yoonho Lee
Jungtaek Kim
Eunho Yang
Sung Ju Hwang
Yee Whye Teh

title: Bootstrapping neural processes

authors:
Ziwei Liu
Ping Luo
Xiaogang Wang
Xiaoou Tang

title: Deep learning face attributes in the wild

venue: In International Conference on Computer Vision, pp

authors:
Christos Louizos
Xiahan Shi
Klamer Schutte
Max Welling

title: The functional neural process

venue: arXiv preprint arXiv:1906.08324,

authors:
Eric Nalisnick
Padhraic Smyth

title: The amortized bootstrap

venue: In International Conference on Machine Learning

authors:
Michael A. Newton
Adrian E. Raftery

title: Approximate Bayesian inference with the weighted likelihood bootstrap

venue: Journal of the Royal Statistical Society: Series B (Methodological),

authors:
Ian Osband
Charles Blundell
Alexander Pritzel
Benjamin Van Roy

title: Deep exploration via bootstrapped DQN

venue: arXiv preprint arXiv:1602.04621,

authors:
Jens Præstgaard
Jon A Wellner

title: Exchangeably weighted bootstraps of the general empirical process

venue: The Annals of Probability,

title: Gaussian processes in machine learning

venue: In Summer School on Machine Learning,

authors:
Sachin Ravi
Alex Beatson

title: Amortized bayesian meta-learning

authors:
Carlos Riquelme
George Tucker
Jasper Snoek

title: Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling

authors:
Adam Santoro
Sergey Bartunov
Matthew Botvinick
Daan Wierstra
Timothy Lillicrap

title: Metalearning with memory-augmented neural networks

authors:
Amr Sharaf
Hal Daumé III

title: Meta-learning for contextual bandit exploration

venue: arXiv preprint arXiv:1901.08159,

authors:
Minsuk Shin
Lu Wang
Jun S Liu

title: Scalable uncertainty quantification via generativebootstrap sampler

venue: arXiv preprint arXiv:2006.00767,

authors:
Minsuk Shin
Hyungjoo Cho
Hyun-seok Min
Sungbin Lim

title: Neural bootstrapper

authors:
Gautam Singh
Jaesik Yoon
Youngsung Son
Sungjin Ahn

title: Sequential neural processes

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Lukasz Kaiser
Illia Polosukhin

title: Attention is all you need

authors:
Lin(dh
do) ◦ (ReLU ◦ Lin(dh
dh))l ◦ ReLU

title: Unlike the case of ANP and BANP, considering that overfitting does not occur in NP and BNP, we can notice that attention modules are the leading cause of overfitting. B.1 SIMPLE LINEAR REGRESSION We experimented using a linear function rather than a function sampled from the GP

heading: 1 INTRODUCTION

text:
Neural Processes (NP) (Garnelo et al., 2018b) define distributions over functions given a set of observations, and are trained via a meta-learning framework so that it can adapt to new functions rapidly. NP learns to model a wide range of stochastic processes and can estimate the uncertainty over the predictions with less computational effort, compared to Gaussian Processes (GP) (Rasmussen, 2003). However, NP frequently suffers from a fundamental drawback of underfitting. As a remedy of the underfitting issue, Attentive Neural Processes (ANP) (Kim et al., 2018) applies the attention modules to the encoder network. Despite this modification, a single Gaussian latent variable of (A)NP has a limitation in inducing functional uncertainty (Louizos et al., 2019), a global uncertainty that decides the distribution over the space of trajectories or functions.
Appropriate modeling of functional uncertainty in stochastic processes improves the predictive performance and diversity in function realizations (Le et al., 2018), thus provides a principled way to guide agents to find optimal candidates in sequential decision-making problems. In these tasks, a model needs to approximate a function and estimate uncertainty correctly to optimize a black-box function whose analytic information is not given. Although GP is widely used for these tasks, these are the promising area for the application of NP because GP is computationally expensive, and it can be hard to choose an appropriate prior. Recently, Bootstrapping Neural Processes (B(A)NP) (Lee et al., 2020) modify (A)NP to induce more robust uncertainty estimation by employing the residual bootstrapping. However, B(A)NP underperforms in capturing a functional uncertainty because the residual bootstrapping works in a homoscedastic way, removing the connection between the feature and the label in its bootstrapped samples. The bootstrap strategy used in B(A)NP demands a higher computational burden compared to (A)NP since it requires multiple computations of the encoder network and additional heuristics, including the adaptation layer and the lower bound on the variance1. Furthermore, ANP and BANP tend to overfit for a simple regression task rather than learning the underlying heteroscedasticity. We observed that this problem is a by-product of the attention modules
1We colored the revised or added sentences in blue only for the rebuttal. This color will not appear in the final version of the manuscript.
used in both models (see Figure 1 and 5). This finding suggests that effective regularization of attention modules is required to prevent overfitting. Additionally, as explained in Section 4, ANP and BANP often tend to estimate homogeneous uncertainty regardless of whether the observation is given.
To resolve the problems of previous NPs, we introduce a novel bootstrapping method for neural processes, Neural Bootstrapping Attentive Neural Processes (NeuBANP). Motivated from the recent work on efficient bootstrapping of the neural network, Neural Bootstrapper (NeuBoots) (Shin et al., 2021), we introduce bootstrapping of the attention in a computationally efficient way by simple modification of the input of attention modules and the loss function, instead of memory-inefficient resampling and contrived heuristics employed in BANP. The simplicity and computational efficiency of NeuBANP directly come from NeuBoots, but it does not guarantee the performance in modeling the functional uncertainty since NeuBANP operates on the meta-learning framework. Thus, we modify the method for training the model to learn the randomness present in the underlying function, allowing the model to estimate random functions generated by any stochastic process. This modification is simple but gives a strong consistency between the bootstrapped samples and the representations from the attention modules, unlike the residual bootstrapping used in BANP. Besides, our bootstrapping method, which implements the concatenation and multiplication of random bootstrap weights, operates as a regularizer on the attention networks, thus preventing overfitted predictions observed in previous attention-based NP methods. NeuBANP is trained to generate a valid predictive distribution by utilizing the uncertainty inherent in observations by bootstrapping instead of the uncertainty that depends on the prior assumption on the latent variable as in (A)NP. This leads to the success in capturing heteroscedasticity of the data and modeling functional uncertainty in stochastic processes. As a result, NeuBANP achieves the best performance in sequential decision-making problems such as Bayesian optimization and contextual multi-armed bandit. The experimental results demonstrate that our model provides promising capabilities as an efficient neural approximation of stochastic processes.
Contributions We propose NeuBANP, an easy-to-implement and computationally efficient method for bootstrapping ANP. The proposed method has a novelty in learning a generator for bootstrapping stochastic processes under a meta-learning framework. NeuBANP resolves overfitting problem of attention modules and shows robust performance on heteroscedastic models without heuristics like the extra adaptation layer in BANP. NeuBANP estimates functional uncertainty better than BANP and achieves the best performance in stochastic optimization problems, including multi-dimensional Bayesian optimization and contextual multi-armed bandit, compared to previous NP methods.

heading: 2 PRELIMINARIES

heading: 2.1 META-LEARNING FRAMEWORK OF NEURAL PROCESSES

text:
Consider dataD = (X,Y ) = {(xi, yi)}ni=1 ⊂ X ×Y , the pairs of inputs xi ∈ X and outputs yi ∈ Y . Let P be a probability distribution over functions f ∈ F ; yi = f(xi) + i where i∼N (0, σ2i ), hence P determines the distribution of D. For disjoint index sets C and T satisfying C ∪ T = [n], define context DC = (XC , YC) = {(xc, yc)}c∈C and target DT = (XT , YT ) = {(xt, yt)}t∈T , so that D = DC ∪ DT . The task is to learn the neural processes pθ that fits the stochastic processes
f ∼ P given DC when C ∼ Pn is a randomly chosen subset of [n], as follows:
θ? = argmin θ
Ef∼P [ EC∼Pn [ − log pθ(Y |X,DC) ]] . (1)
This meta-learning framework allows pθ to learn diverse patterns in F via a prior distribution P; hence NP can predict target points conditioned on contexts adaptively in the inference phase.

heading: 2.2 (BOOTSTRAPPING) ATTENTIVE NEURAL PROCESSES

text:
ANP ANP is a variant of NP equipped with attention modules in the encoder part. See Appendix A for the detailed definition of attention operations. Let the context DC is given, and the model aims to infer y for a given feature x ∈ X . The encoder network of ANP maps (x,DC) into a pair of representation vectors r = (z,h) as follows:
{sc}c∈C = SelfAttn(DC), sC = mean({sc}c∈C), z ∼ N (z|µz(sC), σ2z(sC)) (2) {hc}c∈C = SelfAttn(DC), h = CrossAttn(x,XC , {hc}c∈C) (3)
Here, µz and σz are single linear layers that map sC to mean and standard deviation of the latent variable z, respectively. In ANP, z and h refer to latent path and deterministic path, respectively. The deterministic path models the overall skeleton of the encoder network, while the latent path models the functional uncertainty using a stochastic global latent variable (Garnelo et al., 2018b; Kim et al., 2018). Then the decoder network takes the representations r and target data x as inputs to predict µ(x, r) and σ(x, r), the parameters of the conditional predictive distribution p(y|x,DC) = N (y|µ, σ2).
BANP The global latent variable in ANP potentially limits the flexibility in expressing functional uncertainty. BANP proposes a method that utilizes paired bootstrapping and residual bootstrapping to model stochasticity in a data-driven way. First, they resample pairs of (xc, yc) with replacement to construct D̂(b)C = (X̂ (b) C , Ŷ (b) C ). Here b = 1, . . . , B implies the number of bootstrap samples. The resampled context is encoded into the representation vector r̂(b) = (ẑ(b), ĥ(b)) by replacing DC in (2) and (3) by D̂(b)C . Then BANP conducts the residual bootstrapping to construct bootstrapped contexts again D̃(b)C = (X̃ (b) C , Ỹ (b) C ). By using DC and D̃ (b) C in (2) and (3) separately, BANP gets the representations of contexts (r, r̃(b)). Finally, BANP uses an adaptation layer to merge (r, r̃(b)) and obtain µ(b), σ(b) through the decoder network (see Figure 2).
p(y|x, r, r̃(b)) = N (y|µ(b), (σ(b))2), p(y|x,DC) ≈ 1
B B∑ b=1 p(y|x, r, r̃(b)). (4)
Due to the residual bootstrapping, BANP conducts the encoder computation three times and the decoder computation twice for a single forward propagation. These additional calculations cause the computational bottleneck in the training and inference (see Appendix B.6).

heading: 2.3 NEURAL BOOTSTRAPPER

text:
Repetitions of training restrain the practical use of bootstrap procedures in deep neural networks due to their high computational burden. To alleviate this, Shin et al. (2021) proposes Neural Bootstrapper (NeuBoots). NeuBoots circumvents multiple training of networks by learning a bootstrap generator. Random Weight Bootstrapping Let ∑ c∈C `(f(xc), yc) be the loss function of interest for a neural network f . In standard bootstrap procedures, a bootstrapped neural network can be obtained by minimizing the loss function weighted by a random bootstrap weight wC := {wc : c ∈ C}:
L(wC , f,DC) = ∑ c∈C wc`(f(xc), yc). (5)
According to the choice of the distribution on wC , various bootstrap procedures can be represented under the form of (5); e.g., the paired bootstrap by wC ∼ Multinomial(n; 1/n, . . . , 1/n) and the Random Weight Bootstrapping (RWB) (Præstgaard & Wellner, 1993; Newton & Raftery, 1994) by wC ∼ |C| × Dirichlet(1, . . . , 1). NeuBoots utilizes RWB to avoid the data discard problem which can occur in the standard bootstrapping. We then compute bootstrapped neural networks {f̂ (b) : b = 1, . . . , B} via minimization of (5) for sampled w(1)C , . . . ,w (B) C .
Learning To Generate Bootstrap Distribution The main idea of NeuBoots is to construct a single generative network that models the bootstrapped neural networks with varying bootstrap weights in (5). This formulation modifies the backbone network in a form of f(x,wC) that inputs both feature x and bootstrap weight wC . Shin et al. (2021) show that the minimizer of the following loss generates valid bootstrap evaluations that match the results of the standard bootstrap procedure:
L(f,DC) = EwC∼|C|×Dirichlet(1,...,1) [L(wC , f(·,wC),DC))] , We call this weighted bootstrapping loss. Once this generator is trained via a single optimization procedure, we can efficiently generate bootstrapped predictions by plugging random bootstrap weights in the trained generator; i.e., for a feature of interest x, the trained generator inputs {w(b)C }Bb=1 and produces bootstrapped predictions ŷ(b) = f̂(x,w(b)C ) for b = 1, . . . , B.

heading: 3 NEURAL BOOTSTRAPPING ATTENTIVE NEURAL PROCESSES

text:
We propose a novel class of NP, called Neural Bootstrapping Attentive Neural Processes (NeuBANP). Aligned with the previous formulation of NP families, we can define our model pθ as:
pθ(Y |X,DC) = ˆ pϕ(Y |X,h, z)q(z|DC)dz = ˆ n∏ i=1 pϕ(yi|xi,h, z)q(z|DC)dz (6)
≈ n∏ i=1 pϕ(yi|xi,h, z) where z ∼ q(z|DC). (7)
Here pϕ is the decoder and q denotes the posterior distribution. Since the above integral is intractable, we approximate the predictive distribution by sampling z from the bootstrap distribution q(·|DC), instead of Gaussian distribution as in (A)NP. Precisely, we train a generative encoder network gφ, which outputs bootstrapped representation pairs (h, z) = {(h(b), z(b))}Bb=1 = gφ(X,DC , {w (b) C }Bb=1). Through the meta-learning framework (1), our model learns to generate bootstrapped predictions for an arbitrarily given function f ∼ P . Thus, our approach can be regarded as a learn-to-bootstrap method for stochastic processes. Compared to the fixed bootstrap method used in BANP, our learnable bootstrap method can find the best strategy to appropriately generate the random functions regarding the given context and the general property of the underlying data generating process. This will lead to the better modeling of functional uncertainty of the target stochastic processes. Also, note that NeuBANP can obtain a number of bootstrapped predictions by simply plugging different bootstrap weights into gφ, while BANP needs repetitive data resampling from scratch. Figure 2 shows the difference between the forward computation paths of BANP and NeuBANP in detail.

heading: 3.1 NEURAL BOOTSTRAPPING CONTEXTS WITH ATTENTION MODULES

text:
To train a generative network, which outputs the bootstrapped representations, we modify the encoder in ANP to take both (x,DC) and bootstrap weight wC as inputs. We introduce the posterior and the prediction paths in the encoder, analogous to NP’s latent and deterministic paths. NeuBANP is designed to leverage RWB, which is theoretically proven as a valid bootstrap method. In detail, the posterior and prediction path in NeuBANP take the random bootstrap weight as an auxiliary input. This input provides enough randomness into the network so that our model can successfully capture the functional uncertainty.
Posterior path We tag each context (xc, yc) ∈ DC with a bootstrap weight w(b)c to construct bootstrapped contexts D(b)C := {(xc, yc, w (b) c )}c∈C . Then the posterior path receives D(b)C as input and outputs a latent variable z(b). In detail, we apply self-attention to D(b)C and multiply the resultant representation z(b)C by the bootstrap weight w (b) C before mean aggregation. This path connects the contexts with weighted bootstrapping loss during the training and allows the bootstrap weights to model bootstrapped posterior distribution q by controlling the magnitude of each context representation.
Prediction path The prediction path outputs a target-specific representation h(b) that is relevant for the prediction. We apply self-attention to contexts DC and multiply the bootstrap weight w(b)C
element-wisely to compute another bootstrapped representation. Cross-attention module uses XC and this bootstrapped representation as key-value pairs to which the target query x attends. Consequently, reflecting the bootstrap weights shared with the posterior path, the prediction path models interactions between the given context and the target input.
We summarize the above posterior and prediction paths as follows:
z (b) C = {z (b) c }c∈C = SelfAttn(D (b) C ), z (b) = mean(z(b)C w (b) C ) (8)
hC = {hc}c∈C = SelfAttn(DC), h(b) = CrossAttn(x,XC ,hC w(b)C ) (9)
where denotes element-wise multiplication. We concatenate random bootstrap weights to context data, unlike NeuBoots, which only utilizes random weight multiplication in the final layer. Concatenation of weights and the contexts in the posterior path yields two effects. First, concatenating random information in given contexts provides sufficient randomness to the model, allowing it to cover a wide range of function samples in a space where a true function is likely to exist. Second, we made this modification to maximize the use of random weights and provide bootstrapping information to the model without resampling the data, enabling better uncertainty estimation. Additionally, by multiplying the random weights to the representations from both paths, as in (8) and (9), the representations are consistent with the weights and maximize bootstrapping effect. We think that these multiplications also provide the effect of regularization, which mitigates the overfitting tendency which ANP and BANP show in a simple regression experiment (see Figure 1).
Decoder and Bootsrapped Prediction The decoder takes (z(b),h(b)) from the encoder and target x ∈ X as inputs to generate a prediction ŷ(b) = MLP(x, z(b),h(b)). We can generate bootstrap samples ŷ(1), ŷ(2), ..., ŷ(B) by plugging w(1)C ,w (2) C , ...,w (B) C into the encoder gφ(x,DC , ·), respectively. We estimate the predictive mean and standard deviation using bootstrap samples:
µ = 1
B B∑ b=1 ŷ(b), σ = √√√√ 1 B − 1 B∑ b=1 (ŷ(b) − µ)2 (10)
In the previous NP methods, the decoder directly outputs the parameters of the predictive distribution, but the decoder of NeuBANP outputs the stochastic predictions {ŷ(b)}Bb=1. This design of output is a distinction of our model from other NP methods and allows the nonparametric estimation. Due to this structure, existing NPs set the lower bound of standard deviation for robust performance.
We found that the performance was susceptible to the lower bound value. However, NeuBANP naturally obtains parameters of predictive distribution using bootstrap predictions and shows better performance without such heuristics. It also has the advantage of calculating higher-order statistics without changing the structure.

heading: 3.2 TRAINING

text:
Weighted Bootstrapping Loss We train NeuBANP with weighted loss similar to that of NeuBoots as demonstrated in Section 2.3. Only context data has the corresponding bootstrap weight in our setting, but the model still has to fit target data. Thus we designed loss function as a sum of weighted context loss Lcontext and non-weighted target loss Ltarget as follows:
Ltotal = 1
B B∑ b=1 1 |C| ∑ c∈C ( − w(b)c logN (yc|µc, σ2c ) ) ︸ ︷︷ ︸
Lcontext
+ 1 |T | ∑ t∈T ( − logN (yt|µt, σ2t ) ) ︸ ︷︷ ︸
Ltarget
(11)
where µc, σc, µt, and σt are computed by (10) for context (xc, yc) ∈ DC and target (xt, yt) ∈ DT . Averaging weighted context loss with multiple bootstrap weights improves the robustness of a model by showing various bootstrap samples during training. We trained the model with negative log-likelihood (NLL). One can replace the NLL with a different loss function, such as cross-entropy according to the target tasks.

heading: 4 EXPERIMENTS

text: We conducted experiments to compare NeuBANP with the previous NP methods for regression and sequential decision-making problems. For regression tasks, we conducted one-dimensional (1D) regression experiments on random functions generated from GP prior and image completion tasks as two-dimensional (2D) regression (see Appendix B.5 for image completion). For sequential decision-making problems, we evaluated each method in Bayesian optimization (BO) and Contextual Multi-Armed Bandit (CMAB).

heading: 4.1 NONPARAMETRIC REGRESSION AND UNCERTAINTY ESTIMATION

text:
Settings We followed the settings in Lee et al. (2020). To obtain meta-training datasets, we sampled batches of random functions from GP prior with RBF kernel, and context and target points were chosen randomly from each function. In addition, kernel parameters of GP were randomly sampled so that the models could learn about various functions. NeuBANP was trained with 10 bootstrap samples, and we confirmed that it is robust to the number of samples. Please refer to Appendix B.2 for details.
Results The numerical results are summarized in Table 2. ANP and BANP tend to estimate homogeneous uncertainties for all target points in a situation where context points are sufficiently given (see Figure 3), because they place the heuristic lower bound on the variance. In the case of BANP, homoscedasticity occurs even when the number of context is small. Another problem only occurs in BANP, which is that it does not properly estimate functional uncertainty. Since the uncertainty about the shape of the true function is high in the region where the context point is not given, models should generate a wide range of function samples. However, we can see that BANP generates almost the same functions, not like those of the other methods, so we argue that BANP is not an appropriate method for modeling functional uncertainty. NeuBANP resolves these problems efficiently and achieves the best performance except for target prediction in the Periodic kernel.

heading: 4.2 BAYESIAN OPTIMIZATION

text:
Since NP can approximate a class of arbitrary functions, it can replace GP, the commonly used surrogate model of BO. It is crucial to approximate the objective function from the given observations using the surrogate model, but evaluating the acquisition function and determining the subsequent samples are vital for efficient exploration. We evaluated the proposed method on various black-box functions, which may be unobserved in the meta-training step (see Algorithm 1).
Settings For 1D, we followed the same setting in Lee et al. (2020). We set objective functions generated from GP with RBF, Matérn 5/2, and Periodic kernels and applied the models trained in Section 4.1. Furthermore, we demonstrated the BO performance of NeuBANP for multi-dimensional settings (2D and 3D). We set objective functions as various benchmark functions used in the optimization literature (Kim, 2020; Kim & Choi, 2017). See Appendix B.3 for details. A simple regret measured the performance of each model, and the mean performance over 100 experiments is reported for reliable evaluations.
Results NeuBANP outperformed the other methods in BO experiments (see Figure 4 and 8). Table 3 and 4 shows numerical results. For multi-dimensional BO, NeuBANP achieved the best results for every function except the Hartmann-3D when using Upper Confidence Bound (UCB) as an acquisition function. For the Hartmann-3D, the performance of NeuBANP is statistically comparable to the best method (ANP) when considering the standard deviation. For Goldsteinprice and Rastrigin functions, GP records poor performance due to its numerical errors during the optimization procedure (see Appendix B.3). The rightmost plots of Figure 4 show the time complexity according to the number of observation points. NeuBANP has the fastest decreasing rate of regrets in terms of iterations in both cases. We also conducted BO experiments using Expected Improvement (EI) to test the performance of NeuBANP to be independent of the selection of acquisition functions (Figure 9). Figure 10 demonstrates the better exploration strategy of NeuBANP compared to BANP. The second
and fourth columns show the explored points by BANP and NeuBANP, respectively. In most cases, NeuBANP converged to the optimum faster than BANP. The third and fifth column shows the contour plots of acquisition functions of each method. Note that NeuBANP accurately infers the potential area of the optimum compared to BANP. As explained above, NeuBANP can estimate heterogeneous uncertainty, and thus it is able to handle well for the exploration-exploitation trade-off, which is essential in sequential decision-making problems.

heading: 4.3 CONTEXTUAL MULTI-ARMED BANDIT

text:
We conducted a CMAB experiment, the wheel bandit problem, as in Garnelo et al. (2018b) to show that NeuBANP works efficiently as well as BO based on its performance of uncertainty estimation.
Settings We followed the same environment in Garnelo et al. (2018b), but used UCB policy. For more details, please refer to Appendix B.4. The parameter δ determines the environment of the wheel bandit problem. As δ increases, high-reward observation becomes sparse, which makes the problem more difficult. We set the baselines of CMAB experiment to MAML (Finn et al., 2017), Neural Linear (Riquelme et al., 2018) and NP. We measured cumulative regrets and simple regrets for 2,000 iterations to demonstrate the performance of each model.
Results NeuBANP performed well for various δ values (see Table 1). Note that NeuBANP showed better performance in challenging environments with sparse high rewards. NeuBANP showed the ability to learn various reward distributions based on appropriate functional uncertainty modeling. The result demonstrates that NeuBANP can utilize a small number of context information and make accurate estimations. NeuBANP also has stable results as its small variance shows. On the other hand, ANP has extremely high variance in the performance because it overfits to certain prediction and failed to adapt to the various bandit environments.

heading: 5 RELATED WORK

text:
Neural Processes CNP (Garnelo et al., 2018a) uses a pair of encoder and decoder networks to produce a predictive posterior distribution over functions. NP (Garnelo et al., 2018b) introduces a global latent variable to embed functional uncertainty in the deterministic architecture of CNP and predicts various outputs given the same context data. ANP (Kim et al., 2018) then enhances predictive accuracy by replacing MLP modules in NP with the attention modules. BNP (Lee et al., 2020) proposes the bootstrap method to model uncertainty in stochastic processes without the assumption of a single latent variable on which previous methods rely. In addition to these works, there are many attempts to use NPs in various tasks. Singh et al. (2019) tackles sequential stochastic processes, where the
dynamics of the given system changes as the time being. Leveraging time-variant context points, Singh et al. (2019) models the underlying temporal 3D structures. Gordon et al. (2020) extends NP families to contain translation equivariant functions, providing theoretical formulation to represent translation-invariant functional representations. Louizos et al. (2019) do not assume explicit global latent variables, instead supported by dependency graph among local latent variables, to encode inductive bias for given data easier than NPs that use global latent variables.
Bootstrapping Neural Networks Bootstrap method (Efron, 1987) is a reliable approach to estimate predictive uncertainty (Lakshminarayanan et al., 2017; Osband et al., 2016). However, it is computationally inefficient to go through the feed-forward computation as much as the number of bootstraps; hence, it discourages the practical application of bootstrap in neural networks. There have been several works to circumvent this issue by approximating bootstrapped distribution. Amortized bootstrap (Nalisnick & Smyth, 2017) approximates bootstrap distribution over model parameters by using amortized inference and implicit models. Generative Bootstrap Sampler (Shin et al., 2020) proposes a computational bootstrap procedure that constructs a generator function of bootstrap evaluations for classical statistical models. Neural Bootstrapper (Shin et al., 2021) suggests a simple recipe for generating bootstrapped predictive distributions of MLPs and convolutional neural networks.
Meta-Learning based Stochastic Optimization NPs are trained with a meta-learning framework to solve various tasks related to the data generation process through a single optimization. Santoro et al. (2016); Chen et al. (2017) follow the same training procedure of NP. However, Santoro et al. (2016) proposes an memory-augmented network for robust meta-learning, while Chen et al. (2017) proposes the method to produce an algorithm for black-box optimization using recurrent networks. Sharaf & Daumé III (2019) presents a meta-learning algorithm for learning a good exploration policy in the contextual bandit. Ravi & Beatson (2019) also solves contextual bandits based on the Bayesian framework by inferring a posterior on weights of neural networks. Galashov et al. (2019) introduces a unified framework for applying NP to a wide range of sequential decision-making problems.

heading: 6 CONCLUSION

text: We have proposed NeuBANP, a novel bootstrap method for a family of NP to model functional uncertainty appropriately. Instead of the standard bootstrap, NeuBANP learns to construct a generator function that produces valid bootstrapped distribution without resampling which can be considered as a learn-to-bootstrap method. NeuBANP successfully worked in a meta-learning framework, providing diverse trajectories of underlying data-generating processes, consistent to any given context. In addition, NeuBANP estimates the local uncertainty accurately, resolving overfitted prediction and variance overestimation problems observed in both ANP and BANP. We replace the additional layer and repetitive computations in BANP with the simple attachment of bootstrap weights to the model, which leads to lower computations and smaller memory. NeuBANP shows superior performance to previous NP methods in regression and stochastic optimization tasks, including multi-dimensional setting, which has been the desired application of NP. However, due to the numerical instability of GP, there is a limit in sampling high-dimensional stochastic processes for the meta-learning framework. We suggest that this is a primary task for scalable applications of NP in stochastic optimization, as a challenging research direction.

heading: 7 REPRODUCIBILITY

text: For reproducibility of experimental results, we provide a link to anonymous github that contains our source code in Appendix B. The source code includes the implementation of our model, data generation, and experiments. Additionally, the data generation steps are thoroughly explained in Appendix B.

heading: A MODEL

text:
A.1 BASIC OPERATION
Muti-Layer Perceptron MLP(di, dh, do, nl) denotes the multi-layer perceptron consisting of nl linear transformations and ReLU activations between them. Let parameters di, dh, do the dimension of input, hidden and output feature, repectively. We used same dh for every hidden layers. Lin(p, q) denotes the linear transformation of input with feature dimension p into output with feature dimension q.
MLP(di, dh,do, nl)(X) = Lin(dh, do) ◦ (ReLU ◦ Lin(dh, dh))nl−2 ◦ ReLU ◦ Lin(di, dh)(X) (12)
Dot Product Attention DotProdAttn (Vaswani et al., 2017) denotes the attention operation with attention score based on cosine similarity. Dot product of query (Q) and key (K) tensors calculates the similarity of each vectors in query tensor relative to each vectors in key tensor. Attention score is calculated through softmax operation of normalized similiarity. Let dk the feature dimension of key and query tensors, and dv that of value tensors.
DotProdAttn(Q,K, V ) = softmax ( QTK/ √ dk ) V (13)
Multi-Head Attention MHA denotes multi-head attention with dot product attention. The input is pre-processed with MLP and the output of attention is post-processed with layer normalization. For simplicity, we omit parameters of each MLP layers. split(X,n) denotes splitting of tensor X with respect to feature axis into a tuple of n tensors (X ′i) n i=1 which have same dimension in feature axis. [X1, X2, ..., Xn] denotes the concatenation of tensors X1, X2, ..., Xn with respect to feature axis. LayerNorm denotes the layer normalization introduced in Ba et al. (2016). Let nhead the number of heads in multi-head attention.
(Q′i) nhead i=1 = split(MLPqk(Q), nhead) (14) (K ′i) nhead i=1 = split(MLPqk(K), nhead) (15) (V ′i ) nhead i=1 = split(MLPv(V ), nhead) (16) MHA(Q,K, V ) = LayerNorm([(DotProdAttn(Q′i,K ′ i, V ′ i )) nhead i=1 ]) (17)
Self-Attention We used self-attention to calculate efficient representations of context DC = (XC , YC). We define self-attention based on multi-head attention as follows:
D′C = [XC , YC ] = ([xc, yc])c∈C , SelfAttn(DC) := MHA(D′C ,D′C ,D′C) (18)
Cross-Attention We used cross-attention to calculate representation of context specific to target feature x in interest, when original representations (hc)c∈C is given. We define cross-attention based on multi-head attention as follows:
CrossAttn(x,XC , (hc)c∈C) := MHA(x,XC , [(hc)c∈C ]). (19)
A.2 ARCHITECTURE
For fair comparison, we used the same architecture of all models as in Lee et al. (2020). For NeuBANP, we increased the input dimension of SelfAttn in the posterior path by one to take bootstrap weight as additional input. Please refer to Lee et al. (2020) for detailed model architecture of ANP and BANP.
A.3 NON-ATTENTIVE CASE
As an ablation study, we consider the non-attentive case of NeuBANP, called Neural Bootstrapping Neural Processes (NeuBNP). Like the architecture of NeuBANP is based on ANP, the architecture of this model is based on NP, the non-attentive counterpart of ANP. With D(b)C = {(xc, yc, w (b) c }c∈C as defined in 3, we applied the similar strategy of using random weights in the encoder as follows, and trained with the same loss function:
z (b) C = {z (b) c }c∈C = MLP(D (b) C , z (b) = mean(z(b)C w (b) C ) (20)
h (b) C = {h (b) c }c∈C = MLP(D (b) C , h (b) = mean(h(b)C w (b) C ) (21)
(22)
To make the non-attentive counterpart of NeuBANP, we use random weights in both paths, resulting in two latent variables z(b), h(b). They induce randomness into the decoder output ŷ(b) = MLP(x, z(b),h(b)), and the predictive distribution is construct by (10). The results in Table 2 and Table 3 shows that this model performs worse than NeuBANP, but outperforms BNP and NP, showing the quality of functional uncertainty the model learns with the bootstrap.

heading: B EXPERIMENTS

text:
Implementation2 of NPs other than NeuBANP was borrowed from the source code of Lee et al. (2020)3. Regression and Bayesian optimization experiment was done in single GeForce RTX 2080 Ti GPU with the memory of 11, 019 MiB. Multi-dimensional regression including image completion was done in Tesla V100 GPU with the memory of 32, 480 MiB.
2https://anonymous.4open.science/r/neubanp_initial 3https://github.com/juho-lee/bnp, MIT License.
B.1 SIMPLE LINEAR REGRESSION
We experimented using a linear function rather than a function sampled from the GP to analyze how well the NP models work in a simple regression task. Additionally, we need to examine how well the NP models predict uncertainty in the presence of heterogeneous noise in the data, and we add (x) ∼ N (0, σ2(x)), σ(x) = √ x2 + 10−5 to the linear function. We multiplied the noise by coefficient β to follow the meta-learning framework, and the β value was uniformly randomly sampled from [0.1, 1.0] during training. In this task, NP, BNP, and NeuBANP were able to estimate the underlying true linear function and the uncertainty of heterogeneous noise. On the other hand, in case of ANP and BANP, overfitting occurred and failed to predict the true function and its uncertainty (see Figures 1 and 5). As explained in the main text, we can see that the overfitting phenomenon occurs because of the attention mechanism that has appeared to solve the underfitting issue of the NP. NeuBANP has the advantage of the attention mechanism and achieved better performance through regularization using random weights. For (A)NP and B(A)NP, we used the official code provided by Lee et al. (2020). The remaining training settings are the same as the 1D regression in Appendix B.2, and we changed only the training iteration to 10, 000.
B.2 1D REGRESSION
Training For all models, training dataset consists of randomly sampled context and target from functions following GP with RBF kernel k(x, y) = s2 · exp(−||x − y||2/(2l2)). Parameters of kernel are randomly sampled with s ∼ Uniform(0.1, 1.0) and l ∼ Uniform(0.1, 0.6). Feature values (xi)i∈C∪T is chosen uniformly at random in [−2, 2]. The size of context and target are randomly sampled with |C| ∼ Uniform(3, 47) and |T | ∼ Uniform(3, 50 − |C|). We trained all models for 100,000 iterations and used Adam optimizer (Kingma & Ba, 2015). For stable learning, we used the cosine annealing scheduler with initial learning rate 5× 10−4.
Results Table 2 shows log-likelihood of NPs for various evaluation dataset sampled from GP with RBF, Matérn 5/2, Periodic kernel. As in generation of training dataset, parameters of Matérn 5/2 kernel k(x, y) = s2(1 + √ 5||x − y||2/(3l2)) exp(− √ 5||x − y||/l) and Periodic kernel k(x, y) = s2 exp(−2 sin2(π||x − y||2/p)/l2) was randomly sampled with s ∼ Uniform(0.1, 1.0), l ∼ Uniform(0.1, 0.6) and p ∼ Uniform(0.1, 0.5). For RBF and Matérn 5/2 dataset, NeuBANP showed state-of-the-art performance both in fitting context and predicting target. We added figures showing the predictions of ANP, BANP, and NeuBANP for the Matérn 5/2 and Periodic kernel (see Figure 6 and 7). In the case of the Matérn kernel, the two problems described in the main text occurred identically for ANP and BANP (see Section 4.1). However, in the case of Periodic, we can see that all models failed to approximate the true function correctly. This result came out because we experimented with testing the models’ generalization performance when trained with the RBF kernel. And the quantitative results in Table 2 show that the prediction performance of the attention-based
model on the target data is poor. Among them, NeuBANP has the worst performance, and we think the reason is the lower bound on the predicted variance set by ANP and BANP. Quantitative results show similar predictions for all models, but ANP and BANP achieve numerically more robust performance by setting a lower bound on the variance. We find some cases with the jumps in the function values which do not seem like a smooth function (See Figure 3, 6, and 7). This phenomena occurs in every attentive models including ANP, BANP, and NeuBANP. It looks like a distorted prediction on particular region, however, since our model predicts high variance in such region, this does not raise a problem in predicting the global trend, as we can see in high average log likelihood in 1d regression.
B.3 BAYESIAN OPTIMIZATION
One-dimensional case Table 3 shows the performance of GP and various NPs for 1D Bayesian optimization task. RBF, Matérn 5/2, and Periodic kernels are used to generate evaluation dataset.
Algorithm 1: Neural Process based Bayesian Optimization. Input :Target function f?; Acquisition function U ; Observed data D0 = {(x0, f?(x0))};
Multi-dimensional case In multi-dimensional BO experiments, we used GPyTorch 4 (Gardner et al., 2018) for scalable GP regression, and BoTorch 5 (Balandat et al., 2020) for overall BO process (e.g., optimization of acquisition functions). GP was set to the default setting of BoTorch. In detail, GP model was parameterized with Matérn 5/2 kernel with ARD and constant mean function, and prior distribution for hyperparameters was set as Gamma(3, 6) for length scale l and Gamma(2, 0.15) for output scale s. For three-dimensional BO experiment in Figure 4, the overall time complexity of ANP and BANP is almost the same. This result is seemingly in contrast with the fact that ANP takes a shorter time for prediction than BANP. However, since the BO algorithm contains the optimization of the acquisition function, the qualities of acquisition functions obtained by the model predictions may affect the overall time complexity. We conjecture that BANP gives an acquisition function easier to optimize than that of ANP so that the overall time complexities of both models are similar. Additionally, we did not report the results of GP for the two functions. Specifically, we omitted the Goldstein-Price
4https://github.com/cornellius-gp/gpytorch, MIT License. 5https://github.com/pytorch/botorch, MIT License.

heading: GP (RBF) 0.016 ± 0.052 0.048 ± 0.206 0.104 ± 0.242

text:
since the simple regret value of GP was too large (the performance was poor) compared to other models and omitted the Rastrigin since the 46 errors occurred out of 100 experiments. An error may occur when ill-conditioned data is given during the kernel training process of the GP, and the data recommended by the UCB during the exploration (or exploitation) process seems to correspond to this condition. If we report the performance ignoring the numerical error, the simple regret for the Goldstein-Price was 52049.13, and the simple regret for the Rastrigin was 18.48. For the Rastrigin function, GP with UCB is numerically unstable, but like EI, it achieved the best performance.
B.4 CONTEXTUAL MULTI-ARMED BANDIT
Setting We followed wheel bandit setting introduced in Riquelme et al. (2018). At every step t (< T ), a two-dimensional point (xt, yt) inside the unit circle Runit = {(x, y) ∈ R2 : x2+ y2 ≤ 1} is given as a context ct. The algorithm chooses an action at ∈ {1, 2, 3, 4, 5}. The stochastic reward rt = r(at, ct) is sampled from the reward distribution. Let R1, R2, R3, R4, R5 ⊂ Runit the disjoint
sets (regions) of unit circle as follows:
R1 = {(x, y) : x2 + y2 < δ} (23) R2 = {(x, y) : δ ≤ x2 + y2 ≤ 1, x > 0, y > 0} (24) R3 = {(x, y) : δ ≤ x2 + y2 ≤ 1, x < 0, y > 0} (25) R4 = {(x, y) : δ ≤ x2 + y2 ≤ 1, x < 0, y < 0} (26) R5 = {(x, y) : δ ≤ x2 + y2 ≤ 1, x > 0, y < 0} (27)
where the constant δ determines the size of R1 relative to other regions. Each action results in rewards following different distribution according to the region to which the given context belongs, where N denotes the normal distribution.
r(1, c) ∼ N (1.2, 0.012) (28) r(a, c) ∼ { N (50, 0.012), if c ∈ Ra N (1, 0.012), otherwise ∀a ∈ {2, 3, 4, 5} (29)
Note that action {1} always produces a moderate reward, but the other actions {2, 3, 4, 5} sometimes produce a very high reward when the context is sampled from the corresponding high-reward region. Thus, learning different reward distributions for actions {2, 3, 4, 5} by apprehension of context information is critical to bandit performance. As δ increases, the high-reward regions for each actions become smaller. Then the model should learn from rare observation of such high reward, which means the problem becomes more difficult.
Training and Evaluation When pre-training NeuBANP, as in Garnelo et al. (2018b), 8 training batches of 512 contexts and 50 targets were generated from the environment with hyperparameter randomly sampled; δ ∼ Uniform(0, 1). We consider two-dimensional context point ci as feature xi ∈ R2 and five rewards for actions (r(1, ci), r(2, ci), r(3, ci), r(4, ci), r(5, ci)) as label yi ∈ R5. At evaluation, only rewards for chosen actions are observed by the model. Thus, we replace other unobserved rewards with dummy values randomly sampled fromN (0, 1), following the usual strategy.
B.5 IMAGE COMPLETION
Settings We compared the baseline NPs and NeuBANP on image completion tasks. Following Lee et al. (2020), we trained all models on EMNIST (Cohen et al., 2017) and CelebA (Liu et al., 2015) which was resized to 32 × 32. For EMNIST, we used only 10 classes for training and reported the evaluation results for both seen classes and unseen classes separately. NeuBANP was trained with 10 samples, and the other baselines were trained with 4 samples. For evaluation, we used 50 samples for all methods. Similar to 1D regression experiment, we randomly select the pixels of a given image as context/target, and the number of context/target were drawn randomly from Uniform distribution. However, in this case, we increased the maximum number of given points; i.e., |C| ∼ Uniform(3, 197), |T | ∼ Uniform(3, 200 − |C|). x values were rescaled to [−1, 1] and the corresponding y values were rescaled to [−0.5, 0.5]. We trained all models for 200 epochs and set a initial learning rate of 5× 10−4 using the Adam optimizer (Kingma & Ba, 2015) with cosine annealing scheduler for learning rate decay.
Method Seen classes (0-9) Unseen classes (10-46)context target context target
CNP 0.926 ± 0.007 0.751 ± 0.005 0.766 ± 0.009 0.498 ± 0.012 NP 0.948 ± 0.006 0.806 ± 0.005 0.808 ± 0.005 0.600 ± 0.009 ANP 1.383 ± 0.000 0.993 ± 0.005 1.383 ± 0.000 0.894 ± 0.004 BNP 1.004 ± 0.008 0.880 ± 0.005 0.883 ± 0.010 0.722 ± 0.006
BANP 1.383 ± 0.000 1.010 ± 0.006 1.382 ± 0.000 0.942 ± 0.005 NeuBANP 1.475 ± 0.345 1.337 ± 0.224 1.333 ± 0.516 1.119 ± 0.388
Table 5. Quantitative result of EMNIST image completion. Mean and standard deviationd of likelihoood over 5 experiments.
context target
CNP 2.975 ± 0.013 2.199 ± 0.003 NP 3.066 ± 0.011 2.492 ± 0.014 ANP 4.150 ± 0.000 2.731 ± 0.006 BNP 3.269 ± 0.008 2.788 ± 0.005
BANP 4.149 ± 0.000 3.129 ± 0.005 NeuBANP 13.946 ± 0.590 2.870 ± 0.021
Table 6. Quantitative result of CelebA results. Mean and standard deviationd of likelihoood over 5 experiments.
Results Figure 11 and 12 show the mean prediction and uncertainty estimation of ANP, BANP, and NeuBANP for test images in unseen classes. For both datasets, though our model shows noisy mean prediction due to the random weights, we can demonstrate the advantage of NeuBANP in uncertainty estimation. NeuBANP estimated the uncertainty correctly in the area where the color of the pixel changes and thus possesses significant uncertainty. This leads to the overall improvement in quantitative performance (see Table 5 and 6).
B.6 TIME COMPLEXITY
Settings We measured the time complexity empirically according to the number of context points, target points, and bootstrap samples. We fixed the number of targets to 25 and adjusted the number of contexts to 10, 20, 30, 40, and 50 to see how inference time varies with the number of contexts. Conversely, to see the inference time according to the number of targets, we fixed the number of context points to 25. We fixed the number of bootstrap samples to 50 as in the 1D regression experiment. When conducting experiments with the number of bootstrap samples, the number of context and target points was fixed at 20 and 25. All experiments were conducted with a batch containing 100 tasks.
Results BANP places a remarkably high computational cost in that the approach of bootstrapping the attention module is inefficient, as demonstrated in Figure 13. The inference time becomes noticeably longer as the number of context points increases. NeuBANP, on the other hand, learns to bootstrap efficiently; therefore, its time complexity is comparable to that of BNP, which does not use the attention module.


summaries: In this paper, the authors developed a transformer-based model, GeneBERT to align DNA sequences with regulatory elements. In particular, GeneBERT first applies transformers to learn representations of sequencing data and regulatory regions (e.g., open chromatin), and then aligns the representations of two modalities for identifying region-aligned sequences. The authors applied GeneBERT to recent scATAC-seq data in fetal and use aligned sequences to predict promoters, CTCF binding sites, a disease type and RNA splicing sites.

summaries: The manuscript describes GeneBERT, a self-supervised and multi-modal pre-training approach for genomic data. GeneBERT combines 1D genome sequence data with a 2D representation of regulatory elements in different cell-types in three different pre-training tasks. A large-scale single-cell ATAC dataset is used to identify pseudo-bulk ATAC profiles for 17 cell-types spanning > 1M genomic locations. The sequence of these genomic locations provides the 1D representation of the data which are used in self-supervised training with masked genome modeling and next genome segment prediction loss functions, both inspired by the BERT framework. A 2D representation is derived using accessibility per regulatory element and cell type, and is self-supervised using a infoNCE loss. The pre-training framework is then compared in a variety of biological tasks including promoter prediction, TFBS prediction, splice site prediction. The authors then present ablation studies highlighting the importance of the different components of the loss function.

summaries: The authors describe a pipeline to perform self-supervised learning on genome sequences, guided by accessible chromatin peaks. Their learning procedure is inspired by the NLP method BERT and uses several tasks from that work and its successors. After pre-training transformer layers, they fine-tune on several regulatory sequence classification tasks and demonstrate high performance.

summaries: This work proposes an approach, called GeneBERT, for pre-training genome data in a multi-modal and self-supervised manner. They take the 1d sequence of genomic data and a 2d matrix of (transcription factors × regions) as the input and optimize three pre-training tasks to improve the robustness and generalizability the model. Specifically, they introduce two main objectives for sequence pre-training, Masked Genome Modeling (MGM) and Next Genome-Segment Prediction (NGSP). GeneBERT consists of three main components: sequence pre-training, region pre-training, and sequence-region matching. The main contribution of the paper is that they use transcription factor information in genomic regions which makes the model more generalizable to other cell types than the previous DNABERT that only uses the sequence information.

abstractText: In the genome biology research, regulatory genome modeling is an important topic for many regulatory downstream tasks, such as promoter classification, transaction factor binding sites prediction. The core problem is to model how regulatory elements interact with each other and its variability across different cell types. However, current deep learning methods often focus on modeling genome sequences of a fixed set of cell types and do not account for the interaction between multiple regulatory elements, making them only perform well on the cell types in the training set and lack the generalizability required in biological applications. In this work, we propose a simple yet effective approach for pre-training genome data in a multi-modal and self-supervised manner, which we call GeneBERT. Specifically, we simultaneously take the 1d sequence of genome data and a 2d matrix of (transcription factors × regions) as the input, where three pre-training tasks are proposed to improve the robustness and generalizability of our model. We pre-train our model on the ATAC-seq dataset with 17 million genome sequences. We evaluate our GeneBERT on regulatory downstream tasks across different cell types, including promoter classification, transaction factor binding sites prediction, disease risk estimation, and splicing sites prediction. Extensive experiments demonstrate the effectiveness of multi-modal and self-supervised pre-training for large-scale regulatory genomics data.

authors:
Babak Alipanahi
Andrew Delong
Matthew T Weirauch
Brendan J Frey

title: Predicting the sequence specificities of dna-and rna-binding proteins by deep learning

venue: Nature biotechnology,

authors:
Žiga Avsec
Vikram Agarwal
Daniel Visentin
Joseph R. Ledsam
Agnieszka Grabska-Barwinska
Kyle R. Taylor
Yannis Assael
John Jumper
Pushmeet Kohli
David R. Kelley

title: Effective gene expression prediction from sequence by integrating long-range interactions

venue: Nature Methods,

authors:
Runpu Chen
Le Yang
Steve Goodison
Yijun Sun

title: Deep-learning approach to identifying cancer subtypes using high-dimensional genomic data

authors:
Ting Chen
Simon Kornblith
Mohammad Norouzi
Geoffrey Hinton

title: A simple framework for contrastive learning of visual representations

venue: In Proceedings of International Conference on Machine Learning (ICML),

authors:
Ting Chen
Simon Kornblith
Kevin Swersky
Mohammad Norouzi
Geoffrey Hinton

title: Big selfsupervised models are strong semi-supervised learners

venue: In Proceedings of Advances in Neural Information Processing Systems (NeurIPS),

authors:
Xinlei Chen
Haoqi Fan
Ross Girshick
Kaiming He

title: Improved baselines with momentum contrastive learning

venue: arXiv preprint arXiv:2003.04297,

authors:
Yifei Chen
Yi Li
Rajiv Narayan
Aravind Subramanian
Xiaohui Xie

title: Gene expression inference with deep learning

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: BERT: Pre-training of deep bidirectional transformers for language understanding

venue: arXiv preprint arXiv:1810.04805,

authors:
Silvia Domcke
Andrew J. Hill
Riza M. Daza
Junyue Cao
Diana R. O’Day
Hannah A. Pliner
Kimberly A. Aldinger
Dmitry Pokholok
Fan Zhang
Jennifer H. Milbank
Michael A. Zager
Ian A. Glass
Frank J. Steemers
Dan Doherty
Cole Trapnell
Darren A. Cusanovich
Jay Shendure

title: A human cell atlas of fetal chromatin

venue: accessibility. Science,

authors:
Alexey Dosovitskiy
Lucas Beyer
Alexander Kolesnikov
Dirk Weissenborn
Xiaohua Zhai
Thomas Unterthiner
Mostafa Dehghani
Matthias Minderer
Georg Heigold
Sylvain Gelly
Jakob Uszkoreit
Neil Houlsby

title: An image is worth 16x16 words: Transformers for image recognition at scale

venue: In Proceedings of International Conference on Learning Representations (ICLR),

authors:
Kaiming He
Haoqi Fan
Yuxin Wu
Saining Xie
Ross Girshick

title: Momentum contrast for unsupervised visual representation learning

venue: In Proceedings of IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
Chenyang Hong
Kevin Y Yip

title: Flexible k-mers with variable-length indels for identifying binding sequences of protein dimers

venue: Briefings in Bioinformatics,

title: Bridging vision and language by large-scale multi-modal pre-training

venue: arXiv preprint arXiv:2103.06561,

authors:
Kishore Jaganathan
Jeremy F. McRae
Siavash Fazel Darbandi
David Knowles
Yang I. Li
Jack A. Kosmicki
Juan Arbelaez
Wenwu Cui
Grace B. Schwartz
Eric D. Chow
Efstathios Kanterakis
Hong Gao
Amirali Kia
Serafim Batzoglou
Stephan J. Sanders
Kyle Kai-How Farh

title: Predicting Splicing from Primary Sequence with Deep Learning

authors:
Yanrong Ji
Zhihan Zhou
Han Liu
Ramana V Davuluri

title: Dnabert: pre-trained bidirectional encoder representations from transformers model for dna-language

venue: in genome. Bioinformatics,

authors:
David R Kelley
Jasper Snoek
John Rinn

title: Basset: Learning the regulatory code of the accessible genome with deep convolutional neural networks

venue: bioRxiv,

authors:
Aziz Khan
Oriol Fornes
Arnaud Stigliani
Marius Gheorghe
Jaime A Castro-Mondragon
Robin van der Lee
Adrien Bessy
Jeanne Chèneby
Shubhada R Kulkarni
Ge Tan
Damir Baranasic
David J Arenillas
Albin Sandelin
Klaas Vandepoele
Boris Lenhard
Benoı̂t Ballester
Wyeth W Wasserman
François Parcy
Anthony Mathelier

title: JASPAR 2018: update of the open-access database of transcription factor binding profiles and its web framework

venue: Nucleic Acids Research,

authors:
Yunhai Li
Kee Khoon Lee
Sean Walsh
Caroline Smith
Sophie Hadingham
Karim Sorefan
Gavin C. Cawley
Michael W. Bevan

title: Establishing glucose- and aba-regulated transcription networks in arabidopsis by microarray analysis and promoter classification using a relevance vector machine

venue: Genome research,

authors:
Yinhan Liu
Myle Ott
Naman Goyal
Jingfei Du
Mandar Joshi
Danqi Chen
Omer Levy
Mike Lewis
Luke Zettlemoyer
Veselin Stoyanov

title: Roberta: A robustly optimized bert pretraining approach

authors:
Ze Liu
Yutong Lin
Yue Cao
Han Hu
Yixuan Wei
Zheng Zhang
Stephen Lin
Baining Guo

title: Swin transformer: Hierarchical vision transformer using shifted windows

venue: arXiv preprint arXiv:2103.14030,

authors:
Mhaned Oubounyt
Zakaria Louadi
Hilal Tayara
Kil To Chong

title: Deepromoter: robust promoter predictor using deep learning

venue: Frontiers in genetics,

authors:
Yeping Lina Qiu
Hong Zheng
Olivier Gevaert

title: A deep learning framework for imputing missing values in genomic data

authors:
Alec Radford
Karthik Narasimhan
Tim Salimans
Ilya Sutskever

title: Improving language under-standing by generative pre-training

venue: URL https://s3-us-west-2. amazonaws. com/openaiassets/research- covers/languageunsupervised/language understanding paper.pdf,

authors:
Alec Radford
Jong Wook Kim
Chris Hallacy
Aditya Ramesh
Gabriel Goh
Sandhini Agarwal
Girish Sastry
Amanda Askell
Pamela Mishkin
Jack Clark
Gretchen Krueger
Ilya Sutskever

title: Learning transferable visual models from natural language supervision

venue: arXiv preprint arXiv:2103.00020,

authors:
Adriana Romero
Pierre Luc Carrier
Akram Erraqabi
Tristan Sylvain
Alex Auvolat
Etienne Dejoie
Marc-André Legault
Marie-Pierre Dubé
Julie G. Hussin
Yoshua Bengio

title: Diet networks: Thin parameters for fat genomic

authors:
Douglas T Ross
Uwe Scherf
Michael B Eisen
Charles M Perou
Christian Rees
Paul Spellman
Vishwanath Iyer
Stefanie S Jeffrey
Matt Van de Rijn
Mark Waltham
Alexander Pergamenschikov
Jeffrey C.F. Lee
Deval Lashkari
Dari Shalon
Timothy G Myers
John N Weinstein
David Botstein
Patrick O Brown

title: Systematic variation in gene expression patterns in human cancer cell lines

venue: Nature Genetics,

authors:
Mark Schena
Dari Shalon
Ronald W. Davis
Patrick O. Brown

title: Quantitative monitoring of gene expression patterns with a complementary dna microarray

authors:
Alexander J Stewart
Sridhar Hannenhalli
Joshua B Plotkin

title: Why Transcription Factor Binding Sites Are Ten Nucleotides Long

authors:
Yu Sun
Shuohuan Wang
Yukun Li
Shikun Feng
Xuyi Chen
Han Zhang
Xin Tian
Danxiang Zhu
Hao Tian
Hua Wu

title: Ernie: Enhanced representation through knowledge integration

authors:
Luis Torada
Lucrezia Lorenzon
Alice Beddis
Ulas Isildak
Linda Pattini
Sara Mathieson
Matteo Fumagalli

title: Imagene: a convolutional neural network to quantify natural selection from genomic data

venue: BMC Bioinformatics,

authors:
Hugo Touvron
Matthieu Cord
Matthijs Douze
Francisco Massa
Alexandre Sablayrolles
Hervé Jégou

title: Training data-efficient image transformers & distillation through attention

venue: arXiv preprint arXiv:2012.12877,

authors:
Li Yuan
Yunpeng Chen
Tao Wang
Weihao Yu
Yujun Shi
Francis EH Tay
Jiashi Feng
Shuicheng Yan

title: Tokens-to-token vit: Training vision transformers from scratch on imagenet

heading: 1 INTRODUCTION

text:
In the gene biology research, many important regulatory problems including promoter (Li et al., 2006) classification and transaction factor binding sites (Stewart et al., 2012) prediction requires a well-designed approach for modeling regulatory genome sequences. As a multi-cellular organism, the human body is formed by different cell types, with each has its own gene expression patterns (Schena et al., 1995; Ross et al., 2000). Therefore, modeling regulatory genome, especially across different cell types, is crucial for both the understanding of this fundamental biological process and the development of RNA-level therapeutic intervention of a vast among of diseases. The main challenge is how to model the interaction between regulatory elements and its variability across different cell types.
In recent years, some works (Chen et al., 2016; Kelley et al., 2016; Romero et al., 2017; Qiu et al., 2018; Torada et al., 2019; Chen et al., 2020a; Avsec et al., 2021; Ji et al., 2021) have been proposed to apply deep learning architectures on modeling the genome data. For example, Enformer (Avsec et al., 2021) combines dilated CNN and transformer architecture as well as multi-head output to predict gene expression and epigenomic marks, and gain some performance improvements against traditional methods. Most of these models are based on a supervised learning paradigm, which limits their abilities to learn the general interaction representations between different regulatory elements. As a result, each model is only useful for some specific downstream tasks. More recently, DNABERT (Ji et al., 2021) is introduced to formulate the whole DNA sequence as a sentence of nucleotide k-mers and utilize BERT to model the sequence generatively. In this way, the interactions between different regulatory elements could be well captured, like the language modeling approach in natural language processing (NLP) to capture the co-occur information between different words, and could be employed for modeling several related downstream tasks. However, DNABERT only performs well on the cell types in the training set, and generalizes poorly to unseen cell types. This
is mainly because the whole DNA sequence is common for different cells, and pre-training on such data cannot well reflect various interaction patterns across different cell types.
Integration of genome data modalities across different cell types could help to build a more holistic model of gene expression regulation and benefit downstream applications such as mutation impact evaluation and disease risk prediction, as well as promote our understanding of the cell-type-specific regulatory programs and various development processes and disease etiology. Inspired by this idea, in this work, we present a simple yet effective method called GeneBERT, for pre-training large-scale genome data in a multi-modal and self-supervised manner. Specifically, we simultaneously take the 1D modality (i.e. sequence) and a 2D modality (i.e. regulatory region) of genome data as the input, where three pre-training tasks are proposed to improve the robustness and generalizability of our model. 1) masked genome modeling: we randomly mask some parts of the input k-mers with a special token (i.e., [MASK]), and the model is trained to predict the masked k-mers. 2) next genome segment prediction: we train the model using the embedding [CLS] to classify whether a pair of given sequences are two consecutive sequences in a cell. 3) sequence-region matching: a sequenceregion matching mechanism is proposed to capture the multi-modal alignment between sequence and regulatory region of genome data.
We pre-train the GeneBERT on the ATAC-seq dataset (Domcke et al., 2020) with 17 million gene sequences. Furthermore, we conduct extensive experiments to evaluate our GeneBERT on various downstream tasks, including promoter prediction, transaction factor binding sites prediction, gene mutation localization, and personalized diseases prediction. Comprehensive ablation studies demonstrate the effectiveness of multi-modal self-supervised pre-training for large-scale genome data across different cell types.
The main contributions of this work lie in the proposal of a simple yet effective method named GeneBERT, for large-scale genome data pre-training in a multi-modal and self-supervised manner. To the best of our knowledge, we are the first to incorporate different genome data modalities across various cell types into the pre-training for large-scale genome data, to tackle the regulatory genome modeling problem. Except for meaningful biological improvements, our model makes an important contribution to the machine learning community, by introducing a novel multi-modality construction. Different from existing multi-modal learning tasks, such as visual question answering, image caption, and image-text retrieval, the ‘visual’ modality in our work is constructed based on the regulatory property of the ‘language’ sequential units. This idea brings some inspiration to building a new ‘visual’ modality, based on text matching, part-of-speech tagging, and named entity recognition etc., for the study of NLP.

heading: 2 RELATED WORK

text:
Uni-modal language pre-training. Self-supervised pre-training models such as GPT (Radford et al., 2018), BERT (Devlin et al., 2018), RoBERTa (Liu et al., 2019), and ERNIE (Sun et al., 2019) have led to dramatic improvement on a variety of natural language processing tasks in the past few years, significantly surpassing the traditional context-independent language model such as Word2Vec. Obviously, BERT (Devlin et al., 2018) uses the masked language model (MLM) and next sentence prediction (NSP) for pre-training, which represent the dynamic vector of words according to the context. RoBERTa (Liu et al., 2019) uses dynamic MLM and discards NSP to train the model with a long time. ERNIE (Sun et al., 2019) masks entities and phrases to learn more context relations. These self-supervised pre-training models are capable of performing natural language processing tasks more effectively by improving MLM.
Multi-modal pre-training. Multi-modal pre-training has recently addressed researchers’ attention for learning meaningful representations. Typically, Previous methods (Radford et al., 2021; Huo et al., 2021) learn visual representations from text paired with images in unsupervised, selfsupervised, weakly supervised, and supervised ways. Since language and vision can share a similar semantic meaning, CLIP (Radford et al., 2021) is a commonly-used neural network trained on a variety of (image, text) pairs for learning transferable visual representations from natural language supervision. WenLan (Huo et al., 2021) applies a cross-modal contrastive learning framework called BriVL for image-text pre-training. However, in this work, we leverage the multi-modal self-supervised pre-training on the genome data to improve the robustness and generalizability of pre-trained models used for data-scarce scenarios.
Genome data pre-training. Transformer models have been recently established to better understand the genotype-phenotype relationships (Avsec et al., 2021; Ji et al., 2021). DNABERT used the human genome to pre-train a BERT-based model, trying to decipher the regulatory code related to gene expression (Ji et al., 2021). In order to adapt the DNA scenario, sequences were split into 5 to 510 base-pair long and tokenized to 3- to 6-mer representations. The model was then pre-trained using those k-mers by predicting around 15% randomly masked regions of the sequence. After the pre-training, the model was fine-tuned on three downstream tasks related to gene regulation: prediction of promoters, transcription factor binding sites (TFBSs), and splice sites. Furthermore, by analyzing the attention maps, DNABERT could visualize the important regions contributing to the model decision, which improved the interpretability of the model.

heading: 3 METHOD: GENEBERT

text: In this section, we propose a simple yet effective approach for pre-training genome data in a multimodal and self-supervised manner, as illustrated in Figure 1. Specifically, we introduce two main objectives for sequence pre-training, including Masked Genome Modeling (MGM) and Next GenomeSegment Prediction (NGSP). Next, we extract features of the region for region pre-training. Finally, we present the Sequence-Region Matching mechanism to explicitly align features of the sequence and region in the input genome data in a contrastive self-supervised manner. Therefore, our GeneBERT consists of three main components: sequence pre-training, region pre-training, and sequence-region matching.

heading: 3.1 INPUT DATA

text:
We use potential regulatory elements mapped with ATAC-seq which measures the chromatin accessibility across the genome. The regions with open accessibility could be bound by transcription factors, which together with regulatory elements and RNA polymerase determines the expression level. In one cell type, the list of all regulatory elements can be viewed as a set of non-overlapped sub-regions across the genome. Each region or regulatory element can be viewed as a string like “AATTCCT...” consist of 4 elements: [A, T, C, G]. The binding site of different transcription factors has different sequence patterns, some of which have been determined using experimental methods and have been modeled using probabilistic models like position weight matrix (PWM). Using PWM, one can get the predicted affinity of a sequence segment to a given transcription factor. Each regulatory region can then be associated with a ‘transcription factor binding vector” which contains the in silico prediction of their affinity to a fixed number of transcription factors. By scanning through a regulatory element sequence with PWM of multiple transcription factors, we can get a transcription factor binding affinity vector for that regulatory element.
When the collection of transcription factor binding models is comprehensive and precise enough, we can get a good representation of regulatory potential for a given sequence. In different cell types, the accessibility of regulatory elements is different. Some regulatory elements are shared across different cell types while some are more specific to one cell type. As a result, in a local region of the genome, the accessible transcription factor binding sites are determined by the combination of accessible regulatory regions. This combinatorial interaction between regulatory regions and transcription factors varies across different cell types such that precise control of expression in different cell types is possible. To capture the combinatorial interaction and the implicit cell-type specific information, we then group 10 consecutive regulatory elements into one train sample by concatenating their string (1D sequence) and stacking their binding vector into a 2D region, as shown in Figure 1.

heading: 3.2 SEQUENCE PRE-TRAINING

text:
Similar to DNABERT, we take k-mer tokens as the sequence units. The k-mer refers to a sequence with length k, i.e., for a sequence AGTCAG, the 3-mers are {AGT, GTC, TCA, CAG}, and the 4-mers are {AGTC, GTCA, TCAG}. Therefore for sequence pre-training, we input three types of embeddings: 1) a k-mer embedding xk for each k-mer in a sequence; 2) a segment embedding xs indicating which part of the sequence the k-mer is from; 3) a position embedding xp for the position of the k-mer in the sequence. Then we sum up all three embeddings to obtain a contextual representation en, n ∈ {1, 2, ..., N}, where N denotes the number of k-mers in the sequence. After being fed into a BERT-based transformer fseq(·), those contextual embeddings become Eseq, i.e. the obtained representation for [CLS]. Inspired by BERT (Devlin et al., 2018), we adopt two objectives including Masked Genome Modeling (MGM)Lmgm and Next Genome Segment Prediction (NGSP) Lngsp. For Lmgm, we randomly mask some parts of the input k-mers with a special token (i.e., [MASK]), and the model is trained to predict the masked kmer. The objective is formulated as:
Lmgm = max P (Kmask|K) (1) where K and Kmask denote the unmasked and masked k-mers, respectively. As for Lngsp, we train the model using the embedding [CLS] to classify whether a pair of given sequences are consecutive in a cell. The loss function of Lngsp is calcualted as:
Lngsp = CrossEntropy(P ([CLS]), y) (2) where P ([CLS]) and y represent the predicted probability and target for classification, respectively.

heading: 3.3 REGION PRE-TRAINING

text: For the region features in the pre-training, we consider a strong backbone (i.e. Swin (Liu et al., 2021)) transformer 1 as the region encoder freg(·) to extract representations Ereg. Specifically, we apply the Swin transformer pre-trained on ImageNet to the region input directly to generate Ereg. During the pre-training, we unfreeze the parameters of the Swin transformer and update them for learning better regional representations. In the pre-training setting, each region input corresponds to each sequence such that we can capture the multi-modal alignment between sequence and region of genome data.

heading: 3.4 SEQUENCE-REGION MATCHING

text:
In order to learn the alignments between sequence and region of genome data, we propose a sequence-region matching mechanism based on sequence embeddings Eseq and region embeddings Ereg. Firstly, we calculate the similarity (i.e. inner product) between each pair of ‘linguistic’ embeddings Eiseq and ‘visual’ embeddings E i reg in a batch of size b, where i ∈ 1, 2, ..., b. Then, those similarities are jointly learned for alignments between the whole sequence and each region in a contrastive learning manner, where we maximize the similarity of the sequential and regional embeddings of the b correct pairs in the batch while minimizing the similarity of the embeddings of the b(b− 1) false pairs. Specifically, an info-NCE like loss defined as follows is applied over these similarities scores for optimization:
Lsrm = − log ∑b i=1 E i seq ·Eireg∑b
i=1 ∑b j=1 E i seq ·E j reg
(3)
1For other transformer strcutures, we explore them in the ablation study.
Algorithm 1 GeneBERT multi-modal pre-training algorithm Input: Sequence transformer fseq(·), Region transformer freg(·)
1: 1D-data: k-mer embedding xk, segment embedding xs, position embedding xp 2: 2D-data: Region matrix X 3: Initialize the parameters fseq(·), freg(·) with pre-trained weights 4: for each iteration step do 5: # Sequence Pre-training 6: Feed xk,xs,xp into fseq(·) to generate Eseq 7: calculate Lmgm and Lngsp in Eq. 1 and 2 8: # Region Pre-training 9: Feed 2D region X into freg(·) to generate Ereg
10: # Sequence-Region Matching 11: Apply embeddings Eseq,Ereg to calculate the loss in Eq. 3 12: Compute the total loss in Eq. 4 13: Update the parameters of fseq(·), freg(·)
Output: fseq(·), freg(·)
In this way, alignments between the whole sequence and each region are learned via our GeneBERT in the pre-training process.
Combining the above three losses together, we obtain the overall objective function for GeneBERT:
L = Lmgm + Lngsp + λ · Lsrm (4) where we vary λ in [0.01, 1] to perform the parameter study for λ, and observe that the performance of our model is stable when λ lies in the region of [0.5, 1]. Therefore in our experiments, we directly set λ = 0.5. The overall algorithm is summarized in Algorithm 1.

heading: 4 EXPERIMENTS

heading: 4.1 PRE-TRAINING DATA & SETTINGS

text: For pre-training data, we process public human fetal cerebrum single-cell chromatin accessibility data in the Descartes database (Domcke et al., 2020) to generate pseudo-bulk accessibility tracks for each cell type (Seurat cell clustering provided by the original paper). Specifically, we take the provided “Peak Count Sparse Matrices” and summed up columns (cells) according to cell type definition, producing a regions × cell-types matrix. Then we binarize the matrix and use only non-zero entries (accessible regions) for each cell type. The corresponding sequence for each region is then retrieved from the HG19 human reference genome, a single representation of multiple genomes. While the motif scanning for each region is either retrieved from the Descartes database or scanned following the same approach using JASPAR 2018 (Khan et al., 2018) vertebrate transcription factor binding site motifs. In total, we use 17 cell types and the union of all accessibility track includes 1,000,029 accessible regions across the genome, covering 504,657,456 base pairs. For the 1D modality, we group 10 consecutive accessible regions into one sample, which corresponds to a (10 × number of TFs) matrix for the 2D modality. Following previous works (Ji et al., 2021), we pre-train the model for 120k steps with a warm-up learning rate of 4e-4 and batch size of 2000. 15% of k-mers in each sequence are masked in the first 100k steps and 20% for the last 20k steps.

heading: 4.2 DOWNSTREAM TASKS & RESULTS

text:
Promoter Classification. Promoters are the elements responsible for regulating the initial transcription of the gene, which is located near the transcription start site (TSS). As the promoters play an important role in gene regulation, using machine learning methods to predict promoter sites accurately is one of the most popular problems in bioinformatics. Here we first used the promoter core dataset from (Ji et al., 2021), which are the 70bp sequences centered around TSS. The promoter core is the key part of the promoter flanking region which is sufficient to direct accurate initiation of transcription (Oubounyt et al., 2019). Here we fine-tune our GeneBERT model to predict the promoter core sequences and compared them with DNABERT. We report the experimental results
in Table 1. For a fair comparison with DNABERT, we reproduce the results of DNABERT by finetuning the pre-trained models provided by the authors. During the reproduction process, we observe that it is hard to get the same results reported in the original paper as their model is very sensitive to the hyper-parameter settings. From the comparison results, we can see that our model can predict promoter core more precisely with improved performance when compared to DNABERT. This is due to the inclusion of regulatory genome information in our 2D pre-training data, which guides the pre-training process to focus on those binding site motifs and their interactions with each other.
Transcription Factor Binding Sites (TFBS) Classification. Predicting TFBS is an important step in studying gene regulation. Sequencing technologies like ChIP-seq can provide information on the in vivo binding sequences, which improve the identification of gene regulatory regions. There are several previous studies that tried to predict TFBSs using traditional machine learning (Hong & Yip, 2020) and deep learning methods (Alipanahi et al., 2015). By incorporating the multi-modal pretraining, the prediction of TFBSs can be further improved. Although we utilize the motif information during the region pre-training, we do not provide any matching information to the model, which avoids leaking information about the actual motif of a specific TF. Here we fine-tune our model for predicting TFBSs from the ChIP-seq data, using 497 TF ChIP-seq uniform peak profiles from ENCODE Consortium (Consortium et al., 2012). Note that the data in this task is different from that reported in DNABERT, since they did not provide the data for this task. We reproduce the results of DNABERT on this TFBS data with the same setting as our GeneBERT. We take the peak sequences of each TF as the positive set and generated a corresponding negative set by randomly shuffling the nucleotides in each positive sequence while preserving dinucleotide frequencies. Again, our model can accurately predict the CTCF binding sites in all tested cell lines.
Disease Risks Estimation. GeneBERT could provide more interpretations of complex genetic diseases. On the one hand, while the disease status and genomic mutations were available, by integrating the 2D data, the relationships among regulatory regions of genes could be captured, which allowed us to estimate the disease risk more accurately. As shown in Table 3, GeneBERT can precisely predict Hirschsprung Disease (HSCR), which is known as a genetic disorder with complex patterns of inheritance. Particularly, our GeneBERT outperforms the baseline by a large margin in terms of all metrics (Precision, Recall, AUC), which further demonstrates the effectiveness of our approach in the treatment target sites and proceeded to the medical experimental validation.
RNA-Splicing Sites Prediction. RNA splicing is important for post-transcription processing to remove introns from pre-mRNA sequences and generate mature mRNA for protein translation. Previously, Dilated CNN models have been used to predict splice junction across the genome and evaluate the impact of genomics variants on splicing sites (Jaganathan et al., 2019). In particular, for each nucleotide in a given sequence for splicing site prediction, we follow the previous approach and include a context sequence around the nucleotide, which could potentially capture the sequence specificity features of RNA-binding proteins and splicing machinery. Since open chromatin regions and splicing sites do not always overlap with each other, among all 548,000 splicing sites in the GTEx pre-mRNA transcripts data, our pre-training sequence only fully covers the entire (in the 256nt context setting) sequence of 72,500 sites. The experimental results are reported in Table 4. In total, 26.7% of nucleotides in context and splicing site sequence are included in the open chromatin region that we use for pre-training. Following the same training/testing split scheme and classification metric (Top-k Accuracy and PR-AUC) as in the SpliceAI study (Jaganathan et al., 2019), we achieve comparable or better results in different context settings without an extremely long context sequence. These results clearly demonstrate the generalizability of our GeneBERT pre-trained representations. By integrating sequence binding features of RNA binding proteins, we are able to further expand our model to enable cell-type specific splicing junction prediction in the future.

heading: 5 ABLATION STUDY

text:
In this section, we perform extensive ablation studies on the effect of each module on the final performance of our GeneBERT, the effectiveness of modeling different cell types, the effect of 2D modality on pre-training, and the effect of 2D transformers on the final performance of our approach. Unless specified, we conduct all ablation studies on the promoter classification task and report the mean value of all results with 5 random seeds.
Effect of each loss. In order to explore the effect of each loss on the final performance of our GeneBERT, we ablate each loss in our method and report the comparison results of promoter classification in Table 5. As can be seen, pre-training without three proposed losses achieves the worst results in terms of all metrics. This means that pre-training without any objectives causes embeddings with insignificant genome information. Adding Lmgm first boosts the performance by 0.659, 0.625, 0.672 for precision, recall, and AUC. By combining Lmgm and Lngsp, we achieve the performance gain of 0.028, 0.024, 0.032, which validates the rationality of Lmgm and Lngsp on modeling the meaningful embeddings of genome sequences. In the meanwhile, adding Lsrm to the final objective further improves the results by 0.103, 0.142, and 0.169, in terms of all metrics. This further demonstrates the effectiveness of our sequence-region matching loss in learning the alignments between sequence and region of genome data. In this case, the sequence-region matching loss is helpful for multi-modal self-supervised pre-training on large-scale genome data, which boosts the performance of downstream tasks.
Generalization to different cell/tissue types. To validate the generalizability of our approach to various cell or tissue types, we evaluate our GeneBERT on Transcription Factor Binding Sites (TFBS) classification. Specifically, we choose different proteins from various cells or tissues in the human body, such as skin, lung, heart, brain, etc. The experimental results of TFBS classification on different cell/tissue types are reported in Table 6. In terms of modeling different cell/tissue types, DNABERT (Ji et al., 2021) performs poorly on the generalizability, as we reported in Table 2. Enformer Avsec et al. (2021) only performs well on the gene expression prediction for the single cell. From the results, we can observe that our GeneBERT achieves superior performance on various cell types and tissue types in terms of all metrics, including precision, recall, and AUC. This indeed shows the advantage of our GeneBERT on learning generalizable embeddings for different cell types and tissue types in the human body. Particularly, we achieve the best performance on the brain tissue and neuroblast cell, which might be caused by the pre-training data from the human brain.
Effect of 2D modality on pre-training. In order to further evaluate how the 2D modality data affect the quality of pre-training on large-scale genome data, we visualize the pre-trained representations pre-trained by the uni-modal and multi-modal manner in Figure 2. Typically, we project the embeddings pre-trained from BERT onto 2 dimensions using tSNE tools, and we select randomly 10 cell types in the pre-training data for visualization. We can observe that our GeneBERT pre-trained representations form more separated clusters that are distributed more uniformly on the space in terms of all cell types when compared to the pre-trained representations generated by the uni-modal model. Furthermore, the representations pre-trained by our multi-modal self-supervised method have much smaller distances inside each cluster of cell type. This further demonstrates the effectiveness of our 2D modality in improving the quality of pre-trained representations.
Effect of different transformer strcutures. In order to figure out why we choose the strong Swin transformer to take the 2D modality data as input, we explore the effect of different 2D transformers on the final performance of our GeneBERT. Specifically, we apply four commonly-used backbones
(ViT (Dosovitskiy et al., 2021), T2T-ViT (Yuan et al., 2021), DeiT (Touvron et al., 2020), Swin (Liu et al., 2021)) as the 2D transformer to extract the region features, and report the comparison results in Table 7. As can be seen, our GeneBERT with Swin-B-384 backbone achieves the best performance in terms of all metrics (precision, recall, AUC). In the meanwhile, with the increase of input size, the performance of our GeneBERT improves, which demonstrates the effectiveness of features extraction of the 2D modality data in learning discriminative representations of large-scale genome data. As the 2D transformer gets powerful, our GeneBERT achieves better results in terms of all metrics on the promoter classification task. This further validates the importance of the 2D transformer in multi-modal self-supervised learning. In this case, we choose the Swin transformer as the 2D backbone since we achieve the best results with the Swin transformer.
Effect of batch size. Batch size is always a significant hyper-parameter in the contrastive learning literature. To explore how the batch size affects the final performance of our multi-modal selfsupervised pre-training, we vary the batch size from 16, 32, 64, 128, 256, and 512. The experimental results are reported in Table 8. We can observe that with the increase of the batch size, our GeneBERT achieves better performance due to the effectiveness of more negative samples in the sequence-region matching loss. This demonstrates the importance of a larger batch size and more negative samples on large-scale genome data pre-training, which complies with the conclusion of vision tasks in (Chen et al., 2020b;c; He et al., 2020; Chen et al., 2020d).
Effect of λ. Furthermore, we explore the effect of λ on the final performance of our GeneBERT. Specifically, we vary λ from 0.01, 0.05, 0.1, 0.5, 1.0 and 5.0, and report the comparison results in Table 8. Our GeneBERT achieves stable performance when λ is set to 0.5 and 1. With the decrease of λ, the performance of our approach deteriorates in terms of all metrics, which validates the importance of our sequence-region loss in our GeneBERT. To receive the stable performance of multi-modal self-supervised pre-training, we set λ = 0.5 in all experiments.

heading: 6 CONCLUSION

text: In this work, we present the GeneBERT, a multi-modal self-supervised framework for large-scale genome data pre-training. Specifically, we leverage sequence pre-training, region pre-training, and sequence-region matching together to improve the generalizability of our model across different cell types for regulatory genome modeling. Extensive experiments on four main regulatory downstream tasks demonstrate the effectiveness of our GeneBERT. Besides, comprehensive ablation studies have been conducted to show the important role of multi-modal modeling.

title: REGULATORY GENOME ACROSS CELL TYPES