MINIMAL RANDOM CODE LEARNING: GETTING BITS BACK FROM COMPRESSED MODEL PARAMETERS
While deep neural networks are a highly successful model class, their large memory footprint puts considerable strain on energy consumption, communication bandwidth, and storage requirements. Consequently, model size reduction has become an utmost goal in deep learning. A typical approach is to train a set of deterministic weights, while applying certain techniques such as pruning and quantization, in order that the empirical weight distribution becomes amenable to Shannon-style coding schemes. However, as shown in this paper, relaxing weight determinism and using a full variational distribution over weights allows for more efficient coding schemes and consequently higher compression rates. In particular, following the classical bits-back argument, we encode the network weights using a random sample, requiring only a number of bits corresponding to the KullbackLeibler divergence between the sampled variational distribution and the encoding distribution. By imposing a constraint on the Kullback-Leibler divergence, we are able to explicitly control the compression rate, while optimizing the expected loss on the training set. The employed encoding scheme can be shown to be close to the optimal information-theoretical lower bound, with respect to the employed variational family. Our method sets new state-of-the-art in neural network compression, as it strictly dominates previous approaches in a Pareto sense: On the benchmarks LeNet-5/MNIST and VGG-16/CIFAR-10, our approach yields the best test performance for a fixed memory budget, and vice versa, it achieves the highest compression rates for a fixed test performance.
The authors come up with a surprisingly elegant algorithm ("minimal random coding") which encodes samples from a posterior distribution, only using a number of bits that approximates the KL divergence between posterior and prior, while Shannon-type algorithms can only do this if the posterior is deterministic (a delta distribution). It can also be directly used to sample from continuous distributions, while Shannon-type algorithms require quantization. In my opinion, this is the main contribution of the paper.
The other part of the paper that is specifically concerned with weight compression ("MIRACLE") turns out to be a lot less elegant. It is somewhat ironic that the authors specifically call attention to the their algorithm sending random samples, as opposed to existing algorithms, which quantize and then send deterministic variables. This is clearly true for the basic algorithm, but, if I understand correctly, not for MIRACLE. It seems clear that neural networks are sensitive to random resampling of their weights -- otherwise, the authors would not have to fix the weights in each block and then do further gradient descent for the following blocks. What would happen if the distributions were held constant, and the algorithm would be run again, just with a different (but identical) random seed in both sender and receiver? It seems this would lead to a performance drop, demonstrating that (albeit indirectly), MIRACLE also makes a deterministic choice of weights.
Overall, I find the paper somewhat lacking in terms of evaluation. MIRACLE consists of a lot of parts. It is hard to assess how much of the final coding gain presented in table 1 is due to the basic algorithm. What is the effect of selecting other probability models, possibly different ones than Gaussians? Choosing appropriate distributions can have a significant impact on the value of the KL divergence. Exactly how much is gained by applying the hashing trick? Are the standard deviations of the priors included in the size, and how are they encoded?
This could be assessed more clearly by directly evaluating the basic algorithm. Theorem 3.2 predicts that the approximation error of algorithm 1 asymptotically zero, i.e. one can gain an arbitrarily good approximation to the posterior by spending more bits. But how many more are practically necessary? It would be fantastic to actually see some empirical data quantifying how large the error is for different distributions (even simple toy distributions). What are the worst vs. best cases?