# Bayesian Flow Networks

Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, Faustino Gomez

<EMAIL_ADDRESS>
NNAISENSE )

###### Abstract

This paper introduces _Bayesian Flow Networks_ (BFNs), a new class of
generative model in which the parameters of a set of independent distributions
are modified with Bayesian inference in the light of noisy data samples, then
passed as input to a neural network that outputs a second, interdependent
distribution. Starting from a simple prior and iteratively updating the two
distributions yields a generative procedure similar to the reverse process of
diffusion models; however it is conceptually simpler in that no forward
process is required. Discrete and continuous-time loss functions are derived
for continuous, discretised and discrete data, along with sample generation
procedures. Notably, the network inputs for discrete data lie on the
probability simplex, and are therefore natively differentiable, paving the way
for gradient-based sample guidance and few-step generation in discrete domains
such as language modelling. The loss function directly optimises data
compression and places no restrictions on the network architecture. In our
experiments BFNs achieve competitive log-likelihoods for image modelling on
dynamically binarized MNIST and CIFAR-10, and outperform all known discrete
diffusion models on the text8 character-level language modelling task111Code
and trained models can be found at https://github.com/nnaisense/bayesian-flow-
networks.

##  1 Introduction

Large-scale neural networks have revolutionised generative modelling over the
last few years, with an unprecedented ability to capture complex relationships
among many variables. Building a convincing joint model of all the pixels in a
high resolution image, for example, was impossible before the advent of modern
generative networks.

Key to the expressive power of most of these networks — including
autoregressive models e.g. [46, 9], flow-based models [32], deep VAEs [49] and
diffusion models [41] — is that the joint distribution they encode is broken
down into a series of steps, thereby eluding the “curse of dimensionality”
that would doom any effort to explicitly define all the interactions among so
many variables. In colloquial terms they solve a hard problem by splitting it
into easy pieces.

A general way to view such distributions is as an exchange of messages between
a sender, Alice, who has access to some data, and her friend Bob, who wishes
to receive it in as few bits as possible. At each step Alice sends a message
to Bob that reveals something about the data. Bob attempts to guess what the
message is: the better his guess the fewer bits are needed to transmit it.
After receiving the message, Bob uses the information he has just gained to
improve his guess for the next message. The loss function is the total number
of bits required for all the messages.

In an autoregressive language model, for example, the messages are the word-
pieces the text is divided into. The distribution encoding Bob’s prediction
for the first message is of necessity uninformed: a zero-gram prior based on
the relative frequencies of different word-pieces. The transmission cost is
the negative log-probability under this prior. Bob then uses the first word-
piece to predict the second; on average, the second prediction will be
slightly more informed than the first, and the expected transmission cost will
be slightly lower. The process repeats with the predictions improving at each
step. The sum of the transmission costs is the negative log-probability of the
complete text sequence, which is the loss function minimised by maximum
likelihood training. It is also the minimum number of bits that would be
required for Alice to transmit the pieces to Bob using arithmetic coding [53].
There is therefore a direct correspondence between fitting an autoregressive
model with maximum likelihood and training it for data compression.

Autoregressive networks are currently state-of-the-art for language modelling
[29], and in general perform well on discrete data where a natural ordering
exists. However they have proved less effective in domains such as image
generation, where the data is continuous and no natural order exists among
variables (e.g. there is no reason to generate one pixel before another). They
also have the drawback that generating samples requires as many network
updates as there are variables in the data.

Diffusion models are an alternative framework that has proved particularly
effective for image generation [5, 34]. In this case the transmission
procedure is a little more complex222We are here describing the reverse
process of diffusion models.. Each message Bob receives is a noisy version of
the message before, where the noise is designed so that in expectation the
messages approach the data. The transmission cost at each step is the
Kullback-Leibler divergence between the distribution from which Alice draws
the message and Bob’s prediction of that distribution (which is a
reparameterisation of his prediction of the data, and which is therefore
improved by the information he gained from the previous message). The sum of
the KL divergences is the _evidence lower bound_ minimised by diffusion
training [41]; it is also the expected number of bits needed to transmit the
data using an efficient bits-back coding scheme [51, 11]. Once again there is
an exact equivalence between the loss function used to train the model and the
model’s ability to compress data, as elucidated by previous authors [47].

We posit that the superiority of diffusion over autoregression for image
generation lies in the way diffusion progresses from coarse to fine image
details as the level of noise decreases — a more natural way to construct an
image than one dot at a time. However diffusion has yet to match
autoregression for discrete data, which is unfortunate, as diffusion models
have the advantage of decoupling the number of generation steps from the
number of variables. A fundamental challenge is that when the data is
discrete, the noise in the diffusion process is also discrete, and therefore
discontinuous. To return to the transmission metaphor, if the data is a piece
of text, then Bob begins the process with a totally garbled text, every symbol
of which is either randomly altered or left unchanged by each of Alice’s
messages. A key motivation for this work was our belief that a fully
continuous transmission process — where Alice’s messages smoothly alter Bob’s
beliefs — would be more effective for discrete data. Moreover this should open
the door to gradient-based sample guidance [5] and few-step generation
techniques [37, 52, 43], similar to those that have been developed for
continuous diffusion.

Figure 1: System Overview. The figure represents one step of the modelling
process of a Bayesian Flow Network. The data in this example is a ternary
symbol sequence, of which the first two variables (‘B’ and ‘A’) are shown. At
each step the network emits the parameters of the output distribution based on
the parameters of the previous input distribution. The sender and receiver
distributions (both of which are continuous, even when the data is discrete)
are created by adding random noise to the data and the output distribution
respectively. A sample from the sender distribution is then used to update the
parameters of the input distribution, following the rules of Bayesian
inference. Conceptually, this is the message sent by Alice to Bob, and its
contribution to the loss function is the KL divergence from the receiver to
the sender distribution.

_Bayesian Flow Networks_ (BFNs), the model introduced in this paper, differ
from diffusion models in that the network operates on the parameters of a data
distribution, rather than on a noisy version of the data itself. This ensures
that the generative process is fully continuous and differentiable, even when
the data is discrete. BFNs can be summarised by the following transmission
scheme (Figure 1). Bob has an “input distribution” which is initially a simple
prior: a standard normal for continuous data, a uniform categorical for
discrete data. At each transmission step he feeds the parameters of the input
distribution (e.g. the mean of a normal distribution, the probabilities of a
categorical distribution) into a neural network. The network outputs the
parameters of a second distribution referred to as the “output distribution”.
Alice then creates a “sender distribution” by adding noise to the data
according to a predefined schedule, and Bob creates a “receiver distribution”
by convolving the output distribution with the same noise distribution used by
Alice: intuitively, for every value the data could take on, Bob constructs the
sender distribution Alice would have used if that value was correct, then sums
over all these hypothetical sender distributions, weighted by the probability
of the corresponding value under the output distribution. Alice picks a sample
from the sender distribution and sends it to Bob at a cost equal to the KL
divergence from receiver to sender. Bob then uses the sample to update his
input distribution, following the rules of Bayesian inference. Usefully, the
Bayesian updates are available in closed-form as long as the input
distribution models all the variables in the data independently. Once the
update is complete, Bob again feeds the parameters of the input distribution
to the network which returns the parameters of the output distribution. The
process repeats for $n$ steps, at which point Bob can predict the data
accurately enough that Alice can send it to him without any noise.

Note the key difference between the input and output distributions: the input
distribution receives information about each variable in the data
independently (via the Bayesian updates), and is therefore unable to exploit
contextual information, such as neighbouring pixels in an image or related
words in a text; the output distribution, on the other hand, is produced by a
neural network that jointly processes all the parameters in the input
distribution, giving it access to all available context. Intuitively, the
combination of the input and output distributions represents a division of
labour between Bayesian inference and deep learning that plays to both of
their strengths: the former provides a mathematically optimal and finely
controllable way to collect and summarise information about individual
variables, while the latter excels at integrating information over many
interrelated variables.

The above transmission process defines an $n$-step loss function that can be
generalised to continuous time by sending $n$ to $\infty$. In continuous time
the Bayesian updates become a _Bayesian flow_ of information from the data to
the network. As well as removing the need to predefine the number of steps
during training, the continuous-time loss function is mathematically simpler
and easier to compute than the discrete-time loss. A BFN trained with
continuous-time loss can be run for any number of discrete steps during
inference and sampling, with performance improving as the number of steps
increases.

The rest of the paper is structured as follows. A short summary of related
work is given in Section 2. The basic framework of BFNs, along with a general
derivation of the discrete and continuous time loss functions is provided in
Section 3. Specialisations of the framework to continuous, discretised and
discrete data are provided in Sections 4–6, along with pseudocode for
training, evaluating and sampling from the network. Experimental results on
the CIFAR-10, dynamically binarized MNIST and text8 datasets are provided in
Section 7 and concluding remarks are given in Section 8.

##  2 Related Work

Of existing methods, Bayesian Flow Networks are most closely related to
diffusion models. However the two differ in some crucial aspects. Most
obviously BFNs embody a function from one distribution to another — rather
than from data to a distribution, like diffusion models and most other
probabilistic networks. One advantage of this approach is that, because the
parameters of a categorical distribution are real-valued probabilities, the
inputs to the network are continuous even when the data is discrete. This
contrasts with discrete diffusion, which natively uses discrete samples as
input [41, 14, 1].

Numerous authors have proposed continuous variants of discrete diffusion.
Typically these rely either on mapping to and from a continuous embedding
space [44, 21, 6, 2], or on restricting continuous diffusion to the
probability simplex [33, 24, 23]. While we do not directly compare against the
above methods, we note that continuity is an inherent property of the Bayesian
Flow framework (the network inputs automatically lie on the probability
simplex by virtue of being the parameters of a categorical distribution),
rather than a constraint added to an existing system. As well as reducing the
number of free parameters and design choices (e.g. the continuous embedding
space, the mapping functions), this ensures that BFNs directly optimise the
negative log-likelihood of discrete data, unlike continuous diffusion methods
for discrete data, which typically require either simplified loss functions
[24] or auxiliary loss terms [21] to make learning stable.

For continuous data, BFNs are most closely related to variational diffusion
models [17], with a very similar continuous-time loss function. The main
difference in this case is that the network inputs are considerably less noisy
in BFNs than in variational diffusion and other continuous diffusion models.
This is because the generative process of BFNs begins with the parameters of a
fixed prior, whereas that of diffusion models begins with pure noise. We
hypothesise that the reduction in noise could lead to faster learning on large
datasets where the model underfits; however we have yet to test this
hypothesis experimentally.

Another key difference from diffusion models is that there is no need to
define and invert a forward process for BFNs, which arguably makes it easier
to adapt them to different distributions and data types. We showcase this
flexibility by adapting BFNs to continuous, discretised and discrete data,
with minimal changes to the training procedure. This contrasts with e.g.
discretised diffusion, which requires carefully defined transition matrices
[1].

##  3 Bayesian Flow Networks

This section covers the basic mathematical formalism of Bayesian Flow
Networks, laying out the structure of the various functions and distributions
required by the model, along with the discrete and continuous-time loss
functions used for training. Specific instantiations of the general framework
for continuous, discretised and discrete data are given in Sections 4–6.

###  3.1 Input and Sender Distributions

Given $D$-dimensional data
$\operatorname{\mathbf{x}}=\left(x^{(1)},\dots,x^{(D)}\right)\in\operatorname{\mathcal{X}}^{D}$,
let $\boldsymbol{\theta}=\left(\theta^{(1)},\dots,\theta^{(D)}\right)$ be the
parameters of a factorised _input distribution_
$p_{{}_{I}}(\cdot\mid\boldsymbol{\theta})$, with

$\displaystyle
p_{{}_{I}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta})=\prod_{d=1}^{D}p_{{}_{I}}(x^{(d)}\mid\theta^{(d)}).$
(1)

For example, $\theta^{(d)}$ may consist of the probabilities of a categorical
distribution. Let
$p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha\right)$ be a
similarly factorised _sender distribution_ with
$\operatorname{\mathbf{y}}=\left(y^{(1)},\dots,y^{(D)}\right)\in\operatorname{\mathcal{Y}}^{D}$
and

$\displaystyle
p_{{}_{S}}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}};\alpha\right)=\prod_{d=1}^{D}p_{{}_{S}}\left(y^{(d)}\mid
x^{(d)};\alpha\right),$ (2)

where $\alpha\in\operatorname{\mathbb{R}}^{+}$ is an _accuracy_ parameter
defined such that when $\alpha=0$, the sender samples are entirely
uninformative about $\operatorname{\mathbf{x}}$ and as $\alpha$ increases the
samples become progressively more informative.

###  3.2 Output Distribution $p_{{}_{O}}(\cdot\mid\boldsymbol{\theta},t)$

During the data transmission process, the input parameters
$\boldsymbol{\theta}$ are passed along with the process time $t$ as input to a
neural network $\Psi$. The network then emits an output vector
$\Psi(\boldsymbol{\theta},t)=\left(\Psi^{(1)}(\boldsymbol{\theta},t),\dots,\Psi^{(D)}(\boldsymbol{\theta},t)\right)$
which is used to parameterise an output distribution factorised in the same
way as the input and sender distributions:

$\displaystyle
p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta},t)=\prod_{d=1}^{D}p_{{}_{O}}(x^{(d)}\mid\Psi^{(d)}(\boldsymbol{\theta},t)).$
(3)

As discussed in the introduction, the key difference between the input and
output distributions is that while each $p_{{}_{I}}(x^{(d)}\mid\theta^{(d)})$
depends only on information gathered via $p_{{}_{S}}\left(y^{(d)}\mid
x^{(d)};\alpha\right)$ about $x^{(d)}$, each
$p_{{}_{O}}(x^{(d)}\mid\Psi^{(d)}(\boldsymbol{\theta},t))$ depends (via the
network) on all of $\boldsymbol{\theta}$ and hence all of
$\operatorname{\mathbf{x}}$. The output distribution, unlike the input
distribution, can therefore exploit context information, such as surrounding
pixels in an image or related words in a text.

###  3.3 Receiver Distribution
$p_{{}_{R}}(\cdot\mid\boldsymbol{\theta};t,\alpha)$

Given sender distribution
$p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha\right)$ and output
distribution $p_{{}_{O}}(\cdot\mid\boldsymbol{\theta},t)$ the _receiver
distribution_ over $\operatorname{\mathcal{Y}}^{D}$ is defined as

$\displaystyle
p_{{}_{R}}(\operatorname{\mathbf{y}}\mid\boldsymbol{\theta};t,\alpha)$
$\displaystyle=\operatorname*{\mathbb{E}}_{p_{{}_{O}}(\operatorname{\mathbf{x}}^{\prime}\mid\boldsymbol{\theta};t)}p_{{}_{S}}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}}^{\prime};\alpha\right).$
(4)

Intuitively this can be understood as a receiver who knows the form of the
sender distribution
$p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha\right)$ but does
not know $\operatorname{\mathbf{x}}$, and therefore integrates over all
$\operatorname{\mathbf{x}}^{\prime}\in\operatorname{\mathcal{X}}^{D}$, and
hence all possible sender distributions, weighted by the probability given to
$\operatorname{\mathbf{x}}^{\prime}$ by the output distribution
$p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta},t)$. The receiver
distribution therefore combines two sources of uncertainty: the “known
unknown” of the sender distribution entropy (which is a function of $\alpha$),
and the “unknown unknown” of the output distribution entropy.

###  3.4 Bayesian Updates

Given parameters $\boldsymbol{\theta}$ and sender sample
$\operatorname{\mathbf{y}}$ drawn with accuracy $\alpha$ the _Bayesian update
function_ $h$ is derived by applying the rules of Bayesian inference to
compute the updated parameters $\boldsymbol{\theta}^{\prime}$:

$\displaystyle\boldsymbol{\theta}^{\prime}\leftarrow
h(\boldsymbol{\theta},\operatorname{\mathbf{y}},\alpha).$ (5)

The _Bayesian update distribution_
$p_{{}_{U}}(\cdot\mid\boldsymbol{\theta},\operatorname{\mathbf{x}};\alpha)$ is
then defined by marginalizing out $\operatorname{\mathbf{y}}$:

$\displaystyle
p_{{}_{U}}(\boldsymbol{\theta}^{\prime}\mid\boldsymbol{\theta},\operatorname{\mathbf{x}};\alpha)=\operatorname*{\mathbb{E}}_{p_{{}_{S}}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}};\alpha\right)}\delta\left(\boldsymbol{\theta}^{\prime}-h(\boldsymbol{\theta},\operatorname{\mathbf{y}},\alpha)\right),$
(6)

where $\delta\left(\cdot-\boldsymbol{a}\right)$ is the multivariate Dirac
delta distribution centred on the vector $\boldsymbol{a}$. In Sections 4.4 and
6.7 we will prove that both forms of
$p_{{}_{U}}(\cdot\mid\boldsymbol{\theta},\operatorname{\mathbf{x}};\alpha)$
considered in this paper have the following property: the accuracies are
additive in the sense that if $\alpha=\alpha_{a}+\alpha_{b}$ then

$\displaystyle
p_{{}_{U}}(\boldsymbol{\theta}^{\prime\prime}\mid\boldsymbol{\theta},\operatorname{\mathbf{x}};\alpha)=\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}^{\prime}\mid\boldsymbol{\theta},\operatorname{\mathbf{x}};\alpha_{a})}p_{{}_{U}}(\boldsymbol{\theta}^{\prime\prime}\mid\boldsymbol{\theta}^{\prime},\operatorname{\mathbf{x}};\alpha_{b}).$
(7)

It follows from this property that given prior input parameters
$\boldsymbol{\theta}_{0}$, the probability of observing parameters
$\boldsymbol{\theta}_{n}$ after drawing a sequence of $n$ sender samples
$\operatorname{\mathbf{y}}_{1},\dots,\operatorname{\mathbf{y}}_{n}$ with
accuracies $\alpha_{1},\dots,\alpha_{n}$ is

$\displaystyle\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}_{1}\mid\boldsymbol{\theta}_{0},\operatorname{\mathbf{x}};\alpha_{1})}\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}_{2}\mid\boldsymbol{\theta}_{1},\operatorname{\mathbf{x}};\alpha_{2})}\dots\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}_{n-1}\mid\boldsymbol{\theta}_{n-2},\operatorname{\mathbf{x}};\alpha_{n-1})}p_{{}_{U}}(\boldsymbol{\theta}_{n}\mid\boldsymbol{\theta}_{n-1},\operatorname{\mathbf{x}};\alpha_{n})=p_{{}_{U}}\left(\boldsymbol{\theta}_{n}\mid\boldsymbol{\theta}_{0},\operatorname{\mathbf{x}};\sum_{i=1}^{n}\alpha_{i}\right).$
(8)

###  3.5 Accuracy Schedule $\beta(t)$

By performing an infinite number of transmission steps, the Bayesian update
process can be generalized to continuous time. Let $t\in[0,1]$ be the process
time and let $\alpha(t)>0$ be the _accuracy rate_ at time $t$. Now define the
_accuracy schedule_ $\beta(t)$ as

$\displaystyle\beta(t)=\int_{t^{\prime}=0}^{t}\alpha(t^{\prime})dt^{\prime}.$
(9)

It follows from the above definitions that $\beta(t)$ is a monotonically
increasing function of $t$, that $\beta(0)=0$, and that
$\frac{d\beta(t)}{dt}=\alpha(t)$.

Specific forms of $\beta(t)$ for continuous and discrete data are provided in
Sections 4.5 and 6.8. Both are derived using simple heuristics, with a deeper
investigation left for future work.

###  3.6 Bayesian Flow Distribution
$p_{{}_{F}}(\cdot\mid\operatorname{\mathbf{x}};t)$

Given prior parameters $\boldsymbol{\theta}_{0}$, Bayesian update distribution
$p_{{}_{U}}(\cdot\mid\boldsymbol{\theta},\operatorname{\mathbf{x}};\alpha)$
and accuracy schedule $\beta(t)$, the _Bayesian flow distribution_
$p_{{}_{F}}(\cdot\mid\operatorname{\mathbf{x}};t)$ is the marginal
distribution over input parameters at time $t$, defined by

$\displaystyle
p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t)=p_{{}_{U}}(\boldsymbol{\theta}\mid\boldsymbol{\theta}_{0},\operatorname{\mathbf{x}};\beta(t)).$
(10)

###  3.7 Loss Function $L(\operatorname{\mathbf{x}})$

Given prior parameters $\boldsymbol{\theta}_{0}$ and accuracy schedule
$\beta(t)$, consider a sequence of $n$ sender samples
$\operatorname{\mathbf{y}}_{1},\dots,\operatorname{\mathbf{y}}_{n}$ sampled at
times $t_{1},\dots,t_{n}$ where $t_{i}=i/n$. The sender distribution at step
$i$ is $p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha_{i}\right)$
where

$\displaystyle\alpha_{i}$ $\displaystyle=\beta(t_{i})-\beta(t_{i-1}),$ (11)

the receiver distribution at step $i$ is
$p_{{}_{R}}(\cdot\mid\boldsymbol{\theta}_{i-1};t_{i-1},\alpha_{i})$, and the
input parameter sequence
$\boldsymbol{\theta}_{1},\dots,\boldsymbol{\theta}_{n}$ is recursively
calculated from

$\displaystyle\boldsymbol{\theta}_{i}=h(\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{y}},\alpha_{i}).$
(12)

Define the $n$-step discrete-time loss $L^{n}(\operatorname{\mathbf{x}})$ as
the expected number of nats required to first transmit
$\operatorname{\mathbf{y}}_{1},\dots,\operatorname{\mathbf{y}}_{n}$, and the
reconstruction loss $L^{r}(\operatorname{\mathbf{x}})$ as the expected number
of nats required to then transmit $\operatorname{\mathbf{x}}$. Since — using a
bits-back coding scheme [11, 7] — it requires $D_{KL}\left(p_{{}_{S}}\parallel
p_{{}_{R}}\right)$ nats to transmit a sample from $p_{{}_{S}}$ to a receiver
with $p_{{}_{R}}$,

$\displaystyle
L^{n}(\operatorname{\mathbf{x}})=\operatorname*{\mathbb{E}}_{p(\boldsymbol{\theta}_{1},\dots,\boldsymbol{\theta}_{n-1})}\sum_{i=1}^{n}D_{KL}\left(p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha_{i}\right)\parallel
p_{{}_{R}}(\cdot\mid\boldsymbol{\theta}_{i-1};t_{i-1},\alpha_{i})\right),$
(13)

where

$\displaystyle
p(\boldsymbol{\theta}_{1},\dots,\boldsymbol{\theta}_{n})=\prod_{i=1}^{n}p_{{}_{U}}(\boldsymbol{\theta}_{i}\mid\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{x}};\alpha_{i}),$
(14)

and since the number of nats needed to transmit $x$ using an arithmetic coding
scheme [53] based on $p(x)$ is $-\ln p(x)$, and the marginal probability of
$\boldsymbol{\theta}_{n}$ is given by
$p_{{}_{F}}(\cdot\mid\operatorname{\mathbf{x}},1)$,

$\displaystyle
L^{r}(\operatorname{\mathbf{x}})=-\operatorname*{\mathbb{E}}_{p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}},1)}\ln
p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta};1).$ (15)

Note that $L^{r}(\operatorname{\mathbf{x}})$ is not directly optimised in this
paper; however it is indirectly trained by optimising
$L^{n}(\operatorname{\mathbf{x}})$ since both are minimised by matching the
output distribution to the data. Furthermore, as long as $\beta(1)$ is high
enough, the input distribution at $t=1$ will be very close to
$\operatorname{\mathbf{x}}$, making it trivial for the network to fit
$p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta};1)$.

The loss function $L(\operatorname{\mathbf{x}})$ is defined as the total
number of nats required to transmit the data, which is the sum of the n-step
and reconstruction losses:

$\displaystyle
L(\operatorname{\mathbf{x}})=L^{n}(\operatorname{\mathbf{x}})+L^{r}(\operatorname{\mathbf{x}})$
(16)

Alternatively $L(\operatorname{\mathbf{x}})$ can be derived as the loss
function of a variational autoencoder (VAE; [18]). Consider the sequence
$\operatorname{\mathbf{y}}_{1},\dots,\operatorname{\mathbf{y}}_{n}$ as a
latent code with posterior probability given by

$\displaystyle
q(\operatorname{\mathbf{y}}_{1},\dots,\operatorname{\mathbf{y}}_{n})=\prod_{i=1}^{n}p_{{}_{S}}\left(\operatorname{\mathbf{y}}_{i}\mid\operatorname{\mathbf{x}};\alpha_{i}\right),$
(17)

and autoregressive prior probability given by

$\displaystyle
p(\operatorname{\mathbf{y}}_{1},\dots,\operatorname{\mathbf{y}}_{n})=\prod_{i=1}^{n}p_{{}_{R}}(\operatorname{\mathbf{y}}_{i}\mid\boldsymbol{\theta}_{i-1};t_{i-1},\alpha_{i}).$
(18)

Then, noting that the decoder probability
$p(\operatorname{\mathbf{x}}\mid\operatorname{\mathbf{y}}_{1},\dots,\operatorname{\mathbf{y}}_{n})=p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta}_{n};1)$,
the complete transmission process defines a VAE with loss function given by
the negative variational lower bound (VLB)

$\displaystyle
L(\operatorname{\mathbf{x}})=-\text{VLB}(\operatorname{\mathbf{x}})$
$\displaystyle=D_{KL}\left(q\parallel
p\right)-\operatorname*{\mathbb{E}}_{\operatorname{\mathbf{y}}_{1},\dots,\operatorname{\mathbf{y}}_{n}\sim
q}\ln
p(\operatorname{\mathbf{x}}\mid\operatorname{\mathbf{y}}_{1},\dots,\operatorname{\mathbf{y}}_{n})$
(19)
$\displaystyle=L^{n}(\operatorname{\mathbf{x}})+L^{r}(\operatorname{\mathbf{x}}).$
(20)

###  3.8 Discrete-Time Loss $L^{n}(\operatorname{\mathbf{x}})$

Eq. 13 can be rewritten as

$\displaystyle
L^{n}(\operatorname{\mathbf{x}})=n\operatorname*{\mathbb{E}}_{i\sim
U\\{1,n\\}}\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}_{1}\mid\boldsymbol{\theta}_{0},\operatorname{\mathbf{x}};\alpha_{1})}\dots\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}\mid\boldsymbol{\theta}_{i-2},\operatorname{\mathbf{x}};\alpha_{i-1})}D_{KL}\left(p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha_{i}\right)\parallel
p_{{}_{R}}(\cdot\mid\boldsymbol{\theta};t_{i-1},\alpha_{i})\right),$ (21)

where $U\\{1,n\\}$ is the uniform distribution over the integers from 1 to
$n$. Furthermore, it follows from Eqs. 8 and 10 that

$\displaystyle\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}_{1}\mid\boldsymbol{\theta}_{0},\operatorname{\mathbf{x}};\alpha_{1})}\dots\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}\mid\boldsymbol{\theta}_{i-2},\operatorname{\mathbf{x}};\alpha_{i-1})}$
$\displaystyle=\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}\mid\boldsymbol{\theta}_{0},\operatorname{\mathbf{x}};\beta(t_{i-1}))}$
(22)
$\displaystyle=\operatorname*{\mathbb{E}}_{p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t_{i-1})},$
(23)

and hence

$\displaystyle
L^{n}(\operatorname{\mathbf{x}})=n\operatorname*{\mathbb{E}}_{i\sim
U\\{1,n\\},p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t_{i-1})}D_{KL}\left(p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha_{i}\right)\parallel
p_{{}_{R}}(\cdot\mid\boldsymbol{\theta};t_{i-1},\alpha_{i})\right),$ (24)

which allows us approximate $L^{n}(\operatorname{\mathbf{x}})$ via Monte-Carlo
sampling without computing the $n$-step sum.

###  3.9 Continuous-Time Loss $L^{\infty}(\operatorname{\mathbf{x}})$

Eq. 24 can be used to train the network directly. However this presupposes
that $n$ is fixed during training. Furthermore, for discrete and discretised
data the KL terms do not have analytic solutions, leading to noisy gradient
estimates.

Inspired by Variational Diffusion Models [17] we derive a continuous-time loss
function $L^{\infty}(\operatorname{\mathbf{x}})$ by taking the limit of
$L^{n}(\operatorname{\mathbf{x}})$ as $n\rightarrow\infty$. This turns out to
be mathematically simpler than the discrete-time loss, as well as removing
both the noisy gradients for the discrete and discretised KL terms and the
need to fix $n$ during training.

Let

$\displaystyle\epsilon$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{1}{n},$ (25)
$\displaystyle\alpha(t,\epsilon)$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\beta(t)-\beta(t-\epsilon),$
(26) $\displaystyle L^{\infty}(\operatorname{\mathbf{x}})$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\lim_{n\rightarrow\infty}L^{n}(\operatorname{\mathbf{x}}).$
(27)

Then, from the definition of $L^{n}(\operatorname{\mathbf{x}})$ in Eq. 24,

$\displaystyle L^{\infty}(\operatorname{\mathbf{x}})=\lim_{\epsilon\rightarrow
0}\frac{1}{\epsilon}\operatorname*{\mathbb{E}}_{t\sim
U(\epsilon,1),p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}},t-\epsilon)}D_{KL}\left(p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha(t,\epsilon)\right)\parallel
p_{{}_{R}}(\cdot\mid\boldsymbol{\theta};t-\epsilon,\alpha(t,\epsilon))\right),$
(28)

where $U(a,b)$ is the continuous uniform distribution over the interval
$[a,b]$. As we will see, for all the sender, receiver distribution pairs in
this paper,

$\displaystyle
D_{KL}\left(p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha\right)\parallel
p_{{}_{R}}(\cdot\mid\boldsymbol{\theta};\alpha,t)\right)=\sum_{d=1}^{D}D_{KL}\left(\mathcal{N}\left(g(x^{(d)}),C\alpha^{-1}\right)\parallel
P^{(d)}(\boldsymbol{\theta},t)\ast\mathcal{N}\left(0,C\alpha^{-1}\right)\right),$
(29)

where $g:\operatorname{\mathcal{X}}\rightarrow\operatorname{\mathcal{Y}}$ is a
function from data space to sender space, $P^{(d)}(\boldsymbol{\theta},t)$ is
a distribution over $\operatorname{\mathcal{Y}}$ with finite expectation and
variance, $\ast$ denotes the convolution of two probability distributions and
$C$ is a scalar constant.

The following proposition is now required:

######  Proposition 3.1.

For a continuous univariate probability distribution $P$ with finite
expectation $E[P]$ and variance $Var[P]$, the convolution
$P\ast\mathcal{N}\left(0,\sigma^{2}\right)\rightarrow\mathcal{N}\left(E[P],\sigma^{2}\right)$
as $\sigma^{2}\rightarrow\infty$.

###### Proof.

Let $\epsilon^{2}$ be some variance in the interval
$\left(0,\frac{\pi}{8}\right)$ and consider the sequence of random variables
$X_{0},X_{1},\dots,X_{n}$ where $X_{0}\sim P$ and
$X_{j}\sim\mathcal{N}\left(0,\epsilon^{2}\right)$ for $j>0$. Define

$\displaystyle Y_{j}$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\begin{cases}X_{0}-E[P]&\text{if
}j=0,\\\ X_{j}&\text{ otherwise.}\end{cases}$ (30) $\displaystyle R_{n}$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\sum_{j=0}^{n}Y_{j},$
(31) $\displaystyle S^{2}_{n}$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\sum_{j=1}^{n}Var[Y_{j}]=n\epsilon^{2},$
(32) $\displaystyle T^{2}_{n}$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}Var[P]+S^{2}_{n}.$ (33)

It follows from the definition of convolution that $\sum_{j=0}^{n}X_{j}\sim
P\ast\mathcal{N}\left(0,n\epsilon^{2}\right)$. Since
$n\epsilon^{2}\rightarrow\infty$ as $n\rightarrow\infty$, and
$\sum_{j=0}^{n}X_{j}=R_{n}+E[P]$, the result is proved if it can be shown that
as $n\rightarrow\infty$,
$R_{n}\rightarrow\mathcal{N}\left(0,n\epsilon^{2}\right)$ or equivalently
$R_{n}/(\epsilon\sqrt{n})\rightarrow\mathcal{N}\left(0,1\right)$.

The Lyapunov central limit theorem [8] states that if there exists $\lambda>0$
such that
$\lim_{n\rightarrow\infty}\frac{1}{T_{n}^{2+\lambda}}\sum_{j=0}^{n}E\left(|Y_{j}|^{2+\lambda}\right)=0$
then $R_{n}/T_{n}\rightarrow\mathcal{N}\left(0,1\right)$. First note that
$T_{n}^{2}\rightarrow S_{n}^{2}=n\epsilon^{2}$ as $n\rightarrow\infty$. Hence
if $R_{n}/T_{n}\rightarrow\mathcal{N}\left(0,1\right)$ then
$R_{n}/(\epsilon\sqrt{n})\rightarrow\mathcal{N}\left(0,1\right)$. Now set
$\lambda=1$ and observe that for
$Y_{j}\sim\mathcal{N}\left(0,\epsilon^{2}\right)$,
$\operatorname*{\mathbb{E}}\left(|Y_{j}|^{3}\right)$ is the third moment of
the half-normal distribution, which is $\epsilon^{3}\sqrt{\frac{8}{\pi}}$. Our
choice of $\epsilon^{2}$ therefore ensures that
$E\left(|Y_{j}|^{3}\right)<\epsilon^{2}$ for $j>0$. Also note that
$T_{n}^{3}>S_{n}^{3}$ and, since $E[P]$ and $Var[P]$ are finite,
$E\left(|Y_{0}|^{3}\right)<C$ for some constant $C$. Hence

$\displaystyle\frac{1}{T_{n}^{3}}\sum_{j=0}^{n}E\left(|Y_{j}|^{3}\right)$
$\displaystyle<\frac{1}{S_{n}^{3}}\left(C+n\epsilon^{2}\right)=\frac{C}{\epsilon^{3}n^{3/2}}+\frac{1}{\epsilon\sqrt{n}}\xrightarrow[]{n\rightarrow\infty}0.$
(34)

∎

It follows from the continuity of $\beta(t)$ and Eq. 26 that
$\alpha(t,\epsilon)^{-1}\rightarrow\infty$ as $\epsilon\rightarrow 0$.
Therefore, Proposition 3.1 can be applied to Eq. 29 to yield

$\displaystyle\lim_{\epsilon\rightarrow
0}D_{KL}\left(p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}},\alpha_{t}\right)\parallel
p_{{}_{R}}(\cdot\mid\boldsymbol{\theta},\alpha_{t},t)\right)$
$\displaystyle=\sum_{d=1}^{D}D_{KL}\left(\mathcal{N}\left(g(x^{(d)}),\frac{C}{\alpha(t,\epsilon)}\right)\parallel\mathcal{N}\left(E[P^{(d)}(\boldsymbol{\theta},t)],\frac{C}{\alpha(t,\epsilon)}\right)\right)$
(35)
$\displaystyle=\frac{\alpha(t,\epsilon)}{2C}\left\|g(\operatorname{\mathbf{x}})-E[P(\boldsymbol{\theta},t)]\right\|^{2},$
(36)

where

$\displaystyle
g(\operatorname{\mathbf{x}})=\left(g(x^{(1)}),\dots,g(x^{(D)})\right),$ (37)
$\displaystyle
E[P(\boldsymbol{\theta},t)]=\left(E[P^{(1)}(\boldsymbol{\theta},t)],\dots,E[P^{(D)}(\boldsymbol{\theta},t)]\right).$
(38)

Therefore,

$\displaystyle
L^{\infty}(\operatorname{\mathbf{x}})=\operatorname*{\mathbb{E}}_{t\sim
U(0,1),p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}},t)}\lim_{\epsilon\rightarrow
0}\frac{\alpha(t,\epsilon)}{\epsilon}\frac{\left\|g(\operatorname{\mathbf{x}})-E[P(\boldsymbol{\theta},t)]\right\|^{2}}{2C}.$
(39)

Substituting from Eq. 26,

$\displaystyle\lim_{\epsilon\rightarrow
0}\frac{\alpha(t,\epsilon)}{\epsilon}=\lim_{\epsilon\rightarrow
0}\frac{\beta(t)-\beta(t-\epsilon)}{\epsilon}=\frac{d\beta(t)}{dt}=\alpha(t),$
(40)

and hence

$\displaystyle L^{\infty}(\operatorname{\mathbf{x}})$
$\displaystyle=\operatorname*{\mathbb{E}}_{t\sim
U(0,1),p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}},t)}\alpha(t)\frac{\left\|g(\operatorname{\mathbf{x}})-E[P(\boldsymbol{\theta},t)]\right\|^{2}}{2C}.$
(41)

###  3.10 Sample Generation

Given prior parameters $\boldsymbol{\theta}_{0}$, accuracies
$\alpha_{1},\dots,\alpha_{n}$ and corresponding times $t_{i}=i/n$, the n-step
sampling procedure recursively generates
$\boldsymbol{\theta}_{1},\dots,\boldsymbol{\theta}_{n}$ by sampling
$\operatorname{\mathbf{x}}^{\prime}$ from
$p_{{}_{O}}(\cdot\mid\boldsymbol{\theta}_{i-1},t_{i-1})$,
$\operatorname{\mathbf{y}}$ from
$p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}}^{\prime},\alpha_{i}\right)$
(meaning that $\operatorname{\mathbf{y}}\sim
p_{{}_{R}}(\cdot\mid\boldsymbol{\theta}_{i-1};t_{i-1},\alpha_{i})$ — see Eq.
4), then setting
$\boldsymbol{\theta}_{i}=h(\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{y}})$.
Given $\boldsymbol{\theta}_{n}$ the network is run one more time and the final
sample is drawn from $p_{{}_{O}}(\cdot\mid\boldsymbol{\theta}_{n},1)$.

##  4 Continuous Data

For continuous data $\operatorname{\mathcal{X}}=\operatorname{\mathbb{R}}$ and
hence $\operatorname{\mathbf{x}}\in\operatorname{\mathbb{R}}^{D}$. In our
experiments, $\operatorname{\mathbf{x}}$ is normalised to lie in $[-1,1]^{D}$
to ensure that the network inputs remain in a reasonable range; however this
is not essential for the mathematical framework.

###  4.1 Input Distribution $p_{{}_{I}}(\cdot\mid\boldsymbol{\theta})$

The input distribution for continuous data is a diagonal normal:

$\displaystyle\boldsymbol{\theta}$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\\{\operatorname{\boldsymbol{\mu}},\rho\\}$
(42) $\displaystyle
p_{{}_{I}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta})$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\mathcal{N}\left(\operatorname{\mathbf{x}}\mid\operatorname{\boldsymbol{\mu}},\rho^{-1}\boldsymbol{I}\right),$
(43)

where $\boldsymbol{I}$ is the $D\times D$ identity matrix. We define the prior
parameters as

$\displaystyle\boldsymbol{\theta}_{0}\stackrel{{\scriptstyle\text{def}}}{{=}}\\{\boldsymbol{0},1\\},$
(44)

where $\boldsymbol{0}$ is the length $D$ vectors of zeros. Hence the input
prior is a standard multivariate normal:

$p_{{}_{I}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta}_{0})=\mathcal{N}\left(\operatorname{\mathbf{x}}\mid\boldsymbol{0},\boldsymbol{I}\right).$
(45)

The usual Bayesian approach would be to fit the prior mean and variance to the
training data. However we found that a standard prior worked better in
practice, as well as simplifying the equations. It is important to remember
that the distributions
$p_{{}_{I}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta}_{0})$ are never
used directly to make predictions, but rather to inform the network’s
predictions. All that matters is that the parameters fed into the network
accurately and accessibly encode the information received so far about
$\operatorname{\mathbf{x}}$. The network can easily learn the empirical prior
of the training set and use that to correct its predictions.

###  4.2 Bayesian Update Function
$h(\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{y}},\alpha)$

Given a univariate Gaussian prior
$\mathcal{N}\left(\mu_{a},\rho_{a}^{-1}\right)$ over some unknown data $x$ it
can be shown [27] that the Bayesian posterior after observing a noisy sample
$y$ from a normal distribution $\mathcal{N}\left(x,\alpha^{-1}\right)$ with
known precision $\alpha$ is $\mathcal{N}\left(\mu_{b},\rho_{b}^{-1}\right)$,
where

$\displaystyle\rho_{b}$ $\displaystyle=\rho_{a}+\alpha,$ (46)
$\displaystyle\mu_{b}$
$\displaystyle=\frac{\mu_{a}\rho_{a}+y\alpha}{\rho_{b}}.$ (47)

Since both $p_{{}_{I}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta})$ and
$p_{{}_{S}}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}};\alpha\right)$
distributions are normal with diagonal covariance, Eqs. 46 and 47 can be
applied to obtain the following Bayesian update function for parameters
$\boldsymbol{\theta}_{i-1}=\\{\boldsymbol{\mu}_{i-1},\rho_{i-1}\\}$ and sender
sample $\operatorname{\mathbf{y}}$ drawn from
$p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha\boldsymbol{I}\right)=\mathcal{N}\left(\operatorname{\mathbf{x}},\alpha^{-1}\boldsymbol{I}\right)$:

$\displaystyle
h(\\{\boldsymbol{\mu}_{i-1},\rho_{i-1}\\},\operatorname{\mathbf{y}},\alpha)=\\{\boldsymbol{\mu}_{i},\rho_{i}\\},$
(48)

with

$\displaystyle\rho_{i}$ $\displaystyle=\rho_{i-1}+\alpha,$ (49)
$\displaystyle\boldsymbol{\mu}_{i}$
$\displaystyle=\frac{\boldsymbol{\mu}_{i-1}\rho_{i-1}+\operatorname{\mathbf{y}}\alpha}{\rho_{i}}.$
(50)

Figure 2: Bayesian updates for continuous data. For univariate data $x=0.7$,
the initial input distribution parameters
$\theta_{0}=\\{\mu_{0}=0,\rho_{0}=1\\}$ are updated to
$\theta_{1}=\\{\mu_{1},\rho_{1}\\}$, $\theta_{2}=\\{\mu_{2},\rho_{2}\\}$,
$\theta_{3}=\\{\mu_{3},\rho_{3}\\}$ by iterating Eqs. 49 and 50 with sender
samples $y_{1}$, $y_{2}$, $y_{3}$ drawn with accuracies $2$, $4$, $6$
respectively. Note how the input mean ($\mu_{1}$, $\mu_{2}$, $\mu_{3}$)
stochastically approaches the data, while the input precision smoothly
increases.

###  4.3 Bayesian Update Distribution
$p_{{}_{U}}(\cdot\mid\boldsymbol{\theta},\operatorname{\mathbf{x}};\alpha)$

Eq. 50 computes $\boldsymbol{\mu}_{i}$ given a single sample
$\operatorname{\mathbf{y}}$ from the sender distribution. To marginalise over
$\operatorname{\mathbf{y}}\sim\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}},\alpha^{-1}\boldsymbol{I}\right)$
as defined in Eq. 6, the following standard identity for normal distributions
can be applied:

$\displaystyle X\sim\mathcal{N}\left(\mu_{X},\sigma_{X}^{2}\right)\implies
aX+b\sim\mathcal{N}\left(a\mu_{X}+b,a^{2}\sigma_{X}^{2}\right)\ \forall
a,b\in\operatorname{\mathbb{R}}.$ (51)

Substituting $X=\operatorname{\mathbf{y}}$,
$\mu_{X}=\operatorname{\mathbf{x}}$,
$\sigma^{2}_{X}=\alpha^{-1}\boldsymbol{I}$, $a=\frac{\alpha}{\rho_{i}}$ and
$b=\frac{\boldsymbol{\mu}_{i-1}\rho_{i-1}}{\rho_{i}}$, Eq. 50 gives:

$\displaystyle\boldsymbol{\mu}_{i}\sim\mathcal{N}\left(\frac{\alpha\operatorname{\mathbf{x}}+\boldsymbol{\mu}_{i-1}\rho_{i-1}}{\rho_{i}},\frac{\alpha}{\rho_{i}^{2}}\boldsymbol{I}\right),$
(52)

and therefore (since $\boldsymbol{\mu}_{i}$ is the only random part of
$\boldsymbol{\theta}_{i}$)

$\displaystyle
p_{{}_{U}}(\boldsymbol{\theta}_{i}\mid\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{x}};\alpha)=\mathcal{N}\left(\boldsymbol{\mu}_{i}\mid\frac{\alpha\operatorname{\mathbf{x}}+\boldsymbol{\mu}_{i-1}\rho_{i-1}}{\rho_{i}},\frac{\alpha}{\rho_{i}^{2}}\boldsymbol{I}\right).$
(53)

Figure 3: Bayesian update distribution for continuous data. For $x=0.7$, the
plot shows the distribution $p(\mu\mid\theta_{0},x;\alpha)$ over input mean
$\mu$ from Eq. 52 given initial parameters $\mu_{0}=0,\rho_{0}=1$ and 11
$\alpha$ values spaced log-linearly between $e^{-5}$ and $e^{5}$. Note how the
distribution is tightly concentrated around $\mu_{0}$ for very low alpha, then
smoothly progresses to a tight concentration around $x$ for high alpha.

###  4.4 Additive Accuracies

We can check that the sender accuracies are additive in the sense required by
Eq. 7 by first observing that if
$\boldsymbol{\theta}_{i-1}=\\{\boldsymbol{\mu}_{i-1},\rho_{i-1}\\}$ is drawn
from
$p(\cdot\mid\boldsymbol{\theta}_{i-2},\operatorname{\mathbf{x}};\alpha_{a})$
then

$\displaystyle\boldsymbol{\mu}_{i-1}\sim\mathcal{N}\left(\frac{\alpha_{a}\operatorname{\mathbf{x}}+\boldsymbol{\mu}_{i-2}\rho_{i-2}}{\rho_{i-1}},\frac{\alpha_{a}}{\rho_{i-1}^{2}}\boldsymbol{I}\right).$
(54)

Define

$\displaystyle\boldsymbol{\mu}^{\prime}_{i}\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{\alpha_{b}\operatorname{\mathbf{x}}+\boldsymbol{\mu}_{i-1}\rho_{i-1}}{\rho_{i}}=\frac{\rho_{i-1}}{\rho_{i}}\boldsymbol{\mu}_{i-1}+\frac{\alpha_{b}\operatorname{\mathbf{x}}}{\rho_{i}},$
(55)

and apply Identity 51 with $a=\smash{\frac{\rho_{i-1}}{\rho_{i}}}$ and
$b=\smash{\frac{\alpha_{b}\operatorname{\mathbf{x}}}{\rho_{i}}}$ to see that

$\displaystyle\boldsymbol{\mu}^{\prime}_{i}$
$\displaystyle\sim\mathcal{N}\left(\frac{\rho_{i-1}}{\rho_{i}}\frac{\alpha_{a}\operatorname{\mathbf{x}}+\boldsymbol{\mu}_{i-2}\rho_{i-2}}{\rho_{i-1}}+\frac{\alpha_{b}\operatorname{\mathbf{x}}}{\rho_{i}},\frac{\rho_{i-1}^{2}}{\rho_{i}^{2}}\frac{\alpha_{a}}{\rho_{i-1}^{2}}\boldsymbol{I}\right)$
(56)
$\displaystyle=\mathcal{N}\left(\frac{(\alpha_{a}+\alpha_{b})\operatorname{\mathbf{x}}+\boldsymbol{\mu}_{i-2}\rho_{i-2}}{\rho_{i}},\frac{\alpha_{a}}{\rho_{i}^{2}}\boldsymbol{I}\right).$
(57)

Now observe that if
$\boldsymbol{\theta}_{i}=\\{\boldsymbol{\mu}_{i},\rho_{i}\\}$ is drawn from
$p(\cdot\mid\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{x}};\alpha_{b})$
then

$\displaystyle\boldsymbol{\mu}_{i}$
$\displaystyle\sim\mathcal{N}\left(\frac{\alpha_{b}\operatorname{\mathbf{x}}+\boldsymbol{\mu}_{i-1}\rho_{i-1}}{\rho_{i}},\frac{\alpha_{b}}{\rho_{i}^{2}}\boldsymbol{I}\right),$
(58)

and hence

$\displaystyle\boldsymbol{\mu}_{i}$
$\displaystyle\sim\boldsymbol{\mu}^{\prime}_{i}+\boldsymbol{\epsilon},$ (59)

where

$\displaystyle\boldsymbol{\epsilon}\sim\mathcal{N}\left(\boldsymbol{0},\frac{\alpha_{b}}{\rho_{i}^{2}}\boldsymbol{I}\right).$
(60)

Another standard identity for Gaussian variables can now be applied:

$\displaystyle
X\sim\mathcal{N}\left(\mu_{X},\sigma^{2}_{X}\right),Y\sim\mathcal{N}\left(\mu_{Y},\sigma^{2}_{Y}\right)\implies
X+Y\sim\mathcal{N}\left(\mu_{X}+\mu_{Y},\sigma^{2}_{X}+\sigma^{2}_{Y}\right),$
(61)

to see that

$\displaystyle\boldsymbol{\mu}_{i}$
$\displaystyle\sim\mathcal{N}\left(\frac{(\alpha_{a}+\alpha_{b})\operatorname{\mathbf{x}}+\boldsymbol{\mu}_{i-2}\rho_{i-2}}{\rho_{i}},\frac{\alpha_{a}+\alpha_{b}}{\rho_{i}^{2}}\boldsymbol{I}\right),$
(62)

and hence

$\displaystyle\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}_{i-1}\mid\boldsymbol{\theta}_{i-2},\operatorname{\mathbf{x}};\alpha_{a})}p_{{}_{U}}(\boldsymbol{\theta}_{i}\mid\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{x}};\alpha_{b})=p_{{}_{U}}(\boldsymbol{\theta}_{i}\mid\boldsymbol{\theta}_{i-2},\operatorname{\mathbf{x}};\alpha_{a}+\alpha_{b}),$
(63)

as required.

###  4.5 Accuracy Schedule $\beta(t)$

We derive $\beta(t)$ for continuous data by requiring that the expected
entropy of the input distribution linearly decreases with $t$. Intuitively,
this means that information flows into the input distribution at a constant
rate. Define

$\displaystyle H(t)$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\operatorname*{\mathbb{E}}_{p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t)}H(p_{{}_{I}}(\cdot\mid\boldsymbol{\theta}))$
(64) $\displaystyle=\frac{D}{2}\ln\left(\frac{2\pi e}{1+\beta(t)}\right).$
(65)

Then if $H(t)$ linearly decreases with $t$,

$\displaystyle H(t)$ $\displaystyle=(1-t)H(0)+tH(1)$ (66)
$\displaystyle\implies\ln\left(\frac{2\pi}{1+\beta(t)}\right)$
$\displaystyle=(1-t)\ln(2\pi)+t\ln\left(\frac{2\pi}{1+\beta(1)}\right)$ (67)
$\displaystyle\implies-\ln(1+\beta(t))$ $\displaystyle=-t\ln(1+\beta(1))$ (68)
$\displaystyle\implies(1+\beta(t))^{-1}$ $\displaystyle=(1+\beta(1))^{-t}.$
(69)

Define $\sigma_{1}$ to be the standard deviation of the input distribution at
$t=1$. We will choose $\sigma_{1}$ empirically to minimise the loss; in
general it should be small enough to ensure that the reconstruction loss is
low, but not so small as to create unnecessary transmission costs. Recalling
that the precision $\rho$ at time $t$ is $1+\beta(t)$, we see that

$\displaystyle\sigma_{1}^{2}=(1+\beta(1))^{-1}.$ (70)

Therefore

$\displaystyle(1+\beta(t))^{-1}$ $\displaystyle=\sigma_{1}^{2t}$ (71)
$\displaystyle\implies\beta(t)$ $\displaystyle=\sigma_{1}^{-2t}-1$ (72)
$\displaystyle\implies\alpha(t)$
$\displaystyle=\frac{d\left(\sigma_{1}^{-2t}-1\right)}{dt}$ (73)
$\displaystyle=-\frac{2\ln\sigma_{1}}{\sigma_{1}^{2t}}.$ (74)

###  4.6 Bayesian Flow Distribution
$p_{{}_{F}}(\cdot\mid\operatorname{\mathbf{x}};t)$

Recall from Eq. 10 that

$\displaystyle p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t)$
$\displaystyle=p_{{}_{U}}(\boldsymbol{\theta}\mid\boldsymbol{\theta}_{0},\operatorname{\mathbf{x}},\beta(t)).$
(75)

Therefore, setting
$\boldsymbol{\theta}_{i-1}=\boldsymbol{\theta}_{0}=\\{\boldsymbol{0},1\\}$ and
$\alpha=\beta(t)$ in Eq. 53, and recalling that $\rho=1+\beta(t)$,

$\displaystyle p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t)$
$\displaystyle=\mathcal{N}\left(\operatorname{\boldsymbol{\mu}}\mid\frac{\beta(t)}{1+\beta(t)}\operatorname{\mathbf{x}},\frac{\beta(t)}{(1+\beta(t))^{2}}\boldsymbol{I}\right)$
(76)
$\displaystyle=\mathcal{N}\left(\operatorname{\boldsymbol{\mu}}\mid\gamma(t)\operatorname{\mathbf{x}},\gamma(t)(1-\gamma(t))\boldsymbol{I}\right),$
(77)

where

$\displaystyle\gamma(t)$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{\beta(t)}{1+\beta(t)}$
(78) $\displaystyle=\frac{\sigma_{1}^{-2t}-1}{\sigma_{1}^{-2t}}$ (79)
$\displaystyle=1-\sigma_{1}^{2t}.$ (80)

Figure 4: Bayesian flow for continuous data. For $x=0.8$, $\sigma_{1}=0.02$
and $\gamma(t)$ defined as in Eqn. 80, the plot shows stochastic parameter
trajectories for the input distribution mean $\mu$ (white lines) superimposed
on a log-scale heatmap of the Bayesian flow distribution $p(\theta\mid x;t)$.
Note how the trajectories all begin at $\mu_{0}=0$ then fan out before
converging on $x$. Figure 5: Input variance for Bayesian Flow Networks and
diffusion models. For $\sigma_{1}=0.001$ and $\gamma(t)$ defined as in Eqn.
80, the blue line shows the variance $\gamma(t)(1-\gamma(t))$ of the
distribution over the input mean $\mu$ as a function of $t$ (see Eq. 77). Note
that the variance is 0 at $t=0$ (since the input prior $\mu_{0}$ is
deterministic) and becomes small again as $t$ approaches 1 and $\mu$ becomes
increasingly concentrated around the data. The green and red lines show the
equivalent network input variance for two different noise schedules from the
literature (linear [12] and cosine [28]) during the reverse process of a
diffusion model (note that $t$ is reversed relative to diffusion convention).
The input variance is much lower for Bayesian Flow Networks.

###  4.7 Output Distribution $p_{{}_{O}}(\cdot\mid\boldsymbol{\theta};t)$

Following standard practice for diffusion models [42], the output distribution
is defined by reparameterising a prediction of the Gaussian noise vector
$\boldsymbol{\epsilon}\sim\mathcal{N}\left(\boldsymbol{0},\boldsymbol{I}\right)$
used to generate the mean $\operatorname{\boldsymbol{\mu}}$ passed as input to
the network. Recall from Eq. 77 that

$\displaystyle\operatorname{\boldsymbol{\mu}}\sim\mathcal{N}\left(\gamma(t)\operatorname{\mathbf{x}},\gamma(t)(1-\gamma(t))\boldsymbol{I}\right),$
(81)

and hence

$\displaystyle\operatorname{\boldsymbol{\mu}}$
$\displaystyle=\gamma(t)\operatorname{\mathbf{x}}+\sqrt{\gamma(t)(1-\gamma(t))}\boldsymbol{\epsilon}$
(82) $\displaystyle\implies\operatorname{\mathbf{x}}$
$\displaystyle=\frac{\operatorname{\boldsymbol{\mu}}}{\gamma(t)}-\sqrt{\frac{1-\gamma(t)}{\gamma(t)}}\boldsymbol{\epsilon}.$
(83)

The network outputs an estimate
$\boldsymbol{\hat{\epsilon}}(\boldsymbol{\theta},t)$ of
$\boldsymbol{\epsilon}$ and this is transformed into an estimate
$\mathbf{\hat{x}}(\boldsymbol{\theta},t)$ of $\operatorname{\mathbf{x}}$ by

$\displaystyle\mathbf{\hat{x}}(\boldsymbol{\theta},t)=\frac{\operatorname{\boldsymbol{\mu}}}{\gamma(t)}-\sqrt{\frac{1-\gamma(t)}{\gamma(t)}}\boldsymbol{\hat{\epsilon}}(\boldsymbol{\theta},t).$
(84)

Given $\boldsymbol{\hat{x}}(\boldsymbol{\theta},t)$ the output distribution is

$\displaystyle
p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta};t)=\delta(\operatorname{\mathbf{x}}-\mathbf{\hat{x}}(\boldsymbol{\theta},t)),$
(85)

Note that $\gamma(0)=0$, making the transformation from
$\boldsymbol{\hat{\epsilon}}(\boldsymbol{\theta},t)$ to
$p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta};t)$ undefined at
$t=0$. We therefore set
$p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta};t)=\boldsymbol{0}$
for $t$ under some small threshold $t_{min}$. Also,
$\mathbf{\hat{x}}(\boldsymbol{\theta},t)$ is clipped to lie within the allowed
range $[x_{min},x_{max}]$ for $\operatorname{\mathbf{x}}$. In our experiments
$t_{min}=1\mathrm{e}{-6}$ and $[x_{min},x_{max}]=[-1,1]$.

###  4.8 Sender Distribution
$p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha\right)$

The sender space
$\operatorname{\mathcal{Y}}=\operatorname{\mathcal{X}}=\operatorname{\mathbb{R}}$
for continuous data, and the sender distribution is normal with precision
$\alpha$:

$\displaystyle
p_{{}_{S}}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}};\alpha\right)$
$\displaystyle=\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}},\alpha^{-1}\boldsymbol{I}\right).$
(86)

###  4.9 Receiver Distribution
$p_{{}_{R}}(\cdot\mid\boldsymbol{\theta};t,\alpha)$

Substituting Eqs. 85 and 86 into Eq. 4,

$\displaystyle
p_{{}_{R}}(\operatorname{\mathbf{y}}\mid\boldsymbol{\theta};t,\alpha)$
$\displaystyle=\operatorname*{\mathbb{E}}_{\delta(\operatorname{\mathbf{x}}^{\prime}-\mathbf{\hat{x}}(\boldsymbol{\theta},t))}\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}}^{\prime},\alpha^{-1}\boldsymbol{I}\right)$
(87)
$\displaystyle=\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\mathbf{\hat{x}}(\boldsymbol{\theta},t),\alpha^{-1}\boldsymbol{I}\right).$
(88)

Figure 6: Sender, output and receiver distributions for continuous data. Note
that the sender and receiver distributions have identical variance and the
output distribution is a Dirac delta distribution centred on the network
prediction $\hat{x}(\theta,t)$.

###  4.10 Reconstruction Loss $L^{r}(\operatorname{\mathbf{x}})$

Truly continuous data requires infinite precision to reconstruct, which makes
the reconstruction loss problematic. However it would be reasonable to assume
that either the data is finely discretised (as all information is on a digital
computer), or that it contains some noise. The reconstruction loss for
discretised data is presented in Section 5.3. Alternatively, if we assume the
presence of normally distributed measurement noise on
$\operatorname{\mathbf{x}}$, with fixed isotropic variance $\sigma^{2}$, then
a noisy version of the reconstruction loss can be defined as the expected KL
divergence between
$\mathcal{N}\left(\operatorname{\mathbf{x}},\sigma^{2}\boldsymbol{I}\right)$
and the output distribution at $t=1$:

$\displaystyle L^{r}(\operatorname{\mathbf{x}})$
$\displaystyle=\operatorname*{\mathbb{E}}_{p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}},1)}D_{KL}\left(\mathcal{N}\left(\operatorname{\mathbf{x}},\sigma^{2}\boldsymbol{I}\right)\parallel\mathcal{N}\left(\mathbf{\hat{x}}(\boldsymbol{\theta},1),\sigma^{2}\boldsymbol{I}\right)\right)$
(89)
$\displaystyle=\operatorname*{\mathbb{E}}_{p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}},1)}\frac{1}{2\sigma^{2}}\left\|\operatorname{\mathbf{x}}-\mathbf{\hat{x}}(\boldsymbol{\theta},1)\right\|^{2}.$
(90)

The noise does not directly affect training, as the reconstruction loss is not
optimised. However the value of $\sigma$ places a natural upper limit on the
value that should be chosen for $\sigma_{1}$: there is no point transmitting
the data to greater precision than it was originally measured. Empirically, we
find that when $\sigma_{1}<\sigma/2$ the reconstruction loss is very small.

###  4.11 Discrete-Time Loss $L^{n}(\operatorname{\mathbf{x}})$

From Eqs. 86 and 88,

$\displaystyle
D_{KL}\left(p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}},\alpha_{i}\right)\parallel
p_{{}_{R}}(\cdot\mid\boldsymbol{\theta}_{i-1};t_{i-1},\alpha_{i})\right)$
$\displaystyle=D_{KL}\left(\mathcal{N}\left(\operatorname{\mathbf{x}},\alpha_{i}^{-1}\boldsymbol{I}\right)\parallel\mathcal{N}\left(\mathbf{\hat{x}}(\boldsymbol{\theta}_{i-1},t_{i-1}),\alpha_{i}^{-1}\boldsymbol{I}\right)\right)$
(91)
$\displaystyle=\frac{\alpha_{i}}{2}\left\|\operatorname{\mathbf{x}}-\mathbf{\hat{x}}(\boldsymbol{\theta}_{i-1},t_{i-1})\right\|^{2},$
(92)

and from Eqs. 11 and 72,

$\displaystyle\alpha_{i}$ $\displaystyle=\beta(t_{i})-\beta(t_{i-1})$ (93)
$\displaystyle=\sigma_{1}^{-2i/n}-\sigma_{1}^{-2(i-1)/n}$ (94)
$\displaystyle=\sigma_{1}^{-2i/n}\left(1-\sigma_{1}^{2/n}\right).$ (95)

Therefore, substituting into Eq. 24,

$\displaystyle
L^{n}(\operatorname{\mathbf{x}})=\frac{n}{2}\left(1-\sigma_{1}^{2/n}\right)\operatorname*{\mathbb{E}}_{i\sim
U\\{1,n\\},p_{{}_{F}}(\boldsymbol{\theta}_{i-1}\mid\operatorname{\mathbf{x}};t_{i-1})}\frac{\left\|\operatorname{\mathbf{x}}-\mathbf{\hat{x}}(\boldsymbol{\theta}_{i-1},t_{i-1})\right\|^{2}}{\sigma_{1}^{2i/n}},$
(96)

where $t_{i-1}=(i-1)/{n}$.

###  4.12 Continuous-time Loss $L^{\infty}(\operatorname{\mathbf{x}})$

Eq. 29 claimed that

$\displaystyle
D_{KL}\left(p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}},\alpha\right)\parallel
p_{{}_{R}}(\cdot\mid\boldsymbol{\theta},\alpha,t)\right)=D_{KL}\left(\mathcal{N}\left(g(\operatorname{\mathbf{x}}),C\alpha^{-1}\boldsymbol{I}\right)\parallel
P(\boldsymbol{\theta},t)\ast\mathcal{N}\left(\boldsymbol{0},C\alpha^{-1}\boldsymbol{I}\right)\right),$
(97)

for some embedding function
$g:\operatorname{\mathcal{X}}\rightarrow\operatorname{\mathcal{Y}}$, constant
$C$ and distribution $p_{\boldsymbol{\theta}}$ over
$\operatorname{\mathcal{Y}}^{D}$ with finite mean and variance. If $g$ is the
identity function, $C=1$ and

$\displaystyle P(\operatorname{\mathbf{y}}\mid\boldsymbol{\theta},t)$
$\displaystyle=\delta(\operatorname{\mathbf{y}}-\mathbf{\hat{x}}(\boldsymbol{\theta},t)),$
(98)

then $P(\boldsymbol{\theta},t)$ has finite mean and variance and

$\displaystyle\mathcal{N}\left(\operatorname{\mathbf{y}}\mid
g(\operatorname{\mathbf{x}}),C\alpha^{-1}\boldsymbol{I}\right)=\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}},\alpha^{-1}\boldsymbol{I}\right)$
$\displaystyle=p_{{}_{S}}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}};\alpha\right),$
(99) $\displaystyle
P(\operatorname{\mathbf{y}}\mid\boldsymbol{\theta},t)\ast\mathcal{N}\left(\boldsymbol{0},C\alpha^{-1}\boldsymbol{I}\right)=\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\mathbf{\hat{x}}(\boldsymbol{\theta},t),\alpha^{-1}\boldsymbol{I}\right)$
$\displaystyle=p_{{}_{R}}(\operatorname{\mathbf{y}}\mid\boldsymbol{\theta},\alpha,t),$
(100)

so the claim is true and the continuous-time loss from Eq 41 applies, with
$E[P(\boldsymbol{\theta},t)]=\mathbf{\hat{x}}(\boldsymbol{\theta},t)$ and
$\alpha(t)$ as defined in Eq 74, yielding

$\displaystyle L^{\infty}(\operatorname{\mathbf{x}})$
$\displaystyle=-\ln\sigma_{1}\operatorname*{\mathbb{E}}_{t\sim
U(0,1),p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t)}\frac{\left\|\operatorname{\mathbf{x}}-\mathbf{\hat{x}}(\boldsymbol{\theta},t)\right\|^{2}}{\sigma_{1}^{2t}}.$
(101)

###  4.13 Pseudocode

Pseudocode for evaluating the $n$-step loss $L^{n}(\operatorname{\mathbf{x}})$
and continuous-time loss $L^{\infty}(\operatorname{\mathbf{x}})$ for
continuous data is presented in Algorithms 1 and 2, while the sample
generation procedure is presented in Algorithm 3.

# Note that $\boldsymbol{\theta}=\\{\operatorname{\boldsymbol{\mu}},\rho\\}$,
but $\rho$ is fully determined by $t$

# For our experiments $t_{min}=1\mathrm{e}{-6}$, $[x_{min},x_{max}]=[-1,1]$

function
cts_output_prediction($\operatorname{\boldsymbol{\mu}}\in\operatorname{\mathbb{R}}^{D},t\in[0,1],\gamma>\in\operatorname{\mathbb{R}}^{+}$,
$t_{min}\in\operatorname{\mathbb{R}}^{+}$,
$x_{min},x_{max}\in\operatorname{\mathbb{R}}$)

if $t<t_{min}$ then

$\mathbf{\hat{x}}(\boldsymbol{\theta},t)\leftarrow\boldsymbol{0}$

else

Input $(\operatorname{\boldsymbol{\mu}},t)$ to network, receive
$\boldsymbol{\hat{\epsilon}}(\boldsymbol{\theta},t)$ as output

$\mathbf{\hat{x}}(\boldsymbol{\theta},t)\leftarrow\frac{\operatorname{\boldsymbol{\mu}}}{\gamma}-\sqrt{\frac{1-\gamma}{\gamma}}\boldsymbol{\hat{\epsilon}}(\boldsymbol{\theta},t)$

clip $\mathbf{\hat{x}}(\boldsymbol{\theta},t)$ to $[x_{min},x_{max}]$

end if

Return $\mathbf{\hat{x}}(\boldsymbol{\theta},t)$

end function

Algorithm 1 Discrete-Time Loss $L^{n}(\operatorname{\mathbf{x}})$ for
Continuous Data

Require: $\sigma_{1}\in\operatorname{\mathbb{R}}^{+}$, number of steps
$n\in\mathbb{N}$

Input: continuous data
$\operatorname{\mathbf{x}}\in\operatorname{\mathbb{R}}^{D}$

$i\sim U\\{1,n\\}$

$t\leftarrow\frac{i-1}{n}$

$\gamma\leftarrow 1-\sigma_{1}^{2t}$

$\operatorname{\boldsymbol{\mu}}\sim\mathcal{N}\left(\gamma\operatorname{\mathbf{x}},\gamma(1-\gamma)\boldsymbol{I}\right)$

${\mathbf{\hat{x}}(\boldsymbol{\theta},t)\leftarrow\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{cts\textunderscore
output\textunderscore
prediction}}}}}}}(\operatorname{\boldsymbol{\mu}},t,\gamma)$

$L^{n}(\operatorname{\mathbf{x}})\leftarrow\frac{n\left(1-\sigma_{1}^{2/n}\right)}{2\sigma_{1}^{2i/n}}\left\|\operatorname{\mathbf{x}}-\mathbf{\hat{x}}(\boldsymbol{\theta},t)\right\|^{2}$

Algorithm 2 Continuous-Time Loss $L^{\infty}(\operatorname{\mathbf{x}})$ for
Continuous Data

Require: $\sigma_{1}\in\operatorname{\mathbb{R}}^{+}$

Input: continuous data
$\operatorname{\mathbf{x}}\in\operatorname{\mathbb{R}}^{D}$

$t\sim U(0,1)$

$\gamma\leftarrow 1-\sigma_{1}^{2t}$

$\operatorname{\boldsymbol{\mu}}\sim\mathcal{N}\left(\gamma\operatorname{\mathbf{x}},\gamma(1-\gamma)\boldsymbol{I}\right)$

${\mathbf{\hat{x}}(\boldsymbol{\theta},t)\leftarrow\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{cts\textunderscore
output\textunderscore
prediction}}}}}}}(\operatorname{\boldsymbol{\mu}},t,\gamma)$

$L^{\infty}(\operatorname{\mathbf{x}})\leftarrow-\ln\sigma_{1}\sigma_{1}^{-2t}\left\|\mathbf{\operatorname{\mathbf{x}}-\hat{x}}(\boldsymbol{\theta},t)\right\|^{2}$

Algorithm 3 Sample Generation for Continuous Data

Require: $\sigma_{1}\in\operatorname{\mathbb{R}}^{+}$, number of steps
$n\in\mathbb{N}$

$\boldsymbol{\mu}\leftarrow\boldsymbol{0}$

$\rho\leftarrow 1$

for $i=1$ to $n$ do

$t\leftarrow\frac{i-1}{n}$

${\mathbf{\hat{x}}(\boldsymbol{\theta},t)\leftarrow\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{cts\textunderscore
output\textunderscore
prediction}}}}}}}(\operatorname{\boldsymbol{\mu}},t,1-\sigma_{1}^{2t})$

$\alpha\leftarrow\sigma_{1}^{-2i/n}\left(1-\sigma_{1}^{2/n}\right)$

$\operatorname{\mathbf{y}}\sim\mathcal{N}\left(\mathbf{\hat{x}}(\boldsymbol{\theta},t),\alpha^{-1}\boldsymbol{I}\right)$

$\operatorname{\boldsymbol{\mu}}\leftarrow\frac{\rho\boldsymbol{\mu}+\alpha\operatorname{\mathbf{y}}}{\rho+\alpha}$

$\rho\leftarrow\rho+\alpha$

end for

${\mathbf{\hat{x}}(\boldsymbol{\theta},1)\leftarrow\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{cts\textunderscore
output\textunderscore
prediction}}}}}}}(\operatorname{\boldsymbol{\mu}},1,1-\sigma_{1}^{2})$

Return $\mathbf{\hat{x}}(\boldsymbol{\theta},1)$

##  5 Discretised Data

This section considers continuous data that has been discretised into $K$
bins. For example, 8-bit images are discretised into 256 bins, 16-bit audio is
discretised in $2^{16}=65,536$ bins. This data is represented by tiling
$[-1,1]$ into $K$ intervals, each of length $2/K$. Let $k_{l}$, $k_{c}$ and
$k_{r}$ denote respectively the left, centre and right of interval $k$, and
let $\\{1,K\\}$ denote the set of integers from 1 to $K$. Then for
$k\in\\{1,K\\}$,

$\displaystyle k_{c}$ $\displaystyle=\frac{2k-1}{K}-1,$ (102) $\displaystyle
k_{l}$ $\displaystyle=k_{c}-\frac{1}{K},$ (103) $\displaystyle k_{r}$
$\displaystyle=k_{c}+\frac{1}{K}.$ (104)

Let
$k(\operatorname{\mathbf{x}})=\left(k(x^{(1)}),\dots,k(x^{(D)})\right)\in\\{1,K\\}^{D}$
be the vector of the indices of the bins occupied by
$\operatorname{\mathbf{x}}=\left(x^{(1)},\dots,x^{(D)}\right)\in\operatorname{\mathbb{R}}^{D}$,
and let $k_{l}(\operatorname{\mathbf{x}})$, $k_{c}(\operatorname{\mathbf{x}})$
and $k_{r}(\operatorname{\mathbf{x}})$ be the corresponding vectors of left
edges, centres and right edges of the bins. If the data has not already been
discretised, we set
$\operatorname{\mathbf{x}}=k_{c}(\operatorname{\mathbf{x}})$. For example if
the red channel in an 8-bit RGB image has index 110, it will be represented by
the number $\frac{2*(110)-1}{256}-1=-0.14453125$. Note that each $x^{(d)}$
therefore lies in the range $[\frac{1}{K}-1,1-\frac{1}{K}]$ and not $[-1,1]$.

The input distribution
$p_{{}_{I}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta})$, prior
parameters $\boldsymbol{\theta}_{0}$, sender distribution
$p_{{}_{S}}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}};\alpha\right)$,
Bayesian update function
$h(\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{y}},\alpha)$, Bayesian
update distribution
$p_{{}_{U}}(\boldsymbol{\theta}_{i}\mid\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{x}};\alpha)$,
Bayesian flow distribution
$p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t)$ and accuracy
schedule $\beta(t)$ are all identical to the continuous case described in
Section 4. It may surprise the reader that the output distribution is
discretised while the input, sender and receiver distributions are not. We
made this choice partly for mathematical convenience (Bayesian updates are
considerably more complex for discretised distributions; [1]) and partly
because we suspected that it would easier for the network to interpret
continuous means than discrete probabilities as input. In a similar vein to
our argument for standard priors in Sec. 4.1, we remind the reader that the
input distribution only serves to inform the network and not directly to model
the data; all that matters is that the input parameters contain enough
information to allow the network to make accurate predictions.

Section 4.11 noted that the level of measurement noise assumed for continuous
data should inform the choice of standard deviation $\sigma_{1}$ for the input
distribution at $t=1$ (which in turn defines the accuracy schedule
$\beta(t)$). For discretised data a similar role is played by the width of the
discretisation bins, as these place a natural limit on how precisely the data
needs to be transmitted. For example, for $8$-bit data with 256 bins and hence
a bin width of $1/128$, setting $\sigma_{1}=1\mathrm{e}{-3}$ corresponds to a
final input distribution with standard deviation roughly one eighth of the
width of the bin, which should be precise enough for the network to identify
the correct bin with very high probability.

One caveat with discretisation is that calculating the loss has $O(K)$
computational cost, which may be prohibitive for very finely discretised data.
In any case, the benefits of discretisation tend to decrease as the number of
bins increases, as we will see in our experiments.

Figure 7: Output distribution for discretised data. For univariate data $x$
discretised into $K=16$ bins, the green line shows the continuous distribution
$\mathcal{N}\left(\mu_{x},\sigma^{2}_{x}\right)$ that is discretised to yield
the output distribution $p_{{}_{O}}(x\mid\theta,t)$, as described in Section
5.1. Bin boundaries are marked with vertical grey lines. The heights of the
green bars represent the probabilities assigned to the respective bins by
$p_{{}_{O}}(x\mid\theta,t)$. For ease of visualisation these heights are
rescaled relative to the probability density, as indicated on the right axis.
Note the clipping at $\pm 1$: the area under the dotted green line to the left
of $-1$ is added to the probability of the first bin, the area under the
dotted green line to the right of $1$ is added to the probability of the last
bin.

###  5.1 Output Distribution $p_{{}_{O}}(\cdot\mid\boldsymbol{\theta},t)$

Discretised continuous distributions offer a natural and expressive way to
model discretised data with neural networks [38]. As in Section 4.7, the
network outputs $\Psi(\boldsymbol{\theta},t)$ are not used to predict
$\operatorname{\mathbf{x}}$ directly, but rather to model the Gaussian noise
vector $\operatorname{\mathbf{e}}$ used to generate the mean sample
$\operatorname{\boldsymbol{\mu}}$ passed as input to the network.

First $\Psi(\boldsymbol{\theta},t)$ is split into two length $D$ vectors,
$\operatorname{\boldsymbol{\mu}}_{\epsilon}$ and
$\ln\boldsymbol{\sigma}_{\epsilon}$. Then these are transformed to
$\operatorname{\boldsymbol{\mu}}_{x}$ and $\boldsymbol{\sigma}_{x}$ using

$\displaystyle\operatorname{\boldsymbol{\mu}}_{x}$
$\displaystyle=\begin{cases}\boldsymbol{0}&\text{if $t<t_{min}$},\\\
\frac{\operatorname{\boldsymbol{\mu}}}{\gamma(t)}-\sqrt{\frac{1-\gamma(t)}{\gamma(t)}}\operatorname{\boldsymbol{\mu}}_{\epsilon}&\text{otherwise},\end{cases}$
(105) $\displaystyle\boldsymbol{\sigma}_{x}$
$\displaystyle=\begin{cases}\boldsymbol{1}&\text{if $t<t_{min}$},\\\
\sqrt{\frac{1-\gamma(t)}{\gamma(t)}}\exp(\ln\boldsymbol{\sigma}_{\epsilon})&\text{otherwise}.\end{cases}$
(106)

For each $d\in\\{1,D\\}$, define the following univariate Gaussian cdf

$\displaystyle F\left(x\mid\mu_{x}^{(d)},\sigma_{x}^{(d)}\right)$
$\displaystyle=\frac{1}{2}\left[1+\text{erf}\left(\frac{x-\mu_{x}^{(d)}}{\sigma_{x}^{(d)}\sqrt{2}}\right)\right],$
(107)

and clip at $[-1,1]$ to obtain

$\displaystyle
G\left(x\mid\mu_{x}^{(d)},\sigma_{x}^{(d)}\right)=\begin{cases}0&\text{if
$x\leq-1$},\\\ 1&\text{if $x\geq 1$},\\\
F\left(x\mid\mu_{x}^{(d)},\sigma_{x}^{(d)}\right)&\text{otherwise}.\end{cases}$
(108)

Then, for $k\in\\{1,K\\}$,

$\displaystyle
p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\stackrel{{\scriptstyle\text{def}}}{{=}}G(k_{r}\mid\mu^{(d)}_{x},\sigma^{(d)}_{x})-G(k_{l}\mid\mu^{(d)}_{x},\sigma^{(d)}_{x}),$
(109)

and hence

$\displaystyle
p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta}_{,}t)=\prod_{d=1}^{D}p_{{}_{O}}^{(d)}\left(k(x^{(d)})\mid\boldsymbol{\theta};t\right).$
(110)

###  5.2 Receiver Distribution
$p_{{}_{R}}(\cdot\mid\boldsymbol{\theta};t,\alpha)$

Substituting Eq. 110 and Eq. 86 into Eq. 4 gives

$\displaystyle
p_{{}_{R}}(\operatorname{\mathbf{y}}\mid\boldsymbol{\theta};t,\alpha)$
$\displaystyle=\operatorname*{\mathbb{E}}_{p_{{}_{O}}(\operatorname{\mathbf{x}}^{\prime}\mid\boldsymbol{\theta}_{,}t)}\mathcal{N}\left(y^{(d)}\mid
k_{c}(\operatorname{\mathbf{x}}^{\prime}),\alpha^{-1}\boldsymbol{I}\right)$
(111)
$\displaystyle=\prod_{d=1}^{D}\int_{x^{\prime}}dx^{\prime}{p_{{}_{O}}^{(d)}\left(k(x^{\prime})\mid\boldsymbol{\theta};t\right)}\mathcal{N}\left(y^{(d)}\mid
k_{c}(x^{\prime}),\alpha^{-1}\right)$ (112)
$\displaystyle=\prod_{d=1}^{D}\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\mathcal{N}\left(y^{(d)}\mid
k_{c},\alpha^{-1}\right).$ (113)

Figure 8: Sender, output and receiver distributions for discretised data. For
data $x$ discretised into 8 bins, the three plots depict the sender
distribution (red line), the discretised output distribution (green bars;
heights reflect the probabilities assigned to each bin, rescaled as in Figure
7) and receiver distribution (blue line) for progressively increasing values
of $\alpha$, and for progressively more accurate predictions of $x$ (both of
which typically happen as $t$ increases). Also shown are the continuous
distribution $\mathcal{N}(x\mid\mu_{x},\sigma^{2}_{x})$ (dotted green line)
which is discretized to create the output distribution and the continuous
receiver distribution from Section 4 (dashed orange line). Bin boundaries are
marked with vertical grey lines. Note the KL divergences printed in the top
right: taking discretisation into account leads to a lower KL due to the
density “bumps” at the bin centres where $x$ could be. The advantage of
discretisation becomes more pronounced as the prediction gets closer to $x$
and more of the probability mass is concentrated in the correct bin.

###  5.3 Reconstruction Loss $L^{r}(\operatorname{\mathbf{x}})$

The reconstruction loss for discretised data is

$\displaystyle L^{r}(\operatorname{\mathbf{x}})$
$\displaystyle=-\operatorname*{\mathbb{E}}_{p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}},1)}\ln
p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta};1)$ (114)
$\displaystyle=-\operatorname*{\mathbb{E}}_{p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}},1)}\sum_{d=1}^{D}\ln
p_{{}_{O}}^{(d)}\left(k(x^{(d)})\mid\boldsymbol{\theta};1\right).$ (115)

###  5.4 Discrete-time Loss $L^{n}(\operatorname{\mathbf{x}})$

From Eqs. 86 and 113,

$\displaystyle
D_{KL}\left(p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}},\alpha_{i}\right)\parallel
p_{{}_{R}}(\cdot\mid\boldsymbol{\theta}_{i-1};t_{i-1},\alpha_{i})\right)$
(116)
$\displaystyle\qquad\qquad\qquad=D_{KL}\left(\mathcal{N}\left(\operatorname{\mathbf{x}},\alpha_{i}^{-1}\boldsymbol{I}\right)\parallel\prod_{d=1}^{D}\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta}_{i-1},t_{i-1})\mathcal{N}\left(k_{c},\alpha_{i}^{-1}\right)\right),$
(117)

which cannot be calculated in closed form, but can be estimated with Monte-
Carlo sampling. Substituting into Eq. 24,

$\displaystyle
L^{n}(\operatorname{\mathbf{x}})=n\operatorname*{\mathbb{E}}_{i\sim
U\\{1,n\\},p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t_{i-1}),\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}},\alpha_{i}^{-1}\boldsymbol{I}\right)}\ln\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}},\alpha_{i}^{-1}\boldsymbol{I}\right)$
(118)
$\displaystyle\qquad\qquad\qquad\qquad-\sum_{d=1}^{D}\ln\left(\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta},t_{i-1})\mathcal{N}\left(y^{(d)}\mid
k_{c},\alpha_{i}^{-1}\right)\right).$ (119)

###  5.5 Continuous-time Loss $L^{\infty}(\operatorname{\mathbf{x}})$

Justifying the claim made in Eq. 29 follows almost the same reasoning here as
in Section 4.12, with $C=1$ and $g$ the identity function. The only difference
is that

$\displaystyle
P(\operatorname{\mathbf{y}}\mid\boldsymbol{\theta};t)=\prod_{d=1}^{D}\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta},t)\delta(y^{(d)}-k_{c}),$
(120)

which clearly has finite variance and mean. Since

$\displaystyle
P(\operatorname{\mathbf{y}}\mid\boldsymbol{\theta},t)\ast\mathcal{N}\left(\boldsymbol{0},C\alpha^{-1}\boldsymbol{I}\right)=p_{{}_{R}}(\operatorname{\mathbf{y}}\mid\boldsymbol{\theta},\alpha,t),$
(121)

the claim holds and the continuous time loss from Eq 41 can be applied with

$\displaystyle
E[P(\boldsymbol{\theta},t)]=\left(\sum_{k=1}^{K}p^{(1)}(k\mid\boldsymbol{\theta},t)k_{c},\dots,\sum_{k=1}^{K}p^{(D)}(k\mid\boldsymbol{\theta},t)k_{c}\right)\stackrel{{\scriptstyle\text{def}}}{{=}}\mathbf{\hat{k}}(\boldsymbol{\theta},t),$
(122)

and $\alpha(t)$ as defined in Eq 74, yielding

$\displaystyle L^{\infty}(\operatorname{\mathbf{x}})$
$\displaystyle=-\ln\sigma_{1}\operatorname*{\mathbb{E}}_{t\sim
U(0,1),p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t)}\frac{\left\|\operatorname{\mathbf{x}}-\mathbf{\hat{k}}(\boldsymbol{\theta},t)\right\|^{2}}{\sigma_{1}^{2t}}.$
(123)

Note that $\mathbf{\hat{k}}(\boldsymbol{\theta},t)$ is a function of the
complete discretised distribution
$p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta},t)$, hence
$L^{\infty}(\operatorname{\mathbf{x}})$ depends on both
$\operatorname{\boldsymbol{\mu}}_{\operatorname{\mathbf{x}}}$ and
$\boldsymbol{\sigma}_{\operatorname{\mathbf{x}}}$, and not only on
$\operatorname{\boldsymbol{\mu}}_{\operatorname{\mathbf{x}}}$, as for
continuous data. This also means that calculating
$L^{\infty}(\operatorname{\mathbf{x}})$ has $O(K)$ computational cost for
discretised data.

###  5.6 Pseudocode

Pseudocode for evaluating the discrete-time loss
$L^{n}(\operatorname{\mathbf{x}})$ and continuous-time loss
$L^{\infty}(\operatorname{\mathbf{x}})$ for discretised data is presented in
Algorithms 4 and 5, while sample generation is presented in Algorithm 6.

function
discretised_cdf($\mu\in\operatorname{\mathbb{R}},\sigma\in\operatorname{\mathbb{R}}^{+},x\in\operatorname{\mathbb{R}}$)

$F(x)\leftarrow\frac{1}{2}\left[1+\text{erf}\left(\frac{x-\mu}{\sigma\sqrt{2}}\right)\right]$

$G(x)\leftarrow\begin{cases}0&\text{ if }x\leq-1\\\ 1&\text{ if }x\geq 1\\\
F(x)&\text{ otherwise}\end{cases}$

Return $G(x)$

end function

# For our experiments $t_{min}=1\mathrm{e}{-6}$

# $k_{l}=\frac{2(k-1)}{K}-1$, $k_{r}=\frac{2k}{K}-1$

function
discretised_output_distribution($\operatorname{\boldsymbol{\mu}}\in\operatorname{\mathbb{R}}^{D},t\in[0,1],K\in\mathbb{N},\gamma\in\operatorname{\mathbb{R}}^{+}$,
$t_{min}\in\operatorname{\mathbb{R}}^{+}$).

if $t<t_{min}$ then

$\operatorname{\boldsymbol{\mu}}_{x}\leftarrow\boldsymbol{0}$

$\boldsymbol{\sigma}_{x}\leftarrow\boldsymbol{1}$

else

Input $(\operatorname{\boldsymbol{\mu}},t)$ to network, receive
$(\operatorname{\boldsymbol{\mu}}_{\epsilon},\ln\boldsymbol{\sigma}_{\epsilon})$
as output

$\operatorname{\boldsymbol{\mu}}_{x}\leftarrow\frac{\operatorname{\boldsymbol{\mu}}}{\gamma}-\sqrt{\frac{1-\gamma}{\gamma}}\operatorname{\boldsymbol{\mu}}_{\epsilon}$

$\boldsymbol{\sigma}_{x}\leftarrow\sqrt{\frac{1-\gamma}{\gamma}}\exp(\ln\boldsymbol{\sigma}_{\epsilon})$

end if

for $d\in\\{1,D\\}$, $k\in\\{1,K\\}$ do

${{p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\leftarrow\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{discretised\textunderscore
cdf}}}}}}}(\mu_{x}^{(d)},\sigma_{x}^{(d)},k_{r})-\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{discretised\textunderscore
cdf}}}}}}}(\mu_{x}^{(d)},\sigma_{x}^{(d)},k_{l})$

end for

Return $\boldsymbol{p}_{{}_{O}}(\cdot\mid\boldsymbol{\theta};t)$

end function

Algorithm 4 Discrete-Time Loss $L^{n}(\operatorname{\mathbf{x}})$ for
Discretised Data

# $k_{c}=\frac{2k-1}{K}-1$

Require: $\sigma_{1}\in\operatorname{\mathbb{R}}^{+}$, number of steps
$n\in\mathbb{N}$, number of bins $K\in\mathbb{N}$

Input: discretised data
$\operatorname{\mathbf{x}}\in[\frac{1}{K}-1,1-\frac{1}{K}]^{D}$

$i\sim U\\{1,n\\}$

$t\leftarrow\frac{i-1}{n}$

$\gamma\leftarrow 1-\sigma_{1}^{2t}$

$\operatorname{\boldsymbol{\mu}}\sim\mathcal{N}\left(\gamma\operatorname{\mathbf{x}},\gamma(1-\gamma)\boldsymbol{I}\right)$

$\alpha\leftarrow\sigma_{1}^{-2i/n}\left(1-\sigma_{1}^{2/n}\right)$

$\operatorname{\mathbf{y}}\sim\mathcal{N}\left(\operatorname{\mathbf{x}},\alpha^{-1}\boldsymbol{I}\right)$

${\boldsymbol{p}_{{}_{O}}(\cdot\mid\boldsymbol{\theta};t)\leftarrow\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{discretised\textunderscore
output\textunderscore
distribution}}}}}}}(\operatorname{\boldsymbol{\mu}},t,K,\gamma)$

$L^{n}(\operatorname{\mathbf{x}})\leftarrow
n\left[\ln\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}},\alpha^{-1}\boldsymbol{I}\right)-\sum_{d}\ln\left(\sum_{k}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\mathcal{N}\left(y^{(d)}\mid
k_{c},\alpha^{-1}\right)\right)\right]$

Algorithm 5 Continuous-Time Loss $L^{\infty}(\operatorname{\mathbf{x}})$ for
Discretised Data

Require: $\sigma_{1}\in\operatorname{\mathbb{R}}^{+}$, number of bins
$K\in\mathbb{N}$

Input: discretised data
$\operatorname{\mathbf{x}}\in[\frac{1}{K}-1,1-\frac{1}{K}]^{D}$

$t\sim U(0,1)$

$\gamma\leftarrow 1-\sigma_{1}^{2t}$

$\operatorname{\boldsymbol{\mu}}\sim\mathcal{N}\left(\gamma\operatorname{\mathbf{x}},\gamma(1-\gamma)\boldsymbol{I}\right)$

${\boldsymbol{p}_{{}_{O}}(\cdot\mid\boldsymbol{\theta};t)\leftarrow\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{discretised\textunderscore
output\textunderscore
distribution}}}}}}}(\operatorname{\boldsymbol{\mu}},t,K,\gamma)$

$\mathbf{\hat{k}}(\boldsymbol{\theta},t)\leftarrow\left(\sum_{k}p_{{}_{O}}^{(1)}(k\mid\boldsymbol{\theta};t)k_{c},\dots,\sum_{k}p_{{}_{O}}^{(D)}(k\mid\boldsymbol{\theta};t)k_{c}\right)$

$L^{\infty}(\operatorname{\mathbf{x}})\leftarrow-\ln\sigma_{1}\sigma_{1}^{-2t}\left\|\operatorname{\mathbf{x}}-\mathbf{\hat{k}}(\boldsymbol{\theta},t)\right\|^{2}$

Algorithm 6 Sample Generation for Discretised Data

# $\boldsymbol{k}_{c}=\left(k^{(1)}_{c},\dots,k^{(D)}_{c}\right)$

Require: $\sigma_{1}\in\operatorname{\mathbb{R}}^{+}$, number of steps
$n\in\mathbb{N}$, number of bins $K\in\mathbb{N}$

$\boldsymbol{\mu}\leftarrow\boldsymbol{0}$

$\rho\leftarrow 1$

for $i=1$ to $n$ do

$t\leftarrow\frac{i-1}{n}$

${\mathbf{k}\sim\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{discretised\textunderscore
output\textunderscore
distribution}}}}}}}(\operatorname{\boldsymbol{\mu}},t,K,1-\sigma_{1}^{2t})$

$\alpha\leftarrow\sigma_{1}^{-2i/n}\left(1-\sigma_{1}^{2/n}\right)$

$\operatorname{\mathbf{y}}\sim\mathcal{N}\left(\mathbf{k}_{c},\alpha^{-1}\boldsymbol{I}\right)$

$\operatorname{\boldsymbol{\mu}}\leftarrow\frac{\rho\boldsymbol{\mu}+\alpha\operatorname{\mathbf{y}}}{\rho+\alpha}$

$\rho\leftarrow\rho+\alpha$

end for

${\mathbf{k}\sim\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{discretised\textunderscore
output\textunderscore
distribution}}}}}}}(\operatorname{\boldsymbol{\mu}},1,K,1-\sigma_{1}^{2})$

Return $\mathbf{k}_{c}$

##  6 Discrete Data

We now consider discrete data in which no meaningful order or distance exists
between the classes, unlike the discretised continuous data covered in the
previous section. Some obvious examples are text characters, classification
labels or any binary data. In this context the data is represented as a $D$
dimensional vector of class indices:
$\operatorname{\mathbf{x}}=\left(x^{(1)},\dots,x^{(D)}\right)\in\\{1,K\\}^{D}$,
where $\\{1,K\\}$ is the set of integers from $1$ to $K$.

###  6.1 Input Distribution $p_{{}_{I}}(\cdot\mid\boldsymbol{\theta})$

For discrete data, the input distribution is a factorised categorical over the
class indices. Let
$\boldsymbol{\theta}=\left(\theta^{(1)},\dots,\theta^{(D)}\right)\in[0,1]^{KD}$
with
$\theta^{(d)}=\left(\theta_{1}^{(d)},\dots,\theta_{K}^{(d)}\right)\in\Delta^{K-1}$,
where $\theta_{k}^{(d)}$ is the probability assigned to class $k$ for variable
$d$. Then

$\displaystyle
p_{{}_{I}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta})=\prod_{d=1}^{D}\theta_{x^{(d)}}^{(d)}.$
(124)

The input prior is uniform with

$\displaystyle\boldsymbol{\theta}_{0}=\boldsymbol{\frac{1}{K}},$ (125)

where $\boldsymbol{\frac{1}{K}}$ is the length $KD$ vector whose entries are
all $\frac{1}{K}$. We chose a uniform prior—rather than an empirical prior fit
to the training data—for the same reasons we chose a standard normal prior for
continuous data: it’s mathematically simpler, and the disparity between the
true prior and the simple prior can easily be corrected by the network.

###  6.2 Output Distribution $p_{{}_{O}}(\cdot\mid\boldsymbol{\theta};t)$

Given data $\operatorname{\mathbf{x}}$, network inputs $\boldsymbol{\theta},t$
and corresponding network outputs
$\Psi(\boldsymbol{\theta},t)=\left(\Psi^{(1)}(\boldsymbol{\theta},t),\dots,\Psi^{(D)}(\boldsymbol{\theta},t)\right)\\\
\in\operatorname{\mathbb{R}}^{KD}$, the output distribution for discrete data
is as follows:

$\displaystyle p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)$
$\displaystyle=\left(\text{softmax}(\Psi^{(d)}(\boldsymbol{\theta},t))\right)_{k},$
(126) $\displaystyle
p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta};t)$
$\displaystyle=\prod_{d=1}^{D}p_{{}_{O}}^{(d)}(x^{(d)}\mid\boldsymbol{\theta};t).$
(127)

Note that for binary data only the probability $\theta^{(d)}_{1}$ that $k=1$
is fed into the network, on the grounds that the probability of $k=2$ can
easily be inferred from $\theta^{(d)}_{2}=1-\theta^{(d)}_{1}$. The output
distribution for binary data is determined by applying the logistic sigmoid
function elementwise to the length $D$ output vector to get the probability
for $k=1$:

$\displaystyle
p_{{}_{O}}^{(d)}(1\mid\boldsymbol{\theta};t)=\sigma\left(\Psi^{(d)}(\boldsymbol{\theta},t))\right),$
(128)

where

$\displaystyle\sigma(x)=\frac{1}{1-e^{-x}},$ (129)

then inferring the probabilities for $k=2$ from

$\displaystyle
p_{{}_{O}}^{(d)}(2\mid\boldsymbol{\theta};t)=1-p_{{}_{O}}^{(d)}(1\mid\boldsymbol{\theta};t).$
(130)

In principle one class could also be removed from the inputs and outputs when
$K>2$ and inferred from the others. However this would require the network to
internalise a slightly more sophisticated inference procedure that could
potentially slow down learning. We therefore followed deep-learning convention
and included a redundant input and output unit for $K>2$.

All probabilities are rescaled to the range $[-1,1]$ by multiplying by two
then subtracting one before feeding them into the network.

###  6.3 Sender Distribution
$p_{{}_{S}}\left(\cdot\mid\operatorname{\mathbf{x}};\alpha\right)$

Given $\omega\in[0,1]$, and a vector of $D$ class indices
$\mathbf{k}=\left(k^{(1)},\dots,k^{(D)}\right)\in\\{1,K\\}^{D}$, let

$\displaystyle p(k^{(d)}\mid x^{(d)};\omega)$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{1-\omega}{K}+\omega\delta_{k^{(d)}x^{(d)}},$
(131)

where $\delta_{ij}$ is the Kronecker delta function. Clearly $p(k^{(d)}\mid
x^{(d)};\omega)\geq 0\ \forall k$ and $\sum_{k=1}^{K}p(k^{(d)}\mid
x^{(d)};\omega)=1$, so the vector

$\displaystyle
a(x^{(d)},\omega)\stackrel{{\scriptstyle\text{def}}}{{=}}\left(p(1\mid
x^{(d)};\omega),\dots,p(K\mid x^{(d)};\omega)\right),$ (132)

defines a valid distribution over $K$ classes. To simplify notation we will
from now on drop the superscripts and refer to $x^{(d)}$ as $x$,
$p(k^{(d)}\mid x^{(d)};\omega)$ as $p(k\mid x;\omega)$ and so on, except where
necessary to remove ambiguity.

Consider a vector of integer counts $c=(c_{1},\dots,c_{K})\in\\{1,m\\}^{K}$,
corresponding to the number of times each of the $K$ classes is observed among
$m$ independent draws from $a(x,\omega)$. Then the probability of observing
$c$ is given by the following multinomial distribution:

$\displaystyle p(c\mid x,\omega)$ $\displaystyle=\text{Multi}(m,a(x,\omega))$
(133) $\displaystyle=\frac{m!}{c_{1}!\dots c_{K}!}\prod_{k=1}^{K}\left(p(k\mid
x;\omega)\right)^{c_{k}}$ (134) $\displaystyle=\frac{m!}{c_{1}!\dots
c_{K}!}\prod_{k=1}^{K}\left(\frac{1-\omega}{K}+\omega\delta_{kd}\right)^{c_{k}}.$
(135)

Now consider the fraction $c_{k}/m$ of observations of class $k$ in $c$.
Clearly

$\displaystyle\lim_{m\rightarrow\infty}\frac{c_{k}}{m}=p(k\mid x;\omega),$
(136)

meaning that for any finite $\omega$ it would be possible to deduce from $c$
what the value of $x$ is if $m$ is sufficiently large. However as $\omega$
shrinks, $p(k\mid x;\omega)$ becomes closer to uniform, meaning that a larger
$m$ is required to unambigously identify $x$ from $c$. By defining the
accuracy $\alpha\stackrel{{\scriptstyle\text{def}}}{{=}}m\omega^{2}$ and
sending $m\rightarrow\infty$ (and hence $\omega\rightarrow 0$ for any finite
$\alpha$), $p(c\mid x,\omega)$ can therefore be used to define a continuous-
valued sender distribution that smoothly varies from totally uninformative at
$\alpha=0$ to totally informative as $\alpha\rightarrow\infty$, like the
sender distribution for continuous data.

It can be proved from the central limit theorem that for any set of discrete
probabilities $p=\\{p_{1},\dots,p_{K}\\}$, where $0<p_{k}<1$ $\forall k$, that
if $c\sim\text{Multi}(m,p)$ then in the limit $m\rightarrow\infty$ the
following result holds [8]:

$\displaystyle\frac{c-mp}{\sqrt{mp}}\sim\mathcal{N}\left(0,\boldsymbol{I}\right),$
(137)

where $\boldsymbol{I}$ is the $K\times K$ identity matrix. Therefore

$\displaystyle\lim_{m\rightarrow\infty}p(c_{k}\mid x,\omega)$
$\displaystyle=\mathcal{N}\left(c_{k}\mid mp(k\mid x;\omega),mp(k\mid
x;\omega)\right)$ (138) $\displaystyle=\frac{1}{\sqrt{2\pi mp(k\mid
x;\omega)}}\exp\left(\frac{-\left[c_{k}-mp(k\mid
x,\omega)\right]^{2}}{2mp(k\mid x;\omega)}\right).$ (139)

Now define

$\displaystyle\xi$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}1+\frac{\omega
K}{1-\omega}.$ (140)

And the length $K$ sender sample $y=(y_{1},\dots,y_{K})$ as

$\displaystyle y_{k}$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\left(c_{k}-\frac{m}{K}\right)\ln\xi.$
(141)

Note that $y$, unlike $x$, is continuous
($\operatorname{\mathcal{Y}}=\operatorname{\mathbb{R}}^{K},\operatorname{\mathcal{X}}=\\{1,K\\}$),
and that $\left(c-\frac{m}{K}\right)$ measures the number of times each class
is observed, minus the average number of observations per class. Intuitively,
$y$ provides information about the relative concentration of the classes among
the counts, with (since $\ln\xi>0$) positive values for classes observed more
frequently than the mean and negative values for those observed less
frequently than the mean. As $m\omega^{2}$ grows the concentration increases
around the true class, and hence $y$ become more informative about $x$.

Rearranging Eq. 141,

$\displaystyle c_{k}$ $\displaystyle=\frac{y_{k}}{\ln\xi}+\frac{m}{K}$ (142)
$\displaystyle\implies\frac{dc_{k}}{dy_{k}}$ $\displaystyle=\frac{1}{\ln\xi},$
(143)

which we can use for the following change of variables:

$\displaystyle p(y_{k}\mid x,\omega)$
$\displaystyle=\left|\frac{dc_{k}}{dy_{k}}\right|p(c_{k}\mid x,\omega)$ (144)
$\displaystyle=\frac{1}{\ln\xi\sqrt{2\pi mp(k\mid
x,\omega)}}\exp\left(\frac{-\left[\frac{y_{k}}{\ln\xi}+\frac{m}{K}-mp(k\mid
x,\omega)\right]^{2}}{2mp(k\mid x,\omega)}\right),$ (145)

where we have used the fact that $\xi\geq 1$ and hence
$\frac{dc_{k}}{dy_{k}}\geq 0$. Recall that $\alpha=m\omega^{2}$ and hence
$m=\frac{\alpha}{\omega^{2}}$, which can be substituted into the above to
yield

$\displaystyle p(y_{k}\mid x,\omega)$
$\displaystyle=\frac{1}{\frac{1}{\omega}\ln\xi}\frac{1}{\sqrt{2\pi\alpha
p(k\mid
x,\omega)}}\exp\left(\frac{-\left[\frac{y_{k}}{\frac{1}{\omega}\ln\xi}+\frac{\alpha}{\omega}\left(\frac{1}{K}-p(k\mid
x,\omega)\right)\right]^{2}}{2\alpha p(k\mid x,\omega)}\right).$ (146)

Substituting from Eq. 131,

$\displaystyle\frac{1}{K}-p(k\mid
x,\omega)=\omega\left(\frac{1}{K}-\delta_{kx},\right),$ (147)

and hence

$\displaystyle p(y_{k}\mid x,\omega)$
$\displaystyle=\frac{1}{\frac{1}{\omega}\ln\xi}\frac{1}{\sqrt{2\pi\alpha
p(k\mid
x,\omega)}}\exp\left(\frac{-\left[\frac{y_{k}}{\frac{1}{\omega}\ln\xi}-\alpha\left(\delta_{kx}-\frac{1}{K}\right)\right]^{2}}{2\alpha
p(k\mid x,\omega)}\right).$ (148)

Applying the identity $\ln(1+x)=\sum_{n=1}^{\infty}\frac{(-1)^{n-1}}{n}x^{n}$
for $|x|<1$ to $\ln\xi=\ln\left(1+\frac{\omega K}{1-\omega}\right)$ it can be
seen that

$\displaystyle\ln\xi$ $\displaystyle\in\frac{\omega
K}{1-\omega}+O(\omega^{2}),$ (149)

and hence

$\displaystyle\lim_{\omega\rightarrow 0}\frac{1}{\omega}\ln\xi$
$\displaystyle=K.$ (150)

Furthermore, it follows directly from Eq. 131 that

$\displaystyle\lim_{\omega\rightarrow 0}p(k\mid x,\omega)=\frac{1}{K}\ \forall
k\in\\{1,K\\}.$ (151)

Now define

$\displaystyle p_{{}_{S}}\left(y_{k}\mid
x;\alpha\right)\stackrel{{\scriptstyle\text{def}}}{{=}}\lim_{\omega\rightarrow
0}p(y_{k}\mid x,\omega).$ (152)

Plugging Eq. 150 and 151 into Eq. 148,

$\displaystyle p_{{}_{S}}\left(y_{k}\mid x;\alpha\right)$
$\displaystyle=\frac{1}{K\sqrt{2\pi\alpha\frac{1}{K}}}\exp\left(\frac{-\left[\frac{y_{k}}{K}-\alpha\left(\delta_{kx}-\frac{1}{K}\right)\right]^{2}}{2\alpha\frac{1}{K}}\right)$
(153) $\displaystyle=\frac{1}{\sqrt{2\pi\alpha
K}}\exp\left(\frac{-\left[y_{k}-\alpha\left(K\delta_{kx}-1\right)\right]^{2}}{2\alpha
K}\right)$ (154)
$\displaystyle=\mathcal{N}\left(\alpha\left(K\delta_{kx}-1\right),\alpha
K\right).$ (155)

Restoring the superscript,

$\displaystyle p_{{}_{S}}\left(y^{(d)}\mid x^{(d)};\alpha\right)$
$\displaystyle=\mathcal{N}\left(\alpha\left(K\mathbf{e}_{x^{(d)}}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right),$ (156)

where $\boldsymbol{1}$ is a vector of ones, $\boldsymbol{I}$ is the identity
matrix and $\mathbf{e}_{j}\in\operatorname{\mathbb{R}}^{K}$ is the projection
from the class index $j$ to the length $K$ one-hot vector defined by
$(\mathbf{e}_{j})_{k}=\delta_{jk}$, and therefore

$\displaystyle
p_{{}_{S}}\left(\operatorname{\mathbf{y}}\mid\operatorname{\mathbf{x}};\alpha\right)=\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\alpha\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right),$ (157)

where
$\mathbf{e}_{\operatorname{\mathbf{x}}}\stackrel{{\scriptstyle\text{def}}}{{=}}\left(\mathbf{e}_{x^{(1)}},\dots,\mathbf{e}_{x^{(D)}}\right)\in\operatorname{\mathbb{R}}^{KD}$.

###  6.4 Receiver Distribution
$p_{{}_{R}}(\cdot\mid\boldsymbol{\theta};t,\alpha)$

Substituting Eq. 127 and Eq. 157 into Eq. 4 gives the following receiver
distribution for dimension $d$:

$\displaystyle p_{{}_{R}}^{(d)}(y^{(d)}\mid\boldsymbol{\theta};t,\alpha)$
$\displaystyle=\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\mathcal{N}\left(\alpha\left(K\mathbf{e}_{k}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right),$ (158) $\displaystyle
p_{{}_{R}}(\operatorname{\mathbf{y}}\mid\boldsymbol{\theta};t,\alpha)$
$\displaystyle=\prod_{d=1}^{D}p_{{}_{R}}^{(d)}(y^{(d)}\mid\boldsymbol{\theta};t,\alpha).$
(159)

###  6.5 Bayesian Update Function
$h(\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{y}},\alpha)$

Recall from Section 6.1 that $(\theta_{i-1})^{(d)}_{k}$ is the probability
assigned to $x^{(d)}=k$ by $p(x^{(d)}\mid\theta_{i-1})$. Dropping the
superscript and returning to the count distribution $p(c\mid x,\omega)$
defined in Eq. 133, the posterior probability that $x=k$ after observing $c$
is

$\displaystyle p(k\mid c;\omega)$ $\displaystyle=\frac{p(c\mid
k;\omega)(\theta_{i-1})_{k}}{\sum_{k^{\prime}=1}^{K}p(c\mid
k^{\prime};\omega)(\theta_{i-1})_{k^{\prime}}}.$ (160)

Substituting Eq. 135 into Eq. 160 and cancelling terms in the enumerator and
denominator,

$\displaystyle p(k\mid c;\omega)$
$\displaystyle=\frac{\left[\frac{1-\omega}{K}\right]^{m-c_{k}}\left[\frac{1-\omega}{K}+\omega\right]^{c_{k}}(\theta_{i-1})_{k}}{\sum_{k^{\prime}=1}^{K}{\left[\frac{1-\omega}{K}\right]^{m-c_{k^{\prime}}}\left[\frac{1-\omega}{K}+\omega\right]^{c_{k^{\prime}}}(\theta_{i-1})_{k^{\prime}}}}$
(161)
$\displaystyle=\frac{\left[\frac{1-\omega}{K}\right]^{m}\left[1+\frac{\omega
K}{1-\omega}\right]^{c_{k}}(\theta_{i-1})_{k}}{\left[\frac{1-\omega}{K}\right]^{m}\sum_{k^{\prime}=1}^{K}{\left[1+\frac{\omega
K}{1-\omega}\right]^{c_{k^{\prime}}}(\theta_{i-1})_{k^{\prime}}}}$ (162)
$\displaystyle=\frac{\left[1+\frac{\omega
K}{1-\omega}\right]^{c_{k}}(\theta_{i-1})_{k}}{\sum_{k^{\prime}=1}^{K}{\left[1+\frac{\omega
K}{1-\omega}\right]^{c_{k^{\prime}}}(\theta_{i-1})_{k^{\prime}}}}$ (163)
$\displaystyle=\frac{\xi^{c_{k}}(\theta_{i-1})_{k}}{\sum_{k^{\prime}=1}^{K}{\xi^{c_{k^{\prime}}}(\theta_{i-1})_{k^{\prime}}}}.$
(164)

Now define

$\displaystyle h(\theta,y)$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{e^{y}\theta}{\sum_{k=1}^{K}e^{y_{k}}\theta_{k}}.$
(165)

Substituting the definition of $y_{k}$ from Eq. 141 into the definition of
$h(\theta,y)$ from Eq. 165,

$\displaystyle\left(h(\theta_{i-1},y)\right)_{k}$
$\displaystyle=\frac{\exp(-\frac{m}{K}\ln\xi)\exp(c_{k}\ln\xi)(\theta_{i-1})_{k}}{\exp(-\frac{m}{K}\ln\xi)\sum_{k^{\prime}=1}^{K}\exp(c_{k^{\prime}}\ln\xi)(\theta_{i-1})_{k^{\prime}}}$
(166)
$\displaystyle=\frac{\exp(\ln\xi^{c_{k}})(\theta_{i-1})_{k}}{\sum_{k^{\prime}=1}^{K}\exp(\ln\xi^{c_{k^{\prime}}})(\theta_{i-1})_{k^{\prime}}}$
(167)
$\displaystyle=\frac{\xi^{c_{k}}(\theta_{i-1})_{k}}{\sum_{k^{\prime}=1}^{K}\xi^{c_{k^{\prime}}}(\theta_{i-1})_{k^{\prime}}},$
(168)

and hence, from Eq. 164,

$\displaystyle h(\theta_{i-1},y)_{k}=p(k\mid c;\omega).$ (170)

Therefore in the limit $m\rightarrow\infty$ with $m\omega^{2}=\alpha$, the
stochastic parameter update from $\theta_{i-1}$ to $\theta_{i}$ induced by
drawing $c$ from $\text{multi}(m,a(x,\omega))$ can be sampled by first drawing
$y$ from $p_{{}_{S}}\left(\cdot\mid x,\alpha\right)$ then setting
$\theta_{i}=h(\theta_{i-1},y)$. Hence the Bayesian update function is

$\displaystyle
h(\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{y}},\alpha)\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{e^{\operatorname{\mathbf{y}}}\boldsymbol{\theta}_{i-1}}{\sum_{k=1}^{K}e^{\operatorname{\mathbf{y}}_{k}}(\boldsymbol{\theta}_{i-1})_{k}},$
(171)

where the redundant parameter $\alpha$ has been included for consistency with
the update function for continuous data.

###  6.6 Bayesian Update Distribution
$p_{{}_{U}}(\cdot\mid\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{x}};\alpha)$

Substituting Eqs. 157 and 171 into Eq. 6,

$\displaystyle
p_{{}_{U}}(\boldsymbol{\theta}\mid\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{x}};\alpha)$
$\displaystyle=\operatorname*{\mathbb{E}}_{\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\alpha\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right)}\delta\left(\boldsymbol{\theta}-\frac{e^{\operatorname{\mathbf{y}}}\boldsymbol{\theta}_{i-1}}{\sum_{k=1}^{K}e^{\operatorname{\mathbf{y}}_{k}}(\boldsymbol{\theta}_{i-1})_{k}}\right).$
(172)

###  6.7 Additive Accuracies

It follows from the definition of the update distribution that if $y_{a}$ is
drawn from $p_{{}_{S}}\left(\cdot\mid x;\alpha_{a}\right)$ then
$\theta_{i-1}=h(y_{a},\theta_{i-2})$ is drawn from
$p(\cdot\mid\theta_{i-2},x;\alpha_{a})$. Furthermore, if $y_{b}$ is drawn from
$p_{{}_{S}}\left(\cdot\mid x;\alpha_{b}\right)$ then
$\theta_{i}=h(y_{b},\theta_{i-1})=h(y_{b},h(y_{a},\theta_{i-2}))$ is drawn
from
$\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\theta_{i-1}\mid\theta_{i-2},x;\alpha_{a})}p_{{}_{U}}(\theta_{i}\mid\theta_{i-1},x;\alpha_{b})$.
Substituting the definition of $h$ from Eqn 165,

$\displaystyle h(y_{b},h(y_{a},\theta_{i-2}))$
$\displaystyle=\frac{\exp(y_{b})\frac{\exp(y_{a})\theta_{i-2}}{\sum_{k^{\prime}=1}^{K}\exp\left((y_{a})_{k^{\prime}}\right)(\theta_{i-2})_{k^{\prime}}}}{\sum_{k=1}^{K}\exp\left((y_{b})_{k}\right)\frac{\exp\left((y_{a})_{k}\right)(\theta_{i-2})_{k}}{\sum_{k^{\prime}=1}^{K}\exp\left((y_{a})_{k^{\prime}}\right)(\theta_{i-2})_{k^{\prime}}}}$
(173)
$\displaystyle=\frac{\exp(y_{b})\exp(y_{a})\theta_{i-2}}{\sum_{k=1}^{K}\exp\left((y_{b})_{k}\right)\exp\left((y_{a})_{k}\right)(\theta_{i-2})_{k}}$
(174)
$\displaystyle=\frac{\exp(y_{a}+y_{b})\theta_{i-2}}{\sum_{k=1}^{K}\exp\left((y_{a}+y_{b})_{k}\right)(\theta_{i-2})_{k}}$
(175) $\displaystyle=h(y_{a}+y_{b},\theta_{i-2}).$ (176)

From Eqn. 156

$\displaystyle y_{a}$
$\displaystyle\sim\mathcal{N}\left(\alpha_{a}\left(K\mathbf{e}_{x}-\boldsymbol{1}\right),\alpha_{a}K\boldsymbol{I}\right),$
(177) $\displaystyle y_{b}$
$\displaystyle\sim\mathcal{N}\left(\alpha_{b}\left(K\mathbf{e}_{x}-\boldsymbol{1}\right),\alpha_{b}K\boldsymbol{I}\right)$
(178)

and hence, from Identity 61

$\displaystyle y_{a}+y_{b}$
$\displaystyle\sim\mathcal{N}\left((\alpha_{a}+\alpha_{b})\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),(\alpha_{a}+\alpha_{b})K\boldsymbol{I}\right).$
(180)

Therefore, if $y$ is drawn from $p_{{}_{S}}\left(\cdot\mid
x;\alpha_{a}+\alpha_{b}\right)$ and $\theta_{i}=h(y,\theta_{i-2})$ then
$\theta_{i}$ is drawn from
$\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\theta_{i-1}\mid\theta_{i-2},x;\alpha_{a})}p_{{}_{U}}(\theta_{i}\mid\theta_{i-1},x;\alpha_{b})$
and

$\displaystyle\operatorname*{\mathbb{E}}_{p_{{}_{U}}(\boldsymbol{\theta}_{i-1}\mid\boldsymbol{\theta}_{i-2},\operatorname{\mathbf{x}};\alpha_{a})}p_{{}_{U}}(\boldsymbol{\theta}_{i}\mid\boldsymbol{\theta}_{i-1},\operatorname{\mathbf{x}};\alpha_{b})=p_{{}_{U}}(\boldsymbol{\theta}_{i}\mid\boldsymbol{\theta}_{i-2},\operatorname{\mathbf{x}};\alpha_{a}+\alpha_{b}),$
(181)

as required.

###  6.8 Accuracy Schedule $\beta(t)$

As with continuous data, the guiding heuristic for $\beta(t)$ was to decrease
the expected entropy of the input distribution linearly with $t$. In the
continuous case, where the entropy is a deterministic function of
$\sigma^{2}$, applying the heuristic was straightforward; in the discrete case
an explicit computation of
$\operatorname*{\mathbb{E}}_{p_{{}_{F}}(\boldsymbol{\theta}\mid
x;t)}H\left[p_{{}_{I}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta})\right]$
would be needed. We were unable to derive an analytic expression for this
term, but found that

$\displaystyle\beta(t)=t^{2}\beta(1)$ (182)

was a reasonable approximation, with $\beta(1)$ determined empirically for
each experiment. Therefore

$\displaystyle\alpha(t)=\frac{d\beta(t)}{dt}=\beta(1)2t.$ (183)

Figure 9: Accuracy schedule vs. expected entropy for discrete data. The
surface plot shows the expectation over the parameter distribution
$p(\theta\mid x;\beta)$ of the entropy of the categorical input distribution
$p(x\mid\theta)$ for $K=2$ to $30$ and $\sqrt{\beta}=0.01$ to $3$. The red and
cyan lines highlight the entropy curves for 2 and 27 classes, the two values
that occur in our experiments. The red and cyan stars show the corresponding
values we chose for $\sqrt{\beta(1)}$.

###  6.9 Bayesian Flow Distribution
$p_{{}_{F}}(\cdot\mid\operatorname{\mathbf{x}};t)$

Substituting Eq. 172 into Eq. 10,

$\displaystyle p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t)$
$\displaystyle=\operatorname*{\mathbb{E}}_{\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\beta(t)\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\beta(t)K\boldsymbol{I}\right)}\delta\left(\boldsymbol{\theta}-\frac{e^{\operatorname{\mathbf{y}}}\boldsymbol{\theta}_{0}}{\sum_{k=1}^{K}e^{\operatorname{\mathbf{y}}_{k}}(\boldsymbol{\theta}_{0})_{k}}\right).$
(184)

Since the prior is uniform with
$\boldsymbol{\theta}_{0}=\boldsymbol{\frac{1}{K}}$, this reduces to

$\displaystyle p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t)$
$\displaystyle=\operatorname*{\mathbb{E}}_{\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\beta(t)\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\beta(t)K\boldsymbol{I}\right)}\delta\left(\boldsymbol{\theta}-\text{softmax}(\operatorname{\mathbf{y}})\right),$
(185)

which can be sampled by drawing $\operatorname{\mathbf{y}}$ from
$\mathcal{N}\left(\beta(t)\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\beta(t)K\boldsymbol{I}\right)$
then setting $\boldsymbol{\theta}=\text{softmax}(\operatorname{\mathbf{y}})$.

The sender distribution for discrete data can therefore be interpreted as a
source of softmax logits for the Bayesian flow distribution; the higher the
sender accuracy $\alpha$ is, the larger in expectation the logits
corresponding to $\operatorname{\mathbf{x}}$ will be in
$\operatorname{\mathbf{y}}$, hence the closer $\boldsymbol{\theta}$ will be to
$\mathbf{e}_{\operatorname{\mathbf{x}}}$ and the more information the network
will gain about $\operatorname{\mathbf{x}}$.

Figure 10: Bayesian flow for discrete data. For $K=3$, the input distribution
parameters $\boldsymbol{\theta}=(\theta_{1},\theta_{2},\theta_{3})$ can be
visualised as points on the 2-simplex, with the data $x$ corresponding to the
bottom left corner. For the accuracy schedule $\beta(t)$ from Eq. 182, the
white line shows a single input parameter trajectory starting from
$\boldsymbol{\theta}_{0}=\left(\frac{1}{3},\frac{1}{3},\frac{1}{3}\right)$ and
evolving under the Bayesian update distribution
$p_{{}_{U}}(\boldsymbol{\theta}_{i}\mid\boldsymbol{\theta}_{i-1};x,\beta(t_{i})-\beta(t_{i-1}))$
from Eq. 172, superimposed on log-scale heatmaps of the Bayesian flow
distribution $p_{{}_{F}}(\boldsymbol{\theta}\mid x;t)$ from Eq. 185, plotted
at regular intervals from $t=0.02$ to $1$. Figure 11: Bayesian flow for binary
data. For the input probability $p_{1}$ of class one, the plot shows several
parameter trajectories starting from $p_{1}=0.5$ at $t=0$ and evolving under
the Bayesian update distribution to $t=1$, superimposed on a log-scale heatmap
of the Bayesian flow distribution. $\beta(1)=4$ in this plot. Note that both
here and in Figure 10 the convergence towards the data appears slower and
noisier than the equivalent trajectories for continuous data in Figure 4. This
is a fundamental consequence of discreteness: since all points in
$\operatorname{\mathcal{X}}$ are equidistant the input distributions cannot
concentrate on values close to $\operatorname{\mathbf{x}}$ as the trajectories
progress.

###  6.10 Reconstruction Loss $L^{r}(\operatorname{\mathbf{x}})$

The reconstruction loss for discrete data is

$\displaystyle
L^{r}(\operatorname{\mathbf{x}})=-\operatorname*{\mathbb{E}}_{p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}},1)}\ln
p_{{}_{O}}(\operatorname{\mathbf{x}}\mid\boldsymbol{\theta};1).$ (186)

###  6.11 Discrete-time Loss $L^{n}(\operatorname{\mathbf{x}})$

From Eqs. 156 and 158,

$\displaystyle D_{KL}\left(p_{{}_{S}}\left(\cdot\mid
x^{(d)};\alpha\right)\parallel
p_{{}_{R}}^{(d)}(\cdot\mid\boldsymbol{\theta};t,\alpha)\right)$ (187)
$\displaystyle\qquad=D_{KL}\left(\mathcal{N}\left(\alpha\left(K\mathbf{e}_{x^{(d)}}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right)\parallel\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\mathcal{N}\left(\alpha\left(K\mathbf{e}_{k}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right)\right).$ (188)

Therefore, substituting into Eq. 24,

$\displaystyle
L^{n}(\operatorname{\mathbf{x}})=n\operatorname*{\mathbb{E}}_{i\sim
U\\{1,n\\},p(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}};t_{i-1}),\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\alpha_{i}\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\alpha_{i}K\boldsymbol{I}\right)}\ln\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\alpha_{i}\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\alpha_{i}K\boldsymbol{I}\right)$
(189)
$\displaystyle\qquad\qquad\qquad-\sum_{d=1}^{D}\ln\left(\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t_{i-1})\mathcal{N}\left(y^{(d)}\mid\alpha_{i}\left(K\mathbf{e}_{k}-\boldsymbol{1}\right),\alpha_{i}K\boldsymbol{I}\right)\right),$
(190)

where, from Eq. 182,

$\displaystyle\alpha_{i}$ $\displaystyle=\beta(t_{i})-\beta(t_{i-1})$ (191)
$\displaystyle=\beta(1)\left(\left(\frac{i}{n}\right)^{2}-\left(\frac{i-1}{n}\right)^{2}\right)$
(192) $\displaystyle=\beta(1)\left(\frac{2i-1}{n^{2}}\right).$ (193)

###  6.12 Continuous-time Loss $L^{\infty}(\operatorname{\mathbf{x}})$

Let

$\displaystyle\operatorname{\mathbf{v}}\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{\operatorname{\mathbf{y}}}{\alpha}+1,$
(194)

and apply Identity 51 to see that if

$\displaystyle y^{(d)}\sim p_{{}_{S}}\left(\cdot\mid
x^{(d)};\alpha\right)=\mathcal{N}\left(\alpha(K\mathbf{e}_{x^{(d)}}-\boldsymbol{1}),\alpha
K\boldsymbol{I}\right),$ (195)

then

$\displaystyle
v^{(d)}\sim\mathcal{N}\left(K\mathbf{e}_{x^{(d)}},\frac{K}{\alpha}\boldsymbol{I}\right),$
(196)

and similarly if

$\displaystyle y^{(d)}\sim
p_{{}_{R}}^{(d)}(\cdot\mid\boldsymbol{\theta};t,\alpha)=\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\mathcal{N}\left(y^{(d)}\mid\alpha\left(K\mathbf{e}_{k}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right),$ (197)

then

$\displaystyle v^{(d)}$
$\displaystyle\sim\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\mathcal{N}\left(K\mathbf{e}_{k},\frac{K}{\alpha}\boldsymbol{I}\right)$
(198)
$\displaystyle=K\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\delta(\cdot-\mathbf{e}_{k})\ast\mathcal{N}\left(\boldsymbol{0},\frac{K}{\alpha}\boldsymbol{I}\right).$
(199)

The Kullback-Leibler divergence is invariant under affine transformations of
variables, hence

$\displaystyle D_{KL}\left(p_{{}_{S}}\left(\cdot\mid
x^{(d)};\alpha\right)\parallel
p_{{}_{R}}^{(d)}(\cdot\mid\boldsymbol{\theta};t,\alpha_{i})\right)$ (200)
$\displaystyle\qquad=D_{KL}\left(\mathcal{N}\left(K\mathbf{e}_{x^{(d)}},\frac{K}{\alpha}\boldsymbol{I}\right)\parallel\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)K\delta(\cdot-\mathbf{e}_{k})\ast\mathcal{N}\left(\boldsymbol{0},\frac{K}{\alpha}\boldsymbol{I}\right)\right).$
(201)

Now set $C=K$, $g(x^{(d)})=K\mathbf{e}_{x^{(d)}}$ and

$\displaystyle
P^{(d)}(\boldsymbol{\theta},t)=K\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\delta(\cdot-\mathbf{e}_{k}),$
(202)

which has finite variance and the following finite expectation

$\displaystyle
E[P^{(d)}(\boldsymbol{\theta},t)]=K\mathbf{\hat{e}}^{(d)}(\boldsymbol{\theta},t),$
(203)

where

$\displaystyle\mathbf{\hat{e}}^{(d)}(\boldsymbol{\theta},t)\stackrel{{\scriptstyle\text{def}}}{{=}}\sum_{k=1}^{K}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\mathbf{e}_{k}.$
(204)

The conditions in Eq. 29 are therefore satisfied and Eqs. 203 and 183 can be
substituted into Eq. 41 to yield

$\displaystyle
L^{\infty}(\operatorname{\mathbf{x}})=K\beta(1)\operatorname*{\mathbb{E}}_{t\sim
U(0,1),p_{{}_{F}}(\boldsymbol{\theta}\mid\operatorname{\mathbf{x}},t)}t\|\mathbf{e}_{\operatorname{\mathbf{x}}}-\mathbf{\hat{e}}(\boldsymbol{\theta},t)\|^{2},$
(205)

where

$\displaystyle\mathbf{\hat{e}}(\boldsymbol{\theta},t)\stackrel{{\scriptstyle\text{def}}}{{=}}\left(\mathbf{\hat{e}}^{(1)}(\boldsymbol{\theta},t),\dots,\mathbf{\hat{e}}^{(D)}(\boldsymbol{\theta},t)\right).$
(206)

###  6.13 Pseudocode

Pseudocode for evaluating the discrete-time loss
$L^{n}(\operatorname{\mathbf{x}})$ and continuous-time loss
$L^{\infty}(\operatorname{\mathbf{x}})$ for discrete data is presented in
Algorithms 7 and 8, while sample generation is presented in Algorithm 9.

function
discrete_output_distribution($\boldsymbol{\theta}\in[0,1]^{KD},t\in[0,1]$)

Input $(\boldsymbol{\theta},t)$ to network, receive
$\Psi(\boldsymbol{\theta},t)$ as output

for $d\in\\{1,D\\}$ do

if $k=2$ then

$p_{{}_{O}}^{(d)}(1\mid\boldsymbol{\theta};t)\leftarrow\sigma\left(\Psi^{(d)}(\boldsymbol{\theta},t)\right)$

$p_{{}_{O}}^{(d)}(2\mid\boldsymbol{\theta};t)\leftarrow
1-p_{{}_{O}}^{(d)}(1\mid\boldsymbol{\theta};t)$

else

$p_{{}_{O}}^{(d)}(\cdot\mid\boldsymbol{\theta};t)\leftarrow\text{softmax}(\Psi^{(d)}(\boldsymbol{\theta},t))$

end if

end for

Return $\boldsymbol{p}_{{}_{O}}(\cdot\mid\boldsymbol{\theta};t)$

end function

Algorithm 7 Discrete-Time Loss $L^{n}(\operatorname{\mathbf{x}})$ for Discrete
Data

Require: $\beta(1)\in\operatorname{\mathbb{R}}^{+}$, number of steps
$n\in\mathbb{N}$, number of classes $K\in\mathbb{N}$

Input: discrete data $\operatorname{\mathbf{x}}\in\\{1,K\\}^{D}$

$i\sim U\\{1,n\\}$

$t\leftarrow(i-1)/n$

$\beta\leftarrow\beta(1)t^{2}$

$\operatorname{\mathbf{y}}^{\prime}\sim\mathcal{N}\left(\beta\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\beta
K\boldsymbol{I}\right)$

$\boldsymbol{\theta}\leftarrow\text{softmax}(\operatorname{\mathbf{y}}^{\prime})$

${\boldsymbol{p}_{{}_{O}}(\cdot\mid\boldsymbol{\theta};t)\leftarrow\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{discrete\textunderscore
output\textunderscore distribution}}}}}}}(\boldsymbol{\theta},t)$

$\alpha\leftarrow\beta(1)\left(\frac{2i-1}{n^{2}}\right)$

$\operatorname{\mathbf{y}}\sim\mathcal{N}\left(\alpha\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right)$

$L^{n}(\operatorname{\mathbf{x}})\leftarrow
n\left[\ln\mathcal{N}\left(\operatorname{\mathbf{y}}\mid\alpha\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right)-\sum_{d}\ln\left(\sum_{k}p_{{}_{O}}^{(d)}(k\mid\boldsymbol{\theta};t)\mathcal{N}\left(y^{(d)}\mid\alpha\left(K\mathbf{e}_{k}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right)\right)\right]$

Algorithm 8 Continuous-Time Loss $L^{\infty}(\operatorname{\mathbf{x}})$ for
Discrete Data

Require: $\beta(1)\in\operatorname{\mathbb{R}}^{+}$, number of classes
$K\in\mathbb{N}$

Input: discrete data $\operatorname{\mathbf{x}}\in\\{1,K\\}^{D}$

$t\sim U(0,1)$

$\beta\leftarrow\beta(1)t^{2}$

$\operatorname{\mathbf{y}}\sim\mathcal{N}\left(\beta\left(K\mathbf{e}_{\operatorname{\mathbf{x}}}-\boldsymbol{1}\right),\beta
K\boldsymbol{I}\right)$

$\boldsymbol{\theta}\leftarrow\text{softmax}(\operatorname{\mathbf{y}})$

${\boldsymbol{p}_{{}_{O}}(\cdot\mid\boldsymbol{\theta};t)\leftarrow\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{discrete\textunderscore
output\textunderscore distribution}}}}}}}(\boldsymbol{\theta},t)$

$\mathbf{\hat{e}}(\boldsymbol{\theta},t)\leftarrow\left(\sum_{k}p_{{}_{O}}^{(1)}(k\mid\boldsymbol{\theta};t)\mathbf{e}_{k},\dots,\sum_{k}p_{{}_{O}}^{(D)}(k\mid\boldsymbol{\theta};t)\mathbf{e}_{k}\right)$

$L^{\infty}(\operatorname{\mathbf{x}})\leftarrow
K\beta(1)t\left\|\mathbf{e}_{\operatorname{\mathbf{x}}}-\mathbf{\hat{e}}(\boldsymbol{\theta},t)\right\|^{2}$

Algorithm 9 Sample Generation for Discrete Data

Require: $\beta(1)\in\operatorname{\mathbb{R}}^{+}$, number of steps
$n\in\mathbb{N}$, number of classes $K\in\mathbb{N}$

$\boldsymbol{\theta}\leftarrow\left(\boldsymbol{\frac{1}{K}}\right)$

for $i=1$ to $n$ do

$t\leftarrow\frac{i-1}{n}$

${\mathbf{k}\sim\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{discrete\textunderscore
output\textunderscore distribution}}}}}}}(\boldsymbol{\theta},t)$

$\alpha\leftarrow\beta(1)\left(\frac{2i-1}{n^{2}}\right)$

$\operatorname{\mathbf{y}}\sim\mathcal{N}\left(\alpha\left(K\mathbf{e}_{\mathbf{k}}-\boldsymbol{1}\right),\alpha
K\boldsymbol{I}\right)$

$\boldsymbol{\theta}^{\prime}\leftarrow
e^{\operatorname{\mathbf{y}}}\boldsymbol{\theta}$

$\boldsymbol{\theta}\leftarrow\frac{\boldsymbol{\theta}^{\prime}}{\sum_{k}\boldsymbol{\theta}^{\prime}_{k}}$

end for

${\mathbf{k}\sim\text{\sc{\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor{\@listingGroup{ltx_lst_identifier}{discrete\textunderscore
output\textunderscore distribution}}}}}}}(\boldsymbol{\theta},1)$

Return $\mathbf{k}$

##  7 Experiments

We evaluated Bayesian Flow Networks (BFNs) on the following generative
benchmarks: CIFAR-10 (32$\times$32 8-bit color images), dynamically binarized
MNIST (28$\times$28 binarized images of handwritten digits) and text8 (length
256 character sequences with a size 27 alphabet). The continuous (Sec. 4) and
discretised (Sec. 5) versions of the system were compared on CIFAR-10, while
the discrete version (Sec. 6) was applied to the other datasets. In all cases,
the network was trained using the continuous-time loss
$L^{\infty}(\operatorname{\mathbf{x}})$, with the discrete-time loss
$L^{n}(\operatorname{\mathbf{x}})$ evaluated for testing only, with various
values of $n$. Standard network architectures and training algorithms were
used throughout to allow for direct comparison with existing methods. Because
the focus of this paper is on probabilistic modelling rather than image
generation, FID scores were not calculated. However, examples of generated
data are provided for all experiments.

Model | Dynamically Binarized MNIST | CIFAR-10
---|---|---
Improved DDPM [28] |  | 2.94
NVAE [49] | 78.01 | 2.91
PixelVAE++† [35] | 78.00 | 2.90
Locally Masked PixelCNN† [15] | 77.58 | 2.89
Image Transformer† [30] |  | 2.89
DDPM++ [16] |  | 2.88
LSGM [50] |  | 2.87
VDVAE [3] |  | 2.87
Sparse Transformer† [4] |  | 2.80
Reflected Diffusion [23] |  | 2.68
VDM [17] |  | 2.65
ARDM-Upscale 4 [13] |  | 2.64
BFN | 77.87 | 2.66
CR-NVAE* [40] | 76.93 | 2.51
VDM* [17] |  | 2.49
Table 1: Comparison of dynamically binarized MNIST and CIFAR-10 results with
other methods. The best published results for both datasets (*) use data
augmentation for regularization. Results for models marked with (†) are exact
values; all other results are upper bounds.

###  7.1 Dynamically Binarized MNIST

$n$-steps | 10 | 25 | 50 | 100 | 784 | 1000 | $\infty$
---|---|---|---|---|---|---|---
NPI | $95.21$ | $84.40$ | $81.06$ | $79.46$ | $78.02$ | $78.07$ | $77.87$
Table 2: Dynamically binarized MNIST results. NPI is nats per image averaged
over 2,000 passes through the test set with $L^{n}(\operatorname{\mathbf{x}})$
or $L^{\infty}(\operatorname{\mathbf{x}})$ sampled once per test image per
pass. The reconstruction loss $L^{r}(\operatorname{\mathbf{x}})$ (included in
NPI) was $0.46$. 784 is the total number of pixels per image, hence the number
of steps required to generate an image with an autoregressive model.

Data. The binarized MNIST benchmark data was originally created from the MNIST
dataset of handwritten images [20] by treating the grayscale pixel intensities
as Bernoulli probabilities and sampling a particular binarization [36] which
is held fixed during training. In recent years, a variant of the same
benchmark has become more popular, with a new binarization sampled from the
probabilities for every training batch. The two are not comparable, as the
latter, which we refer to as dynamically binarized MNIST, effectively has a
larger training set and hence gives better test set performance. All our
experiments and the results referenced from the literature use dynamically
binarized MNIST.

(a) Test Data

(b) Generated Data

Figure 12: MNIST real and generated data. Samples generated with 100 steps.

Setup. The network architecture was based on a U-Net introduced for diffusion
models [28]. Starting from the hyperparameters used for the CIFAR-10 dataset
(see Appendix A in the above reference), we made the following modifications:
the number of resblocks was reduced from three to two and the layer widths
were reduced from $[C,2C,2C,2C]$ to $[C,2C,2C]$ with $C=128$. Finally, the
input and output of the standard network were concatenated and projected back
to the output size. 600 randomly selected training images (1% of the training
set) were used as a validation set. The optimiser was AdamW [22] with learning
rate $0.0001$, weight decay 0.01 and $(\beta_{1},\beta_{2})=(0.9,0.98)$.
Dropout was used with probability 0.5, the training batch size was 512, and
$\beta(1)$ was set to $3$ (see Sec. 6.8). The network was trained for
$150\,000$ weight updates until early stopping. An exponential moving average
of model parameters with a decay rate of 0.9999 was used for evaluation and
sample generation. The total number of learnable parameters was approximately
25M.

(a) Input Distribution

(b) Output Distribution

Figure 13: MNIST Input and output distributions. For two test set images the
figure shows the white pixel probability at 20 steps evenly spaced between
$t=0$ and $t=1/3$. Note how the input probabilities are initially uniform
whereas the output distribution initially predicts a superposition of multiple
digits, closely matching the per-pixel marginal prior over the training set:
this supports our belief that the network learns to correct for the uniform
prior in the input distribution. Also note that the output distribution is
much less noisy than the input distribution, and that it changes more
dramatically as new information is received (e.g. the network appears to
switch from predicting a $6$ to a $2$ to a $7$ for the first image). This
highlights the network’s use of context to resolve ambiguity and noise in the
input distribution.

Results. As can be seen from Table 1, BFN is close to state-of-the-art for
this task with no data augmentation. Table 2 shows the expected inverse
relationship between loss and number of steps. Direct optimisation of the
$n$-step loss would likely lead to reduced loss for low values of $n$; however
we leave that for future work. One issue is that the reconstruction loss was
relatively high at 0.46 nats per image. The obvious way to decrease this would
be to increase $\beta(1)$, but we found that doing so led to slower learning
and worse performance. Along with the loss curves in Figure 14, this suggests
that the accuracy schedule is suboptimal for binary data.

Figure 14: MNIST losses against time. The left plot shows the mean over the
test set of the cts. time loss $L^{\infty}(\operatorname{\mathbf{x}})$ used
for training for transmission time $t$ between 0 and 1. The right plot shows
the average cumulative value of $L^{\infty}(\operatorname{\mathbf{x}})$ up to
$t$, along with the reconstruction loss $L^{r}(\operatorname{\mathbf{x}})$
evaluated at $t$ and the sum of these two losses, which would be the total
loss if the transmission process halted at $t$. Note the unevenness of
$L^{\infty}(\operatorname{\mathbf{x}})$ against $t$: we speculate that
rescaling $\beta(t)$ to make the loss curve more uniform could improve
performance.

###  7.2 CIFAR-10

$n$-steps | Cts. (256 bins) | Discd. (256 bins) | Cts. (16 bins) | Discd. (16 bins)
---|---|---|---|---
10 | 6.18 | 3.91 | 1.42 | 1.16
25 | 3.65 | 3.16 | 1.11 | 1.02
50 | 3.10 | 2.93 | 1.03 | 0.98
100 | 2.86 | 2.81 | 0.99 | 0.96
250 | 2.73 | 2.73 | 0.97 | 0.94
500 | 2.69 | 2.71 | 0.96 | 0.94
1000 | 2.67 | 2.70 | 0.96 | 0.94
$\infty$ | 2.66 | 2.68 | 0.96 | 0.94
$L^{r}(\operatorname{\mathbf{x}})$ | 0.001 | 0.003 | 0.073 | 0.070
Updates | 5M | 5M | 250K | 1M
Table 3: CIFAR-10 results. All losses are bits per dimension (BPD) averaged
over 100 passes through the test set with $L^{n}(\operatorname{\mathbf{x}})$
or $L^{\infty}(\operatorname{\mathbf{x}})$ sampled once per test image per
pass. The reconstruction losses $L^{r}(\operatorname{\mathbf{x}})$ (included
in BPD) and the number of training updates for each network are shown below.

Data. Two sets of generative modelling experiments were conducted on the
CIFAR-10 database [19], one at the standard bit-depth of 8, corresponding to
256 discretised bins per colour channel, and one at a reduced bit-depth of 4,
corresponding to $16$ bins per channel. In both cases the bins evenly
partitioned the interval $[-1,1]$ and the data was pre-processed by assigning
each channel intensity to the nearest bin centre, as described in Section 5.
The purpose of comparing 16 and 256 bin discretisation was twofold: (1) to
test the hypothesis that the advantage of training with the discretised loss
from Section 5 rather than the continuous loss from Section 4 would be greater
when the number of bins was lower, and (2) to test whether modelling the data
at lower precision would lead to improved perceptual quality. No data
augmentation, such as horizontal flips or random crops, was used on the
training set.

(a) Test Data (256 bins)

(b) Generated Data (256 bins)

(c) Test Data (16 bins)

(d) Generated Data (16 bins)

Figure 15: CIFAR-10 real and generated data. Samples generated with 4,000
steps, using networks trained with discretised loss. The same random seed was
used for both sets of samples. Note the improved image quality of the 16 bin
samples compared to the 256 bin samples.

Setup. The network architecture was essentially the same as that used for
Variational Diffusion Models (VDMs [17]), including the Fourier feature
inputs. The only modification was an extra input-output connection similar to
the network for MNIST. In total there were approximately 31M learnable
parameters. The following hyperparameters were used for all CIFAR-10
experiments: a validation set of 500 randomly selected training images (1% of
the training set), the AdamW [22] optmizer with weight decay 0.01, learning
rate $0.0002$ and $(\beta_{1},\beta_{2})=(0.9,0.99)$, dropout with probability
0.1, training batch size of 128, $t_{min}=1\mathrm{e}{-6}$,
$[x_{min},x_{max}]=[-1,1]$, and an exponential moving average of model
parameters with a decay rate of 0.9999 for evaluation and sample generation.
For the 256 bin experiments $\sigma_{1}=0.001$, while for the 16 bin
experiments $\sigma_{1}=\sqrt{0.001}$. For the networks trained with
continuous loss, the reconstruction loss was measured using the discretised
version of $L^{r}(\operatorname{\mathbf{x}})$ from Section 5.3 rather than the
continuous version from Section 4.10, using a discretised Gaussian with mean
equal to $\hat{x}(\boldsymbol{\theta},1)$ and std. deviation chosen
empirically to be $\sigma_{1}$ for 256 bins and $0.7\sigma_{1}$ for 16 bins.
This ensured the results were comparable between continuous and discretised
training, and consistent with the literature.

(a) Input Mean

(b) Output Mean

Figure 16: CIFAR-10 Input and output distributions. For two test set images
the figure shows the means of the input and output distributions at steps
evenly spaced between $t=0$ and $t=0.25$.

Results. Table 1 shows that the best performing BFN gives 2.66 BPD for the 256
bin data, which is close to the state-of-the-art at 2.64 BPD. The most obvious
performance benchmark (given the shared network architecture and similarity in
loss function) is the VDM result at 2.65 BPD [17]. However this took 10M
weight updates to achieve, and due to time constraints we were only able to
train BFNs for 5M updates. Validation performance was still improving after 5M
updates, and it remains unclear how much performance would improve with 10M
updates.

Table 3 shows that discretised loss gave better performance than continuous
loss for 16 bins, as well as much faster training time (250K updates vs. 1M).
This supports the hypothesis that training with discretised loss is most
beneficial when the number of bins is relatively low. Furthermore, for both 16
and 256 bins, discretised training gave much better results when the number of
steps $n$ was low (e.g. 10 or 25). However continuous loss gave better
performance than discretised loss on 256 bins (2.66 BPC vs 2.68); more
investigation would be needed to understand why.

Figure 15 shows that discretised training with 16 bins gives better sample
quality than training with 256 bins. This is presumably because the loss
function of the former is restricted to the first four bits of the data in
which — as can be seen by comparing the test data at 16 and 256 bins — most of
the perceptually relevant information is contained. An interesting direction
for future work would be to train one BFN to model the lower bits of an image,
and a second BFN to conditionally upscale to higher bits, as has previously
been explored for autoregressive models [26, 13].

Figure 17: CIFAR-10 losses against time. The plot was made using the network trained with discretised loss on 256 bins. Note the high loss at the very start of the process, which we did not observe with discrete data. | Model | BPC
---|---|---
Flow-based models | IAF/SCF† [54] | 1.88
Argmax Coupling Flow† [14] | 1.80
Discrete Flow† [48] | 1.23
Order-agnostic Models | OA-ARDM [13] | 1.43 $\pm$ 0.001
MAC [39] | 1.40
Diffusion models | Multinomial Diffusion [14] | 1.72
D3PM uniform [1] | 1.61 $\pm$ 0.02
D3PM NN [1] | 1.59 $\pm$ 0.03
| D3PM mask [1] | 1.45 $\pm$ 0.02
| BFN | 1.41
Autoregressive baseline | Transformer† [1] | 1.23
Best result* | Adaptive Span Transformer† [45] | 1.07
Table 4: Comparison of text8 results with other methods. The best published
model on this dataset (*) was trained on sequences of length 512. Rest of the
above models were trained on sequences of length 256. Results for models
marked with (†) are exact values; all other results are upper bounds.

###  7.3 text8

$n$-steps | 10 | 25 | 50 | 100 | 256 | 1000 | $\infty$
---|---|---|---|---|---|---|---
BPC | 1.70 | 1.52 | 1.47 | 1.43 | 1.42 | 1.41 | 1.41
Table 5: text8 results. BPC is bits per character averaged over 1M randomly
cropped sequences from the test set with $L^{n}(\operatorname{\mathbf{x}})$ or
$L^{\infty}(\operatorname{\mathbf{x}})$ sampled once per crop. The
reconstruction loss $L^{r}(\operatorname{\mathbf{x}})$ (included in BPC) was
$0.006$.