[Question] Squeeze and Split flows in Tutorial 11

[phlippe]

[Question] Squeeze and Split flows in Tutorial 11

[phlippe]

Hi, thanks for raising this comparison, I didn't expect that there would be these differences! Regarding the split and squeeze: in general, either way (official implementation or notebook) is fine in theory, although I would expect the version in the tutorial to generalize a bit better. Splitting over pixel positions (in the official code) instead of channels (in the tutorial) seems to me much harder to learn, since you force the network to identify which pixels will be mapped to the prior at the early stage already, while the neural networks are convolutional and thus translation invariant (up to padding borders). So in the end, the model might have to map everything to the prior distribution, which can limit the capacity. In comparison, if you split over channels, it is much easier for the model to simply dedicate some dimensions to be mapped to Gaussians, especially since we use a similar splitting strategy in the coupling layers. Nonetheless, I have not explicitly tested it and thought the tutorial version was the more natural way.

For the coupling layers, the difference in channelwise versus rowwise coupling is how the network perceives the input. In channelwise, we have an input of 4x14x14, while the rowwise coupling has 1x28x28. Thus, different network architectures are learned here. Another key difference is when you start combining these layers with more sophisticated NF tricks. For instance, often invertible 1x1 convolutions are used to intermix the channels between coupling layers. This has a different effect if you squeeze or not.

Hope that helps :)