{"kind":"task","effective_mode":"full","benchmark":{"kind":"benchmark","effective_mode":"full","slug":"longbench-v2","formal_name":"LongBench v2","introduction":"LongBench v2 evaluates deep understanding and reasoning over long contexts through multiple-choice questions. Its official description lists 503 questions spanning tasks such as single-document and multi-document QA and code-repository understanding.","introduction_ja":"","introduction_en":"","category":"Category not supplied","task_count":null,"acquisition_status":"Acquisition status not supplied","official_url":"https://huggingface.co/datasets/zai-org/LongBench-v2","indexing_mode":"noindex","profile":{"resources":[],"task_format":"","scoring":"","metric":"","size":"","answer_access":"","license":"","citation":"","maintainer":"","released":"","why_hard":"","related":[]}},"task_id":"e762ab15-c493-5e73-9231-c781da747c34","task_key":"train--66f2ad2b821e116aacb2ac0f","task_revision_id":"3","upstream_id":"66f2ad2b821e116aacb2ac0f","short_description":"Which is not among the improvements of the method in paper GeoBFN compared with…","config":"","split":"train","body":"{\"choice_A\":\"Be able to process discrete input data like atom charges.\",\"choice_B\":\"Considering atom coordinates and types simultaneously.\",\"choice_C\":\"Equivariant to topological rotation of molecules.\",\"choice_D\":\"Adopt an early mode-seeking sampling strategy.\",\"context\":\"Bayesian Flow Networks\\n\\nAbstract\\nThis paper introduces Bayesian Flow Networks (BFNs), a new class of generative model in\\nwhich the parameters of a set of independent distributions are modified with Bayesian inference\\nin the light of noisy data samples, then passed as input to a neural network that outputs a\\nsecond, interdependent distribution. Starting from a simple prior and iteratively updating\\nthe two distributions yields a generative procedure similar to the reverse process of diffusion\\nmodels; however it is conceptually simpler in that no forward process is required. Discrete and\\ncontinuous-time loss functions are derived for continuous, discretised and discrete data, along\\nwith sample generation procedures. Notably, the network inputs for discrete data lie on the\\nprobability simplex, and are therefore natively differentiable, paving the way for gradient-based\\nsample guidance and few-step generation in discrete domains such as language modelling. The\\nloss function directly optimises data compression and places no restrictions on the network\\narchitecture. In our experiments BFNs achieve competitive log-likelihoods for image modelling\\non dynamically binarized MNIST and CIFAR-10, and outperform all known discrete diffusion\\nmodels on the text8 character-level language modelling task1.\\n1\\nIntroduction\\nLarge-scale neural networks have revolutionised generative modelling over the last few years, with an\\nunprecedented ability to capture complex relationships among many variables. Building a convincing\\njoint model of all the pixels in a high resolution image, for example, was impossible before the\\nadvent of modern generative networks.\\nKey to the expressive power of most of these networks — including autoregressive models\\ne.g. [9, 46], flow-based models [32], deep VAEs [49] and diffusion models [41] — is that the joint\\ndistribution they encode is broken down into a series of steps, thereby eluding the “curse of\\ndimensionality” that would doom any effort to explicitly define all the interactions among so many\\nvariables. In colloquial terms they solve a hard problem by splitting it into easy pieces.\\nA general way to view such distributions is as an exchange of messages between a sender, Alice,\\nwho has access to some data, and her friend Bob, who wishes to receive it in as few bits as possible.\\nAt each step Alice sends a message to Bob that reveals something about the data. Bob attempts\\nto guess what the message is: the better his guess the fewer bits are needed to transmit it. After\\nreceiving the message, Bob uses the information he has just gained to improve his guess for the\\nnext message. The loss function is the total number of bits required for all the messages.\\n1Code and trained models can be found at https://github.com/nnaisense/bayesian-flow-networks\\n1\\narXiv:2308.07037v5  [cs.LG]  3 Feb 2024\\n\\n\\nIn an autoregressive language model, for example, the messages are the word-pieces the text\\nis divided into. The distribution encoding Bob’s prediction for the first message is of necessity\\nuninformed: a zero-gram prior based on the relative frequencies of different word-pieces. The\\ntransmission cost is the negative log-probability under this prior. Bob then uses the first word-piece\\nto predict the second; on average, the second prediction will be slightly more informed than the first,\\nand the expected transmission cost will be slightly lower. The process repeats with the predictions\\nimproving at each step. The sum of the transmission costs is the negative log-probability of the\\ncomplete text sequence, which is the loss function minimised by maximum likelihood training. It\\nis also the minimum number of bits that would be required for Alice to transmit the pieces to\\nBob using arithmetic coding [53]. There is therefore a direct correspondence between fitting an\\nautoregressive model with maximum likelihood and training it for data compression.\\nAutoregressive networks are currently state-of-the-art for language modelling [29], and in general\\nperform well on discrete data where a natural ordering exists. However they have proved less\\neffective in domains such as image generation, where the data is continuous and no natural order\\nexists among variables (e.g. there is no reason to generate one pixel before another). They also\\nhave the drawback that generating samples requires as many network updates as there are variables\\nin the data.\\nDiffusion models are an alternative framework that has proved particularly effective for image\\ngeneration [5, 34]. In this case the transmission procedure is a little more complex2. Each message\\nBob receives is a noisy version of the message before, where the noise is designed so that in\\nexpectation the messages approach the data. The transmission cost at each step is the Kullback-\\nLeibler divergence between the distribution from which Alice draws the message and Bob’s prediction\\nof that distribution (which is a reparameterisation of his prediction of the data, and which is therefore\\nimproved by the information he gained from the previous message). The sum of the KL divergences\\nis the evidence lower bound minimised by diffusion training [41]; it is also the expected number of\\nbits needed to transmit the data using an efficient bits-back coding scheme [11, 51]. Once again\\nthere is an exact equivalence between the loss function used to train the model and the model’s\\nability to compress data, as elucidated by previous authors [47].\\nWe posit that the superiority of diffusion over autoregression for image generation lies in the\\nway diffusion progresses from coarse to fine image details as the level of noise decreases — a more\\nnatural way to construct an image than one dot at a time. However diffusion has yet to match\\nautoregression for discrete data, which is unfortunate, as diffusion models have the advantage of\\ndecoupling the number of generation steps from the number of variables. A fundamental challenge\\nis that when the data is discrete, the noise in the diffusion process is also discrete, and therefore\\ndiscontinuous. To return to the transmission metaphor, if the data is a piece of text, then Bob\\nbegins the process with a totally garbled text, every symbol of which is either randomly altered\\nor left unchanged by each of Alice’s messages. A key motivation for this work was our belief that\\na fully continuous transmission process — where Alice’s messages smoothly alter Bob’s beliefs —\\nwould be more effective for discrete data. Moreover this should open the door to gradient-based\\nsample guidance [5] and few-step generation techniques [37, 43, 52], similar to those that have been\\ndeveloped for continuous diffusion.\\nBayesian Flow Networks (BFNs), the model introduced in this paper, differ from diffusion models\\nin that the network operates on the parameters of a data distribution, rather than on a noisy version\\nof the data itself. This ensures that the generative process is fully continuous and differentiable,\\n2We are here describing the reverse process of diffusion models.\\n2\\n\\n\\nOutput\\nData\\nSender\\nReceiver \\nPrev. Input\\n…\\nNetwork \\nVariable 1\\nNetwork …\\n…\\n+\\nInput\\n+\\nBayesian \\nupdate\\nsample\\nOutput\\nData\\nSender\\nReceiver \\nPrev. Input\\nVariable 2\\n+\\nInput\\nALICE\\n+\\nBayesian \\nupdate\\nKL\\nsample\\nBOB\\nALICE\\nBOB\\nNoise\\nNoise\\nKL\\nFigure 1: System Overview. The figure represents one step of the modelling process of a Bayesian Flow\\nNetwork. The data in this example is a ternary symbol sequence, of which the first two variables (‘B’ and ‘A’)\\nare shown. At each step the network emits the parameters of the output distribution based on the parameters\\nof the previous input distribution. The sender and receiver distributions (both of which are continuous,\\neven when the data is discrete) are created by adding random noise to the data and the output distribution\\nrespectively. A sample from the sender distribution is then used to update the parameters of the input\\ndistribution, following the rules of Bayesian inference. Conceptually, this is the message sent by Alice to Bob,\\nand its contribution to the loss function is the KL divergence from the receiver to the sender distribution.\\neven when the data is discrete. BFNs can be summarised by the following transmission scheme\\n(Figure 1). Bob has an “input distribution” which is initially a simple prior: a standard normal for\\ncontinuous data, a uniform categorical for discrete data. At each transmission step he feeds the\\nparameters of the input distribution (e.g. the mean of a normal distribution, the probabilities of a\\ncategorical distribution) into a neural network. The network outputs the parameters of a second\\ndistribution referred to as the “output distribution”. Alice then creates a “sender distribution” by\\nadding noise to the data according to a predefined schedule, and Bob creates a “receiver distribution”\\nby convolving the output distribution with the same noise distribution used by Alice: intuitively, for\\nevery value the data could take on, Bob constructs the sender distribution Alice would have used if\\n3\\n\\n\\nthat value was correct, then sums over all these hypothetical sender distributions, weighted by the\\nprobability of the corresponding value under the output distribution. Alice picks a sample from the\\nsender distribution and sends it to Bob at a cost equal to the KL divergence from receiver to sender.\\nBob then uses the sample to update his input distribution, following the rules of Bayesian inference.\\nUsefully, the Bayesian updates are available in closed-form as long as the input distribution models\\nall the variables in the data independently. Once the update is complete, Bob again feeds the\\nparameters of the input distribution to the network which returns the parameters of the output\\ndistribution. The process repeats for n steps, at which point Bob can predict the data accurately\\nenough that Alice can send it to him without any noise.\\nNote the key difference between the input and output distributions: the input distribution\\nreceives information about each variable in the data independently (via the Bayesian updates), and\\nis therefore unable to exploit contextual information, such as neighbouring pixels in an image or\\nrelated words in a text; the output distribution, on the other hand, is produced by a neural network\\nthat jointly processes all the parameters in the input distribution, giving it access to all available\\ncontext. Intuitively, the combination of the input and output distributions represents a division\\nof labour between Bayesian inference and deep learning that plays to both of their strengths: the\\nformer provides a mathematically optimal and finely controllable way to collect and summarise\\ninformation about individual variables, while the latter excels at integrating information over many\\ninterrelated variables.\\nThe above transmission process defines an n-step loss function that can be generalised to\\ncontinuous time by sending n to ∞. In continuous time the Bayesian updates become a Bayesian\\nflow of information from the data to the network. As well as removing the need to predefine the\\nnumber of steps during training, the continuous-time loss function is mathematically simpler and\\neasier to compute than the discrete-time loss. A BFN trained with continuous-time loss can be run\\nfor any number of discrete steps during inference and sampling, with performance improving as the\\nnumber of steps increases.\\nThe rest of the paper is structured as follows. A short summary of related work is given\\nin Section 2.\\nThe basic framework of BFNs, along with a general derivation of the discrete\\nand continuous time loss functions is provided in Section 3. Specialisations of the framework to\\ncontinuous, discretised and discrete data are provided in Sections 4–6, along with pseudocode\\nfor training, evaluating and sampling from the network. Experimental results on the CIFAR-10,\\ndynamically binarized MNIST and text8 datasets are provided in Section 7 and concluding remarks\\nare given in Section 8.\\n2\\nRelated Work\\nOf existing methods, Bayesian Flow Networks are most closely related to diffusion models. However\\nthe two differ in some crucial aspects. Most obviously BFNs embody a function from one distribution\\nto another — rather than from data to a distribution, like diffusion models and most other\\nprobabilistic networks.\\nOne advantage of this approach is that, because the parameters of a\\ncategorical distribution are real-valued probabilities, the inputs to the network are continuous even\\nwhen the data is discrete. This contrasts with discrete diffusion, which natively uses discrete samples\\nas input [1, 14, 41].\\nNumerous authors have proposed continuous variants of discrete diffusion. Typically these\\nrely either on mapping to and from a continuous embedding space [2, 6, 21, 44], or on restricting\\n4\\n\\n\\ncontinuous diffusion to the probability simplex [23, 24, 33]. While we do not directly compare\\nagainst the above methods, we note that continuity is an inherent property of the Bayesian Flow\\nframework (the network inputs automatically lie on the probability simplex by virtue of being the\\nparameters of a categorical distribution), rather than a constraint added to an existing system. As\\nwell as reducing the number of free parameters and design choices (e.g. the continuous embedding\\nspace, the mapping functions), this ensures that BFNs directly optimise the negative log-likelihood\\nof discrete data, unlike continuous diffusion methods for discrete data, which typically require either\\nsimplified loss functions [24] or auxiliary loss terms [21] to make learning stable.\\nFor continuous data, BFNs are most closely related to variational diffusion models [17], with a\\nvery similar continuous-time loss function. The main difference in this case is that the network inputs\\nare considerably less noisy in BFNs than in variational diffusion and other continuous diffusion\\nmodels. This is because the generative process of BFNs begins with the parameters of a fixed prior,\\nwhereas that of diffusion models begins with pure noise. We hypothesise that the reduction in noise\\ncould lead to faster learning on large datasets where the model underfits; however we have yet to\\ntest this hypothesis experimentally.\\nAnother key difference from diffusion models is that there is no need to define and invert a\\nforward process for BFNs, which arguably makes it easier to adapt them to different distributions\\nand data types. We showcase this flexibility by adapting BFNs to continuous, discretised and\\ndiscrete data, with minimal changes to the training procedure. This contrasts with e.g. discretised\\ndiffusion, which requires carefully defined transition matrices [1].\\n3\\nBayesian Flow Networks\\nThis section covers the basic mathematical formalism of Bayesian Flow Networks, laying out the\\nstructure of the various functions and distributions required by the model, along with the discrete\\nand continuous-time loss functions used for training. Specific instantiations of the general framework\\nfor continuous, discretised and discrete data are given in Sections 4–6.\\n3.1\\nInput and Sender Distributions\\nGiven D-dimensional data x =\\nx(1), . . . , x(D)\\u0001\\n∈X D, let θ =\\nθ(1), . . . , θ(D)\\u0001\\nbe the parameters of\\na factorised input distribution pI(· | θ), with\\npI(x | θ) =\\nD\\nY\\nd=1\\npI(x(d) | θ(d)).\\n(1)\\nFor example, θ(d) may consist of the probabilities of a categorical distribution. Let pS (· | x; α) be a\\nsimilarly factorised sender distribution with y =\\ny(1), . . . , y(D)\\u0001\\n∈YD and\\npS (y | x; α) =\\nD\\nY\\nd=1\\npS\\n\\u0010\\ny(d) | x(d); α\\n\\u0011\\n,\\n(2)\\nwhere α ∈R+ is an accuracy parameter defined such that when α = 0, the sender samples are entirely\\nuninformative about x and as α increases the samples become progressively more informative.\\n5\\n\\n\\n3.2\\nOutput Distribution pO(· | θ, t)\\nDuring the data transmission process, the input parameters θ are passed along with the process\\ntime t as input to a neural network Ψ.\\nThe network then emits an output vector Ψ(θ, t) =\\nΨ(1)(θ, t), . . . , Ψ(D)(θ, t)\\n\\u0001\\nwhich is used to parameterise an output distribution factorised in the\\nsame way as the input and sender distributions:\\npO(x | θ, t) =\\nD\\nY\\nd=1\\npO(x(d) | Ψ(d)(θ, t)).\\n(3)\\nAs discussed in the introduction, the key difference between the input and output distributions\\nis that while each pI(x(d) | θ(d)) depends only on information gathered via pS\\ny(d) | x(d); α\\n\\u0001\\nabout\\nx(d), each pO(x(d) | Ψ(d)(θ, t)) depends (via the network) on all of θ and hence all of x. The\\noutput distribution, unlike the input distribution, can therefore exploit context information, such as\\nsurrounding pixels in an image or related words in a text.\\n3.3\\nReceiver Distribution pR(· | θ; t, α)\\nGiven sender distribution pS (· | x; α) and output distribution pO(· | θ, t) the receiver distribution\\nover YD is defined as\\npR(y | θ; t, α) =\\nE\\npO(x′|θ;t) pS\\ny | x′; α\\n\\u0001\\n.\\n(4)\\nIntuitively this can be understood as a receiver who knows the form of the sender distribution\\npS (· | x; α) but does not know x, and therefore integrates over all x′ ∈X D, and hence all possible\\nsender distributions, weighted by the probability given to x′ by the output distribution pO(x | θ, t).\\nThe receiver distribution therefore combines two sources of uncertainty: the “known unknown” of\\nthe sender distribution entropy (which is a function of α), and the “unknown unknown” of the\\noutput distribution entropy.\\n3.4\\nBayesian Updates\\nGiven parameters θ and sender sample y drawn with accuracy α the Bayesian update function h is\\nderived by applying the rules of Bayesian inference to compute the updated parameters θ′:\\nθ′ ←h(θ, y, α).\\n(5)\\nThe Bayesian update distribution pU (· | θ, x; α) is then defined by marginalizing out y:\\npU (θ′ | θ, x; α) =\\nE\\npS (y|x;α) δ\\nθ′ −h(θ, y, α)\\n\\u0001\\n,\\n(6)\\nwhere δ (· −a) is the multivariate Dirac delta distribution centred on the vector a. In Sections 4.4\\nand 6.7 we will prove that both forms of pU (· | θ, x; α) considered in this paper have the following\\nproperty: the accuracies are additive in the sense that if α = αa + αb then\\npU (θ′′ | θ, x; α) =\\nE\\npU (θ′|θ,x;αa) pU (θ′′ | θ′, x; αb).\\n(7)\\n6\\n\\n\\nIt follows from this property that given prior input parameters θ0, the probability of observing\\nparameters θn after drawing a sequence of n sender samples y1, . . . , yn with accuracies α1, . . . , αn is\\nE\\npU (θ1|θ0,x;α1)\\nE\\npU (θ2|θ1,x;α2) . . .\\nE\\npU (θn−1|θn−2,x;αn−1) pU (θn | θn−1, x; αn) = pU\\n \\nθn | θ0, x;\\nn\\nX\\ni=1\\nαi\\n!\\n. (8)\\n3.5\\nAccuracy Schedule β(t)\\nBy performing an infinite number of transmission steps, the Bayesian update process can be\\ngeneralized to continuous time. Let t ∈[0, 1] be the process time and let α(t) > 0 be the accuracy\\nrate at time t. Now define the accuracy schedule β(t) as\\nβ(t) =\\nZ t\\nt′=0\\nα(t′)dt′.\\n(9)\\nIt follows from the above definitions that β(t) is a monotonically increasing function of t, that\\nβ(0) = 0, and that dβ(t)\\ndt\\n= α(t).\\nSpecific forms of β(t) for continuous and discrete data are provided in Sections 4.5 and 6.8. Both\\nare derived using simple heuristics, with a deeper investigation left for future work.\\n3.6\\nBayesian Flow Distribution pF (· | x; t)\\nGiven prior parameters θ0, Bayesian update distribution pU (· | θ, x; α) and accuracy schedule β(t),\\nthe Bayesian flow distribution pF (· | x; t) is the marginal distribution over input parameters at time\\nt, defined by\\npF (θ | x; t) = pU (θ | θ0, x; β(t)).\\n(10)\\n3.7\\nLoss Function L(x)\\nGiven prior parameters θ0 and accuracy schedule β(t), consider a sequence of n sender samples\\ny1, . . . , yn sampled at times t1, . . . , tn where ti = i/n. The sender distribution at step i is pS (· | x; αi)\\nwhere\\nαi = β(ti) −β(ti−1),\\n(11)\\nthe receiver distribution at step i is pR(· | θi−1; ti−1, αi), and the input parameter sequence θ1, . . . , θn\\nis recursively calculated from\\nθi = h(θi−1, y, αi).\\n(12)\\nDefine the n-step discrete-time loss Ln(x) as the expected number of nats required to first transmit\\ny1, . . . , yn, and the reconstruction loss Lr(x) as the expected number of nats required to then\\ntransmit x. Since — using a bits-back coding scheme [7, 11] — it requires DKL (pS ∥pR) nats to\\ntransmit a sample from pS to a receiver with pR,\\nLn(x) =\\nE\\np(θ1,...,θn−1)\\nn\\nX\\ni=1\\nDKL (pS (· | x; αi) ∥pR(· | θi−1; ti−1, αi)) ,\\n(13)\\n7\\n\\n\\nwhere\\np(θ1, . . . , θn) =\\nn\\nY\\ni=1\\npU (θi | θi−1, x; αi),\\n(14)\\nand since the number of nats needed to transmit x using an arithmetic coding scheme [53] based on\\np(x) is −ln p(x), and the marginal probability of θn is given by pF (· | x, 1),\\nLr(x) = −\\nE\\npF (θ|x,1) ln pO(x | θ; 1).\\n(15)\\nNote that Lr(x) is not directly optimised in this paper; however it is indirectly trained by optimising\\nLn(x) since both are minimised by matching the output distribution to the data. Furthermore, as\\nlong as β(1) is high enough, the input distribution at t = 1 will be very close to x, making it trivial\\nfor the network to fit pO(x | θ; 1).\\nThe loss function L(x) is defined as the total number of nats required to transmit the data,\\nwhich is the sum of the n-step and reconstruction losses:\\nL(x) = Ln(x) + Lr(x)\\n(16)\\nAlternatively L(x) can be derived as the loss function of a variational autoencoder (VAE; [18]).\\nConsider the sequence y1, . . . , yn as a latent code with posterior probability given by\\nq(y1, . . . , yn) =\\nn\\nY\\ni=1\\npS (yi | x; αi) ,\\n(17)\\nand autoregressive prior probability given by\\np(y1, . . . , yn) =\\nn\\nY\\ni=1\\npR(yi | θi−1; ti−1, αi).\\n(18)\\nThen, noting that the decoder probability p(x | y1, . . . , yn) = pO(x | θn; 1), the complete trans-\\nmission process defines a VAE with loss function given by the negative variational lower bound\\n(VLB)\\nL(x) = −VLB(x) = DKL (q ∥p) −\\nE\\ny1,...,yn∼q ln p(x | y1, . . . , yn)\\n(19)\\n= Ln(x) + Lr(x).\\n(20)\\n3.8\\nDiscrete-Time Loss Ln(x)\\nEq. 13 can be rewritten as\\nLn(x) = n\\nE\\ni∼U{1,n}\\nE\\npU (θ1|θ0,x;α1) . . .\\nE\\npU (θ|θi−2,x;αi−1) DKL (pS (· | x; αi) ∥pR(· | θ; ti−1, αi)) ,\\n(21)\\nwhere U{1, n} is the uniform distribution over the integers from 1 to n. Furthermore, it follows\\nfrom Eqs. 8 and 10 that\\nE\\npU (θ1|θ0,x;α1) . . .\\nE\\npU (θ|θi−2,x;αi−1) =\\nE\\npU (θ|θ0,x;β(ti−1))\\n(22)\\n=\\nE\\npF (θ|x;ti−1),\\n(23)\\n8\\n\\n\\nand hence\\nLn(x) = n\\nE\\ni∼U{1,n},pF (θ|x;ti−1) DKL (pS (· | x; αi) ∥pR(· | θ; ti−1, αi)) ,\\n(24)\\nwhich allows us approximate Ln(x) via Monte-Carlo sampling without computing the n-step sum.\\n3.9\\nContinuous-Time Loss L∞(x)\\nEq. 24 can be used to train the network directly. However this presupposes that n is fixed during\\ntraining. Furthermore, for discrete and discretised data the KL terms do not have analytic solutions,\\nleading to noisy gradient estimates.\\nInspired by Variational Diffusion Models [17] we derive a continuous-time loss function L∞(x)\\nby taking the limit of Ln(x) as n →∞. This turns out to be mathematically simpler than the\\ndiscrete-time loss, as well as removing both the noisy gradients for the discrete and discretised KL\\nterms and the need to fix n during training.\\nLet\\nϵ def\\n= 1\\nn,\\n(25)\\nα(t, ϵ) def\\n= β(t) −β(t −ϵ),\\n(26)\\nL∞(x) def\\n= lim\\nn→∞Ln(x).\\n(27)\\nThen, from the definition of Ln(x) in Eq. 24,\\nL∞(x) = lim\\nϵ→0\\n1\\nϵ\\nE\\nt∼U(ϵ,1),pF (θ|x,t−ϵ) DKL (pS (· | x; α(t, ϵ)) ∥pR(· | θ; t −ϵ, α(t, ϵ))) ,\\n(28)\\nwhere U(a, b) is the continuous uniform distribution over the interval [a, b]. As we will see, for all\\nthe sender, receiver distribution pairs in this paper,\\nDKL (pS (· | x; α) ∥pR(· | θ; α, t)) =\\nD\\nX\\nd=1\\nDKL\\n\\u0010\\nN\\n\\u0010\\ng(x(d)), Cα−1\\u0011\\n∥P (d)(θ, t) ∗N\\n0, Cα−1\\u0001\\u0011\\n,\\n(29)\\nwhere g : X →Y is a function from data space to sender space, P (d)(θ, t) is a distribution over Y\\nwith finite expectation and variance, ∗denotes the convolution of two probability distributions and\\nC is a scalar constant.\\nThe following proposition is now required:\\nProposition 3.1. For a continuous univariate probability distribution P with finite expectation\\nE[P] and variance V ar[P], the convolution P ∗N\\n0, σ2\\u0001\\n→N\\nE[P], σ2\\u0001\\nas σ2 →∞.\\nProof. Let ϵ2 be some variance in the interval\\n0, π\\n8\\n\\u0001\\nand consider the sequence of random variables\\n9\\n\\n\\nX0, X1, . . . , Xn where X0 ∼P and Xj ∼N\\n0, ϵ2\\u0001\\nfor j > 0. Define\\nYj\\ndef\\n=\\n(\\nX0 −E[P]\\nif j = 0,\\nXj\\notherwise.\\n(30)\\nRn\\ndef\\n=\\nn\\nX\\nj=0\\nYj,\\n(31)\\nS2\\nn\\ndef\\n=\\nn\\nX\\nj=1\\nV ar[Yj] = nϵ2,\\n(32)\\nT 2\\nn\\ndef\\n= V ar[P] + S2\\nn.\\n(33)\\nIt follows from the definition of convolution that Pn\\nj=0 Xj ∼P ∗N\\n0, nϵ2\\u0001\\n. Since nϵ2 →∞as\\nn →∞, and Pn\\nj=0 Xj = Rn + E[P], the result is proved if it can be shown that as n →∞,\\nRn →N\\n0, nϵ2\\u0001\\nor equivalently Rn/(ϵ√n) →N (0, 1).\\nThe Lyapunov central limit theorem [8] states that if there exists λ > 0 such that\\nlimn→∞\\n1\\nT 2+λ\\nn\\nPn\\nj=0 E\\n|Yj|2+λ\\u0001\\n= 0 then Rn/Tn →N (0, 1).\\nFirst note that T 2\\nn →S2\\nn = nϵ2\\nas n →∞. Hence if Rn/Tn →N (0, 1) then Rn/(ϵ√n) →N (0, 1). Now set λ = 1 and observe that\\nfor Yj ∼N\\n0, ϵ2\\u0001\\n, E\\n|Yj|3\\u0001\\nis the third moment of the half-normal distribution, which is ϵ3\\nq\\n8\\nπ.\\nOur choice of ϵ2 therefore ensures that E\\n|Yj|3\\u0001\\n< ϵ2 for j > 0. Also note that T 3\\nn > S3\\nn and, since\\nE[P] and V ar[P] are finite, E\\n|Y0|3\\u0001\\n< C for some constant C. Hence\\n1\\nT 3\\nn\\nn\\nX\\nj=0\\nE\\n|Yj|3\\u0001\\n< 1\\nS3\\nn\\nC + nϵ2\\u0001\\n=\\nC\\nϵ3n3/2 +\\n1\\nϵ√n\\nn→∞\\n−\\n−\\n−\\n→0.\\n(34)\\nIt follows from the continuity of β(t) and Eq. 26 that α(t, ϵ)−1 →∞as ϵ →0. Therefore,\\nProposition 3.1 can be applied to Eq. 29 to yield\\nlim\\nϵ→0 DKL (pS (· | x, αt) ∥pR(· | θ, αt, t)) =\\nD\\nX\\nd=1\\nDKL\\n\\u0012\\nN\\n\\u0012\\ng(x(d)),\\nC\\nα(t, ϵ)\\n\\u0013\\n∥N\\n\\u0012\\nE[P (d)(θ, t)],\\nC\\nα(t, ϵ)\\n\\u0013\\u0013\\n(35)\\n= α(t, ϵ)\\n2C\\n∥g(x) −E[P(θ, t)]∥2 ,\\n(36)\\nwhere\\ng(x) =\\n\\u0010\\ng(x(1)), . . . , g(x(D))\\n\\u0011\\n,\\n(37)\\nE[P(θ, t)] =\\n\\u0010\\nE[P (1)(θ, t)], . . . , E[P (D)(θ, t)]\\n\\u0011\\n.\\n(38)\\nTherefore,\\nL∞(x) =\\nE\\nt∼U(0,1),pF (θ|x,t) lim\\nϵ→0\\nα(t, ϵ)\\nϵ\\n∥g(x) −E[P(θ, t)]∥2\\n2C\\n.\\n(39)\\n10\\n\\n\\nSubstituting from Eq. 26,\\nlim\\nϵ→0\\nα(t, ϵ)\\nϵ\\n= lim\\nϵ→0\\nβ(t) −β(t −ϵ)\\nϵ\\n= dβ(t)\\ndt\\n= α(t),\\n(40)\\nand hence\\nL∞(x) =\\nE\\nt∼U(0,1),pF (θ|x,t) α(t)∥g(x) −E[P(θ, t)]∥2\\n2C\\n.\\n(41)\\n3.10\\nSample Generation\\nGiven prior parameters θ0, accuracies α1, . . . , αn and corresponding times ti = i/n, the n-step\\nsampling procedure recursively generates θ1, . . . , θn by sampling x′ from pO(· | θi−1, ti−1), y from\\npS (· | x′, αi) (meaning that y ∼pR(· | θi−1; ti−1, αi) — see Eq. 4), then setting θi = h(θi−1, y).\\nGiven θn the network is run one more time and the final sample is drawn from pO(· | θn, 1).\\n4\\nContinuous Data\\nFor continuous data X = R and hence x ∈RD. In our experiments, x is normalised to lie in [−1, 1]D\\nto ensure that the network inputs remain in a reasonable range; however this is not essential for the\\nmathematical framework.\\n4.1\\nInput Distribution pI(· | θ)\\nThe input distribution for continuous data is a diagonal normal:\\nθ def\\n= {µ, ρ}\\n(42)\\npI(x | θ) def\\n= N\\nx | µ, ρ−1I\\n\\u0001\\n,\\n(43)\\nwhere I is the D × D identity matrix. We define the prior parameters as\\nθ0\\ndef\\n= {0, 1},\\n(44)\\nwhere 0 is the length D vectors of zeros. Hence the input prior is a standard multivariate normal:\\npI(x | θ0) = N (x | 0, I) .\\n(45)\\nThe usual Bayesian approach would be to fit the prior mean and variance to the training data.\\nHowever we found that a standard prior worked better in practice, as well as simplifying the\\nequations. It is important to remember that the distributions pI(x | θ0) are never used directly\\nto make predictions, but rather to inform the network’s predictions. All that matters is that the\\nparameters fed into the network accurately and accessibly encode the information received so far\\nabout x. The network can easily learn the empirical prior of the training set and use that to correct\\nits predictions.\\n11\\n\\n\\nFigure 2: Bayesian updates for continuous data. For univariate data x = 0.7, the initial input\\ndistribution parameters θ0 = {µ0 = 0, ρ0 = 1} are updated to θ1 = {µ1, ρ1}, θ2 = {µ2, ρ2}, θ3 = {µ3, ρ3} by\\niterating Eqs. 49 and 50 with sender samples y1, y2, y3 drawn with accuracies 2, 4, 6 respectively. Note how\\nthe input mean (µ1, µ2, µ3) stochastically approaches the data, while the input precision smoothly increases.\\n4.2\\nBayesian Update Function h(θi−1, y, α)\\nGiven a univariate Gaussian prior N\\nµa, ρ−1\\na\\n\\u0001\\nover some unknown data x it can be shown [27] that\\nthe Bayesian posterior after observing a noisy sample y from a normal distribution N\\nx, α−1\\u0001\\nwith\\nknown precision α is N\\nµb, ρ−1\\nb\\n\\u0001\\n, where\\nρb = ρa + α,\\n(46)\\nµb = µaρa + yα\\nρb\\n.\\n(47)\\nSince both pI(x | θ) and pS (y | x; α) distributions are normal with diagonal covariance, Eqs. 46 and\\n47 can be applied to obtain the following Bayesian update function for parameters θi−1 = {µi−1, ρi−1}\\nand sender sample y drawn from pS (· | x; αI) = N\\nx, α−1I\\n\\u0001\\n:\\nh({µi−1, ρi−1}, y, α) = {µi, ρi},\\n(48)\\nwith\\nρi = ρi−1 + α,\\n(49)\\nµi = µi−1ρi−1 + y α\\nρi\\n.\\n(50)\\n4.3\\nBayesian Update Distribution pU(· | θ, x; α)\\nEq. 50 computes µi given a single sample y from the sender distribution. To marginalise over\\ny ∼N\\ny | x, α−1I\\n\\u0001\\nas defined in Eq. 6, the following standard identity for normal distributions\\n12\\n\\n\\nFigure 3: Bayesian update distribution for continuous data. For x = 0.7, the plot shows the\\ndistribution p(µ | θ0, x; α) over input mean µ from Eq. 52 given initial parameters µ0 = 0, ρ0 = 1 and 11 α\\nvalues spaced log-linearly between e−5 and e5. Note how the distribution is tightly concentrated around µ0\\nfor very low alpha, then smoothly progresses to a tight concentration around x for high alpha.\\ncan be applied:\\nX ∼N\\nµX, σ2\\nX\\n\\u0001\\n=\\n⇒aX + b ∼N\\naµX + b, a2σ2\\nX\\n\\u0001\\n∀a, b ∈R .\\n(51)\\nSubstituting X = y, µX = x, σ2\\nX = α−1I, a = α\\nρi and b = µi−1ρi−1\\nρi\\n, Eq. 50 gives:\\nµi ∼N\\n\\u0012α x +µi−1ρi−1\\nρi\\n, α\\nρ2\\ni\\nI\\n\\u0013\\n,\\n(52)\\nand therefore (since µi is the only random part of θi)\\npU (θi | θi−1, x; α) = N\\n\\u0012\\nµi | α x +µi−1ρi−1\\nρi\\n, α\\nρ2\\ni\\nI\\n\\u0013\\n.\\n(53)\\n4.4\\nAdditive Accuracies\\nWe can check that the sender accuracies are additive in the sense required by Eq. 7 by first observing\\nthat if θi−1 = {µi−1, ρi−1} is drawn from p(· | θi−2, x; αa) then\\nµi−1 ∼N\\n\\u0012αa x +µi−2ρi−2\\nρi−1\\n, αa\\nρ2\\ni−1\\nI\\n\\u0013\\n.\\n(54)\\nDefine\\nµ′\\ni\\ndef\\n= αb x +µi−1ρi−1\\nρi\\n= ρi−1\\nρi\\nµi−1 + αb x\\nρi\\n,\\n(55)\\n13\\n\\n\\nand apply Identity 51 with a = ρi−1\\nρi\\nand b = αb x\\nρi\\nto see that\\nµ′\\ni ∼N\\n\\u0012ρi−1\\nρi\\nαa x +µi−2ρi−2\\nρi−1\\n+ αb x\\nρi\\n, ρ2\\ni−1\\nρ2\\ni\\nαa\\nρ2\\ni−1\\nI\\n\\u0013\\n(56)\\n= N\\n\\u0012(αa + αb) x +µi−2ρi−2\\nρi\\n, αa\\nρ2\\ni\\nI\\n\\u0013\\n.\\n(57)\\nNow observe that if θi = {µi, ρi} is drawn from p(· | θi−1, x; αb) then\\nµi ∼N\\n\\u0012αb x +µi−1ρi−1\\nρi\\n, αb\\nρ2\\ni\\nI\\n\\u0013\\n,\\n(58)\\nand hence\\nµi ∼µ′\\ni + ϵ,\\n(59)\\nwhere\\nϵ ∼N\\n\\u0012\\n0, αb\\nρ2\\ni\\nI\\n\\u0013\\n.\\n(60)\\nAnother standard identity for Gaussian variables can now be applied:\\nX ∼N\\nµX, σ2\\nX\\n\\u0001\\n, Y ∼N\\nµY , σ2\\nY\\n\\u0001\\n=\\n⇒X + Y ∼N\\nµX + µY , σ2\\nX + σ2\\nY\\n\\u0001\\n,\\n(61)\\nto see that\\nµi ∼N\\n\\u0012(αa + αb) x +µi−2ρi−2\\nρi\\n, αa + αb\\nρ2\\ni\\nI\\n\\u0013\\n,\\n(62)\\nand hence\\nE\\npU (θi−1|θi−2,x;αa) pU (θi | θi−1, x; αb) = pU (θi | θi−2, x; αa + αb),\\n(63)\\nas required.\\n4.5\\nAccuracy Schedule β(t)\\nWe derive β(t) for continuous data by requiring that the expected entropy of the input distribution\\nlinearly decreases with t. Intuitively, this means that information flows into the input distribution\\nat a constant rate. Define\\nH(t) def\\n=\\nE\\npF (θ|x;t) H(pI(· | θ))\\n(64)\\n= D\\n2 ln\\n\\u0012\\n2πe\\n1 + β(t)\\n\\u0013\\n.\\n(65)\\n14\\n\\n\\nThen if H(t) linearly decreases with t,\\nH(t) = (1 −t)H(0) + tH(1)\\n(66)\\n=\\n⇒ln\\n\\u0012\\n2π\\n1 + β(t)\\n\\u0013\\n= (1 −t) ln(2π) + t ln\\n\\u0012\\n2π\\n1 + β(1)\\n\\u0013\\n(67)\\n=\\n⇒−ln(1 + β(t)) = −t ln(1 + β(1))\\n(68)\\n=\\n⇒(1 + β(t))−1 = (1 + β(1))−t.\\n(69)\\nDefine σ1 to be the standard deviation of the input distribution at t = 1. We will choose σ1 empirically\\nto minimise the loss; in general it should be small enough to ensure that the reconstruction loss is\\nlow, but not so small as to create unnecessary transmission costs. Recalling that the precision ρ at\\ntime t is 1 + β(t), we see that\\nσ2\\n1 = (1 + β(1))−1.\\n(70)\\nTherefore\\n(1 + β(t))−1 = σ2t\\n1\\n(71)\\n=\\n⇒β(t) = σ−2t\\n1\\n−1\\n(72)\\n=\\n⇒α(t) = d\\nσ−2t\\n1\\n−1\\n\\u0001\\ndt\\n(73)\\n= −2 ln σ1\\nσ2t\\n1\\n.\\n(74)\\n4.6\\nBayesian Flow Distribution pF (· | x; t)\\nRecall from Eq. 10 that\\npF (θ | x; t) = pU (θ | θ0, x, β(t)).\\n(75)\\nTherefore, setting θi−1 = θ0 = {0, 1} and α = β(t) in Eq. 53, and recalling that ρ = 1 + β(t),\\npF (θ | x; t) = N\\n\\u0012\\nµ |\\nβ(t)\\n1 + β(t) x,\\nβ(t)\\n(1 + β(t))2 I\\n\\u0013\\n(76)\\n= N (µ | γ(t) x, γ(t)(1 −γ(t))I) ,\\n(77)\\nwhere\\nγ(t) def\\n=\\nβ(t)\\n1 + β(t)\\n(78)\\n= σ−2t\\n1\\n−1\\nσ−2t\\n1\\n(79)\\n= 1 −σ2t\\n1 .\\n(80)\\n15\\n\\n\\nFigure 4: Bayesian flow for continuous data. For x = 0.8, σ1 = 0.02 and γ(t) defined as in Eqn. 80, the\\nplot shows stochastic parameter trajectories for the input distribution mean µ (white lines) superimposed on\\na log-scale heatmap of the Bayesian flow distribution p(θ | x; t). Note how the trajectories all begin at µ0 = 0\\nthen fan out before converging on x.\\n4.7\\nOutput Distribution pO(· | θ; t)\\nFollowing standard practice for diffusion models [42], the output distribution is defined by reparam-\\neterising a prediction of the Gaussian noise vector ϵ ∼N (0, I) used to generate the mean µ passed\\nas input to the network. Recall from Eq. 77 that\\nµ ∼N (γ(t) x, γ(t)(1 −γ(t))I) ,\\n(81)\\nand hence\\nµ = γ(t) x +\\np\\nγ(t)(1 −γ(t))ϵ\\n(82)\\n=\\n⇒x =\\nµ\\nγ(t) −\\ns\\n1 −γ(t)\\nγ(t)\\nϵ.\\n(83)\\nThe network outputs an estimate ˆ\\nϵ(θ, t) of ϵ and this is transformed into an estimate ˆ\\nx(θ, t) of x by\\nˆ\\nx(θ, t) =\\nµ\\nγ(t) −\\ns\\n1 −γ(t)\\nγ(t)\\nˆ\\nϵ(θ, t).\\n(84)\\nGiven ˆ\\nx(θ, t) the output distribution is\\npO(x | θ; t) = δ(x −ˆ\\nx(θ, t)),\\n(85)\\nNote that γ(0) = 0, making the transformation from ˆ\\nϵ(θ, t) to pO(x | θ; t) undefined at t = 0. We\\ntherefore set pO(x | θ; t) = 0 for t under some small threshold tmin. Also, ˆ\\nx(θ, t) is clipped to lie\\nwithin the allowed range [xmin, xmax] for x. In our experiments tmin = 1e−6 and [xmin, xmax] =\\n[−1, 1].\\n16\\n\\n\\nFigure 5: Input variance for Bayesian Flow Networks and diffusion models. For σ1 = 0.001 and\\nγ(t) defined as in Eqn. 80, the blue line shows the variance γ(t)(1 −γ(t)) of the distribution over the input\\nmean µ as a function of t (see Eq. 77). Note that the variance is 0 at t = 0 (since the input prior µ0 is\\ndeterministic) and becomes small again as t approaches 1 and µ becomes increasingly concentrated around\\nthe data. The green and red lines show the equivalent network input variance for two different noise schedules\\nfrom the literature (linear [12] and cosine [28]) during the reverse process of a diffusion model (note that t is\\nreversed relative to diffusion convention). The input variance is much lower for Bayesian Flow Networks.\\n4.8\\nSender Distribution pS (· | x; α)\\nThe sender space Y = X = R for continuous data, and the sender distribution is normal with\\nprecision α:\\npS (y | x; α) = N\\ny | x, α−1I\\n\\u0001\\n.\\n(86)\\n4.9\\nReceiver Distribution pR(· | θ; t, α)\\nSubstituting Eqs. 85 and 86 into Eq. 4,\\npR(y | θ; t, α) =\\nE\\nδ(x′ −ˆ\\nx(θ,t)) N\\ny | x′, α−1I\\n\\u0001\\n(87)\\n= N\\ny | ˆ\\nx(θ, t), α−1I\\n\\u0001\\n.\\n(88)\\n4.10\\nReconstruction Loss Lr(x)\\nTruly continuous data requires infinite precision to reconstruct, which makes the reconstruction loss\\nproblematic. However it would be reasonable to assume that either the data is finely discretised\\n(as all information is on a digital computer), or that it contains some noise. The reconstruction\\nloss for discretised data is presented in Section 5.3. Alternatively, if we assume the presence of\\nnormally distributed measurement noise on x, with fixed isotropic variance σ2, then a noisy version\\nof the reconstruction loss can be defined as the expected KL divergence between N\\nx, σ2I\\n\\u0001\\nand\\nthe output distribution at t = 1:\\nLr(x) =\\nE\\npF (θ|x,1) DKL\\nN\\nx, σ2I\\n\\u0001\\n∥N\\nˆ\\nx(θ, 1), σ2I\\n\\u0001\\u0001\\n(89)\\n=\\nE\\npF (θ|x,1)\\n1\\n2σ2 ∥x −ˆ\\nx(θ, 1)∥2 .\\n(90)\\n17\\n\\n\\nFigure 6: Sender, output and receiver distributions for continuous data. Note that the sender and\\nreceiver distributions have identical variance and the output distribution is a Dirac delta distribution centred\\non the network prediction ˆ\\nx(θ, t).\\nThe noise does not directly affect training, as the reconstruction loss is not optimised. However the\\nvalue of σ places a natural upper limit on the value that should be chosen for σ1: there is no point\\ntransmitting the data to greater precision than it was originally measured. Empirically, we find\\nthat when σ1 < σ/2 the reconstruction loss is very small.\\n4.11\\nDiscrete-Time Loss Ln(x)\\nFrom Eqs. 86 and 88,\\nDKL (pS (· | x, αi) ∥pR(· | θi−1; ti−1, αi)) = DKL\\nN\\nx, α−1\\ni I\\n\\u0001\\n∥N\\nˆ\\nx(θi−1, ti−1), α−1\\ni I\\n\\u0001\\u0001\\n(91)\\n= αi\\n2 ∥x −ˆ\\nx(θi−1, ti−1)∥2 ,\\n(92)\\nand from Eqs. 11 and 72,\\nαi = β(ti) −β(ti−1)\\n(93)\\n= σ−2i/n\\n1\\n−σ−2(i−1)/n\\n1\\n(94)\\n= σ−2i/n\\n1\\n\\u0010\\n1 −σ2/n\\n1\\n\\u0011\\n.\\n(95)\\nTherefore, substituting into Eq. 24,\\nLn(x) = n\\n2\\n\\u0010\\n1 −σ2/n\\n1\\n\\u0011\\nE\\ni∼U{1,n},pF (θi−1|x;ti−1)\\n∥x −ˆ\\nx(θi−1, ti−1)∥2\\nσ2i/n\\n1\\n,\\n(96)\\nwhere ti−1 = (i −1)/n.\\n4.12\\nContinuous-time Loss L∞(x)\\nEq. 29 claimed that\\nDKL (pS (· | x, α) ∥pR(· | θ, α, t)) = DKL\\nN\\ng(x), Cα−1I\\n\\u0001\\n∥P(θ, t) ∗N\\n0, Cα−1I\\n\\u0001\\u0001\\n,\\n(97)\\n18\\n\\n\\nfor some embedding function g : X →Y, constant C and distribution pθ over YD with finite mean\\nand variance. If g is the identity function, C = 1 and\\nP(y | θ, t) = δ(y −ˆ\\nx(θ, t)),\\n(98)\\nthen P(θ, t) has finite mean and variance and\\nN\\ny | g(x), Cα−1I\\n\\u0001\\n= N\\ny | x, α−1I\\n\\u0001\\n= pS (y | x; α) ,\\n(99)\\nP(y | θ, t) ∗N\\n0, Cα−1I\\n\\u0001\\n= N\\ny | ˆ\\nx(θ, t), α−1I\\n\\u0001\\n= pR(y | θ, α, t),\\n(100)\\nso the claim is true and the continuous-time loss from Eq 41 applies, with E[P(θ, t)] = ˆ\\nx(θ, t) and\\nα(t) as defined in Eq 74, yielding\\nL∞(x) = −ln σ1\\nE\\nt∼U(0,1),pF (θ|x;t)\\n∥x −ˆ\\nx(θ, t)∥2\\nσ2t\\n1\\n.\\n(101)\\n4.13\\nPseudocode\\nPseudocode for evaluating the n-step loss Ln(x) and continuous-time loss L∞(x) for continuous\\ndata is presented in Algorithms 1 and 2, while the sample generation procedure is presented in\\nAlgorithm 3.\\n# Note that θ = {µ, ρ}, but ρ is fully determined by t\\n# For our experiments tmin = 1e−6, [xmin, xmax] = [−1, 1]\\nfunction cts output prediction(µ ∈RD, t ∈[0, 1], γ >∈R+, tmin ∈R+, xmin, xmax ∈R)\\nif t < tmin then\\nˆ\\nx(θ, t) ←0\\nelse\\nInput (µ, t) to network, receive ˆ\\nϵ(θ, t) as output\\nˆ\\nx(θ, t) ←µ\\nγ −\\nq\\n1−γ\\nγ ˆ\\nϵ(θ, t)\\nclip ˆ\\nx(θ, t) to [xmin, xmax]\\nend if\\nReturn ˆ\\nx(θ, t)\\nend function\\nAlgorithm 1 Discrete-Time Loss Ln(x) for Continuous Data\\nRequire: σ1 ∈R+, number of steps n ∈N\\nInput: continuous data x ∈RD\\ni ∼U{1, n}\\nt ←i−1\\nn\\nγ ←1 −σ2t\\n1\\nµ ∼N (γ x, γ(1 −γ)I)\\nˆ\\nx(θ, t) ←cts output prediction(µ, t, γ)\\nLn(x) ←\\nn\\n\\u0010\\n1−σ2/n\\n1\\n\\u0011\\n2σ2i/n\\n1\\n∥x −ˆ\\nx(θ, t)∥2\\n19\\n\\n\\nAlgorithm 2 Continuous-Time Loss L∞(x) for Continuous Data\\nRequire: σ1 ∈R+\\nInput: continuous data x ∈RD\\nt ∼U(0, 1)\\nγ ←1 −σ2t\\n1\\nµ ∼N (γ x, γ(1 −γ)I)\\nˆ\\nx(θ, t) ←cts output prediction(µ, t, γ)\\nL∞(x) ←−ln σ1σ−2t\\n1\\n∥x −ˆ\\nx(θ, t)∥2\\nAlgorithm 3 Sample Generation for Continuous Data\\nRequire: σ1 ∈R+, number of steps n ∈N\\nµ ←0\\nρ ←1\\nfor i = 1 to n do\\nt ←i−1\\nn\\nˆ\\nx(θ, t) ←cts output prediction(µ, t, 1 −σ2t\\n1 )\\nα ←σ−2i/n\\n1\\n\\u0010\\n1 −σ2/n\\n1\\n\\u0011\\ny ∼N\\nˆ\\nx(θ, t), α−1I\\n\\u0001\\nµ ←ρµ+α y\\nρ+α\\nρ ←ρ + α\\nend for\\nˆ\\nx(θ, 1) ←cts output prediction(µ, 1, 1 −σ2\\n1)\\nReturn ˆ\\nx(θ, 1)\\n5\\nDiscretised Data\\nThis section considers continuous data that has been discretised into K bins. For example, 8-bit\\nimages are discretised into 256 bins, 16-bit audio is discretised in 216 = 65, 536 bins. This data\\nis represented by tiling [−1, 1] into K intervals, each of length 2/K. Let kl, kc and kr denote\\nrespectively the left, centre and right of interval k, and let {1, K} denote the set of integers from 1\\nto K. Then for k ∈{1, K},\\nkc = 2k −1\\nK\\n−1,\\n(102)\\nkl = kc −1\\nK ,\\n(103)\\nkr = kc + 1\\nK .\\n(104)\\nLet k(x) =\\nk(x(1)), . . . , k(x(D))\\n\\u0001\\n∈{1, K}D be the vector of the indices of the bins occupied by\\nx =\\nx(1), . . . , x(D)\\u0001\\n∈RD, and let kl(x), kc(x) and kr(x) be the corresponding vectors of left edges,\\ncentres and right edges of the bins. If the data has not already been discretised, we set x = kc(x).\\nFor example if the red channel in an 8-bit RGB image has index 110, it will be represented by the\\n20\\n\\n\\nFigure 7: Output distribution for discretised data. For univariate data x discretised into K = 16 bins,\\nthe green line shows the continuous distribution N\\nµx, σ2\\nx\\n\\u0001\\nthat is discretised to yield the output distribution\\npO(x | θ, t), as described in Section 5.1. Bin boundaries are marked with vertical grey lines. The heights of the\\ngreen bars represent the probabilities assigned to the respective bins by pO(x | θ, t). For ease of visualisation\\nthese heights are rescaled relative to the probability density, as indicated on the right axis. Note the clipping\\nat ±1: the area under the dotted green line to the left of −1 is added to the probability of the first bin, the\\narea under the dotted green line to the right of 1 is added to the probability of the last bin.\\nnumber 2∗(110)−1\\n256\\n−1 = −0.14453125. Note that each x(d) therefore lies in the range [ 1\\nK −1, 1 −1\\nK ]\\nand not [−1, 1].\\nThe input distribution pI(x | θ), prior parameters θ0, sender distribution pS (y | x; α), Bayesian\\nupdate function h(θi−1, y, α), Bayesian update distribution pU(θi | θi−1, x; α), Bayesian flow\\ndistribution pF (θ | x; t) and accuracy schedule β(t) are all identical to the continuous case described\\nin Section 4. It may surprise the reader that the output distribution is discretised while the input,\\nsender and receiver distributions are not. We made this choice partly for mathematical convenience\\n(Bayesian updates are considerably more complex for discretised distributions; [1]) and partly\\nbecause we suspected that it would easier for the network to interpret continuous means than\\ndiscrete probabilities as input. In a similar vein to our argument for standard priors in Sec. 4.1, we\\nremind the reader that the input distribution only serves to inform the network and not directly to\\nmodel the data; all that matters is that the input parameters contain enough information to allow\\nthe network to make accurate predictions.\\nSection 4.11 noted that the level of measurement noise assumed for continuous data should\\ninform the choice of standard deviation σ1 for the input distribution at t = 1 (which in turn defines\\nthe accuracy schedule β(t)). For discretised data a similar role is played by the width of the\\ndiscretisation bins, as these place a natural limit on how precisely the data needs to be transmitted.\\nFor example, for 8-bit data with 256 bins and hence a bin width of 1/128, setting σ1 = 1e−3\\ncorresponds to a final input distribution with standard deviation roughly one eighth of the width of\\nthe bin, which should be precise enough for the network to identify the correct bin with very high\\nprobability.\\nOne caveat with discretisation is that calculating the loss has O(K) computational cost, which\\nmay be prohibitive for very finely discretised data. In any case, the benefits of discretisation tend\\nto decrease as the number of bins increases, as we will see in our experiments.\\n5.1\\nOutput Distribution pO(· | θ, t)\\nDiscretised continuous distributions offer a natural and expressive way to model discretised data\\nwith neural networks [38]. As in Section 4.7, the network outputs Ψ(θ, t) are not used to predict\\n21\\n\\n\\nx directly, but rather to model the Gaussian noise vector e used to generate the mean sample µ\\npassed as input to the network.\\nFirst Ψ(θ, t) is split into two length D vectors, µϵ and ln σϵ. Then these are transformed to µx\\nand σx using\\nµx =\\n(\\n0\\nif t < tmin,\\nµ\\nγ(t) −\\nq\\n1−γ(t)\\nγ(t) µϵ\\notherwise,\\n(105)\\nσx =\\n(\\n1\\nif t < tmin,\\nq\\n1−γ(t)\\nγ(t) exp(ln σϵ)\\notherwise.\\n(106)\\nFor each d ∈{1, D}, define the following univariate Gaussian cdf\\nF\\n\\u0010\\nx | µ(d)\\nx , σ(d)\\nx\\n\\u0011\\n= 1\\n2\\n\\\"\\n1 + erf\\n \\nx −µ(d)\\nx\\nσ(d)\\nx\\n√\\n2\\n!#\\n,\\n(107)\\nand clip at [−1, 1] to obtain\\nG\\n\\u0010\\nx | µ(d)\\nx , σ(d)\\nx\\n\\u0011\\n=\\n\\n\\n\\n\\n\\n\\n\\n0\\nif x ≤−1,\\n1\\nif x ≥1,\\nF\\n\\u0010\\nx | µ(d)\\nx , σ(d)\\nx\\n\\u0011\\notherwise.\\n(108)\\nThen, for k ∈{1, K},\\np(d)\\nO (k | θ; t) def\\n= G(kr | µ(d)\\nx , σ(d)\\nx ) −G(kl | µ(d)\\nx , σ(d)\\nx ),\\n(109)\\nand hence\\npO(x | θ,t) =\\nD\\nY\\nd=1\\np(d)\\nO\\n\\u0010\\nk(x(d)) | θ; t\\n\\u0011\\n.\\n(110)\\n5.2\\nReceiver Distribution pR(· | θ; t, α)\\nSubstituting Eq. 110 and Eq. 86 into Eq. 4 gives\\npR(y | θ; t, α) =\\nE\\npO(x′|θ,t) N\\n\\u0010\\ny(d) | kc(x′), α−1I\\n\\u0011\\n(111)\\n=\\nD\\nY\\nd=1\\nZ\\nx′ dx′p(d)\\nO\\nk(x′) | θ; t\\n\\u0001\\nN\\n\\u0010\\ny(d) | kc(x′), α−1\\u0011\\n(112)\\n=\\nD\\nY\\nd=1\\nK\\nX\\nk=1\\np(d)\\nO (k | θ; t)N\\n\\u0010\\ny(d) | kc, α−1\\u0011\\n.\\n(113)\\n22\\n\\n\\nFigure 8: Sender, output and receiver distributions for discretised data. For data x discretised\\ninto 8 bins, the three plots depict the sender distribution (red line), the discretised output distribution (green\\nbars; heights reflect the probabilities assigned to each bin, rescaled as in Figure 7) and receiver distribution\\n(blue line) for progressively increasing values of α, and for progressively more accurate predictions of x (both\\nof which typically happen as t increases). Also shown are the continuous distribution N(x | µx, σ2\\nx) (dotted\\ngreen line) which is discretized to create the output distribution and the continuous receiver distribution from\\nSection 4 (dashed orange line). Bin boundaries are marked with vertical grey lines. Note the KL divergences\\nprinted in the top right: taking discretisation into account leads to a lower KL due to the density “bumps” at\\nthe bin centres where x could be. The advantage of discretisation becomes more pronounced as the prediction\\ngets closer to x and more of the probability mass is concentrated in the correct bin.\\n23\\n\\n\\n5.3\\nReconstruction Loss Lr(x)\\nThe reconstruction loss for discretised data is\\nLr(x) = −\\nE\\npF (θ|x,1) ln pO(x | θ; 1)\\n(114)\\n= −\\nE\\npF (θ|x,1)\\nD\\nX\\nd=1\\nln p(d)\\nO\\n\\u0010\\nk(x(d)) | θ; 1\\n\\u0011\\n.\\n(115)\\n5.4\\nDiscrete-time Loss Ln(x)\\nFrom Eqs. 86 and 113,\\nDKL (pS (· | x, αi) ∥pR(· | θi−1; ti−1, αi))\\n(116)\\n= DKL\\n \\nN\\nx, α−1\\ni I\\n\\u0001\\n∥\\nD\\nY\\nd=1\\nK\\nX\\nk=1\\np(d)\\nO (k | θi−1, ti−1)N\\nkc, α−1\\ni\\n\\u0001\\n!\\n,\\n(117)\\nwhich cannot be calculated in closed form, but can be estimated with Monte-Carlo sampling.\\nSubstituting into Eq. 24,\\nLn(x) = n\\nE\\ni∼U{1,n},pF (θ|x;ti−1),N(y|x,α−1\\ni\\nI)\\nln N\\ny | x, α−1\\ni I\\n\\u0001\\n(118)\\n−\\nD\\nX\\nd=1\\nln\\n K\\nX\\nk=1\\np(d)\\nO (k | θ, ti−1)N\\n\\u0010\\ny(d) | kc, α−1\\ni\\n\\u0011!\\n.\\n(119)\\n5.5\\nContinuous-time Loss L∞(x)\\nJustifying the claim made in Eq. 29 follows almost the same reasoning here as in Section 4.12, with\\nC = 1 and g the identity function. The only difference is that\\nP(y | θ; t) =\\nD\\nY\\nd=1\\nK\\nX\\nk=1\\np(d)\\nO (k | θ, t)δ(y(d) −kc),\\n(120)\\nwhich clearly has finite variance and mean. Since\\nP(y | θ, t) ∗N\\n0, Cα−1I\\n\\u0001\\n= pR(y | θ, α, t),\\n(121)\\nthe claim holds and the continuous time loss from Eq 41 can be applied with\\nE[P(θ, t)] =\\n K\\nX\\nk=1\\np(1)(k | θ, t)kc, . . . ,\\nK\\nX\\nk=1\\np(D)(k | θ, t)kc\\n!\\ndef\\n= ˆ\\nk(θ, t),\\n(122)\\nand α(t) as defined in Eq 74, yielding\\nL∞(x) = −ln σ1\\nE\\nt∼U(0,1),pF (θ|x;t)\\n\\n\\n\\nx −ˆ\\nk(θ, t)\\n\\n\\n\\n2\\nσ2t\\n1\\n.\\n(123)\\nNote that ˆ\\nk(θ, t) is a function of the complete discretised distribution pO(x | θ, t), hence L∞(x)\\ndepends on both µx and σx, and not only on µx, as for continuous data. This also means that\\ncalculating L∞(x) has O(K) computational cost for discretised data.\\n24\\n\\n\\n5.6\\nPseudocode\\nPseudocode for evaluating the discrete-time loss Ln(x) and continuous-time loss L∞(x) for discretised\\ndata is presented in Algorithms 4 and 5, while sample generation is presented in Algorithm 6.\\nfunction discretised cdf(µ ∈R, σ ∈R+, x ∈R)\\nF(x) ←1\\n2\\nh\\n1 + erf\\n\\u0010\\nx−µ\\nσ\\n√\\n2\\n\\u0011i\\nG(x) ←\\n\\n\\n\\n\\n\\n0\\nif x ≤−1\\n1\\nif x ≥1\\nF(x)\\notherwise\\nReturn G(x)\\nend function\\n# For our experiments tmin = 1e−6\\n# kl = 2(k−1)\\nK\\n−1, kr = 2k\\nK −1\\nfunction discretised output distribution(µ ∈RD, t ∈[0, 1], K ∈N, γ ∈R+, tmin ∈R+).\\nif t < tmin then\\nµx ←0\\nσx ←1\\nelse\\nInput (µ, t) to network, receive (µϵ, ln σϵ) as output\\nµx ←µ\\nγ −\\nq\\n1−γ\\nγ\\nµϵ\\nσx ←\\nq\\n1−γ\\nγ\\nexp(ln σϵ)\\nend if\\nfor d ∈{1, D}, k ∈{1, K} do\\np(d)\\nO (k | θ; t) ←discretised cdf(µ(d)\\nx , σ(d)\\nx , kr) −discretised cdf(µ(d)\\nx , σ(d)\\nx , kl)\\nend for\\nReturn pO(· | θ; t)\\nend function\\n25\\n\\n\\nAlgorithm 4 Discrete-Time Loss Ln(x) for Discretised Data\\n# kc = 2k−1\\nK\\n−1\\nRequire: σ1 ∈R+, number of steps n ∈N, number of bins K ∈N\\nInput: discretised data x ∈[ 1\\nK −1, 1 −1\\nK ]D\\ni ∼U{1, n}\\nt ←i−1\\nn\\nγ ←1 −σ2t\\n1\\nµ ∼N (γ x, γ(1 −γ)I)\\nα ←σ−2i/n\\n1\\n\\u0010\\n1 −σ2/n\\n1\\n\\u0011\\ny ∼N\\nx, α−1I\\n\\u0001\\npO(· | θ; t) ←discretised output distribution(µ, t, K, γ)\\nLn(x) ←n\\n\\u0002\\nln N\\ny | x, α−1I\\n\\u0001\\n−P\\nd ln\\nP\\nk p(d)\\nO (k | θ; t)N\\ny(d) | kc, α−1\\u0001\\u0001\\u0003\\nAlgorithm 5 Continuous-Time Loss L∞(x) for Discretised Data\\nRequire: σ1 ∈R+, number of bins K ∈N\\nInput: discretised data x ∈[ 1\\nK −1, 1 −1\\nK ]D\\nt ∼U(0, 1)\\nγ ←1 −σ2t\\n1\\nµ ∼N (γ x, γ(1 −γ)I)\\npO(· | θ; t) ←discretised output distribution(µ, t, K, γ)\\nˆ\\nk(θ, t) ←\\nP\\nk p(1)\\nO (k | θ; t)kc, . . . , P\\nk p(D)\\nO (k | θ; t)kc\\n\\u0001\\nL∞(x) ←−ln σ1σ−2t\\n1\\n\\n\\n\\nx −ˆ\\nk(θ, t)\\n\\n\\n\\n2\\nAlgorithm 6 Sample Generation for Discretised Data\\n# kc =\\n\\u0010\\nk(1)\\nc , . . . , k(D)\\nc\\n\\u0011\\nRequire: σ1 ∈R+, number of steps n ∈N, number of bins K ∈N\\nµ ←0\\nρ ←1\\nfor i = 1 to n do\\nt ←i−1\\nn\\nk ∼discretised output distribution(µ, t, K, 1 −σ2t\\n1 )\\nα ←σ−2i/n\\n1\\n\\u0010\\n1 −σ2/n\\n1\\n\\u0011\\ny ∼N\\nkc, α−1I\\n\\u0001\\nµ ←ρµ+α y\\nρ+α\\nρ ←ρ + α\\nend for\\nk ∼discretised output distribution(µ, 1, K, 1 −σ2\\n1)\\nReturn kc\\n26\\n\\n\\n6\\nDiscrete Data\\nWe now consider discrete data in which no meaningful order or distance exists between the classes,\\nunlike the discretised continuous data covered in the previous section. Some obvious examples are\\ntext characters, classification labels or any binary data. In this context the data is represented as a\\nD dimensional vector of class indices: x =\\nx(1), . . . , x(D)\\u0001\\n∈{1, K}D, where {1, K} is the set of\\nintegers from 1 to K.\\n6.1\\nInput Distribution pI(· | θ)\\nFor discrete data, the input distribution is a factorised categorical over the class indices. Let\\nθ =\\nθ(1), . . . , θ(D)\\u0001\\n∈[0, 1]KD with θ(d) =\\n\\u0010\\nθ(d)\\n1 , . . . , θ(d)\\nK\\n\\u0011\\n∈∆K−1, where θ(d)\\nk\\nis the probability\\nassigned to class k for variable d. Then\\npI(x | θ) =\\nD\\nY\\nd=1\\nθ(d)\\nx(d).\\n(124)\\nThe input prior is uniform with\\nθ0 = 1\\nK ,\\n(125)\\nwhere\\n1\\nK is the length KD vector whose entries are all\\n1\\nK . We chose a uniform prior—rather than\\nan empirical prior fit to the training data—for the same reasons we chose a standard normal prior\\nfor continuous data: it’s mathematically simpler, and the disparity between the true prior and the\\nsimple prior can easily be corrected by the network.\\n6.2\\nOutput Distribution pO(· | θ; t)\\nGiven data x, network inputs θ, t and corresponding network outputs Ψ(θ, t) =\\nΨ(1)(θ, t), . . . , Ψ(D)(θ, t)\\n\\u0001\\n∈RKD, the output distribution for discrete data is as follows:\\np(d)\\nO (k | θ; t) =\\n\\u0010\\nsoftmax(Ψ(d)(θ, t))\\n\\u0011\\nk ,\\n(126)\\npO(x | θ; t) =\\nD\\nY\\nd=1\\np(d)\\nO (x(d) | θ; t).\\n(127)\\nNote that for binary data only the probability θ(d)\\n1\\nthat k = 1 is fed into the network, on the grounds\\nthat the probability of k = 2 can easily be inferred from θ(d)\\n2\\n= 1 −θ(d)\\n1 . The output distribution for\\nbinary data is determined by applying the logistic sigmoid function elementwise to the length D\\noutput vector to get the probability for k = 1:\\np(d)\\nO (1 | θ; t) = σ\\n\\u0010\\nΨ(d)(θ, t))\\n\\u0011\\n,\\n(128)\\nwhere\\nσ(x) =\\n1\\n1 −e−x ,\\n(129)\\n27\\n\\n\\nthen inferring the probabilities for k = 2 from\\np(d)\\nO (2 | θ; t) = 1 −p(d)\\nO (1 | θ; t).\\n(130)\\nIn principle one class could also be removed from the inputs and outputs when K > 2 and inferred\\nfrom the others. However this would require the network to internalise a slightly more sophisticated\\ninference procedure that could potentially slow down learning. We therefore followed deep-learning\\nconvention and included a redundant input and output unit for K > 2.\\nAll probabilities are rescaled to the range [−1, 1] by multiplying by two then subtracting one\\nbefore feeding them into the network.\\n6.3\\nSender Distribution pS (· | x; α)\\nGiven ω ∈[0, 1], and a vector of D class indices k =\\nk(1), . . . , k(D)\\u0001\\n∈{1, K}D, let\\np(k(d) | x(d); ω) def\\n= 1 −ω\\nK\\n+ ωδk(d)x(d),\\n(131)\\nwhere δij is the Kronecker delta function.\\nClearly p(k(d) | x(d); ω) ≥0 ∀k and PK\\nk=1 p(k(d) |\\nx(d); ω) = 1, so the vector\\na(x(d), ω) def\\n=\\n\\u0010\\np(1 | x(d); ω), . . . , p(K | x(d); ω)\\n\\u0011\\n,\\n(132)\\ndefines a valid distribution over K classes. To simplify notation we will from now on drop the\\nsuperscripts and refer to x(d) as x, p(k(d) | x(d); ω) as p(k | x; ω) and so on, except where necessary\\nto remove ambiguity.\\nConsider a vector of integer counts c = (c1, . . . , cK) ∈{1, m}K, corresponding to the number\\nof times each of the K classes is observed among m independent draws from a(x, ω). Then the\\nprobability of observing c is given by the following multinomial distribution:\\np(c | x, ω) = Multi(m, a(x, ω))\\n(133)\\n=\\nm!\\nc1! . . . cK!\\nK\\nY\\nk=1\\n(p(k | x; ω))ck\\n(134)\\n=\\nm!\\nc1! . . . cK!\\nK\\nY\\nk=1\\n\\u00121 −ω\\nK\\n+ ωδkd\\n\\u0013ck\\n.\\n(135)\\nNow consider the fraction ck/m of observations of class k in c. Clearly\\nlim\\nm→∞\\nck\\nm = p(k | x; ω),\\n(136)\\nmeaning that for any finite ω it would be possible to deduce from c what the value of x is if m\\nis sufficiently large. However as ω shrinks, p(k | x; ω) becomes closer to uniform, meaning that a\\nlarger m is required to unambigously identify x from c. By defining the accuracy α def\\n= mω2 and\\nsending m →∞(and hence ω →0 for any finite α), p(c | x, ω) can therefore be used to define a\\ncontinuous-valued sender distribution that smoothly varies from totally uninformative at α = 0 to\\ntotally informative as α →∞, like the sender distribution for continuous data.\\n28\\n\\n\\nIt can be proved from the central limit theorem that for any set of discrete probabilities\\np = {p1, . . . , pK}, where 0 < pk < 1 ∀k, that if c ∼Multi(m, p) then in the limit m →∞the\\nfollowing result holds [8]:\\nc −mp\\n√mp ∼N (0, I) ,\\n(137)\\nwhere I is the K × K identity matrix. Therefore\\nlim\\nm→∞p(ck | x, ω) = N (ck | mp(k | x; ω), mp(k | x; ω))\\n(138)\\n=\\n1\\np\\n2πmp(k | x; ω)\\nexp\\n \\n−[ck −mp(k | x, ω)]2\\n2mp(k | x; ω)\\n!\\n.\\n(139)\\nNow define\\nξ def\\n= 1 + ωK\\n1 −ω.\\n(140)\\nAnd the length K sender sample y = (y1, . . . , yK) as\\nyk\\ndef\\n=\\n\\u0010\\nck −m\\nK\\n\\u0011\\nln ξ.\\n(141)\\nNote that y, unlike x, is continuous (Y = RK, X = {1, K}), and that\\nc −m\\nK\\n\\u0001\\nmeasures the number\\nof times each class is observed, minus the average number of observations per class. Intuitively, y\\nprovides information about the relative concentration of the classes among the counts, with (since\\nln ξ > 0) positive values for classes observed more frequently than the mean and negative values for\\nthose observed less frequently than the mean. As mω2 grows the concentration increases around\\nthe true class, and hence y become more informative about x.\\nRearranging Eq. 141,\\nck = yk\\nln ξ + m\\nK\\n(142)\\n=\\n⇒dck\\ndyk\\n=\\n1\\nln ξ ,\\n(143)\\nwhich we can use for the following change of variables:\\np(yk | x, ω) =\\n\\f\\n\\f\\n\\f\\n\\f\\ndck\\ndyk\\n\\f\\n\\f\\n\\f\\n\\f p(ck | x, ω)\\n(144)\\n=\\n1\\nln ξ\\np\\n2πmp(k | x, ω)\\nexp\\n\\n\\n\\n−\\nh\\nyk\\nln ξ + m\\nK −mp(k | x, ω)\\ni2\\n2mp(k | x, ω)\\n\\n\\n,\\n(145)\\nwhere we have used the fact that ξ ≥1 and hence dck\\ndyk ≥0. Recall that α = mω2 and hence m = α\\nω2 ,\\nwhich can be substituted into the above to yield\\np(yk | x, ω) =\\n1\\n1\\nω ln ξ\\n1\\np\\n2παp(k | x, ω)\\nexp\\n\\n\\n\\n−\\nh\\nyk\\n1\\nω ln ξ + α\\nω\\n 1\\nK −p(k | x, ω)\\n\\u0001i2\\n2αp(k | x, ω)\\n\\n\\n.\\n(146)\\n29\\n\\n\\nSubstituting from Eq. 131,\\n1\\nK −p(k | x, ω) = ω\\n\\u0012 1\\nK −δkx,\\n\\u0013\\n,\\n(147)\\nand hence\\np(yk | x, ω) =\\n1\\n1\\nω ln ξ\\n1\\np\\n2παp(k | x, ω)\\nexp\\n\\n\\n\\n−\\nh\\nyk\\n1\\nω ln ξ −α\\nδkx −1\\nK\\n\\u0001i2\\n2αp(k | x, ω)\\n\\n\\n.\\n(148)\\nApplying the identity ln(1 + x) = P∞\\nn=1\\n(−1)n−1\\nn\\nxn for |x| < 1 to ln ξ = ln\\n\\u0010\\n1 + ωK\\n1−ω\\n\\u0011\\nit can be seen\\nthat\\nln ξ ∈ωK\\n1 −ω + O(ω2),\\n(149)\\nand hence\\nlim\\nω→0\\n1\\nω ln ξ = K.\\n(150)\\nFurthermore, it follows directly from Eq. 131 that\\nlim\\nω→0 p(k | x, ω) = 1\\nK ∀k ∈{1, K}.\\n(151)\\nNow define\\npS (yk | x; α) def\\n= lim\\nω→0 p(yk | x, ω).\\n(152)\\nPlugging Eq. 150 and 151 into Eq. 148,\\npS (yk | x; α) =\\n1\\nK\\nq\\n2πα 1\\nK\\nexp\\n \\n−\\n\\u0002 yk\\nK −α\\nδkx −1\\nK\\n\\u0001\\u00032\\n2α 1\\nK\\n!\\n(153)\\n=\\n1\\n√\\n2παK\\nexp\\n \\n−[yk −α (Kδkx −1)]2\\n2αK\\n!\\n(154)\\n= N (α (Kδkx −1) , αK) .\\n(155)\\nRestoring the superscript,\\npS\\n\\u0010\\ny(d) | x(d); α\\n\\u0011\\n= N (α (Kex(d) −1) , αKI) ,\\n(156)\\nwhere 1 is a vector of ones, I is the identity matrix and ej ∈RK is the projection from the class\\nindex j to the length K one-hot vector defined by (ej)k = δjk, and therefore\\npS (y | x; α) = N (y | α (Kex −1) , αKI) ,\\n(157)\\nwhere ex\\ndef\\n= (ex(1), . . . , ex(D)) ∈RKD.\\n30\\n\\n\\n6.4\\nReceiver Distribution pR(· | θ; t, α)\\nSubstituting Eq. 127 and Eq. 157 into Eq. 4 gives the following receiver distribution for dimension\\nd:\\np(d)\\nR (y(d) | θ; t, α) =\\nK\\nX\\nk=1\\np(d)\\nO (k | θ; t)N (α (Kek −1) , αKI) ,\\n(158)\\npR(y | θ; t, α) =\\nD\\nY\\nd=1\\np(d)\\nR (y(d) | θ; t, α).\\n(159)\\n6.5\\nBayesian Update Function h(θi−1, y, α)\\nRecall from Section 6.1 that (θi−1)(d)\\nk\\nis the probability assigned to x(d) = k by p(x(d) | θi−1).\\nDropping the superscript and returning to the count distribution p(c | x, ω) defined in Eq. 133, the\\nposterior probability that x = k after observing c is\\np(k | c; ω) =\\np(c | k; ω)(θi−1)k\\nPK\\nk′=1 p(c | k′; ω)(θi−1)k′\\n.\\n(160)\\nSubstituting Eq. 135 into Eq. 160 and cancelling terms in the enumerator and denominator,\\np(k | c; ω) =\\n\\u0002 1−ω\\nK\\n\\u0003m−ck \\u0002 1−ω\\nK + ω\\n\\u0003ck (θi−1)k\\nPK\\nk′=1\\n\\u0002 1−ω\\nK\\n\\u0003m−ck′ \\u0002 1−ω\\nK + ω\\n\\u0003ck′ (θi−1)k′\\n(161)\\n=\\n\\u0002 1−ω\\nK\\n\\u0003m h\\n1 + ωK\\n1−ω\\nick (θi−1)k\\n\\u0002 1−ω\\nK\\n\\u0003m PK\\nk′=1\\nh\\n1 + ωK\\n1−ω\\nick′\\n(θi−1)k′\\n(162)\\n=\\nh\\n1 + ωK\\n1−ω\\nick (θi−1)k\\nPK\\nk′=1\\nh\\n1 + ωK\\n1−ω\\nick′\\n(θi−1)k′\\n(163)\\n=\\nξck(θi−1)k\\nPK\\nk′=1 ξck′(θi−1)k′\\n.\\n(164)\\nNow define\\nh(θ, y) def\\n=\\neyθ\\nPK\\nk=1 eykθk\\n.\\n(165)\\nSubstituting the definition of yk from Eq. 141 into the definition of h(θ, y) from Eq. 165,\\n(h(θi−1, y))k =\\nexp(−m\\nK ln ξ) exp(ck ln ξ)(θi−1)k\\nexp(−m\\nK ln ξ) PK\\nk′=1 exp(ck′ ln ξ)(θi−1)k′\\n(166)\\n=\\nexp(ln ξck)(θi−1)k\\nPK\\nk′=1 exp(ln ξck′)(θi−1)k′\\n(167)\\n=\\nξck(θi−1)k\\nPK\\nk′=1 ξck′(θi−1)k′\\n,\\n(168)\\n(169)\\n31\\n\\n\\nand hence, from Eq. 164,\\nh(θi−1, y)k = p(k | c; ω).\\n(170)\\nTherefore in the limit m →∞with mω2 = α, the stochastic parameter update from θi−1 to θi\\ninduced by drawing c from multi(m, a(x, ω)) can be sampled by first drawing y from pS (· | x, α)\\nthen setting θi = h(θi−1, y). Hence the Bayesian update function is\\nh(θi−1, y, α) def\\n=\\neyθi−1\\nPK\\nk=1 eyk(θi−1)k\\n,\\n(171)\\nwhere the redundant parameter α has been included for consistency with the update function for\\ncontinuous data.\\n6.6\\nBayesian Update Distribution pU(· | θi−1, x; α)\\nSubstituting Eqs. 157 and 171 into Eq. 6,\\npU (θ | θi−1, x; α) =\\nE\\nN(y|α(Kex−1),αKI) δ\\n \\nθ −\\neyθi−1\\nPK\\nk=1 eyk(θi−1)k\\n!\\n.\\n(172)\\n6.7\\nAdditive Accuracies\\nIt follows from the definition of the update distribution that if ya is drawn from pS (· | x; αa) then\\nθi−1 = h(ya, θi−2) is drawn from p(· | θi−2, x; αa). Furthermore, if yb is drawn from pS (· | x; αb) then\\nθi = h(yb, θi−1) = h(yb, h(ya, θi−2)) is drawn from EpU (θi−1|θi−2,x;αa) pU (θi | θi−1, x; αb). Substituting\\nthe definition of h from Eqn 165,\\nh(yb, h(ya, θi−2)) =\\nexp(yb)\\nexp(ya)θi−2\\nPK\\nk′=1 exp((ya)k′)(θi−2)k′\\nPK\\nk=1 exp ((yb)k)\\nexp((ya)k)(θi−2)k\\nPK\\nk′=1 exp((ya)k′)(θi−2)k′\\n(173)\\n=\\nexp(yb) exp(ya)θi−2\\nPK\\nk=1 exp ((yb)k) exp ((ya)k) (θi−2)k\\n(174)\\n=\\nexp(ya + yb)θi−2\\nPK\\nk=1 exp ((ya + yb)k) (θi−2)k\\n(175)\\n= h(ya + yb, θi−2).\\n(176)\\nFrom Eqn. 156\\nya ∼N (αa (Kex −1) , αaKI) ,\\n(177)\\nyb ∼N (αb (Kex −1) , αbKI)\\n(178)\\n(179)\\nand hence, from Identity 61\\nya + yb ∼N ((αa + αb) (Kex −1) , (αa + αb)KI) .\\n(180)\\n32\\n\\n\\nTherefore, if y is drawn from pS (· | x; αa + αb) and θi = h(y, θi−2) then θi is drawn from\\nEpU (θi−1|θi−2,x;αa) pU (θi | θi−1, x; αb) and\\nE\\npU (θi−1|θi−2,x;αa) pU (θi | θi−1, x; αb) = pU (θi | θi−2, x; αa + αb),\\n(181)\\nas required.\\n6.8\\nAccuracy Schedule β(t)\\nAs with continuous data, the guiding heuristic for β(t) was to decrease the expected entropy of\\nthe input distribution linearly with t. In the continuous case, where the entropy is a deterministic\\nfunction of σ2, applying the heuristic was straightforward; in the discrete case an explicit computation\\nof EpF (θ|x;t) H [pI(x | θ)] would be needed. We were unable to derive an analytic expression for this\\nterm, but found that\\nβ(t) = t2β(1)\\n(182)\\nwas a reasonable approximation, with β(1) determined empirically for each experiment. Therefore\\nα(t) = dβ(t)\\ndt\\n= β(1)2t.\\n(183)\\n6.9\\nBayesian Flow Distribution pF (· | x; t)\\nSubstituting Eq. 172 into Eq. 10,\\npF (θ | x; t) =\\nE\\nN(y|β(t)(Kex−1),β(t)KI) δ\\n \\nθ −\\neyθ0\\nPK\\nk=1 eyk(θ0)k\\n!\\n.\\n(184)\\nSince the prior is uniform with θ0 = 1\\nK , this reduces to\\npF (θ | x; t) =\\nE\\nN(y|β(t)(Kex−1),β(t)KI) δ (θ −softmax(y)) ,\\n(185)\\nwhich can be sampled by drawing y from N (β(t) (Kex −1) , β(t)KI) then setting θ = softmax(y).\\nThe sender distribution for discrete data can therefore be interpreted as a source of softmax logits\\nfor the Bayesian flow distribution; the higher the sender accuracy α is, the larger in expectation the\\nlogits corresponding to x will be in y, hence the closer θ will be to ex and the more information the\\nnetwork will gain about x.\\n6.10\\nReconstruction Loss Lr(x)\\nThe reconstruction loss for discrete data is\\nLr(x) = −\\nE\\npF (θ|x,1) ln pO(x | θ; 1).\\n(186)\\n33\\n\\n\\nFigure 9: Accuracy schedule vs. expected entropy for discrete data. The surface plot shows the\\nexpectation over the parameter distribution p(θ | x; β) of the entropy of the categorical input distribution\\np(x | θ) for K = 2 to 30 and √β = 0.01 to 3. The red and cyan lines highlight the entropy curves for 2 and\\n27 classes, the two values that occur in our experiments. The red and cyan stars show the corresponding\\nvalues we chose for\\np\\nβ(1).\\n6.11\\nDiscrete-time Loss Ln(x)\\nFrom Eqs. 156 and 158,\\nDKL\\n\\u0010\\npS\\n\\u0010\\n· | x(d); α\\n\\u0011\\n∥p(d)\\nR (· | θ; t, α)\\n\\u0011\\n(187)\\n= DKL\\n \\nN (α (Kex(d) −1) , αKI) ∥\\nK\\nX\\nk=1\\np(d)\\nO (k | θ; t)N (α (Kek −1) , αKI)\\n!\\n.\\n(188)\\nTherefore, substituting into Eq. 24,\\nLn(x) = n\\nE\\ni∼U{1,n},p(θ|x;ti−1),N(y|αi(Kex−1),αiKI) ln N (y | αi (Kex −1) , αiKI)\\n(189)\\n−\\nD\\nX\\nd=1\\nln\\n K\\nX\\nk=1\\np(d)\\nO (k | θ; ti−1)N\\n\\u0010\\ny(d) | αi (Kek −1) , αiKI\\n\\u0011!\\n,\\n(190)\\n34\\n\\n\\nFigure 10: Bayesian flow for discrete data. For K = 3, the input distribution parameters θ = (θ1, θ2, θ3)\\ncan be visualised as points on the 2-simplex, with the data x corresponding to the bottom left corner. For\\nthe accuracy schedule β(t) from Eq. 182, the white line shows a single input parameter trajectory starting\\nfrom θ0 =\\n 1\\n3, 1\\n3, 1\\n3\\n\\u0001\\nand evolving under the Bayesian update distribution pU (θi | θi−1; x, β(ti) −β(ti−1))\\nfrom Eq. 172, superimposed on log-scale heatmaps of the Bayesian flow distribution pF (θ | x; t) from Eq. 185,\\nplotted at regular intervals from t = 0.02 to 1.\\nwhere, from Eq. 182,\\nαi = β(ti) −β(ti−1)\\n(191)\\n= β(1)\\n \\u0012 i\\nn\\n\\u00132\\n−\\n\\u0012i −1\\nn\\n\\u00132!\\n(192)\\n= β(1)\\n\\u00122i −1\\nn2\\n\\u0013\\n.\\n(193)\\n35\\n\\n\\nFigure 11: Bayesian flow for binary data. For the input probability p1 of class one, the plot shows\\nseveral parameter trajectories starting from p1 = 0.5 at t = 0 and evolving under the Bayesian update\\ndistribution to t = 1, superimposed on a log-scale heatmap of the Bayesian flow distribution. β(1) = 4 in this\\nplot. Note that both here and in Figure 10 the convergence towards the data appears slower and noisier than\\nthe equivalent trajectories for continuous data in Figure 4. This is a fundamental consequence of discreteness:\\nsince all points in X are equidistant the input distributions cannot concentrate on values close to x as the\\ntrajectories progress.\\n6.12\\nContinuous-time Loss L∞(x)\\nLet\\nv def\\n= y\\nα + 1,\\n(194)\\nand apply Identity 51 to see that if\\ny(d) ∼pS\\n\\u0010\\n· | x(d); α\\n\\u0011\\n= N (α(Kex(d) −1), αKI) ,\\n(195)\\nthen\\nv(d) ∼N\\n\\u0012\\nKex(d), K\\nα I\\n\\u0013\\n,\\n(196)\\nand similarly if\\ny(d) ∼p(d)\\nR (· | θ; t, α) =\\nK\\nX\\nk=1\\np(d)\\nO (k | θ; t)N\\n\\u0010\\ny(d) | α (Kek −1) , αKI\\n\\u0011\\n,\\n(197)\\nthen\\nv(d) ∼\\nK\\nX\\nk=1\\np(d)\\nO (k | θ; t)N\\n\\u0012\\nKek, K\\nα I\\n\\u0013\\n(198)\\n= K\\nK\\nX\\nk=1\\np(d)\\nO (k | θ; t)δ(· −ek) ∗N\\n\\u0012\\n0, K\\nα I\\n\\u0013\\n.\\n(199)\\n36\\n\\n\\nThe Kullback-Leibler divergence is invariant under affine transformations of variables, hence\\nDKL\\n\\u0010\\npS\\n\\u0010\\n· | x(d); α\\n\\u0011\\n∥p(d)\\nR (· | θ; t, αi)\\n\\u0011\\n(200)\\n= DKL\\n \\nN\\n\\u0012\\nKex(d), K\\nα I\\n\\u0013\\n∥\\nK\\nX\\nk=1\\np(d)\\nO (k | θ; t)Kδ(· −ek) ∗N\\n\\u0012\\n0, K\\nα I\\n\\u0013!\\n.\\n(201)\\nNow set C = K, g(x(d)) = Kex(d) and\\nP (d)(θ, t) = K\\nK\\nX\\nk=1\\np(d)\\nO (k | θ; t)δ(· −ek),\\n(202)\\nwhich has finite variance and the following finite expectation\\nE[P (d)(θ, t)] = Kˆ\\ne(d)(θ, t),\\n(203)\\nwhere\\nˆ\\ne(d)(θ, t) def\\n=\\nK\\nX\\nk=1\\np(d)\\nO (k | θ; t)ek.\\n(204)\\nThe conditions in Eq. 29 are therefore satisfied and Eqs. 203 and 183 can be substituted into Eq. 41\\nto yield\\nL∞(x) = Kβ(1)\\nE\\nt∼U(0,1),pF (θ|x,t) t∥ex −ˆ\\ne(θ, t)∥2,\\n(205)\\nwhere\\nˆ\\ne(θ, t) def\\n=\\n\\u0010\\nˆ\\ne(1)(θ, t), . . . ,ˆ\\ne(D)(θ, t)\\n\\u0011\\n.\\n(206)\\n6.13\\nPseudocode\\nPseudocode for evaluating the discrete-time loss Ln(x) and continuous-time loss L∞(x) for discrete\\ndata is presented in Algorithms 7 and 8, while sample generation is presented in Algorithm 9.\\nfunction discrete output distribution(θ ∈[0, 1]KD, t ∈[0, 1])\\nInput (θ, t) to network, receive Ψ(θ, t) as output\\nfor d ∈{1, D} do\\nif k = 2 then\\np(d)\\nO (1 | θ; t) ←σ\\nΨ(d)(θ, t)\\n\\u0001\\np(d)\\nO (2 | θ; t) ←1 −p(d)\\nO (1 | θ; t)\\nelse\\np(d)\\nO (· | θ; t) ←softmax(Ψ(d)(θ, t))\\nend if\\nend for\\nReturn pO(· | θ; t)\\nend function\\n37\\n\\n\\nAlgorithm 7 Discrete-Time Loss Ln(x) for Discrete Data\\nRequire: β(1) ∈R+, number of steps n ∈N, number of classes K ∈N\\nInput: discrete data x ∈{1, K}D\\ni ∼U{1, n}\\nt ←(i −1)/n\\nβ ←β(1)t2\\ny′ ∼N (β (Kex −1) , βKI)\\nθ ←softmax(y′)\\npO(· | θ; t) ←discrete output distribution(θ, t)\\nα ←β(1)\\n 2i−1\\nn2\\n\\u0001\\ny ∼N (α (Kex −1) , αKI)\\nLn(x) ←n\\n\\u0002\\nln N (y | α (Kex −1) , αKI) −P\\nd ln\\nP\\nk p(d)\\nO (k | θ; t)N\\ny(d) | α (Kek −1) , αKI\\n\\u0001\\u0001\\u0003\\nAlgorithm 8 Continuous-Time Loss L∞(x) for Discrete Data\\nRequire: β(1) ∈R+, number of classes K ∈N\\nInput: discrete data x ∈{1, K}D\\nt ∼U(0, 1)\\nβ ←β(1)t2\\ny ∼N (β (Kex −1) , βKI)\\nθ ←softmax(y)\\npO(· | θ; t) ←discrete output distribution(θ, t)\\nˆ\\ne(θ, t) ←\\nP\\nk p(1)\\nO (k | θ; t)ek, . . . , P\\nk p(D)\\nO (k | θ; t)ek\\n\\u0001\\nL∞(x) ←Kβ(1)t ∥ex −ˆ\\ne(θ, t)∥2\\nAlgorithm 9 Sample Generation for Discrete Data\\nRequire: β(1) ∈R+, number of steps n ∈N, number of classes K ∈N\\nθ ←\\n 1\\nK\\n\\u0001\\nfor i = 1 to n do\\nt ←i−1\\nn\\nk ∼discrete output distribution(θ, t)\\nα ←β(1)\\n 2i−1\\nn2\\n\\u0001\\ny ∼N (α (Kek −1) , αKI)\\nθ′ ←eyθ\\nθ ←\\nθ′\\nP\\nk θ′\\nk\\nend for\\nk ∼discrete output distribution(θ, 1)\\nReturn k\\n38\\n\\n\\nModel\\nDynamically Binarized MNIST\\nCIFAR-10\\nImproved DDPM [28]\\n2.94\\nNVAE [49]\\n78.01\\n2.91\\nPixelVAE++† [35]\\n78.00\\n2.90\\nLocally Masked PixelCNN† [15]\\n77.58\\n2.89\\nImage Transformer† [30]\\n2.89\\nDDPM++ [16]\\n2.88\\nLSGM [50]\\n2.87\\nVDVAE [3]\\n2.87\\nSparse Transformer† [4]\\n2.80\\nReflected Diffusion [23]\\n2.68\\nVDM [17]\\n2.65\\nARDM-Upscale 4 [13]\\n2.64\\nBFN\\n77.87\\n2.66\\nCR-NVAE* [40]\\n76.93\\n2.51\\nVDM* [17]\\n2.49\\nTable 1: Comparison of dynamically binarized MNIST and CIFAR-10 results with other\\nmethods. The best published results for both datasets (*) use data augmentation for regularization. Results\\nfor models marked with (†) are exact values; all other results are upper bounds.\\nn-steps\\n10\\n25\\n50\\n100\\n784\\n1000\\n∞\\nNPI\\n95.21\\n84.40\\n81.06\\n79.46\\n78.02\\n78.07\\n77.87\\nTable 2: Dynamically binarized MNIST results. NPI is nats per image averaged over 2,000 passes\\nthrough the test set with Ln(x) or L∞(x) sampled once per test image per pass. The reconstruction loss\\nLr(x) (included in NPI) was 0.46. 784 is the total number of pixels per image, hence the number of steps\\nrequired to generate an image with an autoregressive model.\\n7\\nExperiments\\nWe evaluated Bayesian Flow Networks (BFNs) on the following generative benchmarks: CIFAR-10\\n(32×32 8-bit color images), dynamically binarized MNIST (28×28 binarized images of handwritten\\ndigits) and text8 (length 256 character sequences with a size 27 alphabet). The continuous (Sec. 4)\\nand discretised (Sec. 5) versions of the system were compared on CIFAR-10, while the discrete\\nversion (Sec. 6) was applied to the other datasets. In all cases, the network was trained using\\nthe continuous-time loss L∞(x), with the discrete-time loss Ln(x) evaluated for testing only, with\\nvarious values of n. Standard network architectures and training algorithms were used throughout to\\nallow for direct comparison with existing methods. Because the focus of this paper is on probabilistic\\nmodelling rather than image generation, FID scores were not calculated. However, examples of\\ngenerated data are provided for all experiments.\\n39\\n\\n\\n(a) Test Data\\n(b) Generated Data\\nFigure 12: MNIST real and generated data. Samples generated with 100 steps.\\n7.1\\nDynamically Binarized MNIST\\nData.\\nThe binarized MNIST benchmark data was originally created from the MNIST dataset\\nof handwritten images [20] by treating the grayscale pixel intensities as Bernoulli probabilities\\nand sampling a particular binarization [36] which is held fixed during training. In recent years, a\\nvariant of the same benchmark has become more popular, with a new binarization sampled from\\nthe probabilities for every training batch. The two are not comparable, as the latter, which we refer\\nto as dynamically binarized MNIST, effectively has a larger training set and hence gives better test\\nset performance. All our experiments and the results referenced from the literature use dynamically\\nbinarized MNIST.\\nSetup.\\nThe network architecture was based on a U-Net introduced for diffusion models [28].\\nStarting from the hyperparameters used for the CIFAR-10 dataset (see Appendix A in the above\\nreference), we made the following modifications: the number of resblocks was reduced from three\\nto two and the layer widths were reduced from [C, 2C, 2C, 2C] to [C, 2C, 2C] with C = 128. Fi-\\nnally, the input and output of the standard network were concatenated and projected back to\\nthe output size. 600 randomly selected training images (1% of the training set) were used as a\\nvalidation set. The optimiser was AdamW [22] with learning rate 0.0001, weight decay 0.01 and\\n(β1, β2) = (0.9, 0.98). Dropout was used with probability 0.5, the training batch size was 512, and\\nβ(1) was set to 3 (see Sec. 6.8). The network was trained for 150 000 weight updates until early stop-\\nping. An exponential moving average of model parameters with a decay rate of 0.9999 was used for\\nevaluation and sample generation. The total number of learnable parameters was approximately 25M.\\nResults.\\nAs can be seen from Table 1, BFN is close to state-of-the-art for this task with no data\\naugmentation. Table 2 shows the expected inverse relationship between loss and number of steps.\\nDirect optimisation of the n-step loss would likely lead to reduced loss for low values of n; however\\n40\\n\\n\\n(a) Input Distribution\\n(b) Output Distribution\\nFigure 13: MNIST Input and output distributions. For two test set images the figure shows the white\\npixel probability at 20 steps evenly spaced between t = 0 and t = 1/3. Note how the input probabilities are\\ninitially uniform whereas the output distribution initially predicts a superposition of multiple digits, closely\\nmatching the per-pixel marginal prior over the training set: this supports our belief that the network learns\\nto correct for the uniform prior in the input distribution. Also note that the output distribution is much\\nless noisy than the input distribution, and that it changes more dramatically as new information is received\\n(e.g. the network appears to switch from predicting a 6 to a 2 to a 7 for the first image). This highlights the\\nnetwork’s use of context to resolve ambiguity and noise in the input distribution.\\nFigure 14: MNIST losses against time. The left plot shows the mean over the test set of the cts. time\\nloss L∞(x) used for training for transmission time t between 0 and 1. The right plot shows the average\\ncumulative value of L∞(x) up to t, along with the reconstruction loss Lr(x) evaluated at t and the sum of\\nthese two losses, which would be the total loss if the transmission process halted at t. Note the unevenness\\nof L∞(x) against t: we speculate that rescaling β(t) to make the loss curve more uniform could improve\\nperformance.\\nwe leave that for future work. One issue is that the reconstruction loss was relatively high at 0.46\\nnats per image. The obvious way to decrease this would be to increase β(1), but we found that\\ndoing so led to slower learning and worse performance. Along with the loss curves in Figure 14, this\\nsuggests that the accuracy schedule is suboptimal for binary data.\\n41\\n\\n\\nn-steps\\nCts. (256 bins)\\nDiscd. (256 bins)\\nCts. (16 bins)\\nDiscd. (16 bins)\\n10\\n6.18\\n3.91\\n1.42\\n1.16\\n25\\n3.65\\n3.16\\n1.11\\n1.02\\n50\\n3.10\\n2.93\\n1.03\\n0.98\\n100\\n2.86\\n2.81\\n0.99\\n0.96\\n250\\n2.73\\n2.73\\n0.97\\n0.94\\n500\\n2.69\\n2.71\\n0.96\\n0.94\\n1000\\n2.67\\n2.70\\n0.96\\n0.94\\n∞\\n2.66\\n2.68\\n0.96\\n0.94\\nLr(x)\\n0.001\\n0.003\\n0.073\\n0.070\\nUpdates\\n5M\\n5M\\n250K\\n1M\\nTable 3: CIFAR-10 results. All losses are bits per dimension (BPD) averaged over 100 passes through\\nthe test set with Ln(x) or L∞(x) sampled once per test image per pass. The reconstruction losses Lr(x)\\n(included in BPD) and the number of training updates for each network are shown below.\\n7.2\\nCIFAR-10\\nData.\\nTwo sets of generative modelling experiments were conducted on the CIFAR-10 database [19],\\none at the standard bit-depth of 8, corresponding to 256 discretised bins per colour channel, and\\none at a reduced bit-depth of 4, corresponding to 16 bins per channel. In both cases the bins\\nevenly partitioned the interval [−1, 1] and the data was pre-processed by assigning each channel\\nintensity to the nearest bin centre, as described in Section 5. The purpose of comparing 16 and 256\\nbin discretisation was twofold: (1) to test the hypothesis that the advantage of training with the\\ndiscretised loss from Section 5 rather than the continuous loss from Section 4 would be greater when\\nthe number of bins was lower, and (2) to test whether modelling the data at lower precision would\\nlead to improved perceptual quality. No data augmentation, such as horizontal flips or random\\ncrops, was used on the training set.\\nSetup.\\nThe network architecture was essentially the same as that used for Variational Diffusion\\nModels (VDMs [17]), including the Fourier feature inputs. The only modification was an extra\\ninput-output connection similar to the network for MNIST. In total there were approximately 31M\\nlearnable parameters. The following hyperparameters were used for all CIFAR-10 experiments: a\\nvalidation set of 500 randomly selected training images (1% of the training set), the AdamW [22]\\noptmizer with weight decay 0.01, learning rate 0.0002 and (β1, β2) = (0.9, 0.99), dropout with\\nprobability 0.1, training batch size of 128, tmin = 1e−6, [xmin, xmax] = [−1, 1], and an exponential\\nmoving average of model parameters with a decay rate of 0.9999 for evaluation and sample generation.\\nFor the 256 bin experiments σ1 = 0.001, while for the 16 bin experiments σ1 =\\n√\\n0.001. For the\\nnetworks trained with continuous loss, the reconstruction loss was measured using the discretised\\nversion of Lr(x) from Section 5.3 rather than the continuous version from Section 4.10, using a\\ndiscretised Gaussian with mean equal to ˆ\\nx(θ, 1) and std. deviation chosen empirically to be σ1 for\\n256 bins and 0.7σ1 for 16 bins. This ensured the results were comparable between continuous and\\ndiscretised training, and consistent with the literature.\\n42\\n\\n\\n(a) Test Data (256 bins)\\n(b) Generated Data (256 bins)\\n(c) Test Data (16 bins)\\n(d) Generated Data (16 bins)\\nFigure 15: CIFAR-10 real and generated data. Samples generated with 4,000 steps, using networks\\ntrained with discretised loss. The same random seed was used for both sets of samples. Note the improved\\nimage quality of the 16 bin samples compared to the 256 bin samples.\\nResults.\\nTable 1 shows that the best performing BFN gives 2.66 BPD for the 256 bin data, which\\nis close to the state-of-the-art at 2.64 BPD. The most obvious performance benchmark (given the\\nshared network architecture and similarity in loss function) is the VDM result at 2.65 BPD [17].\\nHowever this took 10M weight updates to achieve, and due to time constraints we were only able to\\ntrain BFNs for 5M updates. Validation performance was still improving after 5M updates, and it\\nremains unclear how much performance would improve with 10M updates.\\n43\\n\\n\\n(a) Input Mean\\n(b) Output Mean\\nFigure 16: CIFAR-10 Input and output distributions. For two test set images the figure shows the\\nmeans of the input and output distributions at steps evenly spaced between t = 0 and t = 0.25.\\nFigure 17: CIFAR-10 losses against time. The plot was made using the network trained with discretised\\nloss on 256 bins. Note the high loss at the very start of the process, which we did not observe with discrete\\ndata.\\nTable 3 shows that discretised loss gave better performance than continuous loss for 16 bins, as\\nwell as much faster training time (250K updates vs. 1M). This supports the hypothesis that training\\nwith discretised loss is most beneficial when the number of bins is relatively low. Furthermore, for\\nboth 16 and 256 bins, discretised training gave much better results when the number of steps n was\\nlow (e.g. 10 or 25). However continuous loss gave better performance than discretised loss on 256\\nbins (2.66 BPC vs 2.68); more investigation would be needed to understand why.\\nFigure 15 shows that discretised training with 16 bins gives better sample quality than training\\nwith 256 bins. This is presumably because the loss function of the former is restricted to the first\\nfour bits of the data in which — as can be seen by comparing the test data at 16 and 256 bins —\\nmost of the perceptually relevant information is contained. An interesting direction for future work\\nwould be to train one BFN to model the lower bits of an image, and a second BFN to conditionally\\nupscale to higher bits, as has previously been explored for autoregressive models [13, 26].\\n44\\n\\n\\nModel\\nBPC\\nFlow-based models\\nIAF/SCF† [54]\\n1.88\\nArgmax Coupling Flow† [14]\\n1.80\\nDiscrete Flow† [48]\\n1.23\\nOrder-agnostic Models\\nOA-ARDM [13]\\n1.43 ± 0.001\\nMAC [39]\\n1.40\\nDiffusion models\\nMultinomial Diffusion [14]\\n1.72\\nD3PM uniform [1]\\n1.61 ± 0.02\\nD3PM NN [1]\\n1.59 ± 0.03\\nD3PM mask [1]\\n1.45 ± 0.02\\nBFN\\n1.41\\nAutoregressive baseline\\nTransformer† [1]\\n1.23\\nBest result*\\nAdaptive Span Transformer† [45]\\n1.07\\nTable 4: Comparison of text8 results with other methods. The best published model on this dataset\\n(*) was trained on sequences of length 512. Rest of the above models were trained on sequences of length 256.\\nResults for models marked with (†) are exact values; all other results are upper bounds.\\nn-steps\\n10\\n25\\n50\\n100\\n256\\n1000\\n∞\\nBPC\\n1.70\\n1.52\\n1.47\\n1.43\\n1.42\\n1.41\\n1.41\\nTable 5: text8 results. BPC is bits per character averaged over 1M randomly cropped sequences from the\\ntest set with Ln(x) or L∞(x) sampled once per crop. The reconstruction loss Lr(x) (included in BPC) was\\n0.006.\\n7.3\\ntext8\\nData.\\nThe text8 dataset [25] was derived from a subset of the enwik9 Wikipedia dataset by\\nremoving punctuation and restricting the text to lowercase Latin letters and spaces, giving an\\nalphabet of size 27. For clarity, we represent the space character with an underscore in figures.\\nSetup.\\nThe network architecture was a Transformer similar to the small model (dmodel = 768)\\nused by Radford et al. [31] except that it uses the GELU activation function [10] and the depth\\nwas increased to 24 layers. The input and output of the Transformer were concatenated and then\\nprojected back to the output size to produce the final output. The standard training/validation/test\\nsplit of 90M/5M/5M consecutive characters was used, and the network was trained with a batch size\\nof 3328 sequences of length 256, randomly cropped from the training set, for 1.2 M weight updates\\nusing the AdamW optimizer[22]. The learning rate was set to 10−4, weight decay to 0.1 and (β1, β2)\\nto (0.9, 0.98). An exponential moving average of model parameters with a decay rate of 0.9999 was\\nused for evaluation and sample generation. Dropout was not used, but overfitting was observed\\ntowards the end of training indicating that regularization may further improve results. β(1) was\\n0.75. The total number of learnable parameters was approximately 170M. Note that the batch\\nsize and number of layers were larger than prior results from diffusion models. The first choice\\n45\\n\\n\\nnd_philip_melanchthon_in_one_five_six_one_he_copper_etched_a_map\\n_a_single_sheet_cartouche_of_silesia_which_he_published_under_th\\ne_title_silesiae_typus_and_dedicated_to_nicolaus_rhedinger_his_m\\nap_was_later_republished_in_several_versions_of_abraham_ortelius\\nages_middle_school_high_school_or_both_in_a_manner_that_includes\\n_military_traditions_and_training_in_military_subjects_the_vast_\\nmajority_are_in_the_united_states_many_military_schools_are_also\\n_boarding_schools_and_others_are_simply_magnet_schools_in_a_larg\\nberno_norman_bourkes_have_been_in_mayo_since_the_thirteenth_cent\\nury_like_many_who_came_to_ireland_with_the_norman_invasion_it_wa\\ns_said_of_the_bourkes_that_they_ended_up_more_irish_than_the_iri\\nsh_themselves_her_family_had_links_with_many_diverse_political_s\\nonally_eaten_on_the_day_is_goose_according_to_legend_martin_was_\\nreluctant_to_become_bishop_which_is_why_he_hid_in_a_stable_fille\\nd_with_geese_the_noise_made_by_the_geese_betrayed_his_location_t\\no_the_people_who_were_looking_for_him_also_in_the_east_part_of_t\\nide_film_fn_m_two_four_nine_saw_presentation_mpeg_for_the_religi\\nous_order_known_as_the_minimi_minims_order_of_the_minims_see_min\\nim_religious_order_light_machine_guns_modern_firearms_of_the_uni\\nted_states_fabrique_nationale_de_herstal_five_five_six_mm_machin\\nining_to_fly_planes_atta_traveled_to_prague_stayed_overnight_and\\n_then_entered_the_u_s_on_june_three_atta_and_earlier_arrived_hij\\nackers_opened_bank_accounts_and_continue_to_check_on_flight_scho\\nols_in_july_atta_and_marwan_al_shehhi_enrolled_at_huffman_aviati\\n_boosted_by_its_famously_harsh_winters_the_region_is_reportedly_\\nthe_third_largest_theater_market_in_the_country_attracting_major\\n_performances_the_guthrie_theater_is_the_most_famous_theater_in_\\nthe_city_in_order_to_help_revitalize_the_downtown_and_warehouse_\\n(a) Test Data\\nth_earth_it_will_receive_the_smell_of_a_body_known_as_the_postwa\\nnaplast_for_that_reason_two_devonian_rocks_englathus_slimche_rho\\nse_ice_s_stalk_permanently_and_after_the_last_we_drink_meat_of_c\\nourse_the_sessiology_created_during_the_first_half_of_the_fall_t\\nin_moor_was_employed_as_lake_in_among_others_one_nine_two_zero_b\\nut_the_pair_went_out_of_the_fame_a_cell_on_maytown_in_one_nine_s\\nix_six_kells_was_formally_promoted_by_being_presidential_candida\\nte_for_proclamation_moor_was_the_first_author_of_snowman_s_expos\\nrt_unofficially_naming_pepper_s_best_health_cea_resistant_tum_fu\\nlci_ua_starring_yellow_the_negative_campaign_for_the_use_of_poly\\nmap_button_was_sent_on_systems_comparable_to_a_new_type_of_nasca\\nr_historically_gea_involved_a_commission_that_the_government_net\\nir_own_context_although_finally_loyal_and_genial_to_medieval_voi\\nl_the_honour_best_on_amber_can_be_done_by_scrutinous_members_of_\\nthe_knights_templar_and_grand_master_still_called_amber_the_pope\\n_for_four_years_who_had_on_approximated_journeys_he_was_only_kno\\n_as_were_independent_of_jesus_or_throughout_the_time_of_nazarite\\ns_this_is_not_true_versical_times_despite_the_use_of_the_lattian\\ns_as_still_not_readily_adjacent_to_the_text_the_popular_tendency\\n_of_revision_is_current_in_the_three_gospels_some_subscriptions_\\n_counter_strike_by_the_national_terrorists_c_e_ti_core_blooded_a\\nm_forse_for_elsie_by_the_columbia_university_of_washington_stude\\nnt_in_their_own_right_such_views_change_reflected_by_the_fiercel\\ny_scriptly_knit_allegations_came_under_this_a_tendency_that_the_\\ns_that_of_the_related_spirit_of_the_great_sword_and_by_the_spiri\\nt_of_jehudah_as_they_were_attained_in_hebrew_trees_as_spiritual_\\nwitness_to_a_high_angel_of_the_south_thus_identified_pinhim_with\\n_pulim_where_they_were_emphetic_and_assyrian_languages_in_ancien\\n(b) Generated Data\\nFigure 18: text8 real and generated data. Samples generated with 1000 steps.\\nincreases model capacity while the second tends to make overfitting more likely. These choices\\nwere made to maximize the utilization of available resources while achieving results in reasonable time.\\nResults.\\nTable 4 shows that BFN yielded a 1.41 BPC on the text8 test set, which is better than\\nall discrete diffusion models we found in the literature, and close to the best order-agnostic model,\\nMAC at 1.40 BPC. We note however that both a standard autoregressive baseline and a discrete\\nflow model perform substantially better at 1.23 BPC. Table 5 shows that performance is reasonably\\nrobust to decreased n, with only 100 steps required to reach 1.43 BPC. This result could probably\\nbe improved by training with the discrete-time loss.\\n8\\nConclusion\\nThis paper introduced Bayesian Flow Networks, a new class of generative model that combines\\nBayesian inference with neural networks in an iterative modelling process. Discrete and continuous-\\ntime loss functions were derived along with sampling procedures, and the model was succesfully\\napplied to continuous, discretised and discrete data. We hope this work will inspire fresh perspectives\\nand new directions for generative modelling research.\\nAckowledgements\\nWe would like to thank Vojtech Micka for his invaluable engineering and infrastructure support.\\n46\\n\\n\\nt = 0.0\\nt = 0.2\\nt = 0.4\\nt = 0.6\\nt = 0.8\\nt = 1.0\\n(a) Input Distribution\\n(b) Output Distribution\\nFigure 19: text8 Input and Output Distributions. The heatmaps show the character probability\\ndistributions across part of a test sequence at various times during the flow process. Whereas the expected\\nentropy for each letter decreases independently in the input distribution, the entropy of the output distribution\\ntends to chunk into words and phrases — e.g. the date “one five six one” is confidently predicted early in the\\nprocess.\\nReferences\\n[1] Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Struc-\\ntured Denoising Diffusion Models in Discrete State-Spaces. arXiv preprint arXiv:2107.03006,\\nJuly 2021.\\n[2] Ting Chen, Ruixiang Zhang, and Geoffrey Hinton. Analog bits: Generating discrete data using\\ndiffusion models with self-conditioning. arXiv preprint arXiv:2208.04202, 2022.\\n[3] Rewon Child. Very deep vaes generalize autoregressive models and can outperform them on\\nimages. arXiv preprint arXiv:2011.10650, 2020.\\n47\\n\\n\\n(a) Input Distribution\\n(b) Output Distribution\\n0.1\\n0.2\\n0.4\\n0.6\\n0.5\\n0.7\\n0.3\\nt\\nFigure 20: text8 Input and Output Distributions. An alternative visualisation with the character\\nsizes scaled in proportion to their probability.\\n[4] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with\\nsparse transformers. arXiv preprint arXiv:1904.10509, 2019.\\n[5] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis.\\nAdvances in neural information processing systems, 34:8780–8794, 2021.\\n[6] Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin,\\nPierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Contin-\\nuous diffusion for categorical data. arXiv preprint arXiv:2211.15089, 2022.\\n[7] Jarek Duda. Asymmetric numeral systems. arXiv preprint arXiv:0902.0271, 2009.\\n[8] H.O. Georgii. Stochastics: Introduction to Probability and Statistics. De Gruyter textbook.\\nWalter De Gruyter, 2008. ISBN 9783110191455. URL https://books.google.co.uk/books\\n?id=ttJ5xpQX2MgC.\\n[9] Alex Graves.\\nGenerating sequences with recurrent neural networks.\\narXiv preprint\\narXiv:1308.0850, 2013.\\n[10] Dan Hendrycks and Kevin Gimpel.\\nGaussian error linear units (gelus).\\narXiv preprint\\narXiv:1606.08415, 2016.\\n[11] Geoffrey E Hinton and Drew Van Camp. Keeping the neural networks simple by minimizing the\\ndescription length of the weights. In Proceedings of the sixth annual conference on Computational\\nlearning theory, pages 5–13, 1993.\\n[12] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances\\nin neural information processing systems, 33:6840–6851, 2020.\\n48\\n\\n\\n[13] Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg,\\nand Tim Salimans. Autoregressive diffusion models. arXiv preprint arXiv:2110.02037, 2021.\\n[14] Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forr´\\ne, and Max Welling. Argmax\\nFlows and Multinomial Diffusion: Learning Categorical Distributions. In Advances in Neural\\nInformation Processing Systems, volume 34, pages 12454–12465. Curran Associates, Inc., 2021.\\n[15] Ajay Jain, Pieter Abbeel, and Deepak Pathak. Locally masked convolution for autoregressive\\nmodels. In Conference on Uncertainty in Artificial Intelligence, pages 1358–1367. PMLR, 2020.\\n[16] Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang, and Il-Chul Moon.\\nSoft\\ntruncation: A universal training technique of score-based diffusion model for high precision\\nscore estimation. arXiv preprint arXiv:2106.05527, 2021.\\n[17] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models.\\nAdvances in neural information processing systems, 34:21696–21707, 2021.\\n[18] Diederik P Kingma and Max Welling.\\nAuto-encoding variational bayes.\\narXiv preprint\\narXiv:1312.6114, 2013.\\n[19] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report,\\nUniversity of Toronto, 2009.\\n[20] Yann LeCun and Corinna Cortes. MNIST handwritten digit database, 2010. URL http:\\n//yann.lecun.com/exdb/mnist/.\\n[21] Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori B. Hashimoto.\\nDiffusion-lm improves controllable text generation. arXiv preprint arXiv:2205.14217, 2022.\\n[22] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint\\narXiv:1711.05101, 2017.\\n[23] Aaron Lou and Stefano Ermon. Reflected diffusion models. arXiv preprint arXiv:2304.04740,\\n2023.\\n[24] Rabeeh Karimi Mahabadi, Jaesung Tae, Hamish Ivison, James Henderson, Iz Beltagy,\\nMatthew E. Peters, and Arman Cohan. Tess: Text-to-text self-conditioned simplex diffu-\\nsion. arXiv preprint arXiv:2305.08379, 2023.\\n[25] Matt Mahoney. Large text compression benchmark., 2009. URL http://mattmahoney.net/\\ndc/textdata.html.\\n[26] Jacob Menick and Nal Kalchbrenner. Generating high fidelity images with subscale pixel\\nnetworks and multidimensional upscaling. arXiv preprint arXiv:1812.01608, 2018.\\n[27] Kevin Murphy. Conjugate bayesian analysis of the gaussian distribution. Technical report,\\nUniversity of British Columbia, 2007.\\n[28] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic\\nmodels. In International Conference on Machine Learning, pages 8162–8171. PMLR, 2021.\\n49\\n\\n\\n[29] OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.\\n[30] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku,\\nand Dustin Tran. Image transformer. In International conference on machine learning, pages\\n4055–4064. PMLR, 2018.\\n[31] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al.\\nLanguage models are unsupervised multitask learners. Technical report, OpenAI, 2019.\\n[32] Danilo Rezende and Shakir Mohamed.\\nVariational inference with normalizing flows.\\nIn\\nInternational conference on machine learning, pages 1530–1538. PMLR, 2015.\\n[33] Pierre H. Richemond, Sander Dieleman, and Arnaud Doucet. Categorical SDEs with simplex\\ndiffusion. arXiv preprint arXiv:2210.14784, 2022.\\n[34] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bj¨\\norn Ommer.\\nHigh-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF\\nconference on computer vision and pattern recognition, pages 10684–10695, 2022.\\n[35] Hossein Sadeghi, Evgeny Andriyash, Walter Vinci, Lorenzo Buffoni, and Mohammad H Amin.\\nPixelvae++: Improved pixelvae with discrete prior. arXiv preprint arXiv:1908.09948, 2019.\\n[36] Ruslan Salakhutdinov and Iain Murray. On the quantitative analysis of deep belief networks.\\nIn Proceedings of the 25th international conference on Machine learning, pages 872–879. ACM,\\n2008.\\n[37] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models.\\narXiv preprint arXiv:2202.00512, 2022.\\n[38] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelcnn++: Improving the\\nPixelCNN with discretized logistic mixture likelihood and other modifications. arXiv preprint\\narXiv:1701.05517, 2017.\\n[39] Andy Shih, Dorsa Sadigh, and Stefano Ermon. Training and inference on any-order autoregres-\\nsive models the right way. Advances in Neural Information Processing Systems, 35:2762–2775,\\n2022.\\n[40] Samarth Sinha and Adji Bousso Dieng. Consistency regularization for variational auto-encoders.\\nAdvances in Neural Information Processing Systems, 34:12943–12954, 2021.\\n[41] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper-\\nvised learning using nonequilibrium thermodynamics. In International conference on machine\\nlearning, pages 2256–2265. PMLR, 2015.\\n[42] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and\\nBen Poole. Score-based generative modeling through stochastic differential equations. arXiv\\npreprint arXiv:2011.13456, 2020.\\n[43] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv\\npreprint arXiv:2303.01469, 2023.\\n50\\n\\n\\n[44] Robin Strudel, Corentin Tallec, Florent Altch´\\ne, Yilun Du, Yaroslav Ganin, Arthur Mensch, Will\\nGrathwohl, Nikolay Savinov, Sander Dieleman, Laurent Sifre, et al. Self-conditioned embedding\\ndiffusion for text generation. arXiv preprint arXiv:2211.04236, 2022.\\n[45] Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive\\nAttention Span in Transformers. arXiv preprint arXiv:1905.07799, August 2019.\\n[46] Ilya Sutskever, James Martens, and Geoffrey E Hinton. Generating text with recurrent neural\\nnetworks. In Proceedings of the 28th international conference on machine learning (ICML-11),\\npages 1017–1024, 2011.\\n[47] James Townsend, Tom Bird, and David Barber. Practical lossless compression with latent\\nvariables using bits back coding. arXiv preprint arXiv:1901.04866, 2019.\\n[48] Dustin Tran, Keyon Vafa, Kumar Agrawal, Laurent Dinh, and Ben Poole. Discrete flows:\\nInvertible generative models of discrete data. Advances in Neural Information Processing\\nSystems, 32, 2019.\\n[49] Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in\\nneural information processing systems, 33:19667–19679, 2020.\\n[50] Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space.\\nAdvances in Neural Information Processing Systems, 34:11287–11302, 2021.\\n[51] Chris S. Wallace.\\nClassification by minimum-message-length inference.\\nIn International\\nConference on Computing and Information, 1991.\\n[52] Daniel Watson, William Chan, Jonathan Ho, and Mohammad Norouzi. Learning fast samplers\\nfor diffusion models by differentiating through sample quality. arXiv preprint arXiv:2202.05830,\\n2022.\\n[53] Ian H Witten, Radford M Neal, and John G Cleary. Arithmetic coding for data compression.\\nCommunications of the ACM, 30(6):520–540, 1987.\\n[54] Zachary Ziegler and Alexander Rush. Latent Normalizing Flows for Discrete Sequences. In\\nProceedings of the 36th International Conference on Machine Learning, pages 7673–7682. PMLR,\\nMay 2019.\\n51\\n\\n\\nPublished as a conference paper at ICLR 2024\\nUNIFIED GENERATIVE MODELING OF 3D MOLECULES\\nVIA BAYESIAN FLOW NETWORKS\\nYuxuan Song1∗\\nJingjing Gong1∗\\nYanru Qu2\\nHao Zhou1\\nMingyue Zheng3\\nJingjing Liu1\\n&\\nWei-Ying Ma1\\n1 Institute of AI Industry Research (AIR), Tsinghua University\\n2 University of Illinois Urbana-Champaign\\n3 Shanghai Institute of Materia Medica, Chinese Academy of Sciences\\n{songyuxuan,gongjingjing,zhouhao,maweiying}@air.tsinghua.edu\\nABSTRACT\\nAdvanced generative model (e.g., diffusion model) derived from simplified conti-\\nnuity assumptions of data distribution, though showing promising progress, has\\nbeen difficult to apply directly to geometry generation applications due to the multi-\\nmodality and noise-sensitive nature of molecule geometry. This work introduces\\nGeometric Bayesian Flow Networks (GeoBFN), which naturally fits molecule\\ngeometry by modeling diverse modalities in the differentiable parameter space of\\ndistributions. GeoBFN maintains the SE-(3) invariant density modeling property by\\nincorporating equivariant inter-dependency modeling on parameters of distributions\\nand unifying the probabilistic modeling of different modalities. Through optimized\\ntraining and sampling techniques, we demonstrate that GeoBFN achieves state-of-\\nthe-art performance on multiple 3D molecule generation benchmarks in terms of\\ngeneration quality (90.87% molecule stability in QM9 and 85.6% atom stability in\\nGEOM-DRUG1). GeoBFN can also conduct sampling with any number of steps\\nto reach an optimal trade-off between efficiency and quality (e.g., 20× speedup\\nwithout sacrificing performance).\\n1\\nINTRODUCTION\\nSender Distribution\\nBayesian\\nUpdate\\nEGNN\\nReceiver Distribution\\nSender\\nSample\\nKL\\nNoisy channel\\nSender Distribution\\nBayesian\\nUpdate\\nEGNN\\nReceiver Distribution\\nSender\\nSample\\nNoisy channel\\nKL\\nFigure 1: The framework of GeoBFN\\nMolecular geometries can be represented\\nas three-dimensional point clouds, charac-\\nterized by their Cartesian coordinates in\\nspace and enriched with descriptive fea-\\ntures. For example, proteins can be repre-\\nsented as proximity spatial graphs (Jing\\net al., 2021) and molecules as atomic\\ngraphs in 3D (Schütt et al., 2017). Thus,\\nlearning geometric generative models has\\nthe potential to benefit scientific discover-\\nies such as material and drug design. Re-\\ncent progress in deep generative model-\\ning has paved the way for geometric gen-\\nerative modeling. For example, Gebauer\\net al. (2019); Luo & Ji (2021) and Satorras\\net al. (2021a) use autoregressive models\\nand flow-based models, respectively, for\\ngenerating 3D molecules in-silico. Most\\nrecently, inspired by the huge success of\\ndiffusion model (DM) in image genera-\\ntion Meng et al. (2022); Ho et al. (2020)\\n∗Equal Contribution. Correspondence to Hao Zhou(zhouhao@air.tsinghua.edu).\\n1The scores are reported at 1k sampling steps for fair comparison, and our scores could be further improved\\nif sampling sufficiently longer steps\\n1\\narXiv:2403.15441v1  [physics.chem-ph]  17 Mar 2024\\n\\n\\nPublished as a conference paper at ICLR 2024\\nand beyond Li et al. (2022), DM incorporating geometric symmetries has been widely explored in\\nthe field of geometry generation Hoogeboom et al. (2022); Xu et al. (2023).\\nHowever, two major challenges remain in directly applying DM to molecule geometry: multi-modality\\nand noise sensitivity. The multi-modality issue refers to the dependency on diverse data forms to\\neffectively depict the atomic-level geometry of a molecule. For instance, the continuous variable of\\natom coordinates is essential for describing the spatial arrangement, while either the discretised atom\\ncharge or categorical atom types are employed to completely determine the molecule’s composition.\\nNoise sensitivity refers to the fact that applying noise or perturbing the atom coordinates will not only\\nchange the value of the variable but also have a significant impact on the relationship among different\\natoms as the Euclidean distances are also changed. Therefore, a small noise on atom coordinates\\ncould bring a sudden drop of the signal at the molecule level.\\nTo alleviate these issues, Xu et al. (2023) introduces a latent space for alleviating the inconsistency\\nof unified Gaussian diffusion on different modalities. Anand & Achim (2022) propose to use\\ndecomposed modeling of different modalities. Peng et al. (2023) use different noise schedulers for\\ndifferent modalities to accommodate noise sensitivity. However, these methods either depend on the\\nsophisticated and artifact-filled design or lack of guarantee or constraint on the designed space.\\nIn this work, we propose Geometric Bayesian Flow Networks (GeoBFN) to model 3D molecule\\ngeometry in a principally different way. Bayesian Flow Networks Graves et al. (2023) (BFN) is\\na novel generative model developed quite recently. Taking a unique approach by incorporating\\nBayesian inference to modify the parameters of a collection of independent distributions, brings a\\nfresh perspective to geometric generative modeling. Firstly, GeoBFN uses a unified probabilistic\\nmodeling formulation for different modalities in the molecule geometry. Secondly, regarding the\\nvariable of 3D atom coordinates, the input variance for BFNs is considerably lower than DMs,\\nleading to better compatibility with the inherent noise sensitivity. Further, we bring the geometry\\nsymmetries into the Bayesian update procedure through an equivariant inter-dependency modeling\\nmodule. We also demonstrate that the density function of implied generative distribution is SE-(3)\\ninvariant and the generative process of iterative updating is roto-translational equivariant. Thirdly,\\nwith BFN’s powerful probabilistic modeling capacity, 3D molecule geometry representation can\\nbe further optimized into a representation with only two similar modalities: discretised charge and\\ncontinuous atom coordinates. The mode-redundancy issue on discretised variable in the original\\nBFNs is fixed by an early mode-seeking sampling strategy in GeoBFN.\\nWith operating on the space with less variance, GeoBFN could sample with any number of steps\\nwhich provides a superior trade-off between efficiency and quality, which leads to a 20× speedup\\nwith competitive performance. Besides, GeoBFN is a general framework that can be easily extended\\nto other molecular tasks. We conduct thorough evaluations of GeoBFN on multiple benchmarks, in-\\ncluding both unconditional and property-conditioned molecule generation tasks. Results demonstrate\\nthat GeoBFN consistently achieves state-of-the-art generation performance on molecule stability and\\nother metrics. Empirical studies also show a significant improvement in controllable generation and\\ndemonstrate that GeoBFN enjoys a significantly higher modeling capacity and inference efficiency.\\n2\\nPRELIMINARIES\\n2.1\\nSE-(3) INVARIANT DENSITY MODELING\\nTo distinguish geometry representation and the atomic property features, we use the tuple g = ⟨x, h⟩\\nto represent the 3D molecules. Note here x = (x1, . . . , xN) ∈RN×3 is the atom coordinate matrix,\\nand h = (h1, . . . , hN) ∈RN×d is the node feature matrix, e.g., atomic types and charges. Density\\nestimation on the 3D molecules should satisfy specific symmetry conditions of the geometry. In this\\nwork, we focus on the transformations Tg in the Special Euclidean group (SE-(3)), i.e., the group of\\nrotation and translation in 3D space, where transformations Tg can be represented by a translation\\nt and an orthogonal matrix rotation R. Note for a generative model on molecule geometry with\\nunderlying density function pθ(⟨x, h⟩), the likelihood should not be influenced by the rotation or\\ntranslation of the entire molecule, which means the likelihood function should be SE-(3) invariant on\\nthe input coordinates, i.e., pθ(⟨x, h⟩) = pθ(⟨Rx + t, h⟩).\\n2.2\\nBAYESIAN FLOW NETWORKS\\nThe Bayesian Flow Networks (BFNs) are based on the following latent variable models: for learning\\nthe probability distribution pθ over g, a series of noisy versions ⟨y1, · · · , yn⟩of g are introduced as\\n2\\n\\n\\nPublished as a conference paper at ICLR 2024\\nlatent variables. And then the variational lower bound of likelihood is optimized:\\nlog pθ(g) ≥\\nE\\ny1,...,yn∼q\\n\\u0014\\nlog pϕ (g | y1, . . . , yn) pϕ (y1, . . . , yn)\\nq(y1, . . . , yn|g)\\n\\u0015\\n= −DKL(q∥pϕ (y1, . . . , yn)) +\\nE\\ny1,...,yn∼q log [pϕ (g | y1, . . . , yn)]\\n(1)\\nAnd q is namely the variational distribution. The prior distribution of latent variables is usually\\norganized autoregressively, i.e., pϕ(y1, · · · , yn) = pϕ(y1)pϕ(y2 | y1)pϕ(yn | yn−1 · · · y1) which\\nalso implies the data generation procedure, i.e., y1 →· · · →yn →g (Note: this procedure only\\ndemonstrates the generation order, yet does NOT imply Markov property for the following derivative).\\nOne widely adopted intuition for the generation process is that the information of the data samples\\nshould progressively increase along with the above Markov chain, e.g., noisier images to cleaner\\nimages. The key motivation of BFNs is that the information along the latent variables should change\\nas smoothly as possible for all modalities including discretized and discrete variables. To this end,\\nBFNs operate on the distributions in the parameter space, in contrast to the sample space.\\nWe introduce components of BFNs one by one (Fig.2a). Firstly, the variational distribution q is\\ndefined by the following form:\\nq (y1, . . . , yn | g) =\\nn\\nY\\ni=1\\npS (yi | g; αi)\\n(2)\\npS (yi | g; αi) is termed as the sender distribution, which could be seen as adding noise to the data\\naccording to a predefined accuracy αi.\\nSecondly, for the definition of pϕ, BFNs will first transfer the noisy sample y to the parameter\\nspace, obtaining θ, then apply Bayesian update in the parameters space and transfer back to the\\nnoisy sample space at last. To clarify, θ refers to the parameter of distributions in the sample space,\\ne.g., the mean/variance for Gaussian distribution or probabilities for categorical distribution. In the\\nscope of BFNs, the distributions on the sample space are factorized by default, e.g., p(g | θ) =\\nQD\\nd=1 p\\ng(d) | θ(d)\\u0001\\n.\\nThirdly, a neural network Φ takes θ as input and aims to model the dependency among different\\ndimensions hence to recover the distribution of the original sample g. The output of neural network\\nΦ(θ) still lies in the parameter space, and we termed it as the parameter of output distribution pO,\\nwhere pO(y|θ; ϕ) = QD\\nd=1 pO(y(d) | Φ(θ)(d)).\\nTo map the noisy sample y to the input space, Bayesian update is applied to θ:\\nθi ←h(θi−1, yi, αi),\\n(3)\\nh is called Bayesian update function . The distribution over (θ0, . . . , θn−1) is then defined by the\\nBayesian update distribution via marginalizing out y:\\npϕ (θ0, . . . , θn−1) = p(θ0)\\nn\\nY\\ni=1\\npU (θi | θi−1; αi) ,\\n(4)\\nwhere p(θ0) is a simple prior for ease of generation, e.g., standard normal, and pU could be obtained\\nfrom Eq. 3:\\npU (θi | θi−1; αi) =\\nE\\npR(yi|θi−1;αi)δ (θi −h(θi−1, yi, αi)) ,\\n(5)\\nδ being the Dirac delta distribution. pR(yi|θi−1, αi) =\\nE\\npO(x′|θi−1;ϕ)pS(yi|x′; αi) and is also called\\nthe as receiver distribution.\\nAt last we map Φ(θ) back to the noisy sample space by combining the known form, accuracy of PS\\nand marginalizing out y:\\npϕ (y1, . . . , yn) = pϕ(y1)\\nn\\nY\\ni=2\\npϕ(yi | y{1:i−1} =\\nn\\nY\\ni=1\\npϕ(yi | θi−1)\\n=\\nn\\nY\\ni=1\\nE\\npO(x′\\ni|θi−1;ϕ) [pS(yi|x′\\ni; αi)] ,\\n(6)\\n3\\n\\n\\nPublished as a conference paper at ICLR 2024\\na Graphical Model of BFN\\nb Graphical Model of Diffusion\\nFigure 2: Graphical View of Comparison between BFN and Diffusion\\nwhere we use θ0:n−1 to abbreviate (θ0, . . . , θn−1), and y similar. . Till now, we have defined q,\\npϕ(y1, . . . , yn), and pϕ (g | y1, . . . , yn) is simply pO(g | θn) on each sample, thus Eq.1 can be\\nestimated.\\n3\\nMETHODOLOGY\\n3.1\\nSE-(3) INVARIANT GEOMETRY DENSITY MODELING\\nAs discussed in Sec. 2.1, for a generative model on the 3D molecule geometry, it is crucial to hold the\\nSE-(3) invariant conditions. Recall the mathematics formula of the geometries g = ⟨x, h⟩, we denote\\nthe latent variable, e.g., noisy samples, of g as yg. We are interested in applying the SE-(3) invariant\\nconditions to the probabilistic model pϕ. To this end, we need to first reformulate the likelihood\\nfunction:\\npϕ(g) = pϕ(⟨x, h⟩) =\\nZ\\nyg\\n1,··· ,yg\\nn\\npϕ(g | yg\\n1, · · · , yg\\nn)pϕ(yg\\n1, · · · , yg\\nn)dyg\\n1 . . . dyg\\nn.\\n(7)\\nWith θx and yx and all the distributions defined in the same way as above, we focus on the variables\\ncorresponds to x in the geometry yg. Then we have the following theorem:\\nTheorem 3.1. (SE-(3) Invariant Condition)\\n• With the θx, yx, x constrained in the zero Center of Mass(CoM) space (Köhler et al., 2020;\\nXu et al., 2022), the likelihood function pϕ is translational invariant.\\n• When the following properties are satisfied, the likelihood function pϕ is roto-invariant:\\npO\\nx′ | θx\\ni−1; ϕ\\n\\u0001\\n= pO\\nR(x′) | R(θx\\ni−1); ϕ\\n\\u0001\\n; pS (yx | x′; α) = pS (R(yx) | R(x′); α) ;\\nh(R(θx\\ni−1), R(yx\\ni ), αi) = Rh(θx\\ni−1, yx\\ni , αi); p(x′|θx\\n0) = p(R(x′)|θx\\n0), ∀orthogonal R\\nProposition 3.2. With the condition in Theorem. 3.1 satisfied, the evidence lower bound objective in\\nEq. 1, i.e.,\\nLVLB(x) =\\nE\\npϕ(θx\\n0 ,...,θx\\nn)\\n\\\" n\\nX\\ni=1\\nDKL(pS (· | x; αi) ∥pR(· | θx\\ni−1; αi)) −log pϕ (x | θx\\nn)\\n#\\n,\\n(8)\\nwith the Bayesian update distribution pϕ\\nθx\\n0, . . . , θx\\nn−1\\n\\u0001\\n= Qn\\ni=1 pU (θi | θi−1, x; αi) similar to\\nEq. 4. And pR(· | θx\\ni−1; αi) =\\nE\\npO(x′\\ni|θx\\ni−1;ϕ) [pS(yi|x′\\ni; αi)], pϕ(x|θx\\nn) = pO(x|θx\\nn, ϕ). Derivation\\nfrom Eq. 1 to equation 8 is at Appendix C.4. if pU (θi | θi−1, x; αi) = pU (Rθi | Rθi−1, Rx; αi),\\nthen LVLB(x) is also SE-(3) invariant.\\nWe leave the formal proof of Theorem. 3.1 and Proposition. 3.2 in Appendix C.\\n3.2\\nGEOMETRIC BAYESIAN FLOW NETWORKS\\nThen we introduce the detailed formulation of geometric Bayesian flow networks (GeoBFN) based on\\nthe analysis in Sec. 3.1. For describing a 3D molecule geometry g = ⟨x, h⟩, various representations\\ncan be utilized for the node features. The atom types ht and atomic charges hc are commonly\\nemployed, with the former being discrete (categorical) and the latter being discretized (integer).\\nTogether with the continuous variable, e.g., atom coordinates x, the network module in the modeling\\nof the output distribution of GeoBFN could be parameterized with an equivariant graph neural\\nnetwork (EGNN) (Satorras et al., 2021b) Φ:\\nΦ(Rθx + t, [θht, θhc]) = [Rθx′ + t, θh′\\nt, θh′\\nc],\\n∀R, t\\n(9)\\n4\\n\\n\\nPublished as a conference paper at ICLR 2024\\nwhere Φ(θx, [θht, θhc]) = [θx′, θh′\\nt, θh′\\nc]. And then we introduce the necessary components to\\nderive the objective in Eq. 8\\nAtom Coordinates x and Charge hc: For the continuous and discretized variables, the input\\ndistribution is set as the factorized Gaussian distributions, where θ\\ndef\\n:= {µ, ρ} the parameter of\\nN\\n· | µ, ρ−1I\\n\\u0001\\n. For simplicity, we take x as an example to illustrate the common parts of the\\ntwo variables. And θx\\n0 is set as {0, 1}. The sender distribution pS is also an isotropic Gaussian\\ndistribution:\\npS(· | x; αI) = N\\nx, α−1I\\n\\u0001\\n(10)\\nGiven the nice property of isotropic Gaussian (proof given by Graves et al. (2023)), the simple form\\nof Bayesian update function could be derived as:\\nh ({µi−1, ρi−1} , y, α) = {µi, ρi} ,\\nHere\\nρi = ρi−1 + α, µi = µi−1ρi−1 + yα\\nρi\\n(11)\\nAs shown in Eq. 11, the randomness only exists in µ, and the corresponding Bayesian update\\ndistribution in Eq. 8 is as:\\npU (θi | θi−1, x; α) = N\\n\\u0012\\nµi | αx + µi−1ρi−1\\nρi\\n, α\\nρ2\\ni\\nI\\n\\u0013\\n(12)\\nThe above discrete-time Bayesian update could be easily extended to continuous-time, with an\\naccuracy scheduler defined as β(t) =\\nR t\\nt′=0 α (t′) dt′, t ∈[0, 1]. Given the accuracy additive\\nproperty of pU (proof given by Graves et al. (2023)),\\nE\\npU (θi−1|θi−2,x;αa)pU (θi | θi−1, x; αb) =\\npU (θi | θi−2, x; αa + αb), the Bayesian flow distribution could be obtained as:\\npF (θx | x; t) = pU (θx | θ0, x; β(t))\\n(13)\\nThe key difference of atom coordinates x and charges hc lies in the design of the output distribution.\\nFor continuous variable x, the network module Φ directly outputs an estimated ˆ\\nx = Φ(θg, t). Hence\\nfor timestep t, the output distribution is\\npO (x′ | θg, t; ϕ) = δ(x −Φ(θg, t))\\n(14)\\nWhile for discretized variable hc, the network module will output two variables, µhc and ln σhc with\\ndimension equivalent to hc which implies a distribution N(µhc, σ2\\nhcI). With a K-bins discretized\\nvariable, the support is split into K buckets with each bucket k centered as kc = 2k−1\\nK\\n−1 and left\\nboundary as kl = kc −1\\nK and right boundary as kr = kc + 1\\nK . Then for each k, the probability is\\nthe mass from kl to kr, i.e.,\\nR kr\\nkl N(µhc, σ2\\nhcI). And the first and last bins are curated by making sure\\nthe sum of the probability mass is 1. Then the output distribution is:\\npO(hc | θg, t; ϕ) =\\nD\\nY\\nd=1\\np(d)\\nO\\n\\u0010\\nk\\n\\u0010\\nh(d)\\nc\\n\\u0011\\n| θg, t; ϕ\\n\\u0011\\n,\\n(15)\\nwhere the function k(·) maps the variable to the corresponding bucket.\\nAtom Types ht: The atom types ht are discrete variables with K categories, where the corresponding\\nparameter space lies in probability simplex thus the procedure is slightly different from the others.\\nThe input distribution for ht is pI(ht | θ) = QD\\nd=1 θht(d), where D is number if variables. And\\nthe input prior θht\\n0\\n=\\n1\\nK, where 1\\nK is the length KD vector whose entries are all\\n1\\nK . The sender\\ndistribution, could be derived with the central limit theorem, lies in the form of\\npS(y | ht; α) = N (y | α (Keht −1) , αKI)\\n(16)\\nwhere 1 is a vector of ones, I is the identity matrix, and ej ∈RK is a vector defined as the projection\\nfrom the class index j to a length K one-hot vector (proof given by Graves et al. (2023)). In other\\nwords, each element of ej is defined as (ej)k = δjk, where δjk is the Kronecker delta function. And\\neht\\ndef\\n=\\n\\u0010\\neh(1)\\nt , . . . , eh(D)\\nt\\n\\u0011\\n∈RKD.\\n5\\n\\n\\nPublished as a conference paper at ICLR 2024\\nGeoBFN\\nEDM\\nFigure 3: The Bayesian Flow and Diffusion Process of GeoBFN and EDM.\\nThe Bayesian update function could be derived as h (θi−1, y, α) =\\neyθi−1\\nPK\\nk=1 eyk(θi−1)k (proof given by\\nGraves et al. (2023)) . And similar to Eq. 13, the Bayesian flow distribution for ht is as:\\npF (θht | ht; t) =\\nE\\nN(yht|β(t)(Keht−1),β(t)KI)\\nδ(θht −softmax(yht))\\n(17)\\nWith the network module Φ, the output distribution could be obtained as\\np(d)\\nO (k | θg; t) =\\n\\u0010\\nsoftmax\\n\\u0010\\nΦ(d)(θg, t)\\n\\u0011\\u0011\\nk , pO(ht | θ; t) =\\nD\\nY\\nd=1\\np(d)\\nO\\n\\u0010\\nh(d)\\nt\\n| θg; t\\n\\u0011\\n(18)\\nTraining Objective: By combining the different variables together, we could obtain the unified\\ncontinuous-time loss for GeoBFN based on Eq. 25 to Eq. 41 in (Graves et al., 2023) as:\\nL∞(g) = L∞(⟨x, hc, ht⟩) =\\nE\\nt∼U(0,1),pF (θg|g;t)\\n\\u0014αg(t)\\n2\\n∥g −Φ(θg, t)∥2\\n\\u0015\\n=\\nE\\nt∼U(0,1),\\nθg∼pF (·|g;t)\\n\\u0014αx(t)\\n2\\n∥x −Φx∥2 + αhc(t)\\n2\\n∥hc −Φhc∥2 + αht(t)\\n2\\n∥ht −Φht∥2\\n\\u0015\\n(19)\\nWhere Φ· is short for Φ·(θg, t). The joint Bayesian flow distribution is decomposed as:\\npF (θg | g; t) = pF (θx | x; t)pF (θhc | hc; t)pF (θht | ht; t),\\n(20)\\nwith αx, αhc and αht refer to the corresponding accuracy scheduler (details provided by\\nGraves et al. (2023)).\\nAnd Φx is defined the same as in Eq. 14;\\nwhile Φhc is de-\\nfined by the weighted average of different bucket centers with the output distribution in\\nEq. 15 as\\n\\u0010PK\\nk=1 p(1)\\nO (k | θ, t)kc, . . . , PK\\nk=1 p(D)\\nO (k | θ, t)kc\\n\\u0011\\n; And for Φht, it is defined as the\\nPK\\nk=1 p(d)\\nO (k | θ; t)ek based on Eq. 18.\\nRemark 3.3. The GeoBFN defined in the above formulation satisfied the SE(3)-invariant condition in\\nTheorem. 3.1.\\nSampling GeoBFN will generate samples follow the graphical model in the recursive procedure as\\nillustrated in Fig. 2a: e.g., g′ ∼pO(·|θi−1) →y ∼pS(·|g′, α) →θi = h(θi−1, y, α).\\n3.3\\nOVERCOME NOISE SENSITIVITY IN MOLECULE GEOMETRY\\nOne key obstacle of applying diffusion models to 3D molecule generation is the noise sensitivity\\nproperty of the molecule geometry. The property of noise sensitivity seeks to state the fact: When\\nnoise is incorporated into the coordinates and displaces them significantly from their original positions,\\nthe bond distance between certain connected atoms may exceed the bond length range[1]. Under these\\ncircumstances, the point cloud could potentially lose the critical chemical information inherently\\nencoded in the bonded structures; Another perspective stems from the reality that when noise is\\nadded to the coordinates, the relationships (distance) between different atoms could alter at a more\\nrapid pace, e.g. modifying the coordinates of one atom results in altering its distance to all other\\natoms. Thus, the intermediate steps’ structure in the generation procedure of diffusion models the\\nintermediate steps’ structure might be uninformative. And the majority of the information being\\nacquired in the final few steps of generation (as depicted in Fig. 3).\\n6\\n\\n\\nPublished as a conference paper at ICLR 2024\\nA fundamental belief underpinning GeoBFN is that a smoother transformation during the generative\\nprocess could result in a more favorable inductive bias according to (Graves et al., 2023). This\\nprocess occurs within the parameter space of GeoBFN, which is regulated through the Bayesian\\nupdate procedure. Specifically, samples exhibiting higher degrees of noise are assigned lesser weight\\nduring this update (refer to Eq. 11). This approach consequently leads to a significant reduction in\\nvariance within the parameter space as (Graves et al., 2023), which in turn facilitates the smooth\\ntransformation of molecular geometries. As illustrated in Fig. 3, this is evidenced by the gradual\\nconvergence of the structure of the intermediary steps towards the final structure, thus underscoring\\nthe effectiveness of smoother transformation.\\n3.4\\nOPTIMIZED DISCRETISED VARIABLE SAMPLING\\nPrevious research (Hoogeboom et al., 2022; Xu et al., 2023; Wu et al., 2022) utilizes both the atom\\ntypes ht and charges hc to represent the atomic properties. The hc usually serves as an auxiliary loss\\nfor improving training which is not involved in determining the molecule graph during generation due\\nto the insufficient modeling. However, there is redundant information between these two variables,\\nsince the ht and hc variables have a one-to-one mapping, e.g.the charge value 4 could be uniquely\\ndetermined as the Carbon atom. We found that with advanced probabilistic modeling on discretized\\ndata, GeoBFN could conduct training and sampling only with x and hc. However, there exists a\\ncounterexample for the objective in Eq. 19 and the output distribution during sampling as in Eq. 15.\\nAs shown in Fig 5, the boundary condition for clamping the cumulative probability function in\\nthe bucket could cause the mismatch, e.g., the true density should be centered in the center bucket\\nwhile the ouput distribution instead put the most density in the first and last buckets which cause the\\nmode-redundancy as shown in upper-left in Fig. 5. Though the weighted sum in Eq. 19 is optimized,\\nthe sampling procedure will rarely sample the center buckets. And such cases could be non-negligible\\nin our scenarios, especially when the number of bins is small for low dimensional data. To alleviate\\nthis issue, we instead update the output distribution in the sampling procedure to:\\nˆ\\nkc(θ, t) = NEAREST_CENTER(\\n\\\" K\\nX\\nk=1\\np(1)\\nO (k | θ, t)kc, . . . ,\\nK\\nX\\nk=1\\np(D)\\nO (k | θ, t)kc\\n#\\n)\\n(21)\\nFunction NEAREST_CENTER compares inputs to the center bins ⃗\\nkc =\\n\\u0010\\nk(1)\\nc , . . . , k(D)\\nc\\n\\u0011\\n, and return\\nthe nearest center for each input value. The updated distribution is unbiased towards the training\\nobjective and also reduce the variance during generation which could be found in the trajectory of\\nFig.5.\\n4\\nEXPERIMENTS\\n4.1\\nEXPERIMENT SETUP\\nTask and Datasets We focus on the 3D molecule generation task following the setting of prior\\nworks (Gebauer et al., 2019; Luo & Ji, 2021; Satorras et al., 2021a; Hoogeboom et al., 2022; Wu et al.,\\n2022). We consider both Unconditional Molecular Generation which assesses the capability to learn\\nthe underlying molecular data distribution and generate chemically valid and structurally diverse\\nmolecules and the Conditional Molecule Generation tasks which evaluate the capacity of generating\\nmolecules with desired properties. For Conditional Molecule Generation, we implement a conditional\\nversion GeoBFN with the details in the Appendix. The widely adapted QM9 (Ramakrishnan et al.,\\n2014) and the GEOM-DRUG (Gebauer et al., 2019; 2021) with large molecules are used for the\\nexperiments. And the data configurations directly follow previous work(Anderson et al., 2019;\\nHoogeboom et al., 2022; Xu et al., 2023)2.\\nEvaluation Metrics The evaluation configuration follows the prior works (Hoogeboom et al., 2022;\\nWu et al., 2022; Xu et al., 2023). For the Unconditional Molecular Generation, the bond types\\nare first predicted (single, double, triple, or none) based on pair-wise atomic distance and atom\\ntypes in the 10000 generated molecular geometries (Hoogeboom et al., 2022). With the obtained\\nmolecular graph, we evaluate the quality by calculating both atom stability and molecule stability\\nmetrics. Besides, the validity (based on RDKIT) and uniqueness are also reported. Regarding the\\n2The official implementation is at https://github.com/AlgoMole/GeoBFN\\n7\\n\\n\\nPublished as a conference paper at ICLR 2024\\nTable 1: Results of atom stability, molecule stability, validity, validity×uniqueness (V×U), and\\nnovelty. A higher number indicates a better generation quality. The results marked with an asterisk\\nwere obtained from our own tests. And GeoBFNk denote the results of sampling the molecules with\\na specific number of steps k\\nQM9\\nDRUG\\n# Metrics\\nAtom Sta (%)\\nMol Sta (%)\\nValid (%)\\nV×U (%)\\nNovelty (%)\\nAtom Sta (%)\\nValid (%)\\nData\\n99.0\\n95.2\\n97.7\\n97.7\\n-\\n86.5\\n99.9\\nENF\\n85.0\\n4.9\\n40.2\\n39.4\\n-\\n-\\n-\\nG-Schnet\\n95.7\\n68.1\\n85.5\\n80.3\\n-\\n-\\n-\\nGDM-AUG\\n97.6\\n71.6\\n90.4\\n89.5\\n74.6\\n77.7\\n91.8\\nEDM\\n98.7\\n82.0\\n91.9\\n90.7\\n58.0\\n81.3\\n92.6\\nEDM-Bridge\\n98.8\\n84.6\\n92.0\\n90.7\\n-\\n82.4\\n92.8\\nGEOLDM\\n98.9 ± 0.1\\n89.4 ± 0.5\\n93.8 ± 0.4\\n92.7 ± 0.5\\n57.0\\n84.4\\n99.3\\nGEOBFN 50\\n98.28 ± 0.1\\n85.11 ± 0.5\\n92.27 ± 0.4\\n90.72 ± 0.3\\n72.9\\n75.11\\n91.66\\nGEOBFN 100\\n98.64 ± 0.1\\n87.21 ± 0.3\\n93.03 ± 0.3\\n91.53 ± 0.3\\n70.3\\n78.89\\n93.05\\nGEOBFN 500\\n98.78 ± 0.8\\n88.42 ± 0.2\\n93.35 ± 0.2\\n91.78 ± 0.2\\n67.7\\n81.39\\n93.47\\nGEOBFN 1k\\n99.08 ± 0.06\\n90.87 ± 0.2\\n95.31 ± 0.1\\n92.96 ± 0.1\\n66.4\\n85.60\\n92.08\\nGEOBFN 2k\\n99.31 ± 0.03\\n93.32 ± 0.1\\n96.88 ± 0.1\\n92.41 ± 0.1\\n65.3\\n86.17\\n91.66\\nTable 2: Mean Absolute Error for molecular prop-\\nerty prediction with 500 sampling steps. A lower\\nnumber indicates a better controllable generation\\nresult.\\nProperty\\nα\\n∆ε\\nεHOMO\\nεLUMO\\nµ\\nCv\\nUnits\\nBohr3\\nmeV\\nmeV\\nmeV\\nD\\ncal\\nmolK\\nQM9*\\n0.10\\n64\\n39\\n36\\n0.043\\n0.040\\nRandom*\\n9.01\\n1470\\n645\\n1457\\n1.616\\n6.857\\nNatoms\\n3.86\\n866\\n426\\n813\\n1.053\\n1.971\\nEDM\\n2.76\\n655\\n356\\n584\\n1.111\\n1.101\\nGEOLDM\\n2.37\\n587\\n340\\n522\\n1.108\\n1.025\\nGEOBFN\\n2.34\\n577\\n328\\n516\\n0.998\\n0.949\\nTable 3:\\nAblation study, GeoBFN models\\nmolecule charge settings, the sampling step is\\nset to 1,000.\\nCharge Feature\\nAtom Stable (%)\\nMol Stable (%)\\ndiscretised_basis\\n99.08\\n90.87\\ncontinuous_basis\\n98.97\\n89.94\\ndiscrete\\n98.93\\n88.93\\ndiscrete + continuous\\n98.96\\n89.33\\ndiscrete + discretised\\n98.91\\n88.65\\nConditional Molecule Generation, we evaluate our conditional version of GeoBFN on QM9 with\\n6 properties: polarizability α, orbital energies εHOMO, εLUMO and their gap ∆ε, Dipole moment\\nµ, and heat capacity Cv. Following previous work Hoogeboom et al. (2022); Xu et al. (2023), the\\nconditional GeoBFN is fed with a range of property s to generate samples and the same pre-trained\\nclassifier w is utilized to measure the property of generated molecule as ˆ\\ns. The Mean Absolute Error\\n(MAE) between s and ˆ\\ns is calculated to measure whether the generated molecules is related to the\\nconditioned property.\\nBaselines GeoBFN is compared with several advanced baselines including G-Schnet (Gebauer et al.,\\n2019), Equivariant Normalizing Flows (ENF) (Satorras et al., 2021a) and Equivariant Graph Diffusion\\nModels (EDM) with its non-equivariant variant (GDM) (Hoogeboom et al., 2022). Also with recent\\nadvancements, EDM-Bridge (Wu et al., 2022) which improves upon the performance of EDM by\\nincorporating well-designed informative prior bridges and also GeoLDM (Xu et al., 2023) where\\na latent space diffusion model is applied are both included. To yield a fair comparison, all the\\nmethod-agnostic configurations are set as the same. The implementation details could be found in\\nAppendix. B.\\n0\\n1000\\n2000\\n3000\\n4000\\nSampling Step\\n0.65\\n0.70\\n0.75\\n0.80\\n0.85\\n0.90\\n0.95\\n1.00\\nMolecule Stability\\nGeoBFN\\nEDM\\nEDM-Bridge\\nGEOLDM\\nupper bound\\nFigure 4: QM9 Molecule Stability wrt. Sampling\\nSteps\\nOriginal \\nsampling \\nalgorithm\\nImproved \\nsamping \\nalgorithm\\nFigure 5: 2D Synthetic case of optimized\\nsynthetic example. In the left columns,\\ngenerated samples are in orange, and data\\npoints are in blue.\\n8\\n\\n\\nPublished as a conference paper at ICLR 2024\\n4.2\\nMAIN RESULTS\\nThe results of Unconditional Molecular Generation can be found in Tab. 1. We could observe that in\\nboth the QM9 and GEOM-DRUG datasets, GeoBFN achieves a new state-of-the-art performance\\nregarding both the quality and diversity of the generated molecules which demonstrates the huge\\npotential of GeoBFN on geometry generative modeling. The phenomenon demonstrates that the\\nGeoBFN does not hold the tendency to collapse to the subset of training data which could imply\\na probabilistic generalization ability and could be useful for several application scenarios; The\\nConditional Molecule Generation results can be found in Tab. 2. GeoBFN consistently outperforms\\nother baseline models by an obvious margin in all conditional generation tasks. This clearly highlights\\nthe effectiveness and generalization capability of the proposed methods.\\n4.3\\nANY-STEP SAMPLING\\nOne notable property of GeoBFN is that training with the continuous-time loss, e.g., Eq. 19, the\\nsampling could be conducted with any steps without incurring additional training overhead. As shown\\nin Tab. 1, GeoBFN could get superior performance compared to several advanced models with only 50\\nsteps during sampling which brings 20× speed-up during sampling due to the benefit of low variance\\nparameter space. As we could find in Fig 4, with the sampling steps increasing from 50 to 4600, the\\nmolecule stability could be further boosted to approach the upper bound, e.g., 94.25% with 4000 steps.\\n4.4\\nABLATION STUDIES\\nWe conduct ablation studies on the effect of input modalities in Tab. 3. We try different compositions\\nand losses to represent the atom types, discretised basis refers to the case where the charge feature is\\nused with discretised and the Gaussian basis, i.e., ϕj(x) = exp\\n\\u0012\\n−\\n(x−µj)2)\\n2σ2\\n\\u0013\\nis used as functional\\nembedding for charge; continous basis only differ in that the continous loss is utilized. The discrete\\nrefers to including the one-hot type representation; discrete+continuous refers to both the one-hot\\ntype and charge are included while continuous loss is included; Similar is the discrete+continuous.\\nWith only discretised variable utilized, the performance is superior to including the discrete variable\\nwhich implies powerful probabilistic modeling capacity and the benefits of applying similar modality.\\n5\\nRELATED WORK\\nPrevious molecule generation studies have primarily focused on generating molecules as 2D graphs\\n(Jin et al., 2018; Liu et al., 2018; Shi et al., 2020), but there has been increasing interest in 3D\\nmolecule generation. With the increasing interest in 3D molecule generation, G-Schnet and G-\\nSphereNet (Gebauer et al., 2019; Luo & Ji, 2021) respectively, employ autoregressive techniques\\nto create molecules in a step-by-step manner by progressively connecting atoms or molecular\\nfragments. These frameworks have also been extended to structure-based drug design (Li et al.,\\n2021; Peng et al., 2022; Powers et al., 2022). There are approaches use atomic density grids\\nthat generate the entire molecule in a single step by producing a density over the voxelized 3D\\nspace (Masuda et al., 2020). Most recently, the attention has shifted towards using DMs for 3D\\nmolecule generation (Hoogeboom et al., 2022; Wu et al., 2022; Peng et al., 2023; Xu et al., 2023),\\nwith successful applications in target drug generation (Lin et al., 2022), antibody design (Luo et al.,\\n2022), and protein design (Anand & Achim, 2022; Trippe et al., 2022). However, our method is\\nbased on the Bayesian Flow Network (Graves et al., 2023) objective and hence lies in a different\\nmodel family which fundamentally differs from this line of research in both training and generation.\\n6\\nCONCLUSION\\nWe introduce GeoBFN, a new generative framework for molecular geometry. GeoBFN operates\\nin a differentiable parameter space for variables from different modalities. Also, the less variance\\nin parameter space is naturally compatible with the noise sensitivity of molecule geometry. Given\\nthe appealing property, the GeoBFN achieves state-of-the-art performance on several 3D molecule\\ngeneration benchmarks. Besides, GeoBFN can also conduct sampling with an arbitary number\\nof steps to reach an optimal trade-off between efficiency and quality (e.g., 20× speedup without\\nsacrificing performance).\\n9\\n\\n\\nPublished as a conference paper at ICLR 2024\\nACKNOWLEDGMENTS\\nThe authors thank the anonymous reviewers for reviewing the draft. This work is supported by the\\nNational Science and Technology Major Project (2022ZD0117502), Natural Science Foundation\\nof China (62376133) and Guoqiang Research Institute General Project, Tsinghua University (No.\\n2021GQG1012).\\nREFERENCES\\nNamrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant\\ndenoising diffusion probabilistic models. arXiv preprint arXiv:2205.15019, 2022.\\nBrandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural\\nnetworks. Advances in neural information processing systems, 32, 2019.\\nNiklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry-adapted generation of 3d point\\nsets for the targeted discovery of molecules. Advances in neural information processing systems,\\n32, 2019.\\nNiklas WA Gebauer, Michael Gastegger, Stefaan SP Hessmann, Klaus-Robert Müller, and Kristof T\\nSchütt. Inverse design of 3d molecular structures with conditional generative neural networks.\\narXiv preprint arXiv:2109.04824, 2021.\\nAlex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. Bayesian flow\\nnetworks. arXiv preprint arXiv:2308.07037, 2023.\\nJonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. arXiv preprint\\narXiv:2006.11239, 2020.\\nEmiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion\\nfor molecule generation in 3d. In International Conference on Machine Learning, pp. 8867–8887.\\nPMLR, 2022.\\nWengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for\\nmolecular graph generation. In International conference on machine learning, pp. 2323–2332.\\nPMLR, 2018.\\nBowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror.\\nLearning from protein structure with geometric vector perceptrons. In International Conference on\\nLearning Representations, 2021.\\nDiederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3nd International\\nConference on Learning Representations, 2014.\\nJonas Köhler, Leon Klein, and Frank Noe. Equivariant flows: Exact likelihood generative learning for\\nsymmetric densities. In Proceedings of the 37th International Conference on Machine Learning,\\n2020.\\nXiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori Hashimoto. Diffusion-\\nLM improves controllable text generation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave,\\nand Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL\\nhttps://openreview.net/forum?id=3s9IrEsjLyk.\\nYibo Li, Jianfeng Pei, and Luhua Lai. Structure-based de novo drug design using 3d deep generative\\nmodels. Chemical science, 12(41):13664–13675, 2021.\\nHaitao Lin, Yufei Huang, Meng Liu, Xuanjing Li, Shuiwang Ji, and Stan Z Li. Diffbp: Generative\\ndiffusion of 3d molecules for target protein binding. arXiv preprint arXiv:2211.11214, 2022.\\nQi Liu, Miltiadis Allamanis, Marc Brockschmidt, and Alexander Gaunt. Constrained graph variational\\nautoencoders for molecule design. In Advances in neural information processing systems, 2018.\\n10\\n\\n\\nPublished as a conference paper at ICLR 2024\\nShitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigen-specific\\nantibody design and optimization with diffusion-based generative models for protein structures.\\nIn Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in\\nNeural Information Processing Systems, 2022. URL https://openreview.net/forum?\\nid=jSorGn2Tjg.\\nYouzhi Luo and Shuiwang Ji. An autoregressive flow model for 3d molecular geometry generation\\nfrom scratch. In International Conference on Learning Representations, 2021.\\nTomohide Masuda, Matthew Ragoza, and David Ryan Koes. Generating 3d molecular structures con-\\nditional on a receptor binding site with deep generative models. arXiv preprint arXiv:2010.14442,\\n2020.\\nChenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon.\\nSDEdit: Guided image synthesis and editing with stochastic differential equations. In International\\nConference on Learning Representations, 2022. URL https://openreview.net/forum?\\nid=aBsCjcPu_tE.\\nAdam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito,\\nZeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in\\npytorch. In NIPS-W, 2017.\\nXingang Peng, Shitong Luo, Jiaqi Guan, Qi Xie, Jian Peng, and Jianzhu Ma. Pocket2mol: Efficient\\nmolecular sampling based on 3d protein pockets. In International Conference on Machine Learning,\\n2022.\\nXingang Peng, Jiaqi Guan, Qiang Liu, and Jianzhu Ma.\\nMoldiff: Addressing the atom-bond\\ninconsistency problem in 3d molecule diffusion generation. arXiv preprint arXiv:2305.07508,\\n2023.\\nAlexander S. Powers, Helen H. Yu, Patricia Suriana, and Ron O. Dror. Fragment-based ligand\\ngeneration guided by geometric deep learning on protein-ligand structure. bioRxiv, 2022. doi: 10.\\n1101/2022.03.17.484653. URL https://www.biorxiv.org/content/early/2022/\\n03/21/2022.03.17.484653.\\nRaghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum\\nchemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1–7, 2014.\\nVictor Garcia Satorras, Emiel Hoogeboom, Fabian B Fuchs, Ingmar Posner, and Max Welling. E (n)\\nequivariant normalizing flows for molecule generation in 3d. arXiv preprint arXiv:2105.09016,\\n2021a.\\nVıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks.\\nIn International conference on machine learning, pp. 9323–9332. PMLR, 2021b.\\nKristof T Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R Müller, and Alexandre Tkatchenko.\\nQuantum-chemical insights from deep tensor neural networks. Nature communications, 8:13890,\\n2017.\\nChence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. Graphaf: a\\nflow-based autoregressive model for molecular graph generation. arXiv preprint arXiv:2001.09382,\\n2020.\\nBrian L Trippe, Jason Yim, Doug Tischer, Tamara Broderick, David Baker, Regina Barzilay, and\\nTommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motif-\\nscaffolding problem. arXiv preprint arXiv:2206.04119, 2022.\\nLemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and qiang liu. Diffusion-based molecule\\ngeneration with informative prior bridges. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave,\\nand Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL\\nhttps://openreview.net/forum?id=TJUNtiZiTKE.\\nMinkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric\\ndiffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022.\\n11\\n\\n\\nPublished as a conference paper at ICLR 2024\\nMinkai Xu, Alexander Powers, Ron Dror, Stefano Ermon, and Jure Leskovec. Geometric latent\\ndiffusion models for 3d molecule generation. arXiv preprint arXiv:2305.01140, 2023.\\nA\\nEXPLANATION OF THE DATA EXCHANGE PERSPECTIVE OF BAYESIAN\\nFLOW NETWORKS\\nIn this section, we provide a brief overview of the Bayesian Flow Networks Graves et al. (2023) from\\na data exchange perspective. Bayesian Flow Networks (BFNs) is a new class of generative model\\nthat operates on the parameters of a set of independent distributions with Bayesian inference. The\\nmeta elements of BFNs are the input distributions, sender distributions, and output distributions.\\nTo start with, we denote the D-dimensional variable as m =\\nm(1), . . . , m(D)\\u0001\\n∈MD, and\\nθ =\\nθ(1), . . . , θ(D)\\u0001\\nrepresent the parameters of a D-dimensional factorised distribution, i.e.,\\np(m | θ) = QD\\nd=1 p\\nm(d) | θ(d)\\u0001\\n.\\nThe basic logic of BFN could be better explained by the following communication example between\\nthe sender, referred to as Alice, and the receiver Bob. Alice aims to transfer some data to Bod in a\\nprogressive fashion, i.e., at each timestep, Alice corrupts the data according to some channel noise,\\nand then the noisy sample is transferred. The sender distribution is then defined to describe the noise-\\nadding procedure, which is also a factorized distribution, pS(y | m; α) = QD\\nd=1 pS\\ny(d) | m(d); α\\n\\u0001\\n.\\nα refers to the accuracy parameter, α = 0 refers to the information of the sample being totally\\ndestroyed by the noise, and with α increase the noisy sample will contain more information of the\\noriginal sample. Intuitively, the sender distribution could be approximately understood as adding\\nnoise to each dimension of the data independently.\\nAfter receiving the noisy sample y, Bob will first update an initial “guess” of what is the original\\nsample behind the noisy sample, i.e. input distribution. Note that except for the noisy sample, Bob\\nalso knows the accuracy parameter α and noise formulation while not aware of the original sample m,\\ni.e., the noise level to create such a sample. The input distribution is initially a simple prior on the data\\nspace, e.g., a standard Gaussian, lies in the mean-field family, pI(m | θ) = QD\\nd=1 pI\\nx(d) | θ(d)\\u0001\\n.\\nThe parameter of input distribution will be updated through Bayesian inference, noted as θi =\\nh (θi−1, y, αi). This update usually lies in a simple form, e.g. additive or weighted average.\\nAfter updating the parameter of input distribution, Bob has an “assistant” which will help to provide\\na better guess on the original data which generates the observed noisy sample. The assistant aims\\nto exploit more context information between different dimensions, e.g., the relationship between\\ndifferent pixels in an image, in contrast to updating each dimension independently as in the input\\ndistribution. Empirically, the assistant could be implemented by a neural network Ψ which takes\\nall parameters of input distribution for the prediction of parameters of each dimension, i.e, Ψ(θ) =\\nΨ(1)(θ, t), . . . , Ψ(D)(θ, t)\\n\\u0001\\n.\\nThe output distribution is then implied by the predicted parameter which lies in the formulation of\\npO(m | θ, t) = QD\\nd=1 pO\\nm(d) | Ψ(d)(θ, t)\\n\\u0001\\n. Then Bob could construct a distribution to approxi-\\nmate the sender distribution at accuracy α by combining the output distribution with the known noise\\nform, accuracy, i.e., pR (· | θ; t, α) = EpO(x′|θ;t)pS (y | x′; α). Such distribution is called receiver\\ndistribution. The \\\"assistant\\\" of BFNs Ψ is to minimize the KL divergence with a defined accuracy\\nscheduler under different timesteps, i.e., DKL (pS (· | m; αi) ∥pR (· | θi−1; ti−1, αi)), which could\\nalso be interpreted as transmission cost under the bits-back coding scheme.\\nB\\nIMPLEMENTATION DETAILS\\nThe bayesian flow network is implemented with EGNNs Satorras et al. (2021b) by PyTorch (Paszke\\net al., 2017) package. We set the dimension of latent invariant features k to 1 for QM9 and 2 for\\nDRUG, which extremely reduces the atomic feature dimension. For the training of vector field\\nnetwork vθ: on QM9, we train EGNNs with 9 layers and 256 hidden features with a batch size 64;\\nand on DRUG, we train EGNNs with 4 layers and 256 hidden features, with batch size 64. The\\nmodel uses SiLU activations. We train all the modules until convergence. For all the experiments,\\nwe choose the Adam optimizer (Kingma & Ba, 2014) with a constant learning rate of 10−4 as our\\n12\\n\\n\\nPublished as a conference paper at ICLR 2024\\ndefault training configuration. The training on QM9 takes approximately 2000 epochs, and on DRUG\\ntakes 20 epochs.\\nC\\nPROOF OF THEOREMS\\nIn this Section, we provide the formal proof of the Theorem. 3.1 and Proposition. 3.2, as well as the\\ndetailed derivations for Equations.\\nC.1\\nDISCUSSION ON THE TRANSLATIONAL INVARIANCE\\nRemark C.1. It is important to distinguish it from the rotation invariant. The rotational invariant\\nis defined as p(x) = p(Rx), while the translational is not as p(x) = p(x + t) as such distribution\\ncan not integrate into one and hence does not exist. Fortunately, the freedom of translation could\\nbe eliminated by only focusing on learning distribution on the linear subspace where the center of\\ngravity is always zero. This is, for all configurations on Rn×3 space, the density on the zero CoM\\nspace is utilized to represent their density; It’s important to note that the distribution is not defined for\\nconfigurations outside the zero CoM space. However, it remains possible to leverage the distribution\\nto provide a density-evaluation (not probability density) on the configurations outside the zero CoM\\nspace. This is achieved by projecting them back into the subspace. The evaluation procedure for\\nconfigurations out of zero CoM space could only get a quantity defined artificially instead of the true\\ndensity of some real distribution, e.g. It is referred to as \\\"CoM-free density\\\" in (Xu et al., 2022).\\nThus, there does not exist correctness issues.\\nC.1.1\\nZERO CENTER OF MASS(COM) IN THE GEOBFN\\nHere we provide detailed discussions on optimizing the distribution in the zero CoM space. Recall\\nthe training objective in equation 8,\\nLVLB(x) =\\nE\\npϕ(θx\\n0 ,...,θx\\nn)\\n\\\" n\\nX\\ni=1\\nDKL(pS (· | x; αi) ∥pR(· | θx\\ni−1; αi)) −log pϕ (x | θx\\nn)\\n#\\n,\\n(22)\\nwhere pϕ(x|θx\\nn) = pO(x|θx\\nn, ϕ). For learning a distribution on the zero CoM space of Rn×3, the\\npS (· | x; αi), pR(· | θx\\ni−1; αi) and pϕ (x | θx\\nn) are all defined and supported on the zero CoM space,\\nhere Pn\\ni=1 xi = 0 and Pn\\ni=1 θx = 0. In other words, such distribution has no definition for variable\\nv ∈Rn×3 if Pn\\ni=1 vi ̸= 0. We then express the likelihood function of an isotropic diagonal Gaussian\\ndistribution, which is originally defined on the zero CoM space ((n −1) × 3-dimensional), in the\\nambient space (n × 3-dimensional) as (Hoogeboom et al., 2022):\\nNx\\nx | µ, σ2I\\n\\u0001\\n= (\\n√\\n2πσ)−(n−1)×3 exp\\n\\u0012\\n−1\\n2σ2 ∥x −µ∥2\\n\\u0013\\n(23)\\nHere σ2 is the variance which is equivalent for each dimension.\\nRecall the Eq.\\n35 in the\\n(Graves et al., 2023), which shows that DKL(pS (· | x; αi) ∥pR(· | θx\\ni−1; αi)) takes the form of\\nDKL\\nN\\nx, α−1\\ni I\\n\\u0001\\n∥N(Φx(θg, t), α−1\\ni I)\\n\\u0001\\nwhich is the KL divergence between to diagonal Gaus-\\nsian, then we derive the KL divergence for isotropic diagonal normal distributions of zero CoM with\\nmeans represent on the ambient space. If pS = N\\n ˆ\\nµ1, σ2I\\n\\u0001\\nand pR = N\\n ˆ\\nµ2, σ2I\\n\\u0001\\non subspace,\\nwhere ˆ\\nµ1 and ˆ\\nµ2 is (n −1) × 3-dimension. Then the KL between them could be represented as:\\nDKL(q∥p) = 1\\n2\\n\\\"\\n∥ˆ\\nµ1 −ˆ\\nµ2∥2\\nσ2\\n#\\n(24)\\nThere is an orthogonal transformation Q which transforms the ambient space µi ∈Rn×3 where\\nP\\ni µi = 0 to the subspace in the way that\\n\\u0014\\nˆ\\nµ\\n0\\n\\u0015\\n= Qµ. With ∥ˆ\\nµ∥= ∥\\n\\u0014\\nµ\\n0\\n\\u0015\\n∥= ∥µ∥, there is\\n∥ˆ\\nµ1 −ˆ\\nµ2∥2 = ∥µ1 −µ2∥2. Hence we have:\\nDKL\\nN\\nx, α−1\\ni I\\n\\u0001\\n∥N(Φx(θg, t), α−1\\ni I)\\n\\u0001\\n= αi\\n2 ∥x −Φx (θg, t)∥2\\n(25)\\nHence we demonstrate the correctness of our objective in Eq. 19.\\n13\\n\\n\\nPublished as a conference paper at ICLR 2024\\nC.1.2\\nPROOF OF THE TRANSLATIONAL INVARIANT DENSITY EVALUATION PROCEDURE.\\nProof. For an n-atom molecule g = ⟨x, h⟩, the coordinate variable x has the dimension of n × 3.\\nNote that with the zero Center of Mass mapping (Xu et al., 2022; Satorras et al., 2021a), where\\nwe constrain the center of gravity as zero (Pn\\ni=1 xi = 0), then variable x essentially lies in the\\n(n −1) × 3-dimensional linear subspace. The generative distributions pX mentioned in all of the\\nrelated literature (Satorras et al., 2021a; Hoogeboom et al., 2022; Xu et al., 2022; 2023) is constrained\\nin the zero Center of Mass space. This is, for samples in the ambient space, if Pn\\ni=1 xi ̸= 0, pX is not\\ndefined. The translational invariant property of distribution pX mentioned is not referred to the fact\\nthat pX(x) = pX(x + t) for all translation vector t is satisfied in the ambient space with dimension\\nn × 3. Actually, such conditions could not be satisfied in any space (Satorras et al., 2021a). The\\ntranslational invariant condition actually refers to the invariant function f which could evaluate the\\ndensity of all the ambient space based on pX, the evaluated density by f also referred to as \\\"CoM-free\\nstandard density\\\" in (Xu et al., 2022). The function f is defined as\\nf(x) = pX(Qx)\\n(26)\\nwhere Q refers to the operation which maps the x to the zero Center of Mass space, e.g. in our work\\nQ is defined as\\nQ = I3 ⊗\\n\\u0012\\nIN −1\\nN 1N1T\\nN\\n\\u0013\\n,\\ns.t.\\nQx =\\n\\n\\nx1 −\\nPn\\ni=1 xi\\nn\\n· · ·\\nxn −\\nPn\\ni=1 xi\\nn\\n\\n\\n(27)\\nthat subtracting mean\\nPn\\ni=1 xi\\nn\\nfrom each xi, where Ik denotes the k × k identity matrix and 1k\\ndenotes the k-dimensional vector filled with 1s. Then the density evaluation function f is translational\\ninvariant in the ambient space:\\nf(x + t) = pX( ˆ\\nQ(x + t)) = pX(\\n\\n\\nx1 + ti −\\nPn\\ni=1(xi+ti)\\nn\\n· · ·\\nxn + tn −\\nPn\\ni=1(xi+ti)\\nn\\n\\n)\\n(28)\\nNote t stands for a translation vector, which implies that t1 = · · · = ti = tn = C ∈R3. Then we\\nhave:\\nf(x + t) = pX(\\n\\n\\nx1 + C −\\nPn\\ni=1(xi+C)\\nn\\n· · ·\\nxn + C −\\nPn\\ni=1(xi+C)\\nn\\n\\n) = pX(\\n\\n\\nx1 −\\nPn\\ni=1 xi\\nn\\n· · ·\\nxn −\\nPn\\ni=1 xi\\nn\\n\\n) = pX(Qx) = f(x)\\n(29)\\nFurthermore, the above proof has no constraint on the distribution pX. This is, for any distribution on\\nthe zero Center of Mass space, the corresponding evaluation function defined in Eq. 26 is translational\\ninvariant.\\nC.2\\nPROOF OF THEOREM. 3.1.\\nGiven the above discussion on the translational invariance, for simplicity, we could only focus on the\\nrotation transformation.\\nProof. Recall the graphical model in Fig. 2, we could reformulate the density function in Eq. 7 as:\\npϕ(x) =\\nZ\\npϕ(x | θx\\n1, · · · , θx\\nn)pϕ(θx\\n1, · · · , θx\\nn)dθx\\n1:n\\n(definition of marginal)\\n=\\nZ\\npϕ(x | θx\\nn)p(θ0)\\nn\\nY\\ni=1\\npU (θi | θi−1; αi) dθx\\n1:n.\\n(30)\\nNote that pϕ(x | θx\\nn) = pϕ(Rx | Rθx\\nn) = pO (R(x) | R(θx\\nn); ϕ) due to the property of EGNN,\\nand p(θ0) = p(Rθ0) since θ0 = 0. Then we prove that pU (θi | θi−1; αi) satisfies the equiv-\\nariant condition that pU (θi | θi−1; αi) = pU (Rθi | Rθi−1; αi). Recall that pU (θi | θi−1; αi) =\\n14\\n\\n\\nPublished as a conference paper at ICLR 2024\\nE\\npO(yi|θi−1;αi)δ (θi −h (θi−1, yi, αi)), then we have:\\npU (Rθi | Rθi−1; αi) =\\nE\\npO(yi|Rθi−1;αi)δ (Rθi −h (Rθi−1, yi, αi))\\n=\\nZ\\npO (yi | Rθi−1; αi) δ (Rθi −h (Rθi−1, yi, αi)) dyi\\n(31)\\nThen we apply integration-by-substitution and replace the variable yi with a new variable y′\\ni, i.e.\\nyi = Ry′\\ni, into the Eq. 31:\\nZ\\npO (yi | Rθi−1; αi) δ (Rθi −h (Rθi−1, yi, αi)) dyi\\n=\\nZ\\npO (Ry′\\ni | Rθi−1; αi) δ (Rθi −h (Rθi−1, Ry′\\ni, αi)) dRy′\\ni\\n=\\nZ\\npO (Ry′\\ni | Rθi−1; αi) δ (Rθi −h (Rθi−1, Ry′\\ni, αi)) |det(R)|dy′\\ni\\n(32)\\nThe rotation matrix R is a SO(3) matrix, thus the |det(R)| = 1. And for the continuous coordinate\\nvariable, the update function h defined in Eq. 11 is also equivariant:\\nh (Rθi−1, Ryi, αi) = Rθi−1ρi−1 + Ryiαi\\nρi\\n= Rh (θi−1, yi, αi)\\n(33)\\nPutting these conditions back to the Eq. 32, we have that\\npU (Rθi | Rθi−1; αi) =\\nZ\\npO (yi | Rθi−1; αi) δ (Rθi −h (Rθi−1, yi, αi)) dyi\\n=\\nZ\\npO (Ry′\\ni | Rθi−1; αi) δ (Rθi −Rh (θi−1, y′\\ni, αi)) |det(R)|dy′\\ni\\n=\\nZ\\npO (y′\\ni | θi−1; αi) δ (θi −h (θi−1, y′\\ni, αi)) dy′\\ni\\n= pU (θi | θi−1; αi)\\n(34)\\nHence the transitions on the θ space are the Markov and equivariant to rotation as shown in Eq. 30.\\nThe initial state θ0 is a zero vector 0 which is rotation invariant. To derive the rotation-invariant\\nproperty of pϕ, we will use the following Lemma, which is the direct application of Proposition 1 in\\n(Xu et al., 2022). We changed the notation to make it consistent with our literature.\\nLemma C.2. (Xu et al., 2022) Let p (θ0) be an SE(3)-invariant density function, i.e., p (θ0) =\\np (Tg (θ0)).\\nIf Markov transitions p (θi | θi−1) are SE(3)-equivariant, i.e., p (θi | θi−1) =\\np (Tg (θi) | Tg (θi−1)), then we have that the density p (θn) =\\nR\\np (θ0) p (θ1:n | θ0) dθ0:n is also\\nSE(3)-invariant. (Tg stands for the SE(3) transformations.)\\nFor completeness, we also include the derivation of the lemma from (Xu et al., 2022):\\np (Tg(θn)) =\\nZ\\np (Tg(θ0)) p (Tg(θ1:n) | Tg(θ0)) dθ0:n\\n=\\nZ\\np (Tg (θ0)) Πn\\ni=1p (Tg (θi) | Tg (θi−1)) dθ0:n\\n=\\nZ\\np (θ0) Πn\\ni=1pθ (Tg (θi) | Tg (θi−1)) dθ0:n\\n( invariant prior p (θ0))\\n=\\nZ\\np (θ0) Πn\\ni=1pθ (θi | θi−1) dθ0:n\\n(equivariant kernels p (θi | θi−1))\\n=\\nZ\\np (θ0) p (θ1:n | θ0) dθ0:n\\n= p (θn)\\n(35)\\nGiven the invariant property of θ0 and equivariant property of the transition pU (θi | θi−1; αi) and\\nthe pϕ(x | θx\\nn), we could directly get the conclusion in Theorem. 3.1. Now we finish the proof.\\n15\\n\\n\\nPublished as a conference paper at ICLR 2024\\nC.3\\nPROOF OF PROPOSITION. 3.2.\\nProof. Then we derive the invariant property of the variational lower bound in of the variational\\nlower bounds in equation 8:\\nLV LB(x) =\\nE\\npϕ(θx\\n0 ,...,θx\\nn)\\n[\\nn\\nX\\ni=1\\nDKL(pS (· | x; αi) ∥pR(· | θx\\ni−1; αi)) −log pϕ (x | θx\\nn)]\\nTo start with, we consider the first term:\\nE\\npϕ(θx\\n0 ,...,θx\\nn)\\nn\\nX\\ni=1\\nDKL(pS (· | x; αi) ∥pR(· | θx\\ni ; αi))\\n=\\nn−1\\nX\\ni=0\\nE\\npϕ(θx\\ni )\\nDKL(pS (· | x; αi) ∥pR(· | θx\\ni ; αi))\\n(36)\\nNote a natural conclusion from the Theorem. 3.1 is that the SE(3) invariance property is not only\\nsatisfied in the marginal distribution of the last time step variable p(θn), but also for the distribution\\nof any intermediate p(θi). Such property could be justified based on the condition in Lemma. C.2.\\nActually, the proof of Theorem. 3.1 in the above section does not specify the time steps, hence the\\nmarginal distribution of any time step could be proved in exactly the same way. Consider the i-th\\nterm in the KL part of LVLB(Rx):\\nE\\npϕ(θx\\ni )\\nDKL(pS (· | Rx; αi) ∥pR(· | θx\\ni ; αi))\\n=\\nZ\\npϕ (θx\\ni ) DKL(pS (· | Rx; αi) ∥pR(· | θx\\ni ; αi))dθx\\ni\\n(37)\\nwe introduce the variable θ′\\ni similar to Eq. 32, i.e. θi = Rθ′\\ni, and then we extend i-th term in the\\nEq. 36:\\nZ\\npϕ\\n\\u0010\\nRθx′\\ni\\n\\u0011\\nDKL(pS (· | Rx; αi) ∥pR(· | Rθx′\\ni ; αi))|det(R)|dθx′\\ni\\n(38)\\nAs proved in the proof of Theorem. 3.1, pϕ is invariant and hence pϕ\\n\\u0010\\nRθx′\\ni\\n\\u0011\\n= pϕ(θx′\\ni ); Also the\\n|det(R)| = 1 for SO(3) rotation matrix. And then we discuss the KL divergence term:\\nDKL(pS (· | Rx; αi) ∥pR(· | Rθx′\\ni ; αi)) =\\nZ\\npS (y | Rx; αi) log pS (y | Rx; αi)\\npR(y | Rθx′\\ni ; αi)dy\\n=\\nZ\\npS (Ry′ | Rx; αi) log pS (Ry′ | Rx; αi)\\npR(Ry′ | Rθx′\\ni ; αi)det(R)|dy′\\n=\\nZ\\npS (y′ | x; αi) log pS (y′ | x; αi)\\npR(y′ | θx′\\ni ; αi)dy′ = DKL(pS (· | x; αi) ∥pR(· | θx′\\ni ; αi))\\n(39)\\nNote that pS (y′ | x; αi) = pS (Ry′ | Rx; αi) is due to that the sender distribution is isotropic;\\nAnd for receiver distribution, the equivariant property that pR (y′ | x; αi) = pR (Ry′ | Rx; αi) is\\nguaranteed by both the parameterization of pO with Equivariant Graph Neural Network and the\\nisotropic pS. At last, we put the above conclusion back to Eq. 37, and we get that:\\nE\\npϕ(θx\\ni )\\nDKL(pS (· | Rx; αi) ∥pR(· | θx\\ni ; αi))\\n=\\nZ\\npϕ (θx\\ni ) DKL(pS (· | Rx; αi) ∥pR(· | θx\\ni ; αi))dθx\\ni\\n=\\nZ\\npϕ (θx\\ni ) DKL(pS (· | x; αi) ∥pR(· | θx\\ni ; αi))dθx\\ni\\n=\\nE\\npϕ(θx\\ni )\\nDKL(pS (· | x; αi) ∥pR(· | θx\\ni ; αi))\\n(40)\\nAnd here we prove the first term in LV LB(Rx) is equivalent to LV LB(x). The second term could be\\nderived in exactly the same way, and here we finish the proof.\\n16\\n\\n\\nPublished as a conference paper at ICLR 2024\\nC.4\\nDERIVATION OF EQUATION 8\\nNote that the equation 8:\\nLVLB(x) =\\nE\\npϕ(θx\\n0 ,...,θx\\nn)\\n\\\" n\\nX\\ni=1\\nDKL(pS (· | x; αi) ∥pR(· | θx\\ni−1; αi)) −log pϕ (x | θx\\nn)\\n#\\n,\\n(41)\\nis the extension formulation of Eq. 1. To align the notation of Eq. 1 and equation 8, we use x in\\nthe following derivation. We first consider the term −DKL(q∥pϕ (y1, . . . , yn)) in Eq. 1, we put the\\nEq. 2\\nq = q (y1, . . . , yn | x) =\\nn\\nY\\ni=1\\npS (yi | x; αi)\\n(42)\\nAnd the pϕ (y1, . . . , yn)) in Eq. 6 as\\npϕ (y1, . . . , yn) =\\nE\\npϕ(θ0:n−1)\\n\\\" n\\nY\\ni=1\\nE\\npO(x′\\ni|θi−1;ϕ) [pS(yi|x′\\ni; αi)]\\n#\\n=\\nE\\npϕ(θ0:n)\\nn\\nY\\ni=1\\npR(yi|θi−1; αi)\\n(43)\\nPutting them together into the KL divergence term, and then we get the\\nDKL(q∥pϕ (y1, . . . , yn)) =\\nE\\nQn\\ni=1 pS(yi|x;αi) log\\nQn\\ni=1 pS (yi | x; αi)\\nE\\npϕ(θ0:n−1)\\nQn\\ni=1 pR(yi|θi−1; αi)\\n=\\nE\\npϕ(θ0:n−1)\\nE\\nQn\\ni=1 pS(yi|x;αi) log\\nQn\\ni=1 pS (yi | x; αi)\\nQn\\ni=1 pR(yi|θi−1; αi)\\n=\\nE\\npϕ(θ0:n−1)\\nn\\nX\\ni=1\\nDKL(pS (· | g; αi) ∥pR(· | θi−1; αi))\\n(44)\\nAnd we have derived the first term in equation 8. And for the second term,\\nlog pϕ(x|y1, · · · , yn) = log pϕ(x|θ0, · · · , θn)\\n(Graphical Model in Fig. 2)\\n= log pϕ(x|θn)\\n(Markov Property of θ)\\n(45)\\nAnd here we finish the derivation.\\nD\\nDETAILS ON CONDITIONAL GENERATION EXPERIMENTS\\nD.1\\nPARAMETERIZATION AND SAMPLING\\nFor the conditional experiments, we directly follow the conditional setting of previous litera-\\nture (Hoogeboom et al., 2022). We discuss the details of the parameterization and sampling in\\nthe following. For conditional experiments, we add the property c as the extra input for the interde-\\npendency modeling module in Eq. 19. The conditional objective will be as:\\nL∞(g, c)\\n=\\nE\\nt∼U(0,1),pF (θg|g;t)\\n\\u0014αx(t)\\n2\\n∥x −Φx∥2 + αhc(t)\\n2\\n∥hc −Φhc∥2 + αht(t)\\n2\\n∥ht −Φht∥2\\n\\u0015\\n(46)\\nWhere Φ·(θg, t, c) is short for Φ·(θg, t, c).\\nFor the sampling procedure, the property c and node number M will be firstly sampled from a prior\\np(c, M) defined in (Hoogeboom et al., 2022). Here p(c, M) is computed on the training partition as\\na parametrized two-dimensional categorical distribution where the continuous variable c is discretized\\ninto small uniformly distributed intervals. Then we could conduct generation as in Algorithm 3 based\\non the conditional output distribution pO(·|θ, c, t) base on Φ(θ, c, t).\\n17\\n\\n\\nPublished as a conference paper at ICLR 2024\\nD.2\\nEXPLANATIONS ON THE PROPERTIES IN TAB. 2\\nα Polarizability: Tendency of a molecule to acquire an electric dipole moment when subjected to\\nanexternal electric field.\\nεHOMO: Highest occupied molecular orbital energy.\\nεLUMO: Lowest unoccupied molecular orbital energy.\\n∆ε Gap: The energy difference between HOMO and LUMO.\\nµ : Dipole moment.\\nCv : Heat capacity at 298.15 K\\nE\\nDETAILED ALGORITHMS FOR TRAINING AND SAMPLING\\nFor a better understanding of the whole procedure in training and sampling, we involve the detailed\\nalgorithms and implements of functions in Algorithm 1, Algorithm 2 and Algorithm 3.\\nAlgorithm 1 Functions for GeoBFN\\nfunction DISCRETISED_CDF(µ ∈R, σ ∈R+, x ∈R)\\nF(x) ←1\\n2\\nh\\n1 + erf\\n\\u0010\\nx−µ\\nσ\\n√\\n2\\n\\u0011i\\nG(x) ←\\n\\n\\n\\n0\\nif x ≤−1\\n1\\nif x ≥1\\nF(x)\\notherwise\\nReturn G(x)\\nend function\\nfunction OUTPUT_PREDICTION(µx ∈RD×3, µh ∈RD, t ∈[0, 1], γx, γh ∈R+, tmin ∈R+)\\n# tmin set to 0.0001 by default\\nif t < tmin then\\nˆ\\nx(θ, t) ←0\\nˆ\\nµh ←0\\nˆ\\nσh ←1\\nelse\\nInput (µx, µh, t) to network, receive ˆ\\nϵ(θ, t), ˆ\\nµϵ\\nh, ln ˆ\\nσϵ\\nh as output\\nˆ\\nx(θ, t) ←µx\\nγx −\\nq\\n1−γx\\nγx ˆ\\nϵ(θ, t)\\nˆ\\nµh ←ˆ\\nµh\\nγh −\\nq\\n1−γh\\nγh\\nˆ\\nµϵ\\nh\\nˆ\\nσh ←\\nq\\n1−γh\\nγh\\nln ˆ\\nσϵ\\nh\\nend if\\nfor d ∈1, · · · , D, k ∈K do\\np(d)\\nO (k | θ; t) ←DISCRETISED_CDF(ˆ\\nµ(d)\\nh , ˆ\\nσ(d)\\nh , kr) −DISCRETISED_CDF(ˆ\\nµ(d)\\nh , ˆ\\nσ(d)\\nh , kl)\\nend for\\nReturn ˆ\\nx(θ, t), pO(· | θ; t)\\nend function\\n18\\n\\n\\nPublished as a conference paper at ICLR 2024\\nAlgorithm 2 Training with continuous loss\\nRequire: σx, σh ∈R, number of bins K ∈N\\nInput: coordinates x ∈RD×3, normalized charges h ∈[ 1\\nK −1, 1 −1\\nK ]D\\nt ∼U(0, 1)\\nγx ←1 −σ2t\\nx , γh ←1 −σ2t\\nh\\nµx ∼N(γx, γx(1 −γx)I)\\nµh ∼N(γh, γh(1 −γh)I)\\nˆ\\nx(θ, t), pO(· | θ; t) ←OUTPUT_PREDICTION(µx, µh, t, γx, γh)\\nˆ\\nk(θ, t) ←\\n\\u0010P\\nk p(1)\\nO pO(k | θ; t)kc, . . . , P\\nk p(D)\\nO (k | θ; t)kc\\n\\u0011\\nL∞(x) ←−ln σxσ−2t\\nx\\n∥x −ˆ\\nx(θ, t)∥2\\nL∞(h) ←−ln σhσ−2t\\nh\\n\\n\\n\\nh −ˆ\\nk(θ, t)\\n\\n\\n\\n2\\nReturn L∞(x) + L∞(h)\\nAlgorithm 3 Sampling procedure\\n# ⃗\\nkc =\\n\\u0010\\nk(1)\\nc , . . . , k(D)\\nc\\n\\u0011\\n# Function NEAREST_CENTER compares inputs to the center bins ⃗\\nkc,\\n# and return the nearest center for each input value.\\nRequire: σx, σh ∈R+, number of steps N ∈N\\nµx, µh ←0\\nρx, ρh ←1\\nfor i = 1 to N do\\nt ←i−1\\nn\\nγx ←1 −σ2t\\nx , γh ←1 −σ2t\\nh\\nˆ\\nx(θ, t), pO(· | θ; t) ←OUTPUT_PREDICTION(µx, µh, t, γx, γh)\\nαx ←σ−2i/n\\nx\\n\\u0010\\n1 −σ2/n\\nx\\n\\u0011\\nαh ←σ−2i/n\\nh\\n\\u0010\\n1 −σ2/n\\nh\\n\\u0011\\nˆ\\nkc(θ, t) ←NEAREST_CENTER(\\nhP\\nk p(1)\\nO (k | θ; t)kc, . . . , P\\nk p(D)\\nO (k | θ; t)kc\\ni\\n)\\nyh ∼N(ˆ\\nkc(θ, t), α−1\\nh I)\\nyx ∼N(ˆ\\nx(θ, t), α−1\\nx I)\\nµx, µh ←ρxµx+αxyx\\nρx+αx\\n, ρhµh+αhyh\\nρh+αh\\nρx, ρh ←(ρx + αx), (ρh + αh)\\nend for\\nˆ\\nx(θ, 1), pO(· | θ; 1) ←OUTPUT_PREDICTION(µx, µh, 1, 1 −σ2\\nx, 1 −σ2\\nh)\\nˆ\\nkc(θ, 1) ←NEAREST_CENTER(\\nhP\\nk p(1)\\nO (k | θ; 1)kc, . . . , P\\nk p(D)\\nO (k | θ; 1)kc\\ni\\n)\\nReturn ˆ\\nx(θ, 1), ˆ\\nkc(θ, 1)\\n19\",\"difficulty\":\"hard\",\"domain\":\"Multi-Document QA\",\"length\":\"short\",\"question\":\"Which is not among the improvements of the method in paper GeoBFN compared with paper BFN？\",\"sub_domain\":\"Academic\"}","display_format":"text","language":"","answer_status":"published","assets":[],"source_url":"https://huggingface.co/datasets/zai-org/LongBench-v2","history":"initial import","indexing_mode":"noindex","subproblems":[],"grids":[]}