Improved Variational Inference in Discrete VAEs using Error Correcting Codes

## Abstract

Abstract Despite significant advancements in deep probabilistic models, learning low-dimensional discrete latent representations remains a challenging task. In this paper, we introduce a novel method that enhances variational inference in discrete latent variable models by leveraging Error-Correcting Codes to introduce redundancy in the latent representations. This redundancy is then exploited by the variational posterior to yield more accurate estimates, thereby narrowing the variational gap. Inspired by ECCs commonly used in digital communications and data storage, we demonstrate proof-of-concept using a Discrete Variational Autoencoder ( DVAE ) with binary latent variables and block repetition codes. We further extend this idea to a hierarchical structure based on polar codes, where certain latent bits are more robustly protected. Our method improves generation quality, data reconstruction, and uncertainty calibration compared to the uncoded DVAE , even when trained with tighter bounds such as the Importance Weighted Autoencoder ( IWAE ) objective. In particular, we demonstrate superior performance on MNIST, FMNIST, CIFAR10, and Tiny ImageNet datasets. The general approach of integrating ECCs into variational inference is compatible with existing techniques to boost variational inference, such as importance sampling or Hamiltonian Monte Carlo. We also outline the key properties ECCs must have to effectively enhance discrete variational inference. Keywords: deep generative models, variational inference, discrete representations, error correcting codes

## 1 Introduction

Discrete latent space models seek to represent data using a finite set of features. Recent progress in generative models has increasingly favored these representations, as they are well-suited for datasets characterized by naturally discrete hidden states. However, effective learning of low-dimensional discrete latent representations is technically challenging. Vector Quantized-Variational Autoencoders [, ] stand out as solutions for this problem but rely on a non-probabilistic autoencoder, which does not provide uncertainty quantification in the latent space. To fit a fully probabilistic Variational Autoencoder (VAE) model [], a common approach considers either Concrete [] or Gumble-Softmax [] approximations to sample from a discrete latent distribution in a reparameterizable manner [, ]. However, this approach leads to instabilities since the gradient variance is sensitive to the temperature that controls these approximations. The DVAE in [, , ], augments the binary latent representations with a set of continuous random variables, pairing each bit with a continuous counterpart where reparameterization can be done in a more stable manner after marginalizing the latent bits. The key distinction between the DVAE [] and the DVAE++ [] lies in their smoothing transformations: while [] introduces spike-and-exponential transformations, [] uses overlapping exponential distributions. These overlapping transformations are generalized in [], enabling tighter variational bounds. We demonstrate the effectiveness of our method over a simplified version of the DVAE++ [].

This work presents a novel method to improve variational inference and representation learning in generative models with discrete latent variables. In particular, for a latent variable model, we argue that one should use ECCs to introduce redundancy into the latent sample before the reconstruction decoder network processes it to generate the data. The variational approximation to the true posterior distribution can then exploit the added redundancy to provide more accurate estimates, reducing, in turn, the variational gap to the data likelihood.

Our approach is based on well-known digital communications and data storage techniques where information is protected with ECCs before transmission/storage to reduce the overall error rate during recovery. For different datasets, our results demonstrate that, compared to the uncoded DVAE, the DVAE with ECCs (Coded-DVAE) achieves superior generation quality, better data reconstruction, and critically calibrated uncertainty in the latent space. In Fig. 1, we highlight some representative results for both MNIST and FMNIST data sets. We note that the use of ECCs is a general design approach that is perfectly compatible with state-of-the-art techniques for improved variational inference, such as importance sampling [, ] or Hamiltonian Monte Carlo [, ]. In summary, our main contributions are:

- •
We provide proof-of-concept results demonstrating that training deep generative models can be improved by ECC techniques, an idea that, to the best of our knowledge, is completely novel in the literature.
- •
We formulate a coded version of DVAE using block repetition codes. We show that encoding/decoding of the block repetition code can be efficiently done with linear complexity.
- •
We show that Coded-DVAE improves reconstruction, generation, and uncertainty calibration in the latent space when compared to the uncoded case using the same latent dimension, even when the uncoded DVAE is trained with tighter bounds such as the IWAE objective [].
- •
We discuss the generalization of this method to other coding schemes and introduce a hierarchical structure, inspired by polar codes [], that effectively separates high-level information from finer details.
- •
Through an extensive ablation study, we show that the enhancement in performance is not attributed to the increased dimensionality introduced by the redundancy from the ECC.

Figure: (a)
Refer to caption: x1.png

## 2 Our baseline: the uncoded DVAE

This section introduces a simplified version of the *uncoded* DVAE [, , ], serving as the foundational model upon which the subsequent aspects of our work are constructed. Let $\boldsymbol{X}=\{\boldsymbol{x}_{0},\ldots,\boldsymbol{x}_{N}\}$ denote a collection of unlabelled data, where $\boldsymbol{x}_{i}$ represents a K-dimensional feature vector. While [], [] and [] use Boltzmann machine priors, we consider a generative probabilistic model characterized by a simple low-dimensional binary latent variable $\boldsymbol{m}\in\{0,1\}^{M}$ comprising independent and identically distributed (i.i.d.) Bernoulli components $p(\boldsymbol{m})=\prod_{j=1}^{M}p(m_{j})=\prod_{j=1}^{M}\text{Ber}(\nu)$. Since backpropagation through discrete variables is generally not possible, a smoothing transformation of these binary variables is introduced. While the smoothing transformations proposed in [] are limited to spike-and-X type of transformations, [] show better results by using truncated exponential distributions:

$$ $\displaystyle p(\boldsymbol{z}|\boldsymbol{m})=\prod_{j=1}^{M}p(z_{j}|m_{j}), \quad p(z_{j}|m_{j})$ $\displaystyle=\frac{e^{-\beta(z_{j}-m_{j})}}{Z_{\beta}},$ (1) $$

for $m_{j}\in\{0,1\}$, $z_{j}\in[0,1]$, and $Z_{\beta}=(1-e^{-\beta})/\beta$. The parameter $\beta$ serves as an inverse temperature term, similar to the one in the Gumbel-Softmax relaxation []. Given the simplicity of the defined binary prior, the complexity of the model is primarily determined by the likelihood function $p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})=p(f_{\boldsymbol{\theta
}}(\boldsymbol{z}))$, where the likelihood is a Neural Network (NN) (referred to as the decoder) with parameter set $\boldsymbol{\theta}$.

Variational family and inference

Following [], we assume an amortized variational family of the following form:

$$ $\displaystyle q_{\boldsymbol{\eta}}(\boldsymbol{m},\boldsymbol{z}|\boldsymbol{ x})$ $\displaystyle=q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})p( \boldsymbol{z}|\boldsymbol{m}),\quad q_{\boldsymbol{\eta}}(\boldsymbol{m}| \boldsymbol{x})=\prod_{j=1}^{M}\text{Ber}(g_{j,\boldsymbol{\eta}}(\boldsymbol{ x})),$ (2) $$

where $g_{\boldsymbol{\eta}}(\boldsymbol{x})$ represents a parameterized function; here, a NN (referred to as the encoder) with parameter set $\boldsymbol{\eta}$. Inference is achieved by maximizing the Evidence Lower Bound (ELBO), which can be expressed as

$$ $\begin{gathered}\log p(\boldsymbol{x})\geq\int q_{\boldsymbol{\eta}}( \boldsymbol{m},\boldsymbol{z}|\boldsymbol{x})\log\left(\frac{p_{\boldsymbol{ \theta}}(\boldsymbol{x},\boldsymbol{z},\boldsymbol{m})}{q_{\boldsymbol{\eta}}( \boldsymbol{m},\boldsymbol{z}|\boldsymbol{x})}\right)d\boldsymbol{m}d \boldsymbol{z}=\mathbb{E}_{q_{\boldsymbol{\eta}}(\boldsymbol{m},\boldsymbol{z} |\boldsymbol{x})}\log\left(\frac{p_{\boldsymbol{\theta}}(\boldsymbol{x}| \boldsymbol{z})p(\boldsymbol{z}|\boldsymbol{m})p(\boldsymbol{m})}{q_{ \boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})p(\boldsymbol{z}|\boldsymbol{ m})}\right)\\ =\mathbb{E}_{q_{\boldsymbol{\eta}}(\boldsymbol{m},\boldsymbol{z}|\boldsymbol{x })}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})-\mathcal{D}_{KL }\big{(}q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})||p(\boldsymbol{m} )\big{)},\end{gathered}$ (3) $$

where the first term corresponds to the reconstruction of the observed data and the second term is the Kullback-Leibler (KL) Divergence between the variational family and the binary prior distribution, which acts as a regularization term. This can be computed in closed form as $\mathcal{D}_{KL}\big{(}q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})||p
(\boldsymbol{m})\big{)}=\sum_{j=1}^{M}\Big{[}q_{j}\log\frac{q_{j}}{\nu}+(1-q_{
j})\log\frac{1-q_{j}}{1-\nu}\Big{]}$, where $q_{j}=q_{\boldsymbol{\eta}}(m_{j}=1|\boldsymbol{x})$. The reconstruction term needs to be approximated via Monte Carlo. Since $p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})$ does not depend on the binary latent variable $\boldsymbol{m}$, we can marginalize the posterior distribution as

$$ $\displaystyle q_{\boldsymbol{\eta}}(\boldsymbol{z}|\boldsymbol{x})$ $\displaystyle=\prod_{j=1}^{M}q_{\boldsymbol{\eta}}(z_{j}|\boldsymbol{x}),\quad q _{\eta}(z_{j}|\boldsymbol{x})=\sum_{k=0}^{1}q_{\eta}(m_{j}=k|\boldsymbol{x})p( z_{j}|m_{j}=k).$ (4) $$

As shown in [], the corresponding inverse Cumulative Density Function (CDF) is given by

$$ $F^{-1}_{q_{\boldsymbol{\eta}}(z_{j}|\boldsymbol{x})}(\rho)=-\frac{1}{\beta} \log\left(\frac{-b+\sqrt{b^{2}-4c}}{2}\right),$ (5) $$

where $b=\big{(}\rho+e^{-\beta}(q_{j}-\rho)\big{)}/(1-q_{j})-1$ and $c=-[q_{j}e^{-\beta}]/(1-q_{j})$. The (5) is a differentiable function that converts a sample $\rho$ from an independent uniform distribution $\mathcal{U}(0,1)$ into a sample from $q_{\eta}(\boldsymbol{z}|\boldsymbol{x})$. Thus, we can apply the reparameterization trick to sample from the latent variable $\boldsymbol{z}$ and optimize the ELBO with respect to the model’s parameters. Algorithm 1 outlines the training procedure for the uncoded DVAE.

Figure: Algorithm 1 Training the model with uncoded inference.

## 3 Improving inference by adding redundancy to latent vectors

In the DVAE framework [, , ], the authors use Boltzmann machines as priors instead of the independent prior $p(\boldsymbol{m})$ presented in Section 2. While these complex priors increase the model’s flexibility and can produce competitive results, our objective is to enhance inference through model design by maintaining the simpler independent prior. This approach would improve interpretability and encourage the model to learn independent components in the latent space, which is essential for capturing potentially disentangled representations.

VAEs [] are often viewed as lossy compression models, where the goal is to minimize reconstruction error while imposing regularization through a prior distribution. However, our approach is better understood from a generative standpoint. We first sample a latent vector $\boldsymbol{m}$, generate an observation $\boldsymbol{x}$, and focus on minimizing the error rate when recovering $\boldsymbol{m}$ from $\boldsymbol{x}$. Achieving this requires the variational approximation to be sufficiently accurate. In fields where reliable data transmission or storage is important, introducing ECCs is a well-established approach to reduce the error rate when estimating a discrete source $\boldsymbol{m}$ transmitted through a noisy channel with output $\boldsymbol{x}$. Estimating $\boldsymbol{m}$ from $\boldsymbol{x}$ implies approximating the true and unknown posterior distribution $p(\boldsymbol{m}|\boldsymbol{x})$ with a proposed $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$. The gap between $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$ and $p(\boldsymbol{m}|\boldsymbol{x})$ is precisely the variational gap. We propose employing ECCs to safeguard $\boldsymbol{m}$ with controlled and known redundancy that can be leveraged by the variational posterior $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$ by design. This way, it is possible to reduce the mistakes committed when comparing $\boldsymbol{m}$ with samples drawn from $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$, obtaining a tighter approximation to the true posterior $p(\boldsymbol{m}|\boldsymbol{x})$, therefore reducing the gap to optimal inference.

ECCs play a crucial role in information theory and digital communications by enabling reliable data transmission over unreliable channels []. They introduce redundancy into the transmitted data, allowing the receiver to detect errors and, in many cases, correct them without retransmission. In his seminal work, Shannon [] demonstrated that arbitrarily reliable communication is possible through error correction. Our approach builds on the idea that the generative model in Fig. 1(c) can be conceptualized as a communication system, where the bits sampled from $p(\boldsymbol{m})$ undergo continuous modulation into $\boldsymbol{z}$ and are then transmitted through a nonlinear communication channel (in this setting, the decoder NN) characterized by the input/output response $p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})=p(f_{\boldsymbol{\theta
}}(\boldsymbol{z}))$. In this scenario, the complexity of the channel is essential since it is necessary to account for the intricate nature of the data at its output (e.g., complex images). Following this idea, the process of inference via $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$ can be thought of as deciphering the latent variable $\boldsymbol{m}$ given the observed data $\boldsymbol{x}$, where the encoder NN plays the role of the channel equalizer, trying to reverse the channel’s effects without knowing the bit correlations from the ECC.

## 4 Coded DVAE

This section extends the previously described DVAE, introducing an ECC over $\boldsymbol{m}$. We refer to this model as coded DVAE. In ECCs, we augment the dimensionality of the binary latent space from $M$ to $D$ in a controlled and deterministic manner, where $R=M/D$ is the *coding rate*. An ECC is typically designed so that the $2^{M}$ possible codewords are separated as much as possible in the space of binary vectors of $D$ bits. This facilitates algorithms in detecting and/or correcting errors by searching for the nearest code word. A random choice of the codewords brings what is known as a *random block code* []. While they are known to be very robust and amenable to theoretical analysis, their lack of structure makes them computationally intractable since we have to rely on codeword enumeration during the encoding/decoding process. In Appendix J, we include the formulation of a random code’s encoding/decoding process within the DVAE model.

Instead, we adopt a much simpler linear coding scheme, namely repetition codes. In a repetition code, each bit of the original message $\boldsymbol{m}$ is repeated multiple times to create the encoded message $\boldsymbol{c}$. Intuitively, the more times an information bit is repeated, the better it is protected. Our experiments consider uniform $(M,D)$ repetition codes where all bits are repeated $L$ times, resulting in codewords of dimension $D=ML$ and a coding rate of $R=1/L$. Note that repetition codes represent a special case of linear ECCs since each codeword can be deterministically computed by multiplying a binary vector $\boldsymbol{m}$ by an $M\times D$ *generator matrix* $\mathbf{G}$, such that $\boldsymbol{c}=\boldsymbol{m}^{T}\mathbf{G}$, where $u$-th row, with $u=1,\ldots,M$, has entries equal to one at columns $L(u-1)+1,L(u-1)+2,\dots,Lu$, and zero elsewhere. For example, for $M=3$ and $L=2$, the generator matrix of the $(3,6)$ repetition code is

$$ $\mathbf{G}=\begin{bmatrix}1&1&0&0&0&0\\ 0&0&1&1&0&0\\ 0&0&0&0&1&1\end{bmatrix}.$ (6) $$

The generative process of the coded DVAE follows similarly to the uncoded case, and it is represented in Fig. 1(d). We assume the same prior distribution $p(\boldsymbol{m})$, but in this case the samples $\boldsymbol{m}$ are deterministically encoded using $\mathbf{G}$. Now, the smoothing $\boldsymbol{z}$ transformations are defined over $\boldsymbol{c}$

$$ $\displaystyle p(\boldsymbol{z}|\boldsymbol{c})=\prod_{j=1}^{D}p(z_{j}|c_{j}), \quad p(z_{j}|c_{j})$ $\displaystyle=\frac{e^{-\beta(z_{j}-c_{j})}}{Z_{\beta}},$ (7) $$

for $z_{j}\in[0,1]$, $c_{j}\in\{0,1\}$ and $Z_{\beta}=(1-e^{-\beta})/\beta$. The likelihood $p(\boldsymbol{x}|\boldsymbol{z})$ is again of the form $p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})=p(f_{\boldsymbol{\theta
}}(\boldsymbol{z}))$. Note that, compared to the uncoded case, we have a larger input dimensionality to the decoder NN $f_{\boldsymbol{\theta}}(\boldsymbol{z})$. When comparing uncoded vs. coded DVAEs, the structure of the decoder NN $f_{\boldsymbol{\theta}}(\boldsymbol{z})$ (detailed in Appendix A) is equal in both cases except for the first Multilayer Perceptron (MLP) layer that attacks the input $\boldsymbol{z}$. Therefore, if a rate $R=1/L$ repetition code is used, the number of additional parameters of the $f_{\boldsymbol{\theta}}(\boldsymbol{z})$ NN is given by $(L-1)\times h$, where $h$ is the dimension of the first hidden space of $f_{\boldsymbol{\theta}}(\boldsymbol{z})$.

Variational family and inference

The repetition code introduces correlations between the bits in $\boldsymbol{c}$ that we will exploit to obtain an improved variational bound. We again assume a variational family factorizing as

$$ $q_{\boldsymbol{\eta}}(\boldsymbol{m},\boldsymbol{z}|\boldsymbol{x})=q_{ \boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})p(\boldsymbol{z}|\boldsymbol{ c})$ (8) $$

where $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})=\prod_{u=1}^{M}q_{
\boldsymbol{\eta}}(m_{u}|\boldsymbol{x})$ is computed in two steps. First, we construct an encoder NN $\mathbf{g}_{\boldsymbol{\eta}}(\boldsymbol{x})$ similar to that of (2), that retrieves the probabilities of the bits in $\boldsymbol{c}$ from $\boldsymbol{x}$ without exploiting the correlations introduced by the repetition code:

$$ $\displaystyle q^{u}_{\boldsymbol{\eta}}(\boldsymbol{c}|\boldsymbol{x})$ $\displaystyle=\prod_{j=1}^{D}\text{Ber}(g_{j,\boldsymbol{\eta}}(\boldsymbol{x} )),$ (9) $$

where the $u$ superscript serves as a reminder that this posterior does not exploit the redundancy introduced by the ECC.

Now, we utilize the known redundancy introduced by the ECC to constrain the solution of $q_{\boldsymbol{\eta}}^{u}(\boldsymbol{c}|\boldsymbol{x})$, given that each bit from $\boldsymbol{m}$ has been repeated $L$ times to create $\boldsymbol{c}$. To do so, we follow a *soft decoding* approach, where the marginal posteriors of the information bits are derived from the marginal posteriors of the encoded bits, exploiting the repetition code’s known structure. In the case of repetition codes, we compute the all-are-zero and the all-are-ones products of probabilities of the bits in $\boldsymbol{c}$ that are copies of the same message bit and renormalize as

$$ $\displaystyle q(m_{u}=1|\boldsymbol{x})$ $\displaystyle=\frac{1}{Z}\prod_{j=L(u-1)+1}^{Lu}g_{j,\boldsymbol{\eta}}( \boldsymbol{x})\doteq\frac{g^{+}_{u,\boldsymbol{\eta}}(\boldsymbol{x})}{Z},$ (10) $\displaystyle q(m_{u}=0|\boldsymbol{x})$ $\displaystyle=\frac{1}{Z}\prod_{j=L(u-1)+1}^{Lu}(1-g_{j,\boldsymbol{\eta}}( \boldsymbol{x}))\doteq\frac{g^{-}_{u,\boldsymbol{\eta}}(\boldsymbol{x})}{Z},$ (11) $$

for $u=1,\ldots,M$ and $Z=\left(g^{+}_{u,\boldsymbol{\eta}}(\boldsymbol{x})+g^{-}_{u,\boldsymbol{\eta}
}(\boldsymbol{x})\right)^{-1}$. This approach can be seen as a soft majority voting strategy, enabling the recovery of the original information vector even if some bits in the inferred encoded word are corrupted. All operations in (10) preserve the gradients concerning the parameters in the encoder $\mathbf{g}_{\boldsymbol{\eta}}(\boldsymbol{x})$. We implement them in the log domain for stability.

When compared to the uncoded case, as in the likelihood term, we consider the same NN structure for the encoder $\mathbf{g}_{\boldsymbol{\eta}}(\boldsymbol{x})$ where both cases only differ in the last MLP layer. The additional overhead in the coded cases requires $(L-1)\times h^{\prime}$ parameters in the last layer, where $h^{\prime}$ is the dimension at the output of the last layer. In Appendix H.2, we conduct an ablation study on the number of trainable parameters to demonstrate that the improvement in performance does not stem from this increase in the number of parameters, but rather from the incorporation of the ECC in the latent space.

Figure: Algorithm 2 Training the coded DVAE with repetition codes.

Soft encoding for efficient reparameterization

Given the variational family in (8), the ELBO matches the expression in (3). However, the reparameterization trick in (5) requires independent bits, which is not the case in $\boldsymbol{c}$. To efficiently circumvent this issue during training, we employ a soft encoding approach. With soft encoding, a marginal probability is computed for each bit in the codeword $\boldsymbol{c}$, taking into account the structure of the ECC and the marginal probabilities of the information bits. For a repetition code, this involves simply replicating the posterior probabilities $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})=\prod_{u=1}^{M}q_{
\boldsymbol{\eta}}(m_{u}|\boldsymbol{x})$ for each copy of the same information bit. Hence, we treat the bits in $\boldsymbol{c}$ as independent but distributed according to $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$. The Algorithm 2 shows the training pseudo-code.

When marginalizing $\boldsymbol{c}$ using the soft encoding marginals, we disregard potential correlations between the coded bits. For instance, with repetition codes, sampling from the marginals could produce inconsistent bits, leading to an invalid codeword. However, since we do not sample the coded bits during training but instead propagate their marginal probabilities, we consider this approximation to have minimal negative impact. In fact, it can be seen as a form of probabilistic dropout, which enhances robustness during training. It is important to note that when sampling from the generative model in test time, we use hard bits encoded into valid codewords, yielding visually appealing samples, indicating that our training approach is reliable.

Related work

While the use of deep neural networks and generative models in digital communications problems have been profusely reported in recent years (see [], [], [], [], and [] for representative examples), the use of ECC techniques as a design tool in machine learning is scarce. The most prominent example is [], where the authors address multiclass learning problems via ECCs. In [], the authors proposed using linear codes for applications that may require removing the trace of a sample from the system, e.g., a user requests their data to be deleted or corrupted data is discovered. They address a regression problem by introducing a coded learning protocol that employs linear encoders to divide the training data before the learning phase. More recently, in [] the authors introduced ECCs to improve code-to-code translation using transformers.

### 4.1 Reconstruction

We first evaluate the model’s performance of reconstructing data by examining its uncoded and coded versions across different configurations, varying the number of information bits and code rates. The introduction of the repetition code led to improved reconstruction and smaller KL values, indicating that the posterior latent features are disentangled and less correlated. In Appendices B, C, D and E, we show the behavior of the ELBO loss function for all the models and datasets.

Image reconstruction quality.   In the table included in Fig. 2, we first quantify the quality of the reconstructions in FMNIST by measuring the Peak Signal-To-Noise Ratio (PSNR) in the test set. The results for the rest of the datasets are provided in sections B, C, D and E of the Appendix. In all the cases, the coded models yield higher PSNR values than their uncoded counterparts, indicating a superior performance in reconstruction. This improvement is also evident by visual inspection of Fig. 2, where the coded models exhibit a greater ability to capture details in the images for the same latent dimension. We observe a general improvement in PSNR as we increase the number of information bits, i.e., as we augment the latent dimensionality of the model. This increase in the number of available latent vectors provides greater flexibility, enabling the models to capture the underlying structure of the data more effectively. We also observe a general improvement in PSNR as we decrease the code rate, i.e., as we add more redundancy. Note that adding redundancy does not increase the model’s flexibility, since the information bits determine the number of latent vectors. However, coded models yield more accurate and detailed reconstructions.

Figure: Figure 2: Reconstruction performance over the test set in FMNIST. The figure at the left shows an example of reconstructed test images obtained with different model configurations. Observe that more details are visualized as we increase the number of bits in the latent space and decrease the coding rate. The table at the right includes reconstruction metrics. Acc is the semantic accuracy and Conf. Acc the confident semantic accuracy. Entropy is the average entropy of $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$ in the test set.
Refer to caption: x5.png

Semantic accuracy.   As the PSNR operates at the pixel level, it does not account for the semantic errors committed by the model. For example, if the model incorrectly reconstructs a nine instead of a four in the MNIST dataset, the PSNR may still yield a large value due to the similarity between the two images. Nonetheless, this would represent a severe failure in correct reconstruction of the intended class. Therefore, we additionally evaluate the reconstruction accuracy, ensuring that the model successfully reconstructs images within the same class as the original ones. For this purpose, we trained an image classifier for each dataset and compared the reconstructed images’ predicted labels against the originals’ ground truth labels. Additionally, we provide a confident reconstruction accuracy. While the reconstruction accuracy is computed across the entire dataset partitions, for the confident accuracy, we only consider those images projected into a latent vector with a probability exceeding 0.4.(^1^11Namely, we do not count errors when the Maximum a Posteriori (MAP) value of $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$ is below 0.4.). Results for FMNIST are detailed in the table included in Fig.2, and corresponding results for MNIST are available in Table 2 within Appendix C. In light of the results, we can conclude that introducing an ECC in the model allows for latent spaces that better capture the semantics of the images while employing the same number of latent vectors, significantly outperforming the uncoded models in terms of accuracy in all the cases.

Figure: Figure 3: Example of erroneous reconstructions in FMNIST using the 4 most-probable a posteriori latent vectors. The first column in each image shows the original input to the model, while the second column displays the reconstructions. The $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$ probability is indicated in each row.
Refer to caption: x6.png

Posterior uncertainty calibration.   Finally, also in the table included in Fig.2, we report the average entropy of the variational posterior $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$ over the test set. The low entropy observed in the uncoded models suggests a low uncertainty when the model projects data points into the latent space, which could be advantageous if the model consistently assigned high probability to the correct latent vectors. However, the semantic accuracy results demonstrate this is not true in the uncoded model. In other words, the uncoded variational family projects images into the wrong class with high confidence. This indicates the uncertainty of the uncoded case is severely miscalibrated.

Coded models, on the other hand, improve semantic accuracy and present a larger entropy. This suggests that i) the coded DVAE is aware that multiple latent vectors might be related to the image class and ii) that the model posterior shows large uncertainties (high entropy) for certain images for which the model has not properly identified the class. We illustrate this in Fig. 3, where we show some images that were selected so that the MAP latent word from $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$ induces class reconstruction errors. We display the reconstruction of the 4 most probable latent vectors and their corresponding probabilities. Observe that the uncoded model is confident no matter the reconstruction outcome while, in the coded posterior, the uncertainty is much larger. These results are also observed for MNIST, indicating that the posterior distribution in coded models exhibits a better uncertainty calibration. Note also that the increase in the number of latent bits (from 8 to 10) does not result in an excessive increase in the entropy despite the exponential growth of the number of vectors.

### 4.2 Generation

Figure: Figure 4: Evaluation of generation in FMNIST. The figure at the left shows an example of randomly generated, uncurated FMNIST images. The table at the right shows the quantitative results on the evaluation of the BER, WER, and log-likelihood (LL).
Refer to caption: x7.png

In this section, we evaluate the model for the image generation task. In Fig. 4, we show examples of randomly generated images using different model configurations in FMNIST. Results for the rest of the datasets are available in Appendices B, C, D, and E. These results are consistent with the ones obtained in reconstruction since we can observe that the coded models can generate more detailed and diverse images. Both uncoded and coded models generate more intricate and varied images with increased information bits. However, if the number of latent vectors becomes too large for the dataset’s complexity, not all words in the codebook are specialized during model training. This leads to generation artifacts, images where different classes of objects are overlapped. A visual inspection of Fig. 4 suggests these artifacts are more frequent in the uncoded case. Note that, since we are dealing with discrete latent variables, we could simply detect and prune uninformative vectors.

Accuracy metrics in generation.   The improved inference given by the repetition code can also be tested by generating images using the generative model and counting errors using the MAP solution of the variational distribution $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$. The table included in Fig. 4 reports the BER and WER for FMNIST. As expected, at the same number of latent bits, the coded models significantly reduce both the BER and WER w.r.t. the uncoded case. Note also that the error rates grow with the number of latent bits, which is expected due to the increased complexity of the inference process. We may commit more errors by taking the MAP, but errors typically fall in consistent reconstructions (latent words that also reconstruct the same type of image), as the results presented in Section 4.1 indicated.

Log-likelihood.   We additionally estimated the log-likelihood (LL). Results for FMNIST with different model configurations, estimated through importance sampling with 300 samples per observation, are presented in the table included in Fig. 4, please refer to Appendix I for further details. We observe that coded models consistently outperform their uncoded counterparts for both train and test sets, aligning with the results previously presented. Thus, we can argue that the introduction of repetition codes in the definition of the model allows for an improved inference and tighter posterior approximation. We draw similar conclusions for other datasets.

We observe a general improvement in LL values as we increase the number of information bits, i.e., as we augment the latent dimensionality of the model and its flexibility. However, reducing the code rate does not lead to an improvement in log-likelihood. We argue that this might indicate overfitting of the decoder, as the LL deteriorates while reconstruction metrics improve. We must note that we use feed-forward networks at the decoder’s input. However, this may not be appropriate for the correlations we present in our coded words. We might overcome this overfitting tendency by using an architecture that properly leverages these coded bits correlations.

Figure: Figure 5: Reconstruction results for CIFAR10 (left) and Tiny ImageNet (right).
Refer to caption: x8.png

Figure: Figure 6: Generation results for CIFAR10 (left) and Tiny ImageNet (right).
Refer to caption: x10.png

### 4.3 Additional results with CIFAR10 and Tiny ImageNet

Since MNIST-like datasets are rather simple, it is difficult to assess the true gain in performance resulting from the introduction of ECCs proposed in our model. This section presents additional results using CIFAR10 and Tiny ImageNet, which contain colored images with more intricate shapes, patterns, and greater diversity than the previous datasets. We trained uncoded and coded models using different configurations to gain intuition regarding the effect of introducing the ECC. For a reference, in the case of the DVAE++ [], the authors needed 128 binary latent variables to achieve state-of-the-art performance in generation and reconstruction for this dataset. They employed a more intricate model than the one introduced in this study, featuring Boltzmann Machine priors. In Fig. 5 we show examples of reconstruction using different configurations of the model and in Fig. 6 we show examples of randomly generated images. Additional results are provided in Appendices D and E. The results are consistent with those presented in previous sections; but in this case, the difference in performance is even more pronounced. We observe that the uncoded DVAE cannot decouple spatial information from the images and project it in the latent space. Nevertheless, the coded DVAE shows particular promise for learning low-dimensional discrete latent representations in complex datasets. Note that we used a rather simple architecture as we want to focus on the gain obtained only by introducing ECCs in the latent space.

## 5 Beyond repetition codes

We have presented compelling proof-of-concept results that incorporating ECCs, like repetition codes, into DVAEs can improve performance. We believe this opens a new path for designing latent probabilistic models with discrete latent variables. Although a detailed analysis of the joint design of ECC and encoder-decoder networks is beyond the scope of this work, we will outline key properties that any ECCs must satisfy to be integrated within this framework.

- •
Scalable hard encoding $\big{(}\boldsymbol{m}\rightarrow\boldsymbol{c}\big{)}$. Our model requires hard encoding for generation once the model is trained. This process should have linear complexity in $M$.
- •
Scalable soft encoding $\big{(}p(\boldsymbol{m})\rightarrow p(\boldsymbol{c})\big{)}$. Soft encoding is required during training for reparameterization. This process should also have linear complexity in $M$.
- •
Scalable soft decoding $\big{(}p(\boldsymbol{c})\rightarrow p(\boldsymbol{m})\big{)}$. Our model employs soft-in soft-out (SISO) decoding during inference. This process should again be linearly complex in $M$.
- •
Differentiability. Both encoding and decoding processes must be differentiable w.r.t the inputs to enable gradient computation and backpropagation.

Figure: Figure 7: Graphical model of the hierarchical coded DVAE.
Refer to caption: x12.png

Since Shannon’s seminal work [], researchers have developed effective ECC schemes that meet these properties, including state-of-the-art ECCs such as Low Density Parity Check (LDPC) codes [], or polar codes []. Efforts have also focused on developing efficient SISO decoders, such as the sum-product algorithm [].

Inspired by polar codes [], we present a hierarchical coded DVAE with two layers of latent bits. In this model, the latent bits $\boldsymbol{m}_{1}$ are encoded using a repetition code in the first layer, producing $\boldsymbol{c}_{1}$ and $\boldsymbol{z}_{1}$. Simultaneously, the vector $\boldsymbol{m}_{2}$ is linearly combined with $\boldsymbol{m}_{1}$ using modulo 2 operations $\big{(}\boldsymbol{m}_{1}\oplus\boldsymbol{m}_{2}\big{)}$ and then encoded using a repetition code, yielding $\boldsymbol{c}_{2}$ and $\boldsymbol{z}_{2}$. Both soft vectors are concatenated and fed to the decoder NN to generate $\boldsymbol{x}$. The model provides stronger protection for $\boldsymbol{m}_{1}$, as it appears in both branches of the generative model. Inference follows a similar approach to the coded DVAE, incorporating the linear combination of $\boldsymbol{m}_{1}$ and $\boldsymbol{m}_{2}$ used in the second branch. This hierarchical structure allows the model to effectively separate high-level information from finer details, as we show in the results presented in Appendix G.

## 6 Conclusion

This paper presents the first proof-of-concept demonstration that safeguarding latent information with ECCs within deep generative models holds promise for enhancing overall performance. By integrating redundancy into the latent space, the variational family can effectively refine the inference network’s output according to the structure of the ECC. Our findings underscore the efficacy of simple and efficient ECCs, like repetition codes, showcasing remarkable improvements over a lightweight version of the DVAE introduced in [].

Furthermore, our work reveals numerous avenues for future research. Firstly, investigating decoder architectures capable of efficiently utilizing the correlations and structure introduced by the ECCs, in contrast to the feed-forward networks employed in this study. We also contemplate exploring more complex and robust coding schemes, conducting theoretical analyses aligned with Shannon’s channel capacity and mutual information concepts to determine the fundamental parameters of the ECC needed to achieve reliable variational inference, exploring different modulations, and integrating these concepts into state-of-the-art models based on discrete representations.

## 7 Acknowledgements

The work was partially supported by the Ministerio de Ciencia, Innovación y Universidades, Agencia Estatal de Investigación, under grant PID2021-123182OB-I00, MCIN/AEI/10.13039/501100011033, co-financed by European Regional Development Fund (ERDF), “A way of making Europe”. Pablo M. Olmos also acknowledges the support by BBVA Foundation (Beca Leonardo a Investigadores y Creadores Culturales 2024) and by the Comunidad de Madrid under grants IND2022/TIC-23550 and ELLIS Unit Madrid. David Mitchell acknowledges support from the National Science Foundation under Grant No. CCF-2145917 and travel support from the European Union’s Horizon 2020 research and innovation programme under Grant Agreement No 951847. Grace Villacrés acknowledges support from Comunidad de Madrid within the 2023-2026 agreement with Universidad Rey Juan Carlos for the granting of direct subsidies for the promotion and encouragement of research and technology transfer, Line of Action A: Emerging Doctors, under Project OrdeNGN.

## Appendix A Architecture

In this section, we detail the architecture used to obtain the experimental results with FMNIST and MNIST (28x28 gray-scale images). Note that across experiments we only modify the output layer of the encoder and the input layer of the decoder to adapt to the different configurations of the model. This modification leads to a minimal alteration in the total number of parameters. In Section H.2, we conduct an ablation study on the number of trainable parameters to show that the enhancement in performance is not attributed to the increased dimensionality introduced by redundancy. For the additional CIFAR10 experiments, we change the input of the encoder and the output of the decoder to process the 32x32 color images. For the Tiny ImageNet experiments, we do the same to process 64x64 color images. The rest of the architecture remains unchanged.

### A.1 Encoder

The encoder NN consists of 3 convolutional layers followed by two fully connected layers. We employed Leaky ReLU as the intermediate activation function and a Sigmoid as the output activation, as the encoder outputs bit probabilities. The full architecture is detailed in Fig. 8.

Figure: Figure 8: Block diagram of the encoder architecture for FMNIST and MNIST.
Refer to caption: x13.png

### A.2 Decoder

The decoder architecture is inspired by the one proposed in []. It is composed of two fully connected layers, followed by transposed convolutional layers with residual connections and Squeeze-and-Excitation (SE) layers []. We employed Leaky ReLU as the intermediate activation function and a Sigmoid as output activation, given that we consider datasets with gray-scale images. The complete architecture is detailed in Fig. 9.

Figure: Figure 9: Block diagram of the decoder architecture for FMNIST and MNIST.
Refer to caption: x14.png

## Appendix B FMNIST results

In this section, we present supplementary results obtained with the FMNIST dataset.

### B.1 Training

We present the evolution of the ELBO and its terms throughout the training process. The models were trained for 200 epochs using an Adam optimizer with a learning rate of $10^{-4}$, and a batch size of 128. Fig. 10 displays the results for configurations with 5 information bits, Fig. 11 for 8 information bits, and Fig. 12 for 10 information bits. The colors in all plots represent the various code rates.

Across all cases, coded models achieve superior bounds. The main differences in the ELBO come from the different performances in reconstruction. As we have observed in the different experiments, coded models are capable of generating more detailed images and accurate reconstructions. Introducing the repetition code also leads to smaller KL values, indicating that the posterior latent features are disentangled and less correlated.

We observe that, as we decrease the code rate, we obtain better bounds in general. Adding redundancy does not increase the model’s flexibility, since the information bits determine the number of latent vectors. However, the introduction of ECCs in the model allows for latent spaces that better capture the structure of the images while employing the same number of latent vectors.

Figure: Figure 10: Evolution of the ELBO during training with 5 information bits on FMNIST.
Refer to caption: x15.png

Figure: Figure 11: Evolution of the ELBO during training with 8 information bits on FMNIST.
Refer to caption: x16.png

Figure: Figure 12: Evolution of the ELBO during training with 10 information bits on FMNIST.
Refer to caption: x17.png

### B.2 Reconstruction and generation

In this section, we augment the results presented in the main text, including outcomes obtained with the training dataset in Table 1. We include again the results obtained in the test to facilitate comparison. The results remain consistent across the two data partitions, and the analysis conducted for the test set also applies to training data.

In all the cases, the coded models yield higher PSNR values than their uncoded counterparts, indicating a superior performance in reconstruction. We observe a general improvement in PSNR as we increase the number of information bits (i.e., as we augment the latent dimensionality of the model) and decrease the code rate (i.e., as we introduce more redundancy).

As we discussed in the main text, the PSNR does not account for the semantic errors committed by the model. Therefore, we additionally report the semantic accuracy and the confident semantic accuracy. While the reconstruction accuracy is computed across the entire dataset partitions, for the confident accuracy, we only consider those images projected into a latent vector with a probability exceeding $0.4$. We observe that coded models better capture the semantics of the images while employing the same number of latent vectors, significantly outperforming the uncoded models in terms of accuracy in all the cases.

In Fig 13 we include additional examples of randomly generated images using different model configurations. We observe that coded models can generate more detailed and diverse images than their uncoded counterparts.

Figure: Figure 13: Example of randomly generated, uncurated images using different model configurations.
Refer to caption: x18.png

**Table 1: Evaluation of reconstruction performance in FMNIST.**
| Model | PSNR<br>(train) | Acc<br>(train) | Conf. Acc.<br>(train) | PSNR<br>(test) | Acc<br>(test) | Conf. Acc.<br>(test) |
| --- | --- | --- | --- | --- | --- | --- |
| uncoded 5 | 14.490 | 0.541 | 0.541 | 14.477 | 0.536 | 0.536 |
| coded 5/50 | 16.375 | 0.656 | 0.702 | 16.241 | 0.647 | 0.700 |
| coded 5/80 | 16.824 | 0.694 | 0.751 | 16.624 | 0.688 | 0.748 |
| coded 5/100 | 17.001 | 0.708 | 0.760 | 16.702 | 0.700 | 0.757 |
| uncoded 8 | 15.644 | 0.601 | 0.602 | 15.598 | 0.594 | 0.595 |
| coded 8/80 | 17.877 | 0.769 | 0.842 | 17.318 | 0.750 | 0.816 |
| coded 8/160 | 18.828 | 0.807 | 0.878 | 17.713 | 0.783 | 0.831 |
| coded 8/240 | 19.345 | 0.831 | 0.921 | 17.861 | 0.799 | $\mathbf{0.893}$ |
| uncoded 10 | 16.053 | 0.650 | 0.652 | 16.000 | 0.644 | 0.648 |
| coded 10/100 | 18.827 | 0.813 | 0.885 | 17.694 | 0.790 | 0.850 |
| coded 10/200 | 19.937 | 0.846 | 0.897 | 18.009 | 0.814 | 0.871 |
| coded 10/300 | $\mathbf{20.529}$ | $\mathbf{0.855}$ | $\mathbf{0.907}$ | $\mathbf{18.111}$ | $\mathbf{0.817}$ | 0.870 |

## Appendix C MNIST results

In this section, we report the results obtained with the MNIST dataset.

### C.1 Training

We present the evolution of the ELBO and its terms throughout the training process. The models were trained for 100 epochs using an Adam optimizer with a learning rate of $10^{-4}$, and a batch size of 128. Fig. 14 displays the results for configurations with 5 information bits, Fig. 15 for 8 information bits, and Fig. 16 for 10 information bits. The colors in all plots represent the various code rates.

The results are consistent with the ones obtained for FMNIST. Across all the configurations, coded models achieve superior bounds. The main differences in the ELBO come from the different performances in reconstruction. As we have observed across the different experiments, coded models are capable of better capturing the structure of the data, generating more detailed images and accurate reconstructions.

We observe that, as we decrease the code rate, we obtain better bounds in general. Adding redundancy does not increase the model’s flexibility, since the information bits determine the number of latent vectors. However, the introduction of ECCs in the model allows for latent spaces that better capture the structure of the images while employing the same number of latent vectors.

Figure: Figure 14: Evolution of the ELBO during training with 5 information bits on MNIST.
Refer to caption: x19.png

Figure: Figure 15: Evolution of the ELBO during training with 8 information bits on MNIST.
Refer to caption: x20.png

Figure: Figure 16: Evolution of the ELBO during training with 10 information bits on MNIST.
Refer to caption: x21.png

### C.2 Reconstruction

We first evaluate the model’s performance in reconstructing data by examining its uncoded and coded versions across different configurations, varying the number of information bits and code rates. All the results obtained with MNIST are consistent with those presented in the main text for FMNIST.

In Table 2 we quantify the quality of the reconstructions measuring the PSNR in both training and test sets. In all the cases, coded models yield higher PSNR values, indicating a superior performance in reconstruction. This improvement is also evident through visual inspection of Fig. 17, where the coded models better capture the details in the images. As in FMNIST, we observe a general improvement of the PSNR as we increase the number of information bits and decrease the code rate.

As we discussed in the main text, the PSNR does not account for the semantic errors committed by the model. Therefore, we additionally evaluate the reconstruction accuracy, ensuring that the model successfully reconstructs images within the same class as the original ones. We also provide a confident reconstruction accuracy, for which we do not count errors when the MAP value of $q(\boldsymbol{m}|\boldsymbol{x})$ is below 0.4. In light of the results, we argue that introducing ECCs in the model allows for latent spaces that better capture the semantics of the images while employing the same number of latent vectors, outperforming the uncoded models in all the cases.

Figure: Figure 17: Example of reconstructed test images obtained with different model configurations. Observe that more details are visualized as we increase bits in the latent space and decrease the coding rate.
Refer to caption: x22.png

Figure: Figure 18: Example of erroneous reconstructions in MNIST using the 4 most-probable words (a posterior). The first column in each image shows the original input to the model, while the second column displays the reconstructions. The a posteriori probability of the word used for reconstruction is indicated in each row.
Refer to caption: x23.png

**Table 2: Evaluation of reconstruction performance in MNIST.**
| Model | PSNR<br>(train) | Acc<br>(train) | Conf. Acc.<br>(train) | PSNR<br>(test) | Acc<br>(test) | Conf. Acc.<br>(test) | Entropy |
| --- | --- | --- | --- | --- | --- | --- | --- |
| uncoded 5 | 13.483 | 0.702 | 0.703 | 13.483 | 0.701 | 0.702 | 0.277 |
| coded 5/50 | 14.983 | 0.887 | 0.923 | 14.888 | 0.887 | 0.920 | 2.073 |
| coded 5/80 | 15.436 | 0.899 | 0.936 | 15.263 | 0.895 | 0.929 | 2.237 |
| coded 5/100 | 15.590 | 0.905 | 0.931 | 15.352 | 0.898 | 0.924 | 2.382 |
| uncoded 8 | 14.530 | 0.860 | 0.864 | 14.490 | 0.860 | 0.868 | 0.513 |
| coded 8/80 | 16.878 | 0.937 | 0.964 | 16.042 | 0.912 | 0.947 | 3.105 |
| coded 8/160 | 18.108 | 0.957 | 0.974 | 16.497 | 0.927 | 0.951 | 3.645 |
| coded 8/240 | 19.984 | 0.967 | 0.978 | 16.688 | 0.936 | 0.957 | 3.881 |
| uncoded 10 | 14.879 | 0.888 | 0.891 | 14.816 | 0.887 | 0.890 | 0.636 |
| coded 10/100 | 17.584 | 0.945 | 0.972 | 16.795 | 0.928 | $\mathbf{0.968}$ | 4.080 |
| coded 10/200 | 20.060 | 0.973 | 0.977 | 16.863 | 0.932 | 0.944 | 4.411 |
| coded 10/300 | $\mathbf{21.083}$ | $\mathbf{0.979}$ | $\mathbf{0.984}$ | $\mathbf{17.114}$ | $\mathbf{0.941}$ | 0.945 | 4.810 |

We also report the average entropy of the variational posterior over the test set in Table 2. If we analyze the entropy together with the semantic accuracy, we can argue that coded VAE is aware that multiple vectors might be related to the same image class, and that the posterior shows larger uncertainties for images for which the model has not properly identified the class. We illustrate this argument in Fig. 18, where we show some images selected so that the MAP latent word of $q(\boldsymbol{m}|\boldsymbol{x})$ induces class reconstruction errors. We show the reconstruction of the 4 most probable latent vectors and their corresponding probabilities. Observe that the uncoded model is confident no matter the reconstruction outcome, while in the coded posterior, the uncertainty is much larger.

Table 4 shows the log-likelihood values obtained for the MNIST dataset with various model configurations. Coded models consistently outperform their uncoded counterparts for both the training and test sets, consistent with the findings observed using the FMNIST dataset.

### C.3 Generation

In this section, we evaluate the model in the image generation task. In Fig. 19, we show examples of randomly generated images using different model configurations in MNIST. These results are consistent with the ones obtained in reconstruction, and with the ones obtained for FMNIST, as we observe that the coded models can generate more detailed and diverse images.

The improved inference provided by the repetition code can also be tested by generating images using the generative model and counting errors using the MAP solution of the variational posterior distribution. Table 4 reports the BER and WER for MNIST. Remarkably, for the same number of latent bits, coded models reduce both the BER and WER w.r.t. the uncoded case. Note also that the error rates grow with the number of latent bits, but this is expected due to the increased complexity of the inference process.

Figure: Figure 19: Example of randomly generated, uncurated images using different model configurations.
Refer to caption: x24.png

**Table 3: Evaluation of the BER, WER in MNIST.**
| Model | BER | WER |
| --- | --- | --- |
| uncoded 5 | 0.002 | 0.008 |
| coded 5/50 | 0.007 | 0.034 |
| coded 5/80 | 0.004 | 0.021 |
| coded 5/100 | 0.009 | 0.045 |
| uncoded 8 | 0.015 | 0.071 |
| coded 8/80 | 0.020 | 0.147 |
| coded 8/160 | 0.021 | 0.160 |
| coded 8/240 | 0.023 | 0.167 |
| uncoded 10 | 0.057 | 0.373 |
| coded 10/100 | 0.030 | 0.258 |
| coded 10/200 | 0.034 | 0.282 |
| coded 10/300 | 0.041 | 0.331 |

## Appendix D CIFAR10 results

In this section, we provide additional results using the CIFAR10 dataset with different model configurations.

### D.1 Training

We present the evolution of the ELBO and its terms throughout the training process. The models were trained for 300 epochs using Adam optimizer with a learning rate of $10^{-4}$, and a batch size of 128. Fig. 20 displays the results for configurations with 70 information bits, Fig. 21 for 100 information bits, and Fig. 22 for 130 information bits. The colors in all plots represent the various code rates.

Across all the configurations, coded models achieve superior bounds. The main differences in the ELBO come from the different performances in reconstruction. As we have observed across the different experiments, coded models are capable of better capturing the structure of the data, generating more detailed images and accurate reconstructions.

Figure: Figure 20: Evolution of the ELBO during training with 70 information bits on CIFAR10.
Refer to caption: x25.png

Figure: Figure 21: Evolution of the ELBO during training with 100 information bits on CIFAR10.
Refer to caption: x26.png

Figure: Figure 22: Evolution of the ELBO during training with 130 information bits on CIFAR10.
Refer to caption: x27.png

In the coded case, we do not observe significant differences in the obtained bounds as we increase the number of information bits and reduce the code rate. However, the difference is notable if we compare the coded and uncoded models. Adding redundancy does not increase the model’s flexibility, since the information bits determine the number of latent vectors. However, the introduction of ECCs in the model allows for latent spaces that better capture the structure of the images while employing the same number of latent vectors.

It’s important to note that we are currently using feed-forward networks at the decoder’s input. However, this approach may not be suitable for the correlations present in our coded words. Utilizing an architecture capable of effectively leveraging these correlations among the coded bits could potentially enable us to better exploit the introduced redundancy.

### D.2 Reconstruction and generation

We first evaluate the model’s performance in reconstructing data by examining its uncoded and coded versions across different configurations, varying the number of information bits and code rates.

In Table 6 we quantify the quality of the reconstructions measuring the PSNR in both training and test sets. Coded models yield higher PSNR values in train, and similar values in test, although the coded models with lower rates outperform the rest of the configurations. However, the improvement in reconstruction is evident through visual inspection of Fig. 23, where the coded models better capture the details in the images. We observe that the coded model yields images that better resemble the structure of the dataset, while the uncoded DVAE cannot decouple spatial information from the images and project it in the latent space.

We hypothesize that to adequately model complex images, transitioning to a hierarchical structure may be necessary. This would allow for the explicit modeling of both global and local information. However, despite employing this rather simple model, we observe that coded configurations outperform their uncoded counterparts in capturing colors and textures.

We also evaluate the model in the image generation task. In Fig. 24, we show examples of randomly generated images using different model configurations in CIFAR10. These results are consistent with the ones obtained in reconstruction, as we observe that the coded models can generate more detailed and diverse images. Additionally, we obtained the Fréchet Inception Distance (FID) score using the test set and 10k generated samples. For this, we used the implementation available at [https://github.com/mseitzer/pytorch-fid](https://github.com/mseitzer/pytorch-fid). We can observe that the coded models significantly reduced the FID score in all the cases compared to their uncoded counterparts.For the coded models, we do not observe a clear influence of the code rate on the quality of the generations.

Figure: Figure 23: Example of reconstructed test images obtained with different model configurations. Observe that more details are visualized as we increase bits in the latent space and introduce redundancy.
Refer to caption: x28.png

Figure: Figure 24: Example of randomly generated, uncurated images using different model configurations.
Refer to caption: x29.png

**Table 5: Evaluation of reconstruction performance in CIFAR10 with different model configurations.**
| Model | PSNR (train) | PSNR (test) |
| --- | --- | --- |
| uncoded 70 | 17.985 | 17.596 |
| coded 70/700 | 23.790 | 17.731 |
| coded 70/1400 | 24.555 | 18.008 |
| coded 70/2100 | 25.551 | 18.401 |
| uncoded 100 | 18.509 | 18.334 |
| coded 100/1000 | 24.754 | 18.229 |
| coded 100/2000 | 24.866 | 18.927 |
| coded 100/3000 | 25.646 | 18.920 |
| uncoded 130 | 18.951 | 18.758 |
| coded 130/1300 | 25.007 | 18.887 |
| coded 130/2600 | 25.460 | 19.416 |
| coded 130/3900 | $\mathbf{25.515}$ | $\mathbf{19.292}$ |

## Appendix E Tiny ImageNet results

In this section, we provide additional results using the Tiny ImageNet dataset with different model configurations.

### E.1 Training

We present the evolution of the ELBO and its terms throughout the training process. The models were trained for 300 epochs using an Adam optimizer with a learning rate of $10^{-4}$, and a batch size of 128. Fig. 25 displays the results for configurations with 70 information bits, Fig. 26 for 100 information bits, and Fig. 27 for 130 information bits. The colors in all plots represent the various code rates.

As in the rest of the datasets, coded models achieve superior bounds across all the configurations. The main differences in the ELBO come from the different performances in reconstruction. As we have observed across the different experiments, coded models are capable of better capturing the structure of the data, generating more detailed images and accurate reconstructions.

In the coded case, we do not observe significant differences in the obtained bounds as we increase the number of information bits and reduce the code rate. However, the difference is notable if we compare the coded and uncoded models. Adding redundancy does not increase the model’s flexibility, since the information bits determine the number of latent vectors. However, the introduction of ECCs in the model allows for latent spaces that better capture the structure of the images while employing the same number of latent vectors.

Figure: Figure 25: Evolution of the ELBO during training for the configurations with 70 information bits.
Refer to caption: x30.png

Figure: Figure 26: Evolution of the ELBO during training for the configurations with 100 information bits.
Refer to caption: x31.png

Figure: Figure 27: Evolution of the ELBO during training for the configurations with 130 information bits.
Refer to caption: x32.png

It’s important to note that we are currently using feed-forward networks at the decoder’s input. However, this approach may not be suitable for the correlations present in our coded words. Utilizing an architecture capable of effectively leveraging these correlations among the coded bits could potentially enable us to better exploit the introduced redundancy.

### E.2 Reconstruction and generation

We first evaluate the model’s performance in reconstructing data by examining its uncoded and coded versions across different configurations, varying the number of information bits and code rates.

In Table 8 we quantify the quality of the reconstructions measuring the PSNR in both training and test sets. Coded models yield higher PSNR values in train, and similar values in test, although the coded models with lower rates outperform the rest of the configurations. However, the improvement in reconstruction is evident through visual inspection of Fig. 28, where the coded models better capture the details in the images. We observe that the coded model yields images that better resemble the structure of the dataset, while the uncoded DVAE cannot decouple spatial information from the images and project it in the latent space.

We hypothesize that to adequately model complex images, transitioning to a hierarchical structure may be necessary. This would allow for the explicit modeling of both global and local information. However, despite employing this rather simple model, we observe that coded configurations outperform their uncoded counterparts in capturing colors and textures.

Figure: Figure 28: Example of reconstructed test images obtained with different model configurations. Observe that more details are visualized as we increase the bits in the latent space and introduce redundancy.
Refer to caption: x33.png

Figure: Figure 29: Example of randomly generated, uncurated images using different model configurations.
Refer to caption: x34.png

**Table 7: Evaluation of reconstruction performance in Tiny ImageNet with different model configurations.**
| Model | PSNR (train) | PSNR (test) |
| --- | --- | --- |
| uncoded 70 | 15.598 | 15.402 |
| coded 70/700 | 18.156 | 15.158 |
| coded 70/1400 | 18.396 | 15.419 |
| coded 70/2100 | 18.228 | 15.789 |
| uncoded 100 | 16.012 | 15.774 |
| coded 100/1000 | 18.298 | 15.677 |
| coded 100/2000 | 18.647 | 15.892 |
| coded 100/3000 | 18.729 | 16.167 |
| uncoded 130 | 16.278 | 16.009 |
| coded 130/1300 | 18.719 | 15.901 |
| coded 130/2600 | 18.818 | $\mathbf{16.329}$ |
| coded 130/3900 | $\mathbf{19.020}$ | 16.288 |

We also evaluate the model in the image generation task. In Fig. 29, we show examples of randomly generated images using different model configurations in Tiny ImageNet. These results are consistent with the ones obtained in reconstruction, as we observe that the coded models can generate more detailed and diverse images. Additionally, we obtained the FID score using the test set and 10k generated samples. For this, we used the implementation available at [https://github.com/mseitzer/pytorch-fid](https://github.com/mseitzer/pytorch-fid). We can observe that the coded models significantly reduced the FID score in all the cases compared to their uncoded counterparts.

For the coded models, we do not observe a clear influence of the code rate on the quality of the generations in CIFAR-10. However, in Tiny ImageNet, smaller code rates produce worse FID scores. We hypothesize that this may be due to the presence of artifacts in the generated images. Our experiments indicate that coded models with lower rates attempt to model fine details in images, which can lead to artifacts in generation.

## Appendix F IWAE results

One could draw a parallel between the coded DVAE with repetition codes and the well-known IWAE [], but the two approaches are fundamentally different. In the IWAE, independent samples are drawn from the variational posterior and propagated independently through the generative model to obtain a tighter variational bound on the marginal log-likelihood. In our method, we jointly propagate the output of the ECC encoder through the generative model, obtaining a single prediction and exploiting the introduced known correlations in the variational approximation of the posterior. In the case of repetition codes, the ECC encoder outputs are repeated bits, or repeated probabilities in the case of soft encoding. However, our approach extends beyond repetition codes, opening a new field for improved inference in discrete latent variable models.

In this work, we specifically utilize the redundancy introduced by the repetition code to correct potential errors made by the encoder through a soft decoding approach, leading to a more accurate approximation of $p(\boldsymbol{m}|\boldsymbol{x})$ and an improved proposal for sampling. The results obtained in the coded DVAE case cannot be achieved by training the uncoded DVAE with the IWAE objective. The following results compare the uncoded IWAE model with the coded DVAE trained on FMNIST. While the uncoded model shows slight performance gains with an increasing number of IWAE samples (which improves the evidence lower bound), it still underperforms compared to the coded model. Furthermore, when using 20 and 30 IWAE samples, the metrics slightly declined compared to using 10 samples, likely due to overfitting, as we applied a common early stopping point.

Figure: Figure 30: Evolution of the ELBO and the IWAE objectives for various configurations. Observe that the IWAE provides a tighter bound than the ELBO in the uncoded setting. However, coded models obtain even better bounds using the same number of samples/repetitions.
Refer to caption: extracted/5915767/Figures/iwae_loss.png

**Table 9: Comparison of the metrics obtained with our method and the uncoded DVAE trained with the IWAE objective.**
| Model | BER | WER | Entropy | Acc. | Conf. Acc. | PSNR |
| --- | --- | --- | --- | --- | --- | --- |
| uncoded 8 | 0.089 | 0.384 | 0.467 | 0.594 | 0.595 | 15.598 |
| uncoded 8 IWAE 10 samples | 0.063 | 0.372 | 1.309 | 0.617 | 0.640 | 14.282 |
| uncoded 8 IWAE 20 samples | 0.075 | 0.447 | 1.391 | 0.634 | 0.651 | 14.237 |
| uncoded 8 IWAE 30 samples | 0.074 | 0.438 | 1.564 | 0.619 | 0.641 | 13.757 |
| coded 8/80 | 0.021 | 0.144 | 2.905 | 0.750 | 0.816 | 17.318 |
| coded 8/160 | 0.027 | 0.189 | 3.637 | 0.783 | 0.831 | 17.713 |
| coded 8/240 | 0.037 | 0.231 | 4.000 | 0.799 | 0.893 | 17.861 |

## Appendix G Hierarchical coded DVAE results

Inspired by polar codes [], we present a hierarchical coded DVAE with two layers of latent bits, as illustrated in Figure 31. In this model, the latent bits $\boldsymbol{m}_{1}$ are encoded using a repetition code in the first layer, producing $\boldsymbol{c}_{1}$ and $\boldsymbol{z}_{1}$. Concurrently, the bits in the second layer, $\boldsymbol{m}_{2}$, are linearly combined with $\boldsymbol{m}_{1}$ following $\boldsymbol{m}_{1,2}=\boldsymbol{m}_{1}\oplus\boldsymbol{m}_{2}$, considering a binary field or Galois field. The resulting vector is then encoded with another repetition code to produce $\boldsymbol{c}_{2}$, which is subsequently modulated into $\boldsymbol{z}_{2}$. Finally, both $\boldsymbol{z}_{1}$ and $\boldsymbol{z}_{2}$ are concatenated and passed through the decoder network to generate $\boldsymbol{x}$. The model provides stronger protection for $\boldsymbol{m}_{1}$, as it appears in both branches of the generative model. Inference follows a similar approach to the one employed in the coded DVAE, incorporating the linear combination of $\boldsymbol{m}_{1}$ and $\boldsymbol{m}_{2}$ used in the second branch.

Figure: Figure 31: Graphical model of the hierarchical Coded VAE with two layers.
Refer to caption: x35.png

We adopt the same variational family as used in the standard Coded VAE; however, in this case, we incorporate both hierarchical levels, leading to

$$ $q_{\boldsymbol{\eta}}(\boldsymbol{m},\boldsymbol{z}|\boldsymbol{x})=q_{ \boldsymbol{\eta}}(\boldsymbol{m}_{1}|\boldsymbol{x})q_{\boldsymbol{\eta}}( \boldsymbol{m}_{2}|\boldsymbol{x})p(\boldsymbol{z}_{1}|\boldsymbol{c}_{1})p( \boldsymbol{z}_{2}|\boldsymbol{c}_{2}),$ (12) $$

where $q_{\boldsymbol{\eta}}(\boldsymbol{m}_{1}|\boldsymbol{x})$ is calculated following the same approach as in the coded DVAE with repetition codes, computing the all-are-zero and all-are-ones products of probabilities of the bits in $\boldsymbol{c}_{1}$ that are copies of the same message bit. The posterior $q_{\boldsymbol{\eta}}(\boldsymbol{m}_{2}|\boldsymbol{x})$ considering both the encoder’s output and the inferred posterior distribution $q_{\boldsymbol{\eta}}(\boldsymbol{m}_{1}|\boldsymbol{x})$. Note that in this case, the decoder outputs the probabilities for both $\boldsymbol{c}_{1}$ and $\boldsymbol{c}_{2}$, with $\boldsymbol{c}_{2}$ being the encoded version of the linear combination $\boldsymbol{m}_{1,2}=\boldsymbol{m}_{1}\oplus\boldsymbol{m}_{2}$. Consequently, we first obtain $q_{\boldsymbol{\eta}}(\boldsymbol{m}_{1,2}|\boldsymbol{x})$ following the same approach as in the coded DVAE with repetition codes, and determine $q_{\boldsymbol{\eta}}(\boldsymbol{m}_{2}|\boldsymbol{x})$ as

$$ $q_{\boldsymbol{\eta}}(\boldsymbol{m}_{2}|\boldsymbol{x})=\prod_{u=1}^{M}\text{ Ber}(p_{u}),$ (13) $$

$$ $p_{u}=q_{\boldsymbol{\eta}}(m_{1,2,u}=1|\boldsymbol{x})q_{\boldsymbol{\eta}}(m _{1,u}=0|\boldsymbol{x})+q_{\boldsymbol{\eta}}(m_{1,2,u}=0|\boldsymbol{x})q_{ \boldsymbol{\eta}}(m_{1,u}=1|\boldsymbol{x}).$ (14) $$

After obtaining $q_{\boldsymbol{\eta}}(\boldsymbol{m}_{1}|\boldsymbol{x})$ and $q_{\boldsymbol{\eta}}(\boldsymbol{m}_{2}|\boldsymbol{x})$, we recalculate the posterior bit probabilities for the linear combination $q^{\prime}_{\boldsymbol{\eta}}(\boldsymbol{m}_{1,2}|\boldsymbol{x})$ as

$$ $q^{\prime}_{\boldsymbol{\eta}}(\boldsymbol{m}_{1,2}|\boldsymbol{x})=\prod_{u=1 }^{M}\text{Ber}(q_{u}),$ (15) $$

$$ $q_{u}=q_{\boldsymbol{\eta}}(m_{1,u}=1|\boldsymbol{x})q_{\boldsymbol{\eta}}(m_{ 2,u}=0|\boldsymbol{x})+q_{\boldsymbol{\eta}}(m_{1,u}=0|\boldsymbol{x})q_{ \boldsymbol{\eta}}(m_{2,u}=1|\boldsymbol{x}).$ (16) $$

Next, we apply the soft encoding approach to incorporate the repetition code structure at both levels of the hierarchy.The posterior probabilities $q_{\boldsymbol{\eta}}(\boldsymbol{m}_{1}|\boldsymbol{x})$ are repeated to obtain $q_{\boldsymbol{\eta}}(\boldsymbol{c}_{1}|\boldsymbol{x})$, and the posterior probabilities $q_{\boldsymbol{\eta}}(\boldsymbol{m}_{1,2}|\boldsymbol{x})$ are repeated to produce $q_{\boldsymbol{\eta}}(\boldsymbol{c}_{2}|\boldsymbol{x})$. Utilizing the reparameterization trick from Eq. 5, we sample $\boldsymbol{z}_{1}$ and $\boldsymbol{z}_{2}$, concatenate them to form $\boldsymbol{z}$, and pass this through the decoder to generate $p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})$. The model is trained by maximizing the ELBO, given by

$$ $\text{ELBO}=\mathbb{E}_{q_{\boldsymbol{\eta}}(\boldsymbol{m},\boldsymbol{z}| \boldsymbol{x})}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})- \mathcal{D}_{KL}\big{(}q_{\boldsymbol{\eta}}(\boldsymbol{m}_{1}|\boldsymbol{x} )||p(\boldsymbol{m}_{1})\big{)}-\mathcal{D}_{KL}\big{(}q_{\boldsymbol{\eta}}( \boldsymbol{m}_{2}|\boldsymbol{x})||p(\boldsymbol{m}_{2})\big{)},$ (17) $$

where both $p(\boldsymbol{m}_{1})$ and $p(\boldsymbol{m}_{2})$ are assumed to be independent Bernoulli distributions with bit probabilities of $0.5$, consistent with the other scenarios.

We obtained results on the FMNIST dataset using a model with 5 information bits per branch and repetition rates of $R=1/10$ and $R=1/20$. In this case, we applied the same code rate to both branches, although varying code rates could be used to control the level of protection at each hierarchy level. Tables 10 and 11 present the metrics obtained for the different configurations. Specifically, Table 10 shows the overall metrics obtained with this structure, and Table 11 compares the error metrics across the two hierarchy levels. As expected, $\mathbf{m}_{2}$ shows poorer error metrics compared to $\mathbf{m}_{1}$, since the model provides more redundancy to $\mathbf{m}_{1}$ incorporating it in both branches. Although the overall metrics and generation quality are somewhat similar to those of the coded DVAE with 10 information bits (see tables in Figures 2 and 4), the introduced hierarchy results in a more interpretable latent space. In this setup, $\mathbf{m}_{1}$ captures global features (such as clothing types in the FMNIST dataset), while $\mathbf{m}_{2}$ controls individual features, as we can observe in Figure 32, where we show examples of the model’s generative outputs for fixed $\mathbf{m}_{1}$ and random samples of $\mathbf{m}_{2}$.

Figure: (a)
Refer to caption: x36.png

**Table 10: Comparison of the obtained metrics for the coded DVAE with polar codes with different configurations, which we refer to as ’hierarchical coded DVAE’.**
| Model | BER | WER | Acc | Conf. Acc | PSNR |
| --- | --- | --- | --- | --- | --- |
| hier. 5/50 | 0.099 | 0.400 | 0.753 | 0.800 | 17.130 |
| hier. 5/100 | 0.050 | 0.330 | 0.784 | 0.870 | 17.513 |

**Table 11: Comparison of the obtained error metrics in the different hierarchy levels.**
| Model | BER $\boldsymbol{m}_{1}$ | WER $\boldsymbol{m}_{1}$ | BER $\boldsymbol{m}_{2}$ | WER $\boldsymbol{m}_{2}$ |
| --- | --- | --- | --- | --- |
| hier. 5/50 | 0.079 | 0.259 | 0.119 | 0.362 |
| hier. 5/100 | 0.026 | 0.110 | 0.075 | 0.287 |

## Appendix H Ablation study

In this section, we conduct ablation studies on the hyperparameter $\beta$ of the model, responsible for regulating the decay of exponentials in the smoothing transformation, as well as on the number of trainable parameters in the models.

### H.1 Ablation study on the hyperparameter β 𝛽 \beta italic_β

Across all experiments, we have consistently configured the hyperparameter $\beta$, which controls the decay of exponentials in the smoothing transformation, to a value of 15. To illustrate its impact on the overall performance of the model, we conducted an ablation study on the value of this hyperparameter for both uncoded and coded cases.

The smoothing distribution employed for the reparameterization trick consists of two overlapping exponentials. The hyperparameter $\beta$ functions as a temperature term, regulating the decay of the distributions and, consequently, influencing the degree of overlapping. A lower $\beta$ value results in more overlapped tails, while a higher value leads to less overlapped distributions. A priori, we would like these distributions to be separated, allowing us to retrieve the true value of the bit and effectively use the latent structure of the model.

#### H.1.1 Coded model

We first evaluate the influence of the parameter $\beta$ in coded models. We take as a reference the coded model with 8 information bits and a rate $R=1/30$, and train it using $\beta=5,10,15,20$. We assess the performance of the model in reconstruction and generation tasks. We observe the model is fairly robust, achieving similar performance across configurations in most metrics.

In Fig. 33 we show examples of reconstructed images using the different configurations to assess reconstruction through visual examination, and Table 12 contains the associated reconstruction metrics. All the configurations achieve similar performances, although the models trained with $\beta=10$ and $\beta=15$ seem to be the best configurations for this scenario. Larger values may result in unstable training and inferior performance.

Figure: Figure 33: Example of reconstructed images obtained with different values of $\beta$ using the coded model with an 8/240 code.
Refer to caption: x38.png

**Table 12: Evaluation of reconstruction performance in FMNIST with different values of $\beta$ using the coded model with an 8/240 code.**
| Model | PSNR<br>(train) | Acc<br>(train) | Conf. Acc.<br>(train) | PSNR<br>(test) | Acc<br>(test) | Conf. Acc.<br>(test) | Entropy |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 5 | 18.344 | 0.791 | 0.895 | 17.614 | 0.766 | 0.849 | 4.025 |
| 10 | 19.106 | 0.822 | 0.904 | 17.737 | 0.793 | 0.872 | 4.023 |
| 15 | 19.345 | 0.831 | 0.921 | 17.861 | 0.799 | 0.893 | 4.000 |
| 20 | 18.797 | 0.809 | 0.887 | 17.837 | 0.787 | 0.877 | 3.810 |

Figure: Figure 34: Example of randomly generated, uncurated images with different values of $\beta$ using the coded model with an 8/240 code.
Refer to caption: x39.png

**Table 13: Evaluation of the BER and WER in FMNIST with different values of $\beta$ using the coded model with an 8/240 code.**
| Beta | BER | WER |
| --- | --- | --- |
| 5 | 0.150 | 0.726 |
| 10 | 0.080 | 0.480 |
| 15 | 0.037 | 0.231 |
| 20 | 0.065 | 0.399 |

Next, we evaluate the model in the image generation task. Fig. 34 contains examples of randomly generated images using the different configurations. Table 14 reports the obtained BER and WER, and Table 14 the estimated log-likelihood of the different values of $\beta$. The model trained with $\beta=15$ stands out in terms of error metrics, although achieves similar log-likelihood values as the model trained with $\beta=10$. Again, these two configurations appear to be the most suitable in this scenario.

#### H.1.2 Uncoded model

We first evaluate the influence of the parameter $\beta$ in uncoded models. We take as a reference the coded model with 8 information bits and train it using $\beta=5,10,15,20$. We assess the performance of the model in reconstruction and generation tasks. We observe that the uncoded model is also robust, achieving similar performance across configurations.

In Fig. 35 we show examples of reconstructed images using the different configurations to assess reconstruction through visual examination, and Table 15 contains the associated reconstruction metrics. All the configurations achieve similar performances, although the models trained with $\beta=10$ and $\beta=15$ seem to be the best configurations for this scenario. Larger values may result in unstable training and inferior performance, as we can clearly observe in this case.

Figure: Figure 35: Example of reconstructed images obtained with different values of $\beta$ using an uncoded model 8 information bits.
Refer to caption: x40.png

Figure: Figure 36: Example of randomly generated, uncurated images with different values of $\beta$ using an uncoded model 8 information bits.
Refer to caption: x41.png

**Table 15: Evaluation of reconstruction performance in FMNIST with different values of $\beta$ using an uncoded model 8 information bits.**
| Model | PSNR<br>(train) | Acc<br>(train) | Conf. Acc.<br>(train) | PSNR<br>(test) | Acc<br>(test) | Conf. Acc.<br>(test) | Entropy |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 5 | 14.239 | 0.503 | 0.501 | 14.237 | 0.503 | 0.491 | 0.231 |
| 10 | 15.624 | 0.606 | 0.603 | 15.571 | 0.598 | 0.598 | 0.357 |
| 15 | 15.644 | 0.601 | 0.602 | 15.598 | 0.594 | 0.595 | 0.467 |
| 20 | 13.717 | 0.464 | 0.466 | 13.743 | 0.460 | 0.462 | 0.383 |

**Table 16: Evaluation of the BER and WER in FMNIST with different values of $\beta$ using an uncoded model 8 information bits.**
| Beta | BER | WER |
| --- | --- | --- |
| 5 | 0.203 | 0.852 |
| 10 | 0.086 | 0.384 |
| 15 | 0.089 | 0.384 |
| 20 | 0.278 | 0.939 |

Next, we evaluate the model in the image generation task. Fig. 36 contains examples of randomly generated images using the different configurations. Table 17 reports the obtained BER and WER, and Table 17 the estimated log-likelihood of the different values of $\beta$. The models trained with $\beta=15$ and $\beta=10$ clearly outperform the other two in this task, generating more diverse and detailed images, and obtaining better error metrics and log-likelihood values.

### H.2 Ablation study on the number of trainable parameters

A consistent architecture was employed across all experiments, which is detailed in Section A. However, since the introduction of the code alters the dimensionality of the latent space, it is necessary to adjust the encoder’s output and the decoder’s input. This results in an augmentation of the trainable parameters in the coded cases compared to their uncoded counterparts.

Given that a higher number of parameters usually results in better performance, we conducted an ablation study on the model’s trainable parameters to confirm that the improved performance introduced by the coded models is not due to this factor. We adjusted the hidden dimensions of the encoder and decoder architectures to ensure both configurations (coded and uncoded) have roughly the same number of trainable parameters. We have conducted the ablation study using the uncoded model with 8 bits and the coded 8/240 model trained on FMNIST.

We adjusted the encoder’s last hidden dimension and the decoder’s first hidden dimension to equalize the parameter count between the uncoded and coded models. This adjustment was straightforward since the last layers of the encoder and the first layers of the decoder are feed-forward layers. We kept the latent dimension of the model unchanged, ensuring that the modification solely pertained to the neural network architecture.

**Table 18: Parameter count.**
| Model | # encoder parameters | # decoder parameters |
| --- | --- | --- |
| uncoded 8 | 6,592,008 | 19,341,185 |
| uncoded 8 adjusted | 6,717,538 | 19,581,035 |
| coded 8/240 | 6,711,024 | 19,578,753 |
| coded 8/240 adjusted | 6,583,174 | 19,332,871 |

#### H.2.1 Evaluation

This section provides an empirical evaluation of the models trained with the adjusted parameter count, demonstrating that the enhanced performance observed in coded models does not result from an augmented number of trainable parameters. We found that the performance of the original and adjusted models is very similar, meaning that the conclusions drawn in the main text hold even in this scenario.

We first evaluate the reconstruction performance, measuring the PSNR and reconstruction accuracy in both train and test sets, which are included in Table 19. Then, we assess generation measuring the BER and WER, reported in Table 21. Finally, we compute the log-likelihood for train and test sets, shown in Table 21.

In terms of reconstruction, both the original and adjusted models exhibit very similar performance, observed in both reconstruction quality and accuracy. However, the adjusted models show slightly inferior results than the original ones. For the coded model, this might be attributed to reduced flexibility when decreasing the number of parameters. As for the uncoded model, the increased complexity while maintaining the same low-dimensional latent vector may not provide enough expressiveness to leverage the added flexibility in the architecture, potentially causing the observed decrease in performance. We observe the same behavior when we evaluate the BER and WER.

Analyzing the results, especially the log-likelihood values shown in Table 21, we can argue that increasing the flexibility in the architecture does not necessarily lead to improved performance in this scenario. The coded models exhibit similar performance with both the original and adjusted parameter counts, consistently outperforming the uncoded models. These results indicate that the performance enhancement is attributed to the introduction of ECCs in the latent space, rather than differences in the architecture required to handle the introduced redundancy.

Figure: Figure 37: Example of reconstructed images obtained with different configurations.
Refer to caption: x42.png

Figure: Figure 38: Example of randomly generated, uncurated images using different model configurations.
Refer to caption: x43.png

**Table 19: Evaluation of reconstruction performance in FMNIST with the adjusted parameter count.**
| Model | PSNR<br>(train) | Acc<br>(train) | Conf. Acc.<br>(train) | PSNR<br>(test) | Acc<br>(test) | Conf. Acc.<br>(test) | Entropy |
| --- | --- | --- | --- | --- | --- | --- | --- |
| uncoded 8 | 15.644 | 0.601 | 0.602 | 15.598 | 0.594 | 0.595 | 0.659 |
| uncoded 8 adj. | 15.530 | 0.586 | 0.586 | 15.491 | 0.581 | 0.580 | 0.449 |
| coded 8/240 | 19.345 | 0.831 | 0.921 | 17.861 | 0.799 | 0.893 | 4.609 |
| coded 8/240 adj. | 19.383 | 0.828 | 0.883 | 17.771 | 0.792 | 0.828 | 3.952 |

**Table 20: Evaluation of the BER and WER in FMNIST with the adjusted parameter count.**
| Model | BER | WER |
| --- | --- | --- |
| uncoded 8 | 0.089 | 0.384 |
| uncoded 8 adj. | 0.125 | 0.561 |
| coded 8/240 | 0.037 | 0.231 |
| coded 8/240 adj. | 0.064 | 0.399 |

## Appendix I Evaluating log-likelihood using the soft-encoding model

The reparameterization trick introduced in (5) requires that bits are independent, a condition not met in $\boldsymbol{c}$ once the code’s structure is introduced. To address this issue, during training, we employ a *soft encoding strategy*. We assume the bits in $\boldsymbol{c}$ are independent and equally distributed according to $q_{\boldsymbol{\eta}}(\boldsymbol{m}|\boldsymbol{x})$. Therefore, instead of directly repeating sampled bits in $\boldsymbol{m}$ to obtain $\boldsymbol{c}$ following $\boldsymbol{c}=\boldsymbol{m}^{T}\mathbf{G}$, we repeat the posterior probabilities for the copies of the same bit and sample $\boldsymbol{z}$ using the reparameterization trick in (5). Remark that, despite the soft encoding assumption during training, the generative results in Fig. 4, 6, 13, 19, 24, 24, and 29 are obtained through *hard encoding*. Namely, we sampled an information word $\boldsymbol{m}$ and obtained $\boldsymbol{c}$ by repeating its bits.

While the hard-encoding images are visually appealing, to evaluate the coded DVAE LL for a given image $\boldsymbol{x}$, we have to leverage the soft-encoding model since it is unrealistic for a sample from the soft-encoding model to produce equally repeated bits. In the soft-coded model, we sample bit probabilities from a prior distribution that we model through the product of $M$ independent Beta distributions, and we use a proposal distribution model similar to the Vamp-prior in []. Namely, a mixture model with components given by $q(\boldsymbol{m}|\boldsymbol{x})$ for different training points.

In the main text and Section H, we report the log-likelihood values for different model configurations trained on the FMNIST dataset. Table 4 presents the results obtained for the MNIST dataset. Due to their simplicity, these datasets do not require high-dimensional latent spaces, and competitive results can be achieved with just 8 or 10 information bits. However, for more complex datasets like CIFAR10 or Tiny ImageNet, high-dimensional latent spaces are necessary to capture spatial information, colors, and textures. For these high-dimensional datasets, we were unable to obtain valid log-likelihood estimates because the number of samples needed for importance sampling to converge was too large. However, given the difference in the level of detail in the reconstructed and generated images between coded and uncoded models (see Sections D and E), coded models are expected to better approximate the true marginal likelihood of the data.

## Appendix J Variational inference at codeword level

Here, we present an alternative variational family that is valid for any ECC, including random codes. We assume that we have a deterministic mapping of the form $\boldsymbol{c}=\mathcal{C}(\boldsymbol{m})$. We assume a variational family of the form

$$ $q_{\boldsymbol{\eta}}(\boldsymbol{c},\boldsymbol{z}|\boldsymbol{x})=q_{ \boldsymbol{\eta}}(\boldsymbol{c}|\boldsymbol{x})p(\boldsymbol{z}|\boldsymbol{ c}),$ (18) $$

$$ $q_{\boldsymbol{\eta}}(\boldsymbol{c}|\boldsymbol{x})\propto p(\boldsymbol{c})q _{\boldsymbol{\eta}}^{u}(\boldsymbol{c}|\boldsymbol{x}),$ (19) $$

$$ $q_{\boldsymbol{\eta}}^{u}(\boldsymbol{c}|\boldsymbol{x})=\prod_{j=1}^{M/R} \text{Ber}(g_{j,\boldsymbol{\eta}}(\boldsymbol{x})),$ (20) $$

where $g_{\boldsymbol{\eta}}(\boldsymbol{x})$ represents the output of the encoder with a parameter set denoted as $\boldsymbol{\eta}$. Note that $q_{\boldsymbol{\eta}}^{u}(\boldsymbol{c}|\boldsymbol{x})$ corresponds to the uncoded posterior, which we subsequently constrain using the prior distribution $p(\boldsymbol{c})$ over the code words to obtain the coded posterior $q_{\boldsymbol{\eta}}(\boldsymbol{c}|\boldsymbol{x})$. Then, the coded posterior distribution can be defined as a categorical distribution over the set of code words $\mathcal{C}(\boldsymbol{m})$, which is given by

$$ $\begin{gathered}q_{\boldsymbol{\eta}}(\boldsymbol{c}|\boldsymbol{x})=\text{Cat }\Big{(}\Big{[}\frac{1}{W}q_{\boldsymbol{\eta}}^{u}(\boldsymbol{c}_{1}| \boldsymbol{x}),\ldots,\frac{1}{W}q_{\boldsymbol{\eta}}^{u}(\boldsymbol{c}_{2^ {M}}|\boldsymbol{x})\Big{]}\Big{)}\\ =\frac{1}{W}\prod_{\boldsymbol{c}_{i}\in\mathcal{C}(\boldsymbol{m})}\prod_{j=1 }^{M/R}g_{j,\boldsymbol{\eta}}(\boldsymbol{x})^{c_{i,j}}(1-g_{j,\boldsymbol{ \eta}}(\boldsymbol{x}))^{(1-c_{i,j})},\end{gathered}$ (21) $$

where $W=\sum_{\boldsymbol{c}_{i}\in\mathcal{C}(\boldsymbol{m})}q_{\boldsymbol{\eta}}
^{u}(\boldsymbol{c}_{i}|\boldsymbol{x})$ is a constant for normalization.

Inference is done by maximizing the ELBO, which can be expressed as

$$ $\begin{gathered}\text{ELBO}=\int q_{\boldsymbol{\eta}}(\boldsymbol{c}, \boldsymbol{z}|\boldsymbol{x})\log\left(\frac{p_{\boldsymbol{\theta}}( \boldsymbol{x},\boldsymbol{z},\boldsymbol{c})}{q_{\boldsymbol{\eta}}( \boldsymbol{c},\boldsymbol{z}|\boldsymbol{x})}\right)d\boldsymbol{c}d \boldsymbol{z}\\ =\mathbb{E}_{q_{\boldsymbol{\eta}}(\boldsymbol{c},\boldsymbol{z}|\boldsymbol{x })}\log\left(\frac{p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})p( \boldsymbol{z}|\boldsymbol{c})p(\boldsymbol{c})}{q_{\boldsymbol{\eta}}( \boldsymbol{c}|\boldsymbol{x})p(\boldsymbol{z}|\boldsymbol{c})}\right)\\ =\mathbb{E}_{q_{\boldsymbol{\eta}}(\boldsymbol{c},\boldsymbol{z}|\boldsymbol{x })}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})-\mathcal{D}_{KL }\big{(}q_{\boldsymbol{\eta}}(\boldsymbol{c}|\boldsymbol{x})||p(\boldsymbol{c} )\big{)}.\end{gathered}$ (22) $$

In this case, due to the inability to compute the KL Divergence in closed form, we approximate it via Monte Carlo, sampling from the categorical distribution $q_{\boldsymbol{\eta}}(\boldsymbol{c}|\boldsymbol{x})$. The reconstruction term also needs to be approximated via Monte Carlo. Since the use of channel coding introduces structural dependencies among the components of the vectors $\boldsymbol{c}$, we can no longer assume their independence. Consequently, the formulation of the smoothing transformation as independent mixtures introduced in the DVAE is no longer applicable. Hence, this approach involves sampling $\boldsymbol{c}^{\prime}$ from the categorical distribution $q_{\boldsymbol{\eta}}(\boldsymbol{c}|\boldsymbol{x})$ and subsequently applying the transformation over the sampled word. Thus, we obtain a smooth transformation for each sample $\boldsymbol{c}^{\prime}$ using the inverse CDFs of $p(z_{j}|c_{j}=0)$ and $p(z_{j}|c_{j}=1)$, which are given by

$$ $F^{-1}_{p(z_{j}|\boldsymbol{c}^{\prime}_{j}=0)}(\rho)=-\frac{1}{\beta}\log\big {(}1-\rho(1-e^{-\beta})\big{)},$ (23) $$

$$ $F^{-1}_{p(z_{j}|\boldsymbol{c}^{\prime}_{j}=1)}(\rho)=\frac{1}{\beta}\log\big{ (}\rho(1-e^{-\beta})+e^{-\beta}\big{)}+1.$ (24) $$

These are differentiable functions that convert samples $\rho$ from a uniform distribution $\mathcal{U}(0,1)$ into a sample from $q_{\eta}(\boldsymbol{z},\boldsymbol{c}=\boldsymbol{c}^{\prime}|\boldsymbol{x})$ following

$$ $\begin{gathered}q_{\eta}(\boldsymbol{z},\boldsymbol{c}=\boldsymbol{c}^{\prime} |\boldsymbol{x})=\\ =\prod_{j=1}^{M/R}\Big{[}F^{-1}_{p(z_{j}|\boldsymbol{c}^{\prime}_{j}=1)}(\rho) ^{c^{\prime}_{j}}+F^{-1}_{p(z_{j}|\boldsymbol{c}^{\prime}_{j}=0)}(\rho)^{(1-c^ {\prime}_{j})}\Big{]}.\end{gathered}$ (25) $$

Thus, we can apply the reparameterization trick to obtain samples from the latent variable $\boldsymbol{z}$ and optimize the reconstruction term of the ELBO with respect to the parameters $\boldsymbol{\theta}$ of the decoder.

The KL Divergence term is approximated via Monte Carlo, drawing samples from $q_{\boldsymbol{\eta}}(\boldsymbol{c}|\boldsymbol{x})$. As it is not possible to backpropagate through discrete variables, we approximate the gradients with respect to the parameters of the encoder using the REINFORCE leave-one-out estimator [, ], given by

$$ $\begin{gathered}\widehat{g}_{LOO}=\\ =\frac{1}{S-1}\Bigg{[}\sum_{s=1}^{S}f_{\boldsymbol{\eta}}\big{(}\boldsymbol{z} ^{(s)},\boldsymbol{c}^{(s)}\big{)}\bigtriangledown_{\boldsymbol{\eta}}\log q_{ \boldsymbol{\eta}}\big{(}\boldsymbol{c}^{(s)}|\boldsymbol{x}\big{)}-\overline{ f}_{\boldsymbol{\eta}}\sum_{s=1}^{S}\bigtriangledown_{\boldsymbol{\eta}}\log q _{\boldsymbol{\eta}}\big{(}\boldsymbol{c}^{(s)}|\boldsymbol{x}\big{)}\Bigg{]}, \end{gathered}$ (26) $$

where

$$ $f_{\boldsymbol{\eta}}(\boldsymbol{z}^{(s)},\boldsymbol{c}^{(s)})=\log\Bigg{(} \frac{q_{\boldsymbol{\eta}}(\boldsymbol{c}^{(s)}|\boldsymbol{x})}{p_{ \boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z}^{(s)})p(\boldsymbol{c}^{(s) })}\Bigg{)},$ (27) $$

$$ $\overline{f_{\boldsymbol{\eta}}}=\frac{1}{S}\sum_{s=1}^{S}f_{\boldsymbol{\eta} }(\boldsymbol{z}^{(s)},\boldsymbol{c}^{(s)}).$ (28) $$

Defining a distribution over the codebook can seem intuitive, but scalability becomes challenging as the size of the codebook increases. The reason is that the posterior distribution must be evaluated for all codewords during both inference and test time. However, it can still provide a bound, enabling the utilization of more complex codes with theoretical guarantees that can outperform the previously proposed repetition codes.

Figure: Algorithm 3 Training the model with inference at codeword level.

## Appendix K Computational resources

All the experiments in this paper were conducted on a single GPU. Depending on availability, we used either a Titan X Pascal with 10GB of RAM, a Nvidia GeForce GTX with 10GB of RAM, or a Nvidia GeForce RTX 4090 with 24GB of RAM. Since training times varied significantly based on the hardware used, we were unable to provide comparable training times.

## Appendix L Coded DVAE scheme

Figure: Figure 39: Graphic representation of the coded DVAE, illustrating both the generative and inference paths.
Refer to caption: x44.png

## Appendix M Hierarchical coded DVAE scheme

Figure: Figure 40: Graphic representation of the hierarchical coded DVAE, illustrating both the generative and inference paths.
Refer to caption: x45.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. Please clarify what is the gain of introducing ECC into the encoding of $m$ is beneficial intuitively. Where does the noise come from (analogously to communication channels), and why ECC helps reducing the reconstruction error through this encoding. 

2. If possible, please use theoretical support to prove that the encoded DVAE can achieve a better SNR (signal to noise ratio) through the proposed encoding of latent variable. 

3. I appreciate if any algorithmic flow is further provided in the appendix or somewhere to better clarify the sampling process of c and how it is backpropogated.
soundness: 2
strengths: Novelty: The paper is novel in it borrows the method of Error-Correcting Codes (ECC) and introduces it into the problem of DVAE. The motivation is to safeguard the latent variables by adding redundancy to the latent representations through the error correction method in information theory.  In this way, the proposed method claims that it is able to utilizes the redundancy introduced by the ECC to constrain the solution space of $q(c|x)$, therefore leading to a better reconstruction signal to noise ratio of the DVAE than that without ECC. To me, this seems to be the first paper that introduces this method.  

Motivation and Significance: The paper aims to improve the reconstruction of the DVAE by introducing the ECC method from information theory. This is a good example of how the techniques in information theory and communication techniques can help improve generative models. The motivation of the work is good and interesting.
summary: The paper introduces the method of Error-Correcting Codes (ECC) into the problem of DVAE. The goal is to safeguard the latent variables to add redundancy to the latent representations.  In this way, the proposed method utilizes the redundancy introduced by the ECC to constrain the solution space of $q(c|x)$, given that each bit from $m$ is repeated L times (through the known G) to create c. Empirical result demonstrated the reconstruction signal to noise ratio of the VAE with ECC is improved than that without ECC.
weaknesses: Clarity and rigour: I am a lost in section 4, where how the $m$ is sampled and encoded into the required $c$ is introduced. It seems to me that c is now a linear transformation of $m$ through the known $G$. However, how c is sampled and backpropogated through the DVAE is unclear to me. This section also says: `When comparing uncoded vs. coded DVAEs, the structure of the decoder NN is equal in both cases except for the first Multilayer Perceptron (MLP) layer that attacks the input z'. This is very confusing to me. What does the ``attack'' mean here? Where did the attack come from? In the earlier section of the paper, I did not assume there is any noise injection or attacks existing in the generative model. Also, as the c is linear transformation of the original m. it is unclear to me why such redundancy can improve the signal to noise ratio of the reconstruction. For communication channels, yes, because there is loss during the signal transmission in the channels, but I do not understand why the c=mG can help better constrain the sampling space of $m$ and why this is relevant in achieving a better generative model. I guess better instantiation of the the method with further theoretical support and an algorithmic flow can better support the idea of the paper.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: My suggestion to the authors is to read the papers referenced in 'Weaknesses' and write a small section in the introduction describing the relationship between variational inference and Shannon's rate-distortion theory of lossy compression. This will substantially reinforce the validity of the proposed approach, as there is theoretical justification for utilizing a tool originally used in digital communications. Alternatively, the overlap between digital communications and variational inference is made more explicit, so that the unfamiliar reader can make sense of why such a tool is useful.
soundness: 3
strengths: The authors present a theoretically sound and practically useful enhancement to discrete VAEs, noting its practical utility in digital communication. The paper is well-presented; it is easy to follow the contribution from the problems and weaknesses of contemporary approaches to the proposed solution. The experiments validate the theoretical claims, including the claim that the proposed method can be used alongside existing enhancements like IWAEs.
summary: In this paper, the authors propose to improve variational inference in discrete variational autoencoders by utilizing error correcting codes (ECCs) in the latent space. Despite significant advances in VAEs, the ability to learn low-dimensional discrete latent-variable models remains a challenging task. To that end, the authors employ ECCs as borrowed from digital communications in order to add redundance in the posterior samples. During training, the VAE is able to capture meaning in the latent variables, resulting in significantly improved generation quality and more accurate estimates. The authors emphasize that this technique is general and can be used alongside existing VAE techniques, such as tighter importance weighted autoencoders and Hamiltonian Monte Carlo.
weaknesses: Although the particular enhancement presented is novel, the motivation behind has been thoroughly studied within the context of VAEs in previous papers, specifically in relation to lossy/lossless compression and communication [1, 2]. The paper mentions Vector Quantized VAEs as an existing approach only briefly, failing to discuss the significance of it in relation to the aforementioned overlapping subjects. It would be helpful to discuss the general relationship between compression and variational inference and how certain VAE implementations exploit this close relationship to their advantage for certain applications, such as lossy compression [2, 3] and lossless compression [4].

References

[1] Aaron van den Oord, Oriol Vinyals, & Koray Kavukcuoglu. (2018). Neural Discrete Representation Learning.

[2] Will Williams, Sam Ringer, Tom Ash, John Hughes, David MacLeod, & Jamie Dougherty. (2020). Hierarchical Quantized Autoencoders.

[3] Xi Chen, Diederik P. Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, & Pieter Abbeel. (2017). Variational Lossy Autoencoder.

[4] James Townsend, Tom Bird, & David Barber. (2019). Practical Lossless Compression with Latent Variables using Bits Back Coding.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: - Can the proposed method be extended to VAE with continuous latent space?
soundness: 3
strengths: - The contribution is novel and interesting.
- The paper is well-written, and mathematics is easy to follow.
- Introducing ECCs results in significant performance improvement in DVAEs across several datasets as demonstrated in the paper.
summary: The authors propose to introduce redundancy into the latent representation in a Discrete Variational Autoencoder (DVAE) by employing error-correcting codes (ECC). This would reduce errors in samples drawn from the variational posterior, resulting in a better estimation of the true posterior.
weaknesses: While I appreciate the contribution and the presented results in the paper, there are few concerns I have which are listed below:

- Comparison baselines: The baseline used for comparison is just DVAE with a simple independent prior. It would better highlight the effectiveness of the proposed approach if the comparison is also done with DVAEs with Boltzmann machine priors and other discrete VAE models such as VQ-VAE (discussed in the introduction section).
- The ablation study is conducted by adjusting the number of parameters of a neural network. I think it’s more convincing to compare the results of the uncoded and coded DVAEs with the same dimensionality of latent space (let's say D). (The case where the coded DVAE has redundancy included in it (D = ML with only M original message bits) while the uncoded VAE has D original message bits.)
- To directly verify the improvement in variational inference, I believe measuring the variational gap (inference gap) between the uncoded and the coded DVAEs (as done in [1]) would be helpful. This would provide a direct verification of the main claim of the paper.
- Minor: The resolution of Figure 1(b) is poor. Also, the metric used for measuring reconstruction loss in the vertical axis is not stated.

[1] Shu et. al., Amortized Inference Regularization. NeurIPS, 2018.

If the concerns are adequately addressed, I’d be ready to increase my rating for the paper.

---

confidence: 3
contribution: 4
initial_rating: 5
presentation: 3
questions: - Why is the parameter $c$ absent on the left-hand side of Equation (8)? Could you clarify the reasoning behind this formulation?
- How does the dependency on $\beta$ manifest in Equation (7)? It would be helpful to have a more detailed explanation of its role.
- If the transformations in Equations (10) and (11) are applied, can it still be demonstrated that the Evidence Lower Bound (ELBO) serves as a lower bound for the log-likelihood?
soundness: 3
strengths: - The integration of ECC into variational inference is innovative, and the concept shows promise for applicability across various latent space models.
- The method is validated on diverse datasets, including MNIST, Fashion MNIST, CIFAR-10, and Tiny ImageNet, providing evidence of its generalizability.
- The paper comprehensively evaluates coded and uncoded models using various metrics, such as PSNR, log-likelihood, reconstruction quality, and uncertainty calibration.
summary: This paper proposes a novel approach to enhancing variational inference for discrete latent variable models. The method introduces redundancy in the latent space by leveraging error-correcting codes (ECC) from coding theory, resulting in improved performance. The proposed technique is applied to discrete variational autoencoders (DVAEs) with binary latent variables. It demonstrates improvements in generative quality, data reconstruction, and uncertainty calibration across datasets such as MNIST, Fashion MNIST, CIFAR-10, and Tiny ImageNet.
weaknesses: - The paper appears to mix theoretical claims with empirical results. For instance, the statement, "This way, it is possible to reduce the mistakes committed when comparing $m$ with samples drawn from $q(m|x)$, obtaining a tighter approximation to the true posterior $ p(m|x)$," suggests a theoretical result. However, it is unclear if this has been formally proven. Additionally, there is no explicit theoretical proof showing that the inclusion of ECC improves generalization performance. Based on my review, such proof does not seem present in the paper.
- As Equation (9) indicates, using ECC increases the number of parameters due to the added redundancy. This corresponds to an effective increase in the dimensionality of the latent variables compared to the original model without ECC. In the numerical experiments, comparisons should be made with uncoded models with the same total number of latent dimensions. For example, in the case where $ c=10$, it would be more appropriate to compare against an uncoded model with $z=10$. Whether the observed performance improvements are due to the added parameters or the ECC itself remains unclear. The authors should clarify the number of parameters in the models and specify the dimensions used for $c$ and $z$. Based on my understanding, the notation "CODED 5/100" in Figures 4, 5, and 6 likely corresponds to $z=5$ and $c=100$, meaning the model has parameters associated with $c=100$. In this case, wouldn't the appropriate baseline be an uncoded model with $z=100$?
- The paper appears to exceed the page limit, which I understand to be 10 pages. Could you confirm whether the page count adheres to the conference guidelines?

Meta-review: This paper proposes a novel variational inference framework using Error-Correcting Codes (ECC) to improve the performance of VAEs with discrete latent variables. Specifically, it models discrete variables with repetition codes based on ECC and employs an appropriate variational posterior distribution, demonstrating significant performance improvements across various tasks.
The main weakness of this study is the unclear scope of its applicability. As pointed out by multiple reviewers, the validity of the proposed method has been evaluated solely through numerical experiments. While performance improvements are observed under a simple prior using independent Bernoulli distributions, it remains uncertain whether the method is effective under other, more realistic model settings. Thus, as several reviewers have noted, it would be desirable to strengthen the numerical evidence supporting the study's contributions and to provide numerical or theoretical insights into the settings where this method is most applicable or less effective.
On the other hand, all reviewers have recognized the novelty of the idea, suggesting that the research direction holds sufficient significance. However, in its current form, the study is not sufficient for publication, and I recommend a rejection at this time.

Decision: Reject