Patent Description:
<NPL>, describes a way of improving disentangling in a β-VAE (variational autoencoder) which involves gradually increasing the encoding capacity of the VAE's information bottleneck. <NPL>, addresses the problem of powerful decoders in a VAE by using a loss function with a KL penalty term with a multiplier β which is less than <NUM>. <NPL>, addresses the problem of powerful decoders by including skip connections in the generative model. <NPL>, describes a transformer-based method of encoding a target sequence into a shorter sequence of discrete latent variables to facilitate parallel processing. <NPL>, describes a VampPrior that consists of a mixture distribution, e.g. a mixture of Gaussians, and describes a model that avoids the usual local optima issues related to useless latent dimensions that plague VAEs.

In broad terms a variational autoencoder (VAE) determines a distribution for a set of latent variables representing an input data item, x. Thus the encoder determines parameters of a posterior distribution q(z|x) over the latent variables z. The VAE is trained with an objective which encourages the system to keep the posterior distribution close to a prior p(z), generally a standard Gaussian. A sample can be drawn from this distribution to generate an output data item. The VAE may be trained using unlabeled data. The objective may include a term, such as a KL (Kullback-Leibler) divergence, which measures the difference between the posterior and prior distributions. However a problem with some VAE implementations, in particular those with powerful decoders, is that the decoder may be able to generate the output data item without relying on the latent variables.

The invention is set out in independent claims <NUM>, <NUM> and <NUM>; further aspects are defined in the remaining dependent claims.

The specification describes a variational autoencoder neural network system. The system is implemented as computer programs on one or more computers in one or more locations. The system has an input to receive an input data item that consists of an image e.g. from a camera or LIDAR system. The system comprises an encoder neural network configured to encode the input data item to determine a set of parameters defining a first, posterior distribution of a set of latent variables. The system comprises a subsystem to sample from the posterior distribution to determine values of the set of latent variables. The system comprises a decoder neural network configured to receive the values of the set of latent variables and to generate an output data item representing the values of the set of latent variables. The variational autoencoder neural network system is configured for training with an objective function. This has a first term, such as a cross-entropy term, dependent upon a difference between the input data item and the output data item and a second term, such as a KL divergence term, dependent upon a difference between the posterior distribution and a second, prior distribution of the set of latent variables. The prior distribution is different to the posterior distribution; more particularly a structure of the prior distribution is different to a structure of the posterior distribution so that the posterior distribution cannot be matched to the prior distribution.

In implementations information which is transferred from the encoder to the decoder is manifest as a non-zero divergence between the posterior and prior distributions. The divergence may be a KL divergence, a Jensen-Shannon divergence, or some other difference metric.

By imposing a different structure on the posterior and prior distributions the second term of the objective function is guaranteed to be non-zero, and hence the decoder is forced to rely on the latent variables when generating an output data item.

The structure of a distribution may be determined by its inherent shape which may be determined for example by its mathematical form, and/or it may be determined by parameters of the distribution. Thus the posterior and prior distributions may be defined to have different structures by fixing one or more parameters of the distributions to be different to one another. For example where the distributions have the same or a similar mathematical form the distributions may be constrained to be structurally to be different to one another by constraining one or more parameters of each distribution to have a different relative or absolute value.

For example the posterior distribution and the prior distribution may each comprise a multivariate Gaussian distribution. Then a variance of the posterior distribution may be a factor of α different to a variance of the prior distribution (where α ≠ <NUM>). This structure for the posterior and prior distributions has an advantage of simplicity; it also facilitates determination of the KL divergence term in closed form, for example during training. In some implementations the parameter α can be determined (in closed form) from a desired committed, i.e. minimum, information rate transfer from the encoder to the decoder via the latent variables.

It is desirable to be able to use an autoregressive neural network as the decoder because such decoders can be very powerful, that is they can be capable of generating very accurate samples of data items. One example of such a system, which can be used, for example for generating images, is described in <NPL>; and in <NPL>. An example of such a system, which can be used, for example for generating sound (waveforms), is described in <NPL>. In this context an autoregressive neural network is a neural network which is configured to generate a sequence of output data item values, xt each conditioned upon previously generated output data item values x<t and conditioning variable <MAT>. The data item values are pixel values or, in embodiments not covered by the claims, sound signal values. When generating an image the autoregressive neural network generates the pixels in sequence, for example in a raster scan row by row and pixel by pixel.

However when such a powerful decoder is incorporated in a VAE the latent variables may be ignored. The above described variational autoencoder neural network system facilitates the use of such autoregressive neural network systems as the decoder whilst still using information from the latent variables. More particularly the autoregressive structure of the decoder may provide local structure for the output data item whilst the decoder output may also be conditioned on the information provided by sampling from the latent variable distribution, which may provide longer-range structure and/or global structure for the output data item.

Thus the encoder is configured to determine a sequence of sets of parameters defining a sequence of distributions for a sequence of sets of latent variables, one for each of a plurality of time steps. Here the time steps for the encoder may be different to the sample generating steps of the decoder. For example the interval between determining successive sets of latent variables may be longer than the decoder time steps i.e. more than one output data item value may be generated for each set of latent variables.

Thus a sequence of sets of latent variables is generated. The prior distribution of latent variables comprises an autoregressive distribution such that at each time step the prior distribution depends on the prior distribution at a previous time step. The posterior distribution may, however, be determined independently at each time step. This approach can help to capture correlation in the latent space from one time step to another (where the time steps may e.g. correspond to spatial locations in an image). The autoregressive process defining the evolution of the prior distribution of latent variables over time may be a linear autoregressive process. Values of the set of latent variables at a time step t, zt are defined by a sum of α times the values of the set of latent variables at a previous time step zt-<NUM> and a noise component, e.g. a Gaussian noise component, where |α| < <NUM>. The parameter α defines a degree of temporal correlation in the latent variables, with less correlation as α approaches zero.

Where the decoder neural network is an autoregressive neural network the VAE may further comprise a system to restrict the values of the set of latent variables passed to the decoder at each time step to those which encode information about in the sequence of output data values yet to be generated, i.e. about future values of the sequence of output data values. Thus the values of the set of latent variables passed to the decoder may be derived from input data item values for x>t. This is because the autoregressive neural network effectively already has access to information about past, i.e. previously generated values of the output data item. In implementations the encoder may be configured to restrict the values of the set of latent variables passed to the decoder in this way. Thus the encoder may be configured to have an "anti-causal" structure. When the VAE system is being trained the system will typically have access to a complete data item, and thus will have access to values of the data item which are later than those being generated by the decoder at any particular decoder time step, so causality is not violated. This can facilitate computational efficiency and can also allow for an increased learning rate.

The decoder may comprise a convolutional autoregressive neural network configured to implement causal convolutions i.e. where the generated data item values depend on previously generated data item values but not on future data item values. A causal convolution may be implemented using a mask to mask the input from data item values in a sequence following those at a current time step, or by shifting the convolution location (filter length - <NUM>) time steps.

Although the system may be trained to match the posterior and prior distributions, there is a built-in mismatch which encodes information. The VAE system may thus include a system, in particular an auxiliary prior neural network, configured to learn the sequence of distributions for the sequence of sets of latent variables i.e. configured to learn an approximate (aggregate) posterior distribution. The auxiliary prior neural network may be an autoregressive neural network, and may be trained concurrently with the encoder neural network and decoder neural network.

There is also described a method of training a variational autoencoder neural network system as described above, which is unsupervised. The method comprises receiving training data, the training data comprising training data items; providing each training data item to the input of the variational autoencoder neural network system to generate a corresponding output data item; and determining a gradient of the objective function from a difference between the training data item and the corresponding output data item and from a difference between the posterior distribution and the prior distribution of the set of latent variables. The training data items may be processed in batches. The method further comprises backpropagating the gradient through the variational autoencoder neural network system to adjust parameters of the encoder neural network and of the decoder neural network to optimize the objective function.

In broad terms the training may employ stochastic gradient descent (SGD) with an objective which includes a reconstruction cost and a closed-form KL divergence term. The gradients may be back-propagated through the decoder into the encoder using the "reparameterization trick" (see, e. g<NPL>), in which a sampling node is replaced by a deterministic operation with a noise input to allow a gradient to flow through the node. The objective function, which defines a built-in difference between the posterior and prior distributions ensures a (minimum) rate of information flow via the latent variables from the encoder to the decoder.

The VAE system during training includes both the encoder and decoder. However, once trained, each of these may have independent utility.

For example because the latent variable distribution effectively defines a compressed version of the input data item the encoder may be used to compress data items of the same type as used to train the VAE system. In another example, the encoder may be used as a front end for another machine learning system, for example a classifier. Because the encoder has learned the distribution of the training data items a classifier trained/operating on the latent variables may perform better than a classifier trained on the raw data items. In another example the encoder may be used as a front end for a reinforcement learning (RL) system in which the learned latent variable distribution is used to represent an image of an environment in which the RL system operates and/or to encode other sensor data such as data representing state of a mechanical agent such as the configuration of a robot arm. Although examples have been described using image data, the VAE system may also be trained on video data and thus the trained encoder may encode or compress video data.

The decoder of the trained system may also have independent utility. For example a sample may be drawn from the prior and provided to the decoder to generate a sample output data item. In a system with an auxiliary neural network a sample may be provided to the auxiliary neural network to generate a sequence of latent variables which may then be provided to the decoder to generate a sample output data item. In embodiments not covered by the claims, a generated data item may be a one dimensional data item such as a sound, in which case the latent variables may have a 1D feature space and the data item values may comprise values defining a sound waveform. A generated data item is at least a two dimensional image data item, in which case the latent variables may have a 2D feature space and the data item values comprise pixel values for the image such as brightness and/or color values. A generated data item may be a three dimensional data item such as an image sequence (video), in which case the latent variables may have a 3D feature space and the data item values may comprise pixel values for the image sequence (video). In the case of video the decoder may be an image decoder. For example, the VAE system could be trained to generate images conditional upon an additional data input defining movement/change of a viewpoint and then a sequence of sets of latent variables could be generated by applying a succession of such additional data inputs, each image being generated independently from a respective set of latent variables. In a further example, not covered by the claims, the data items may represent the structure of a molecule or other biological entity, e.g. a protein molecule, and the decoder may be used to generate output data items with similar properties to the training data items e.g. candidate drug molecules.

More generally the encoder and decoder may be configured to be conditioned on an additional data input, such as a label and/or text for generating an image or, in embodiments not covered by the claims, a label and/or text and/or a speaker identifier for generating audio, for example in a text-to-speech system. In embodiments not covered by the claims, the trained decoder may be used to generate an output data item of a desired type, for example for generating a waveform representing speech from natural language text.

The trained VAE system and/or encoder/decoder is used for image data item processing tasks such as an image data item completion task in which missing parts of a data item are generated or filled in by the system.

The VAE system, and in particular the trained decoder, may be used to generate further examples of data items for training another machine learning system. For example the VAE system may be trained on a set of data items and then a set of latent variables may be determined and used generate new data items similar to those in the training data set. The set of latent variables may be determined by sampling from the (prior) distribution of latent variables and/or using the auxiliary neural network. Where the VAE system has been trained conditioned on additional data, e.g. labels, new data items may be generated conditioned on additional data e.g. a label provided to the decoder. In this way additional labelled data items may be generated, for example to supplement a dearth of unlabeled training data items.

Thus a method of obtaining an encoder/decoder comprises training a variational autoencoder neural network system as described above and then using the trained encoder/decoder neural network as the encoder/decoder.

There is further described a method of encoding and/or decoding data using a trained encoder and/or decoder, as described above.

For example, in one implementation an autoregressive decoder comprises a causal convolutional neural network configured to generate a data item by, at each of a plurality of iterations, generating a value of the data item conditioned upon values of the data item previously generated at previous iterations, wherein the generating uses a soft attention query vector dependent upon the previously generated values of the data item to query a memory comprising values derived from the set of latent variables at each time step.

For example, in one implementation an anti-causal autoregressive encoder comprises a causal convolutional neural network configured to input a data item, reverse an order of values of the data item, and generate a representation of the data item from the reverse ordered values.

The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following further advantages.

Some implementations of the system are able to use powerful decoders, such as autoregressive decoders, to generate examples of output data items with improved fidelity, that is improved accuracy and more detail. In particular, implementations of the VAE system can be trained even when the decoder is implemented using an autoregressive technique. The system is potentially also able to learn latent variable distributions which provide an improved representation of the training data items. This in turn facilitates advantages such as reduced memory usage and better control over the generated data items. Where, for example, the system is used for data compression this may facilitate a greater degree of data compression. Some implementations of the system reduce the computational power needed for training because they are able to train faster and better than previous systems. Some implementations of the system are able to effectively model time series data such as video. Some implementations of the system allow efficient computation of the objective function with closed for divergence, as described above. Some implementations of the system allow a minimum information rate between the encoder and decoder to be tuned with one parameter; this parameter can be adjusted according to the type of data items processed, for example a degree of correlation in the data items, to tune the effectiveness of the system in computational resource and memory usage.

This specification generally describes a variational autoencoder (VAE) neural network system implemented as computer programs on one or more computers in one or more locations, and methods of training the system.

When trained the VAE neural network system comprises a trained encoder neural network and a trained decoder neural network. The encoder neural network learns to compress data from a training data distribution into a simpler distribution, represented by a set of latent variables. The decoder neural network learns to decode the set of latent variables into an example drawn from a distribution which approximates the training distribution. Thus trained VAE neural network system operates as a data compression/decompression system, with the encoder neural network acting as a data compressor and the decoder neural network acting as a complementary data decompressor.

During training the VAE learns the structure in the training data and can thus perform efficient data compression/decompression. The data is image data (including video).

The latent variable representation from the encoder neural network can also be useful, for example, for classification and reinforcement learning tasks. The data from the decoder neural network can be used to generate examples from the training data distribution in particular image data, conditioned on a set of (latent) variables provided as an input, optionally further conditioned on a label defining the example type, e.g. image content.

A problem with VAE neural networks is that they can fail to train effectively due to posterior collapse, explained below. Another problem is that whilst they can be good at capturing global structure in the training data they can fail to capture more complex local structure e.g. a generated image may appear blurred. A solution to the latter problem is to use a more powerful generative model for the decoder i.e. one which is better able to reproduce details in the training data distribution such as image detail, but this exacerbates the posterior collapse problem.

One VAE objective function aims to maximize the probability p(x) of obtaining the training data examples x from the decoder neural network (a generative model) by maximizing a lower bound on this probability <MAT> DKL(q(z|x) ∥ p(z)). Here DKL is the Kullback-Leibler divergence; p(z) defines a prior distribution for z, which should be continuous but may be otherwise arbitrary and may be e.g. a standard Gaussian, <IMG>(<NUM>,<NUM>); and the encoder neural network is represented by the function q(z|x) i.e. it defines parameters of an approximate posterior distribution for each component of z. Thus the encoder neural network defines parameters of a multivariate distribution for z from which a sample z~q(z|x) is taken and provided to the decoder neural network, represented by p(x|z). During training the second term can go to zero, that is a probability of the decoder producing examples which match the training data distribution examples can be maximized without using the latent variables, so-called posterior collapse. This is particularly a problem if the decoder is powerful, for example an autoregressive decoder from the PixelCNN family (ibid), which can generate output examples with fine detail, modelling the full data distribution without any conditioning input.

<FIG> shows an example variational autoencoder (VAE) neural network system <NUM> which addresses this problem. The VAE neural network system <NUM> can be implemented as computer programs on one or more computers in one or more locations.

The VAE neural network system <NUM> is provided with training data items e.g. from a data store <NUM>. These comprise image data items including e.g. video and/or LIDAR data items.

The data items are provided to an encoder neural network <NUM> which outputs a set of parameters <NUM> defining a posterior distribution of a set of latent variables e.g. defining the mean and variance of a multivariate Gaussian distribution. The system is configured to sample values for a set of latent variables <NUM> from the posterior distribution. The set of latent variables may define values for a latent variable data structure such as a latent variable vector z.

The latent variables are processed using a decoder neural network <NUM> which generates a data item output <NUM>. In some implementations the decoder neural network <NUM> generates the data item directly; in others it generates parameters of an output data item distribution which is sampled to obtain an example output data item. For example the decoder output may specify parameters of a distribution of the intensity of each pixel (or color sub-pixel) of an image.

A training engine <NUM> is configured to train the VAE neural network system <NUM> by back-propagating gradients of an objective function, in order to update neural network parameters <NUM> of the encoder neural network <NUM> and decoder neural network <NUM>. The training engine uses prior distribution parameters <NUM> of a prior distribution of the set of latent variables in conjunction with the posterior distribution parameters to determine a divergence loss term of the objective function. The training engine <NUM> also determines, from an input data item and a corresponding output data item, a reconstruction loss term of the objective function which aims to match a distribution of the output data items to a distribution of the training data items, e.g. log p(x|z).

In some implementations an auxiliary prior neural network <NUM> is provided to determine parameters of an auxiliary prior distribution, as described later. The auxiliary prior neural network <NUM>, where present, may also be trained by training engine <NUM>.

The log p(x|z) term of the previously described objective function is a reconstruction term and the DKL term encourages the system to keep the posterior distribution close to the prior. The DKL term measures information flow from the encoder to the decoder, more particularly a number of nats which are required, on average, to send through the latent variables from the encoder to the decoder. In implementations the training engine <NUM> is configured to use an objective function for which there is a guaranteed, non-zero minimum rate of information flow from the encoder to the decoder, here termed a "committed rate", δ. The rate of information flow δ is constrained to be equal to or greater than a committed or minimum rate to inhibit posterior collapse. For example a posterior-prior divergence term of the objective function, e.g. DKL(q(z|x) ∥ p(z)), ≥ δ.

In implementations this may be achieved by defining a structure of the (continuous) prior distribution p(z) which is different to a structure of the posterior distribution q(z|x) such that the posterior distribution cannot be matched to the prior distribution. The structure of a distribution may be defined by an inherent shape of the distribution e.g. whether or not it is a Gaussian distribution, and/or by constraining one or more parameters defining a shape the distribution e.g. by defining the variance of a Gaussian distribution. In implementations the committed rate of information flow, δ, is not fixed but may be different for each data item, allowing the system to allocate more bits to more complex input data items. Metrics other than Kullback-Leibler divergence may be used for the difference between p(z) and q(z|x), e.g. the Jensen-Shannon divergence. Some example distribution structures are described later.

The application of these techniques are not dependent on any particular form or architecture of the encoder or decoder neural networks. In particular the techniques avoid posterior collapse when using a powerful decoder such as an autoregressive decoder i.e. a decoder which is configured to generate a sequence of output data item values, xt each conditioned upon previously generated output data item values x<t.

The latent variables in the <FIG> system can capture and represent global characteristics of the data items. However data items may include spatial, temporal, or other sequential variation. For example images typically have some spatial continuity. It can be advantageous also to capture such finer, often shifting attributes variation in the latent variables, such as image texture or pose.

Thus the VAE neural network system <NUM> is configured to a model data item using a sequence of latent variables drawn from a sequence of prior distributions. Thus each data item may be subdivided into a sequential set of data item parts. By way of example for an image data item the parts may comprise rows, columns or regions of the image; for a video data item the parts may comprise individual images. Steps in the sequence may be referred to as time steps, e.g. when the data item parts are processed sequentially.

The same posterior distribution model, e.g. a (diagonal) multivariate Gaussian distribution, may be used for each step of the sequence. That is, the same encoder neural network may process each of the data item parts. However the prior distribution may change with each step in a correlated manner. More specifically the sequence of prior distributions is an autoregressive sequence, e.g. a first order and/or linear autoregressive sequence. Thus at each time step the prior distribution may depend on the prior distribution at a previous time step. This creates a mismatch in the correlation structure between the prior and posterior distributions which results in a positive lower bound on the divergence between these distributions.

For example a latent variable vector at a time step t, zt, may be defined by zt = αzt-<NUM> + ∈t where ∈t is a noise component e.g. Gaussian noise with zero mean and constant variance <MAT>, and |α| < <NUM> (so that the prior distribution has constant sufficient statistics through its time evolution). Then zt has zero mean and variance <MAT> and the choice <MAT> can be made so that <MAT>, facilitating determining an analytical form for the committed minimum rate of information flow e.g. the lower bound on DKL(q(z|x) || p(z)). Thus during training the optimization process is allowed to settle on a rate higher than the minimum but is restricted from going below this. The different dimensions of the latent variable vector zt may, but need not, have different values of α.

In implementations the committed information rate is given by <MAT> where n is the length of the sequence, d is the dimension of the (multidimensional) latent variable vector, and αk is the value of α for each dimension. This allows choices to be made for α or αk to achieve a target value for δ. In broad terms the value of α determines a degree of correlation between the prior distributions from one step to the next - with α = <NUM> there is no correlation and as α approaches one the correlation increases. This can be viewed as variable rate feature analysis, with different values of alpha corresponding to different speeds of variation.

In implementations where the decoder comprises an autoregressive neural network, in principle the decoder neural network can accurately estimate a part of an output data item, xt, given previous parts of the output data item, x<t, which it has already generated. The set of latent variables need not, therefore, transmit this information to the decoder neural network. Thus in some implementations the encoder neural network may have an anti-causal structure i.e. one in which the set of latent variables do not encode information about parts of the output data item which have already been generated. For example the set of latent variables may encode information about a current part of a data item in a sequence and about those parts which have yet to be generated.

This is illustrated schematically in <FIG>, which shows the operation of an anti-causal inference model i.e. an anti-causal encoder neural network: At a first time step z<NUM> is generated using information from x<NUM>, x<NUM>, and x<NUM>; at a second time step z<NUM> is generated using information from x<NUM>, and x<NUM>; and at a third time step z<NUM> is generated using information from only x<NUM>. In practice x<NUM>, x<NUM> and x<NUM> are available in parallel and the anti-causal dependence of the encoder neural network may be defined by appropriate selection of the parts of a data item e.g. via a mask and/or by reversing a sequential order of the parts of a data item. <FIG> illustrates, schematically, the autoregressive nature of a corresponding generative model i.e. of an autoregressive decoder neural network.

Use of an anti-causal encoder neural network may further help to avoid posterior collapse, and may also increase the data compression of the system by removing potentially redundant information from the set of latent variables.

Once trained the encoder neural network <NUM> may be used as a data compression engine, because the set of latent variables encodes meaningful information about the global and finer features of an input data item. Fine detail may, however, be produced by the autoregressive decoder neural network <NUM>. This approach can achieve a high data compression ratio.

The trained decoder neural network <NUM> may be used to an generate example data item by sampling a set of latent variables from the prior distribution and providing the set to the decoder neural network <NUM>.

Where there is a significant mismatch between the prior distribution p(z) and an aggregate posterior distribution <MAT> for an ensemble of the training data items <IMG>, there may be regions of the prior distribution which the decoder neural network does not see during training. Thus the VAE neural network system <NUM> may include an auxiliary prior neural network <NUM> which is trained to output parameters for an auxiliary distribution paux which matches the aggregate posterior distribution. The auxiliary prior neural network <NUM> may be trained at the same time as the encoder and decoder neural networks but does not take part in the training of the encoder and decoder neural networks. However after training a set of latent variables may be sampled from the auxiliary prior distribution rather than from the prior distribution.

An autoregressive model may be used to estimate paux. For example a single layer LSTM (Long Short-Term Memory) neural network may be used to estimate parameters of a posterior distribution of the i-th latent variable q(zi|x) conditioned on previous latent variable samples, <MAT> where paux(zi|z<i) may be a Gaussian distribution with mean and variance parameters for step i output by the LSTM after processing the previous latent variable samples. The LSTM neural network may be trained by minimizing DKL(q(z|x) ∥ paux(z)).

In some implementations the VAE neural network system <NUM> of <FIG> may be modified so that the decoder neural network <NUM> generates output data items from a labelled class e.g. an example image of a particular digit. This can be achieved by adding a conditioning input to both the encoder and decoder neural networks during training. The conditioning input nay be e.g. a one hot conditioning vector c identifying a class to which a training data item belongs. A corresponding conditioning input may then be provided to the decoder neural network <NUM> to generate an example output data item of a specified class.

<FIG> is a flow diagram of an example training process for the VAE neural network system <NUM>. The encoder and decoder neural network parameters are initialized e.g. randomly (step <NUM>), and a training data item is obtained (step <NUM>) e.g. from data store <NUM>. The training data item is processed using the encoder neural network <NUM> to obtain parameters defining the posterior distribution (step <NUM>), and a set of latent variables is sampled from this distribution (step <NUM>). The set of latent variables is then processed by the decoder neural network <NUM> to obtain an output data item (step <NUM>), either directly or e.g. by sampling from a multivariate distribution parameterized by an output of the decoder neural network. The process then backpropagates gradients of an objective function of the type previously described to update the parameters of the encoder and decoder neural networks. Any suitable backpropagation method may be used e.g. Adam. To backpropagate through the latent variable sampling the "reparameterization trick" may be used, rewriting the sampling operation for each latent variable as z = µ + σε where ε is standard Gaussian noise. The encoder neural network generates an output defining the mean (µ) and variance (σ<NUM>) of a distribution for the latent variable but the stochastic (sampling) element is provided by the noise so that gradients can be backpropagated through the "sampling" step. Gradients may be averaged over a minibatch; the process of <FIG> may be repeated until convergence of the neural network parameters.

Optionally gradients of an objective function for the auxiliary prior neural network are backpropagated through the auxiliary prior neural network to train this neural network in parallel with the main VAE system.

The objective function may have the general form log p(x|z) - DKL(q(z|x) || p(z)). The reconstruction loss term log p(x|z) term can be evaluated from training data item and an output data item e.g. by determining a cross-entropy loss or MSE (Mean Square Error) loss between the training data item and output data item.

The second, divergence term of the objective function, DKL(q(z|x) || p(z)), may be calculated from the parameters of the prior and posterior distributions. As previously described, the structures of the prior and posterior distributions are different such that they cannot be matched.

For example in one implementation the encoder neural network <NUM> may output a set of parameters <NUM> defining a multivariate Gaussian posterior distribution (with diagonal covariance) for the set of latent variables and the prior distribution may be defined by a standard Gaussian distribution (i.e. zero mean, unit variance). The KL divergence of two Gaussians <IMG>(µq, σq), <IMG>(µp, σp) is given by: <MAT>.

To calculate the KL divergence of a d-dimensional latent variable the KL divergence for each of the d-dimensions may be summed. In a simple implementation the prior and posterior Gaussian distributions may be defined to have different but fixed variance, e.g. different by a factor of α where α ≠ <NUM>, to constrain δ as non-zero.

In some implementations the prior distribution is similarly defined by a standard Gaussian distribution and the posterior distribution comprises a multivariate Gaussian distribution with a mean µq and variance <MAT> and diagonal covariance i.e. assuming that the latent variables in the set of latent variables are independent of one another. The posterior distribution is constrained by the committed rate of information flow, δ, whilst allowing the rate to go above δ. This can be achieved by constraining the mean µq and variance <MAT> for each latent variable (i.e. for each component of the posterior distribution) according to <MAT>.

This can be solved numerically for µq and δ to obtain a feasible interval <MAT> where <MAT> define respective lower and upper values for the variance. This in turn defines mean µq and variance <MAT> for the posterior distribution in terms of corresponding outputs µ(x) and σ(x) from the encoder neural network: <MAT> <MAT>.

Where the latent variables drawn are from an autoregressive sequence of prior distributions the divergence term is given by <MAT>.

Where, as previously described, p(z<NUM>) = <IMG>(µ<NUM> , σ<NUM>) is e.g. a standard Gaussian, and for t > <NUM> zt is defined such that <MAT>, and the posterior distribution q(zt|x) at time step t comprises a multivariate Gaussian distribution with a mean µt and variance <MAT> and diagonal covariance <MAT> where f(a) = α - ln(a) - <NUM>.

<FIG> shows a flow diagram of a process for using the trained encoder neural network to encode a data item, e.g. to generate a compressed representation of the data item. At step <NUM> the process obtains the data item e.g. directly or indirectly from a sensor, from storage, or via a computer network or other communications link. The encoder neural network <NUM> then processes the data item and outputs a set of parameters defining the posterior distribution (step <NUM>). A set of latent variables is then sampled from the posterior distribution (step <NUM>), the set of latent variables representing an encoded, compressed representation of the data item.

<FIG> shows a flow diagram of a process for using the trained decoder neural network to generate a data item. At step <NUM> the process samples a set of latent variables from the prior distribution, and then processes the set of latent variables using the decoder neural network <NUM> to generate the data item (step <NUM>). The training process allows a valid example data item to be generated by randomly sampling from the prior, which may be a multivariate Gaussian distribution or which may be the auxiliary prior.

<FIG> shows an example implementation of the VAE neural network system <NUM> in which both the encoder and decoder neural networks are implemented using an autoregressive approach. The techniques described in this specification are not limited to a particular form of the encoder and decoder but facilitate the use of a powerful, autoregressive decoder; it is not necessary for the encoder to be autoregressive when the decoder is autoregressive.

In the example of <FIG> the decoder neural network <NUM> comprises a PixelCNN neural network (arXiv:<NUM>) or a variant thereof such as PixelCNN++ (arXiv:<NUM>) or PixelSNAIL (arXiv:<NUM>). Thus, for example, an output of the neural network may define parameters of a mixture of logistics distributions (shown schematically) which are sampled to obtain pixel values, and the neural network may incorporate attention layers e.g. attention over an output of the encoder neural network. The encoder neural network may have the same architecture as the decoder neural network. VAE neural network system <NUM> may be trained using dropout and layer normalization (arXiv: <NUM>).

As previously mentioned an anti-causal architecture can be useful for reducing posterior collapse and encouraging the latent variables to be meaningful. <FIG> shows one example of how an anti-causal context can be implemented when processing an image. In this particular example a pixel order of the input image is reversed and each spatial dimension is padded by one before providing the image to the encoder. An output of the encoder is cropped and reversed again, giving each pixel anti-causal context i.e. pooling information from its own and future values (as shown in curly brackets). Each encoded pixel value may be used as a latent variable but for improved computational efficiency pooling, e.g. average pooling, may then be applied row-wise to encode each row as a multidimensional latent variable.

In implementations the VAE neural network system can be trained to produce example data items of a particular type by conditioning each of the encoder neural network and decoder neural network on an additional conditioning input during training. The conditioning input may identify the type of data item and may comprise e.g. a one-hot vector labelling the type of data item or an embedding of the type of data item. Thus the conditioning input may identify an image label, description, viewpoint and/or pose for generating an image.

Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non transitory program carrier for execution by, or to control the operation of, data processing apparatus. The computer storage medium is not, however, a propagated signal.

The elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data.

While this specification contains many specific implementation details, these should not be construed as limitations on the scope of protection which is defined by the claims, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Moreover, although features may be described above as acting in certain combinations, one or more features from a combination can in some cases be excised from the combination, and the combination may be directed to a subcombination or variation of a subcombination.

Claim 1:
A computer-implemented variational autoencoder neural network system (<NUM>), comprising:
an input to receive an input image data item consisting of pixel values;
an encoder neural network (<NUM>) configured to encode the input image data item to determine a set of parameters defining a first, posterior distribution of a set of latent variables (<NUM>);
a subsystem to sample from the posterior distribution to determine values of the set of latent variables;
a decoder neural network (<NUM>) configured to receive the values of the set of latent variables and to generate an output image data item representing the values of the set of latent variables, the output image data item consisting of pixel values for an output image;
wherein the variational autoencoder neural network system is configured (<NUM>) for training the encoder neural network and the decoder neural network with an objective function which has a first term dependent upon a difference between the input image data item and the output image data item and a second term dependent upon a difference between the posterior distribution and a second, prior distribution of the set of latent variables, and wherein a structure of the prior distribution is different to a structure of the posterior distribution such that the posterior distribution cannot be matched to the prior distribution;
wherein the encoder is configured to determine a sequence of sets of parameters defining a sequence of distributions for a sequence of sets of latent variables, one for each of a plurality of time steps, wherein the prior distribution comprises an autoregressive distribution such that at each time step the prior distribution depends on the prior distribution at a previous time step, and wherein the values of the set of latent variables at a time step t, are defined by a sum of α times the values of the set of latent variables at a previous time step and a noise component, where |α| < <NUM>.