Count the code before calling it a bottleneck

STATEMENT

Track shapes and parameters through a dense and a convolutional autoencoder for the same image, and find the two places where the convolutional one goes wrong.

GIVEN

One input is a $1\times28\times28$ image, $784$ numbers.

The dense autoencoder has layer widths $784$, $128$, $32$, $128$, $784$, with a bias on every layer.

The convolutional encoder is two layers with $3\times3$ kernels, stride $2$ and padding $1$, taking $1$ channel to $8$ and then $8$ to $16$. The decoder is two transposed convolutions with the same kernel, stride and padding, taking $16$ channels to $8$ and then $8$ to $1$. All four layers have biases.

FIND

(a) The dense parameter count, untied and with tied weights. (b) The spatial sizes through the convolutional encoder, and how many numbers its code holds. (c) The decoder’s output size with output padding $0$ and with $1$. (d) The convolutional parameter count. (e) The code channel counts for which the convolutional code is narrower than its input.

STRATEGY

Count each dense layer as inputs times outputs plus outputs. Use (I.11.4) for the encoder’s sizes and (I.11.5) for the decoder’s.

SOLUTION

(a) The four dense layers hold

for a total of $209{,}968$. With tied weights the two decoder matrices are transposes of the encoder’s. The weights are then $100{,}352+4{,}096=104{,}448$ and the four bias vectors add $1{,}072$, for $105{,}520$. The code is $32$ numbers, $24.5$ times fewer than the input.

(b) By (I.11.4), with no dilation,

The code has shape $16\times7\times7$ and holds $784$ numbers, exactly as many as the input. This is the first fault. Each layer halved the height and the width and the picture looks like an hourglass, but the channel count grew as fast as the grid shrank. By count, nothing forces this code to discard anything.

(c) By (I.11.5) a transposed layer returns $(n-1)\cdot2-2+2+o+1=2n-1+o$ for output padding $o$. With $o=0$ the sizes run $7\to13\to25$. The reconstruction is $25\times25$, with $625$ numbers, and cannot be compared with a $28\times28$ target. This is the second fault. With $o=1$ they run $7\to14\to28$.

The ambiguity is real. Under the encoder’s stride-two layer an input of size $27$ and one of size $28$ both give $14$. The transposed layer cannot know which it is undoing, and output padding tells it.

(d) A transposed convolution stores its weight with shape $(C_{\mathrm{in}},C_{\mathrm{out}},k_h,k_w)$ and one bias per output channel. The four layers hold

a total of $2{,}481$, against $209{,}968$ for the dense stack.

(e) A code of $C$ channels on the $7\times7$ grid holds $49C$ numbers. It is narrower than the input when $49C<784$, that is, $C<16$. With $C=4$ the code holds $196$ numbers, a quarter of the input.

Check — sanity

The tied count is half the untied weights plus all the biases: $208{,}896/2+1{,}072=105{,}520$. The transposed sizes $14$ and $28$ are the encoder’s sizes read backwards.