Compression and inference in high dimensions: An approximate message passing perspective

Inferring signal from large quantities of data is one of the most important problems in today’s world, with neural networks emerging as the most powerful solution. While there have been massive breakthroughs in the ML space, there are two main drawbacks. The first is that most of these models are “black boxes”, with very little understanding of why they work, which makes it hard to find the correct architecture for a given problem. The architecture is crucial, as the widely used training algorithms are “local” and thus rely on the loss landscape being benign enough to find a global optimum. This brings us to the second problem: while there is a lot of evidence that this can be overcome with large enough data, this is less clear for datasets which are noisy, do not have enough samples, or have a non-standard structure. Historically, on a smaller scale, the alternative has been simple (Bayesian) statistical methods, which also have the advantage of being interpretable and adaptable. However, many of the classical results cannot be easily extended to the high-dimensional regime in which modern data lives. In this work we want to consider two different settings relating to the above. The first is approximate message passing (AMP), an iterative algorithm for high-dimensional (generalized) linear models that easily lends itself to Bayesian setups, and which in many cases is provably Bayes-optimal. The second is an autoencoder with a shallow encoder, for which we can analyze the weights along the training trajectory of gradient descent. Surprisingly, we will find a connection between the two, as we recover certain structures of AMP. We will first introduce the AMP framework by studying an extension that can be applied to a wider variety of data: rotationally invariant generalized approximate message passing (RI-GAMP). This shows that one can extend the AMP framework to a much wider class of algorithms, relevant to more practical settings. The results are then also confirmed on real data that go beyond the theoretical setting. We then move to the setup of an autoencoder with a shallow encoder and a Bayesian prior in the high-dimensional regime, in which the input and compression dimensions are large and proportional; this is also a first step towards opening one of the black boxes above, as we can say precisely when and why the network works. While autoencoders are widely used in modern ML architectures, most fundamental research in this direction has focused on extreme asymptotic regimes, where the autoencoder is either severely over- or under-parameterized, and which do not extend to the proportional high-dimensional regime. We show that for Gaussian priors the network finds the global optimum, and these theoretical findings extrapolate strikingly well to real image data. Next we extend this analysis to more structured priors and a deeper decoder. For sparse Gaussian inputs in the training regime studied here, the shallow decoder does not exploit sparsity and yields the same dynamics as in the Gaussian setting. We then show how a deeper decoder can break this behavior, and interestingly, if we treat the problem as an AMP problem, the network eventually learns the same decoding as AMP. This too extrapolates to real image data. Lastly, we will show that the AMP setup has real practical use, by using a practical variant, genetic vector approximate message passing (gVAMP), for predicting plasma protein levels from genetic information using the UKB GWAS data. This is a perfect example of the problems in the high-dimensional regime: approximately 8.4 million dimensions and 40,000 samples, where the high noise, low sample count, and large size make it hard for ML methods to work. To get the optimal prediction accuracy we consider a two-step inference. We first consider a pre-processed version of the data, of about 880,000 dimensions, on which we run the gVAMP algorithm; this allows us to find predictive regions, taking into account joint effects between different regions, and for most proteins we end up selecting only about 0.1% of the dimensions. We then consider the analogue of this small subset, on which we can now run a standard tree algorithm (XGB) that no longer suffers from the aforementioned problems. This pipeline yields better prediction accuracy than recent nonlinear deep learning approaches applied to this data, as we are able to handle the pre-processing in a principled Bayesian way.