Convert categorical probabilities to soft binary representation

Question

Let's say you have an artificial neural network (ANN) where the last layer is a softmax, so the output could be interpreted to be a categorical distribution, i.e. output neuron $i \in \{0, 1, \ldots, N-1 \}$ represents the probability the ANN "thinks" the input sample belongs to class $i$.

Let's say that that we restrict $N$ to be a power of two, i.e. $N = 2^K$. In this case you could represent each class as a binary codeword comprised of $K$ bits $\{ B_k \}_{k=1}^K$. By convention we will write the codewords as $B_{K-1} B_{K-2} \ldots B_1 B_0$, where $B_0$ is the least significant bit (LSB).

Now let's say we want to quantify our certainty in ANN output $i$ using a soft binary representation $\{ b_k \}_{k=1}^K$, where $b_k = \operatorname{P}[B_k = 1]$.

If we denote the categorical probability outputs $p_0, p_1, \ldots, p_{N-1}$, my intuition would be that we could express these in terms of $b_k$'s as: $$ p_0 = \operatorname{P}[(B_{K-1} = 0) \cap (B_{K-2} = 0) \cap \ldots \cap (B_{1} = 0) \cap (B_{0} = 0) ] \\\\ p_1 = \operatorname{P}[(B_{K-1} = 0) \cap (B_{K-2} = 0) \cap \ldots \cap (B_{1} = 0) \cap (B_{0} = 1) ] \\\\ p_2 = \operatorname{P}[(B_{K-1} = 0) \cap (B_{K-2} = 0) \cap \ldots \cap (B_{1} = 1) \cap (B_{0} = 0) ] $$ etc.

If we assume the bits are independent and substitute in the $b_k$'s, we obtain: $$ p_0 = (1 - b_{K-1}) (1 - b_{K-2}) \ldots (1 - b_{1}) (1 - b_{0}) \\\\ p_1 = (1 - b_{K-1}) (1 - b_{K-2}) \ldots (1 - b_{1}) b_{0} \\\\ p_2 = (1 - b_{K-1}) (1 - b_{K-2}) \ldots b_1 (1-b_{0}) \\\\ p_3 = (1 - b_{K-1}) (1 - b_{K-2}) \ldots b_1 b_0 \\\\ \vdots \\\\ p_{N-1} = b_{K-1} b_{K-2} \ldots b_1 b_0 $$

The goal then is to solve for the $b_k$'s from the $p_i$'s. There are $N$ equations with $K < N$ unknowns, so I know it should be possible, but I don't know if the solution is guaranteed.

Let's say $K=3$. Naively starting at $p_7$ and working backwards, we obtain: $$ p_7 = b_2 b_1 b_0 \implies b_2 = \frac{p_7}{b_1 b_0} \\\\ p_6 = b_2 b_1 (1 - b_0) = \frac{p_7}{b_1 b_0} b_1 (1 - b_0) \implies \frac{p_7}{p_6} = \frac{b_0}{1 - b_0} $$ This intuitively makes sense if we look at the corresponding code words: $$ 6 \rightarrow 110 \\\\ 7 \rightarrow 111 $$ In that the ratio of the probabilities of being in class 7 to class 6 is the same as the ratio of the probabilities to $B_0 = 1$ and $B_0 = 0$.

Thus $$ b_0 = \frac{1}{1 + p_6/p_7} $$ By similar arguments, we could see that $$ 5 \rightarrow 101 \\\\ 7 \rightarrow 111 $$ so $$ \implies b_1 = \frac{1}{1 + p_5/p_7} $$ and $$ 3 \rightarrow 011 \\\\ 7 \rightarrow 111 $$ $$ \implies b_2 = \frac{1}{1 + p_3/p_7} $$

Minimal Broken Example:

import numpy as np

# generate random logits and class probs
N = 8
np.random.seed(3) # repeatable results
logits = 3*np.random.randn(N) # coefficient to make one class much higher prob
p = np.exp(logits)/np.sum(np.exp(logits)) # softmax

# estimate soft bits
b0 = 1/(1 + p[6]/p[7])
b1 = 1/(1 + p[5]/p[7])
b2 = 1/(1 + p[3]/p[7])

np.set_printoptions(suppress=True, precision=5)
print('logits =', logits)
print('p =', p)
print('b2,b1,b0 = %.5f, %.5f, %.5f'%(b2, b1, b0))

output:

logits = [ 5.36589  1.30953  0.28949 -5.59048 -0.83216 -1.06428 -0.24822 -1.881  ]
p = [0.96939 0.01678 0.00605 0.00002 0.00197 0.00156 0.00353 0.00069]
b2,b1,b0 = 0.97610, 0.30646, 0.16345

In this toy example, class 0 has the highest probability, which implies the codeword should be $000$. But when I use the above formulae to calculate the soft bits, we get that the most likely codeword is $100$.

Is there an error in the assumptions used to derive these expressions?

Answer 1

Is there an error in the assumptions used to derive these expressions?

Yes there is: you assume there is always a solution, which isn't guaranteed (as you already said in your question).

In your example, while it's true that

$$b_0=\frac{1}{1+p_6/p_7},$$

there are 3 other possibilities for $b_0$, namely $\frac{1}{1+p_0/p_1}$, $\frac{1}{1+p_2/p_3}$, and $\frac{1}{1+p_4/p_5}$, which can be derived by using the same argument as you did for the first one. All these possibilities give rise to 4 equations that must be satisfied simultaneously by $b_0$. The same can be done for $b_1$ and $b_2$ for a total of 12 equations. These equations define constraints on the categorical distribution that can be expressed using the soft binary representation. Therefore, you can't express an arbitrary categorical distribution in this way.

To obtain one distribution that works with your $b_0,b_1,b_2$, you can run

p_corr = []
for i2 in range(2):
  for i1 in range(2):
    for i0 in range(2):
      p_corr.append((b2 if i2 else 1-b2)*(b1 if i1 else 1-b1)*(b0 if i0 else 1-b0))
print(p_corr)

which gives

[0.013869183567376016,
 0.0027098459367482265,
 0.006128464791425466,
 0.001197416944758982,
 0.5663121619290953,
 0.11064953488282403,
 0.2502399747236977,
 0.04889341722407433]

You can verify that this distribution satisfies all the constraints, and its mode/argmax (4) is equal to what you get with the modes of the soft binary representation (100).