Explicit Bits-back Coding (a.k.a. Free Energy Coding) applied to Gaussian mixtures

Question

I've been trying to understand Bits-back coding (Frey, B. J., and G. E. Hinton. 1997.) a bit more (pun intended), which can be used to encode data with latent variable models. This tutorial by Pieter Abbeel et al. sums up the procedure with an example which I present below, but with an explicit auxiliary variable (i.e. I added the variable $k$).

Consider a mixture of Gaussians where $i$ indexes the Gaussian from which the data $x$ is sampled; $p(x\vert i) = \mathcal{N}(x|\mu_i, \sigma^2_i)$. Assume, although possibly intractable, the true posterior $p(i \vert x, k)$ is available from which it is possible to sample and encode values $i$, where $k \sim p(k)$ is some auxiliary information used to sample $i$. Note that $k$ is used solely in the process of selecting $i$, therefore $k \rightarrow i \rightarrow x$ forms a Markov chain and $p(x \vert i, k) = p(x \vert i)$

Given a data point $x$ to compress, bits-back coding then proceeds to

0. Sample $k$ from $p(k)$
1. Sample $i$ from $p(i \vert x, k)$
2. Encode $i$ with a prior $p(i)$, resulting in code-length $\log_2(1/p(i))$
3. Encode $x$ with $p(x \vert i)$, resulting in code-length $\log_2(1/p(x \vert i))$
4. Although $i$ was encoded with $p(i)$, given that we know $p(i \vert x, k)$, it is possible to get "bits back" (hence the name) by reconstructing $k$.

My goal is to understand this concretely, since most tutorials I've found stop here and never actually give an example. Hence, consider the following:

Say there are 4 modes (Gaussians) of the mixture represented by $i \in \{0, 1, 2, 3\}$ and $k \in \{0, 1\}$. For a given point $x_o$, assume the given values for $p(i \vert x_o, k)$ and $p(i)$, as well as the code used to represent them from a binary alphabet, are the following:

$$\begin{array}{c|c|c} i & p(i) & \text{code-word} \\ \hline 0 & 1/4 & 00 \\ 1 & 1/4 & 01 \\ 2 & 1/4 & 10 \\ 3 & 1/4 & 11 \\ \end{array}$$

$$\begin{array}{c|c|c} k &i & p(i \vert x_o, k) \\ \hline 0 & 0 & 1/2 \\ 0 & 1 & 1/4 \\ 0 & 2 & 1/4 \\ 0 & 3 & 0 \\ 1 & 0 & 0 \\ 1 & 1 & 1/2 \\ 1 & 2 & 1/4 \\ 1 & 3 & 1/4 \end{array}$$

I've conveniently made the symbols I.I.D. with dyadic probability mass functions, therefore making traditional Huffman Coding optimal.

On the receiving end, if $i=0$, surely $k=0$ since $p(i=0 | x_o, k=1) = 0$. Similarly, $i=3$ implies $k=1$. In general:

$$\begin{array}{c|c|c} i & p(k=0 \vert x_o, i) & p(k=1 \vert x_o, i) \\ \hline 0 & 1 & 0 \\ 1 & 1/3 & 2/3 \\ 2 & 1/2 & 1/2 \\ 3 & 0 & 1 \end{array}$$

Therefore, we are able to recover $k$ to some extent.

I have 2 questions:

1. is this the bits-back coding scheme or am I missing something?
2. For $i=0$ and $i=3$ we can recover $k$, but for $i \in \{1, 2\}$ all we have are the probabilities over $k$. How do we get these bits-back?

Answer 1

Apparently, no. This is a generalization where the relationship between $k$ and $i$ is stochastic (i.e. defined by a distribution and not deterministic).

To recover the original bits-back argument, we would need to pick a value $i_k$ for each possible value of $k$ and force $p(i=i_k \vert x, k)=1$ and 0 if $i \neq i_k$. Naturally, we would need to extend the support of $k$ to the same of $i$. In this particular example, $k \in \{0, 1, 2, 3\}$