Bayesian Attention Networks for Data Compression

Data compression is a very important area of study for both practical applications as well as for fundamental research in information theory (Salomon, 2006),(MacKay, 2002), (Mahoney, 2013b), (Lossless Compression, 2020) and subject of many competitions, like (Hutter Prize, 2021). See also (Mahoney, 2013a). The neural compression is an active research area (Kingma et al., 2019).

The fundamental approach to compressing a sequence of data elements is to use a probabilistic model that is able to reproducibly compute a probability of the sequence via conditional probabilities of all elements and then encode the sequence probability with an entropy encoder (Shannon, 1948),(Rissanen, 1976),(Pasco, 1976).

The direct consequence of that approach is that the finding of a good predictive model becomes critical for the development of the high performance compression algorithm.

For a sequence of elements

$$s=a_{1}a_{2}a_{3}..a_{n}..a_{N}$$

(1)

the probability of the sequence is a product of conditional probabilities of each element

$$P(s)=P(a_{1})P(a_{2}|a_{1})P(a_{3}|a_{1}a_{2})..P(a_{n}|s_{n})..P(a_{N}|s_{N}),$$

(2)

where sub-sequences $s_{n}$ are defined as $s_{n}=a_{1}a_{2}..a_{n-1},n=2..N$. To simplify we will consider conditional dependencies for all probabilities as functions of fixed-length sequences $s_{n,n+l}$ of length $l$ with left-side zero padding as needed. Then the probability of the sequence $s$ will be

$$P(s)=\prod_{n=1}^{N}P(a_{n}|s_{n-l,n}),\;s_{n-l,n}=a_{n-l}a_{n-l+1}..a_{n-1}.$$

(3)

By defining a model $P(a|s,w)$ for probability of an element $a$ given a previous sequence $s$ and parametrized by weights $w$ we obtain a parametrized probability of any sequence $s$

$$P(s|w)=\prod_{n=1}^{N}P(a_{n}|s_{n-l,n},w),\;s_{n-l,n}=a_{n-l}a_{n-l+1}..a_{n-1}.$$

(4)

For a known training dataset of independent sequences $\{s_{i},i=1..M\}$ the probability of the data for a selected model is given by Bayesian integral over weights as it follows from the Bayes theorem (Bayes & Price, 1763)

$$Prob(\{s_{i}\}|H)=\int_{w}\prod_{i=1}^{M}P(s_{i}|w)P_{0}(w|H)dw,$$

(5)

where $P_{0}(w|H)$ is a prior of weights conditional on hyperparameters $H$.

The prediction of the probability $P(s|\{s_{i}\},H)$ of a new sequence $s$ then will be given by an average

$$P(s|\{s_{i}\},H)=\langle P(s|w)\rangle=\int_{w}P(s|w)\prod_{i=1}^{M}P(s_{i}|w)P_{0}(w|H)dw\;/\;Prob(\{s_{i}\}|H).$$

(6)

We will call the new sequence $s$ a test sequence.

The prediction probability $P(s|\{s_{i}\},H)$ is an average over weights with the most significant contribution coming from training sequences that maximize both corresponding probabilities of training data and the test sequence $P(s|w)$ for the same weights in the Eq.(6). The resulting average probability is high when the test sequence $s$ is similar to the most typical training sequences. However, the average will be low when $s$ is very different from the typical training data. The problem is that all training sequences contributing equally to the average and a good prediction is possible only for a test sequence that is very similar to training data.

To improve the prediction of the probability for the given sequence $s$ we will introduce the importance weights for training data samples in the Eq.(6) that will increase contributions of samples with sequences $s_{i}$ similar to the test sequence $s$ and reduce contributions of other samples.

The idea is an extension of a very well known approach used for compression that maps first a prediction sample to a certain context and then a prediction for the sample is made using a context specific predictor.

The approach with importance weights described in this paper leads to a definition of Bayesian Attention Networks (BAN) considered next.

The idea to account for long-range correlations between data samples for achieving better compression has a long history. The attention mechanism recently proposed in (Vaswani et al., 2017) was a great success for sequence prediction tasks as well as for other problems including image recognition (Dosovitskiy et al., 2020).

We will define BAN in a few steps. First, we introduce the importance factors $\rho_{i}(s)$ for sequences by modifying contribution of each training data sample in the Bayesian integrals

	$\displaystyle P(s|\{s_{i}\},H)$	$\displaystyle=$	$\displaystyle\int_{w}P(s|w)\prod_{i=1}^{M}e^{-\rho_{i}(s)l(s_{i}|w)}P_{0}(w|H)dw\;/\;Prob(\{s_{i}\}|H),$		(7)
	$\displaystyle Prob(\{s_{i}\}|H)$	$\displaystyle=$	$\displaystyle\int_{w}\prod_{i=1}^{M}e^{-\rho_{i}(s)l(s_{i}|w)}P_{0}(w|H)dw.$		(8)

Here $l(s_{i}|w)$ are loss functions for sequences defined as logs of probabilities $l(s|w)=-\log(P(s|w))$. The importance factors $\rho_{i}(s)$ are assigned to each training sequence $s_{i}$ and are functions of predicted sequence $s$. The importance factors must be constrained by normalization condition $\sum_{i}\rho_{i}(s)=M$, so the total contribution of all training sequences is preserved.

The meaning of the importance factors is very simple - it defines the importance of a given training sample $s_{i}$ to influence a solution for the prediction probability of the test sequence $s$ with higher importance $\rho_{i}(s)$ increasing contribution of the sample $s_{i}$ to predicting $s$.

In the limit when unnormalized importance is only zero or one it is equivalent to clustering when the solution for predicting sample $s$ is defined by a few training samples $s_{i}$ correlated with the prediction sample.

We will clarify the definition of importance weights in the following steps.

Next, we will approximate the average of probability of a sequence as product of averages for predictions of probabilities of elements conditional on corresponding immediate sub-sequences

$$P(s)=P(a_{1}a_{2}..a_{N})=\langle P(a_{1}|s_{1},w)\rangle\langle P(a_{2}|s_{2},w)\rangle..\langle P(a_{N}|s_{N},w)\rangle.$$

(9)

Then the prediction of the probability of an element is given by

$$\langle P(a|s,w)\rangle=\int_{w}P(a|s,w)\prod_{i=1}^{K}e^{-\rho_{i}(s)l(a_{i}|s_{i},w)}P_{0}(w|H)dw\;/\;Prob(\{a_{i}\}|H)$$

(10)

with the normalization factor $Prob(\{a_{i}\}|H)$ that keeps $\langle P(a|s,w)\rangle$ normalized defined as follows

$$Prob(\{a_{i}\}|H)=\int_{w}\prod_{i=1}^{K}e^{-\rho_{i}(s)l(a_{i}|s_{i},w)}P_{0}(w|H)dw.$$

(11)

Here the importance factors depend only on input sequences in each sample and not on predicted elements. The index $i$ runs over all $K$ sequence elements $a_{i}$ in training data in Eqs.(10,11). The losses $l(a_{i}|s_{i},w)$ are per element in training sequences and the importance factors $\rho_{i}(s)$ control the contributions of each training sample $i$ and depend on the sub-sequence $s$ of the element $a$ in the prediction probability $P(a|s)$. The importance factors are normalized by condition $\sum_{i}\rho_{i}(s)=K$ where $K$ is a total number of training data samples.

The Eqs.(10,11) are defining the Bayesian Attention Network with attention variable $\rho_{i}(s)$ which is a function of two sequences. The attention or importance factors allow to take into account long range interactions between different data samples.

It is important to note that for any approximation of the integrals in Eqs.(10,11) the solution is dependent on the input sequence $s$ in the data sample we are making prediction for $P(a|s)$, which requires computing the integrals for every test sample.

To resolve this problem and find the practical computing approach we will use a latent variable $z$ representation with the encoder-decoder network that is considered in the next section.

Let’s define encoder $\rho(z|s)$ and decoder $\rho(s|z)$. To do a prediction for a sample (a,s) the encoder $\rho(z|s)$ gives a probability of $z$ for input $s$ and the decoder $\rho(s_{i}|z)$ gives a probability of a training sequence $s_{i}$ for a given $z$.

The decoder allows to define the normalized importance factor

$$\rho_{i}(z)=K\rho(s_{i}|z)/\sum_{j}\rho(s_{j}|z).$$

(12)

With that we can reformulate Eqs.(10,11) as follows

	$\displaystyle\langle P(a|s,w)\rangle$	$\displaystyle=$	$\displaystyle\int_{w,z}P(a|s,w)\prod_{i=1}^{K}e^{-\rho_{i}(z)l(a_{i}|s_{i},w)}P_{0}(w|H)\rho(z|s)dwdz\;/\;Prob(\{a_{i}\}|H)$
	$\displaystyle Prob(\{a_{i}\}|H)$	$\displaystyle=$	$\displaystyle\int_{w,z}\prod_{i=1}^{K}e^{-\rho_{i}(z)l(a_{i}|s_{i},w)}P_{0}(w|H)\rho(z|s)dwdz$		(13)

Due to Bayesian theorem (Bayes & Price, 1763) the encoder probability related to the decoder probability and can be expressed via decoder probability and priors for $z$ and $s$

$$\rho(z|s)=\frac{\rho(s|z)\rho(z)}{\rho(s)}.$$

(14)

By introducing an approximation for the encoder $q(z|s)$ we can use a variational method for simplifying the integral over $z$ following (Kingma & Welling, 2014)

$$\int_{z}e^{-L(z)}\rho(s|z)\rho(z)dz\geqslant\exp\left(-\int_{z}q(z|s)L(z,w)dz+\int_{z}q(z|s)\log\frac{\rho(s|z)\rho(z)}{q(z|s)}dz\right),$$

(15)

where total training loss is $L(z,w)=\sum_{i}\rho_{i}(z)l(a_{i}|s_{i},w)$. Maximizing the right side of Eq.(15) w.r.t. parameters of $q(z|s)$ allows to find the encoder $q$.

However, the variational approximation that is based on Jensen’s inequality, (Jensen, 1906) $\langle\exp(x)\rangle\geqslant\exp(\langle x\rangle)$, results in the loss of dependency on importance factors in the prediction probability in the Eqs.(3). To retain that dependency we have to use a sharpened Jensen’s inequality, (see Appendix A for the proof)

$$\langle e^{x}\rangle\geqslant e^{\langle x\rangle+\frac{1}{2}\langle(x-\langle x\rangle)^{2}\rangle}.$$

(16)

Then the probability of prediction can be maximized w.r.t. importance weights. Now we can find that the variation of the log of the prediction probability is equal to a correlation function of test and train losses

	$\displaystyle\frac{\delta\log{\langle P(a|s,w)\rangle}}{\delta\rho_{i}(z)}=\langle l(a_{i}|s_{i},w)l(a|s,w)\rangle_{w}-\langle l(a_{i}|s_{i},w)\rangle_{w}\langle l(a|s,w)\rangle_{w}\approx$
	$\displaystyle\sum_{j}\sigma_{j}^{2}\frac{\partial l(a_{i}|s_{i},w)}{\partial w_{j}}\frac{\partial l(a|s,w)}{\partial w_{j}},$		(17)

where averages $\langle(..)\rangle_{w}$ are computed with a posterior distribution of weights

$$Post(w|H_{0},z)=\prod_{i=1}^{K}e^{-\rho_{i}(z)l(a_{i}|s_{i},w)}P_{0}(w|H_{0}),$$

(18)

by expanding the weights around the means up to the second order and averaging over it.

Due to the sharpened Jensen’s inequality in Eq.(16) the complete loss includes a cross-correlation term between the test sample loss and training sample losses

$$\int_{w,z}P_{0}(w|H_{0})q(z|s)\left(L(z,w)-l(a|s,w)L(z,w)\right)dwdz-\int_{z}q(z|s)\log\frac{\rho(s|z)\rho(z)}{q(z|s)}dz,$$

(19)

where the training loss $L(z,w)=\sum_{i}\rho_{i}(z)l(a_{i}|s_{i},w)$.

Finding the importance factors by maximizing the prediction probability for test sample requires to minimize only the cross-correlation loss due to normalization factor in Eq.(3). All other parameters could be found by minimizing the loss in the Eq.(19).

Let’s use a posterior distribution for weights that depends on initial hyperparameters $H_{0}$ and the latent variable $z$ in the Eq.(18). The method developed in (Tetelman, 2020) allows to compute approximations of the Bayesian integrals by re-parametrizing a prior of weights $P_{0}(w|H)$ to represent a posterior

$$Post(w|H^{\prime},z)=\prod_{i=1}^{K}e^{-\rho_{i}(z)l(a_{i}|s_{i},w)}P_{0}(w|H_{0})\sim P_{0}(w|H(z,H_{0})).$$

(20)

The re-parametrization is computed recursively. With Gaussian priors for all network weights and Gaussian encoder $q(z|s)$ we can find the following update rules for hyperparameters: (mean, variance) pairs $(\mu,\sigma^{2})$ for all weights of all networks. The training update steps are as follows

	$\displaystyle i$	$\displaystyle\sim$	$\displaystyle\mathbf{Uniform}[1..K]$		(21)
	$\displaystyle z(s,v)$	$\displaystyle\sim$	$\displaystyle q(z|s,v)$		(22)
	$\displaystyle w(z,v)$	$\displaystyle\sim$	$\displaystyle P_{0}(w|\mu_{w}(z|v),\sigma_{w}(z|v))$		(23)
	$\displaystyle\rho_{i}(z,u)$	$\displaystyle=$	$\displaystyle K\rho(s_{i}|z,u)/\sum_{j}\rho(s_{j}|z,u)$		(24)
	$\displaystyle\mathbf{grad_{v}}$	$\displaystyle=$	$\displaystyle\frac{\partial}{\partial v}\bigg{[}\rho_{i}(z,u)l(a_{i}|s_{i},w(z,v))-\log{\frac{{\textbullet}\rho(s|z,u)\rho(z)}{q(z|s,v)}}\bigg{]}$		(25)
	$\displaystyle\Delta\mu_{v}$	$\displaystyle+=$	$\displaystyle-\varepsilon\sigma_{v}^{2}\mathbf{grad_{v}}$		(26)
	$\displaystyle\Delta\frac{1}{\sigma_{v}^{2}}$	$\displaystyle+=$	$\displaystyle\varepsilon(\mathbf{grad_{v}})^{2}$		(27)
	$\displaystyle\mathbf{grad_{u}}$	$\displaystyle=$	$\displaystyle\sum_{j}\bigg{(}\frac{\sigma_{w_{j}}^{2}}{K}\frac{\partial l(a_{i}|s_{i},w)}{\partial w_{j}}\frac{\partial l(a|s,w)}{\partial w_{j}}\bigg{)}\frac{\partial\rho_{i}(z,u)}{\partial u}$		(28)
	$\displaystyle\Delta\mu_{u}$	$\displaystyle+=$	$\displaystyle-\varepsilon\sigma_{u}^{2}\mathbf{grad_{u}}$		(29)
	$\displaystyle\Delta\frac{1}{\sigma_{u}^{2}}$	$\displaystyle+=$	$\displaystyle\varepsilon(\mathbf{grad_{u}})^{2}.$		(30)

Here the learning rate $\varepsilon$ is equal to the inverse number of epochs $T$: $\varepsilon=1/T$.

The prediction for a sample with input $s$ is done by sampling $z\sim q(z|s,v)$, computing a mean for weight $w$ equal to $\mu_{w}(z,v)$ and finally computing a prediction probability $P(a|s,\mu_{w})$.

The compression is defined by computing a probability of a prediction sample that is then encoded with the entropy coder. The contribution of this paper is to derive the compression algorithm by defining the Bayesian Attention Networks for predicting a probability of a new sample by introducing an attention or importance factor per the training sample that is dependent on the prediction sample. The attention factors allow to find a specific solution for each prediction sample that is influenced by a few training samples most correlated with the prediction sample.

Because the solution becomes specific to a given prediction sample the training and the prediction processes are coupled together. To make this approach practical we introduce a latent space representation with encoder-decoder networks that allows to separate training and prediction by mapping any prediction sample to a latent space and learning all solutions as a function of the latent space along with learning attention as a function of the latent space and a training sample.

The latent space plays a role of the context representation with a prediction sample defining a context and a learned context dependent solution used for the prediction.