On the Out-of-distribution Generalization of Probabilistic Image Modelling

1 Introduction . Probabilistic modeling has achieved great success in the modeling of images . Likelihood based models , e.g . Variational Auto-Encoders ( VAE ) [ 22 ] , Flow [ 21 ] , Pixel CNN [ 49 , 41 ] are shown to successfully generate high quality images , in addition to estimating underlying densities . The goal of probabilistic modelling is to use a model pθ to approximate the unknown data distribution pd using the training data { x1 , . . . , xN } ∼ pd . A common method to learn parameters θ is to minimize some divergence between pd and pθ , for example , a popular choice is the KL divergence KL ( pd||pθ ) = −H ( pd ) − ∫ log pθ ( x ) pd ( x ) dx , ( 1 ) where the entropy of the data distribution is a constant . Since we only have access to finite pd data samples x1 , . . . , xN , the second term is typically approximated using a Monte Carlo estimation KL ( pd||pθ ) ≈ − 1 N N∑ n=1 log pθ ( xn ) − const .. ( 2 ) Minimizing the KL divergence is therefore equivalent to Maximum Likelihood Estimation . A typical evaluation criterion for the learned models is the test likelihood 1M ∑M m=1 log pθ ( xm ) , with test set { x1 , . . . , xM } ∼ pd . We refer to this evaluation as in-distribution ( ID ) generalization , since both training and test data are sampled from the same distribution pd . However , in this work , we are interested in out-of-distribution ( OOD ) generalization such that the test data are drawn from po , where po 6= pd . To motivate our study of this topic , we firstly introduce two applications : lossless compression and OOD detection , that both can make use of the OOD generalization property . ∗Co-first author , the work was done during an internship in Huawei Noah ’ s Ark Lab . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 1.1 Model-based Lossless Compression . Lossless compression has a strong connection to probabilistic models [ 33 ] . Let { x1 , . . . , xM } be test data to compress , where xm is sampled from some underlying distribution with PMF pd . If pd is known , we can design a compression scheme to compress each data xm to a bit string with length approximately − log2 pd ( xm ) 2 . As M → ∞ , the averaged compression length will approach the entropy of the distribution pd , that is ; 1M ∑M m=1− log2 pd ( xm ) → H ( pd ) where H ( · ) denotes entropy . In this case , the compression length is optimal under Shannon ’ s source coding theorem [ 45 ] , i.e . we can not find another compression scheme with compression length less than H ( pd ) . In practice , however , the true data distribution pd is unknown and we must use a model pθ ≈ pd to build the lossless compressor . Recent work successfully apply deep generative models such as VAEs [ 47 , 48 , 23 ] , Flow [ 19 , 3 ] to conduct lossless compression . We note that the underlying models are designed to focus on test data that follows the same distribution as the training data , resulting in test compression rates that depend on the ID generalization ability of the model . However , in practical compression applications , the distribution of the incoming test data is unknown , and is usually different from the training distribution pd : X otest = { x′1 , . . . , x′M } where x′ ∼ po 6= pd . To achieve good compression performance , a lossless compressor model should still be able to assign high likelihood for these OOD data . This practical consideration motivates encouragement of the OOD generalization ability of the model . Empirical results [ 48 , 19 ] have shown that we can still use model pθ ( trained to approximate pd ) in order to compress test data X otest with reasonable compression rates . However , the phenomenon that these models can generalize to OOD data lacks intuition and key components , affecting this generalization ability , remain underexplored . Consideration of recent advances in likelihood-based OOD detection next allows us to further investigate these questions and lead to our proposal of a new model that can encourage OOD generalization . 1.2 Likelihood-based OOD Detection . Given a set of unlabeled data , sampled from pd , and a test data x′ then the goal of OOD detection is to distinguish whether or not x′ originates from pd . A natural approach [ 5 ] involves fitting a model pθ to approximate pd and treat sample x′ as in-distribution if its ( log ) likelihood is larger than a threshold ; log pθ ( x′ ) > . Therefore , a good OOD detector model should assign low likelihood to the OOD data . In contrast to lossless compression , this motivates discouragement of the OOD generalization ability of the model . Surprisingly , recent work [ 36 ] report results showing that popular deep generative models , including e.g . VAE [ 22 ] , Flow [ 21 ] and PixelCNN [ 41 ] , can assign OOD data higher density values than in-distribution samples , where such OOD data may contain differing semantics , c.f . the samples used for maximum likelihood training . We demonstrate this phenomenon using PixelCNN models , trained on Fashion MNIST ( CIFAR10 ) and tested using MNIST ( SVHN ) . Figure 1 provides histograms of model evaluation using negative bits-per-dimension ( BPD ) , that is ; the log2 likelihood normalized by data sample dimension ( larger negative BPD corresponds to larger likelihood ) . We corroborate previous work and observe that tested models assign higher likelihood to the OOD data , in both cases . This counter intuitive phenomenon suggests that likelihood-based approaches may not make for good OOD image detection criterion , yet encouragingly also illustrates that a probabilistic model , trained using one dataset , may be employed to compress data originating from a different distribution with a potentially higher compression rate . This intuition builds a connection between OOD detection and lossless compression . Inspired by this link , we next investigate the underlying latent causes of image model generalizability , towards improving both lossless compression and OOD detection . 2 OOD Generalizations of Probabilistic Image Models . Previous work studies the potential causes of the surprising OOD detection phenomenon : OOD data may have higher model likelihood than ID data . For example , [ 37 ] used a typical set to reason about the source of the effect , while the work of [ 40 ] argues that likelihoods are significantly affected by 2For a compression method like Arithmetic Coding [ 50 ] , the message length is always within two bits of − log2 pd ( xm ) [ 33 ] , also see Section 4.1. image background statistics or by the size and smoothness of the background [ 26 ] . In this work , we alternatively consider a recent hypothesis proposed by [ 43 ] ( also implicitly discussed in [ 24 ] for a flow-based model ) : low-level local features , learned by ( CNN-based ) probabilistic models , are common to all images and dominate the likelihood . From the perspective of OOD generalization , this allows formulation of the following related conjectures : • Models can generalize to OOD images as local features are shared between image distributions . • Models can generalize well to OOD images since local features dominate the likelihood . In the work of [ 43 ] , the authors investigated their original hypothesis by studying the differences between individual pixel values and neighbourhood mean values and additionally considered the correlation between models trained on small image patches and trained , alternatively , on full images . To further investigate this hypothesis , we rather propose to directly model the in-distribution dataset , using only local feature information . If the hypothesis is true , then the proposed local model alone should generalize well to OOD images . By contrasting such an approach with a standard full model , that considers both local and non-local features , we are also able to study the contribution that local features make to the full model likelihood . We first discuss how to build a local model for the image distribution and then use the proposed model to study generalization on OOD datasets . 2.1 Local Model Design . Autoregressive models have proven popular for modeling image datasets and common instantiations include PixelCNN [ 49 , 41 ] , PixelRNN [ 49 ] and Transformer based models [ 7 ] . Assuming data dimension D , the Autoregressive model pf ( x ) can be written as pf ( x ) = p ( x1 ) D∏ d=2 p ( xd|x1 , . . . , xd−1 ) , ( 3 ) informally we refer to this type of model as a “ full model ” since it can capture all dependencies between each dimension ( pixel ) . Similarly , we can define a local autoregressive model pl ( x ) where pixel xij , at image row i column j , depends on previous pixels with fixed horizon h : pl ( x ) = ∏ ij p ( xij |x [ i−h : i−1 , j−h : j+h ] , x [ i , j−h : j−1 ] ) , ( 4 ) with zero-padding used in cases where i or j are smaller than h. Figure 2b illustrates the resulting local autoregressive model dependency relationships . We implement this model using a masked CNN [ 49 ] , with kernel size k=2×h+1 in our first network layer , to mask out future pixel dependency . A full PixelCNN model would then proceed to stack multiple masked CNN layers , where increasing kernel depth affords receptive field increases . In contrast , we employ masked CNN with 1×1 convolutions in subsequent layers . Such 1×1 convolutions can model the correlation between channels , as in [ 21 ] , and additionally prevent our local model obtaining information from pixels outwith the local feature region , defined by h. Pixel dependencies are therefore defined solely using the kernel size of the first masked CNN layer , allowing for easy control over model local feature size . We note that the proposed local autoregressive model can also be implemented using alternative backbones e.g . Pixel RNN [ 49 ] or Transformers [ 7 ] . We plot local model samples in Figure 4 . Unlike full autoregressive models [ 49 , 41 ] , which can be used to generate semantically coherent samples , we find the samples from the local model are locally consistent yet have no clear semantic meaning . 2.2 Local Model Generalization . To investigate the generalization ability of our local autoregressive model , we fit the model to Fashion MNIST ( grayscale ) and CIFAR10 ( color ) training datasets and test using in-distribution ( ID ) images ( respective dataset test images ) and additional out-of-distribution ( OOD ) datasets : MNIST , KMNIST ( grayscale ) and SVHN , CelebA3 ( color ) . Both models use the discretized mixture of logistic distributions [ 41 ] with 10 mixtures for the predictive distribution and a ResNet architecture [ 49 , 14 ] . We use a local horizon length h=3 ( kernel size k=7 ) for both grayscale and color image data . We compare our local autoregressive model to a standard full autoregressive model ( i.e . a standard PixelCNN ) , with additional network architecture and training details found in Appendix A . Tables 1 , 2 report comparisons in terms of BPD ( where lower values entail higher likelihood ) for Fasion MNIST and CIFAR10 , respectively . We observe that for in-distribution ( ID ) data , the full model has better generalization ability c.f . the local model ( 0.11 and 0.13 BPD , respectively ) . This is unsurprising as training and test data originate from the same distribution ; both local and non-local features , as learned by the full model , help ID generalization . For OOD data , we observe that the local model has generalization ability similar to the full model , exhibiting very small empirical gaps ( only ≈ 0.02 BPD on average ) , showing that the local model alone can generalize well to OOD distributions . We thus verify the hypothesis considered at the start of Section 2 . For simple datasets containing gray-scale images , the PixelCNN model is flexible enough to capture both local and global features . We notice that , in Table 1 , our local model exhibits even better OOD generalization than the full model . This drives us to further study the role of non-local features for generalization . When the local horizon size increases , the model will be able to learn features with greater non-locality . We thus vary the local horizon size to study generalization ability under this property , see Table 3 . We find the model has poor generalization performance when local features are too small and increasing the horizon size helps ID generalization but decreases the OOD generalization . A consistent phenomenon is observed when considering color images , see Appendix B.1 . We can thus conclude : non-local features are not shared between images distributions , overfitting to non-local features will hurt generalization . Table 3 : Generalization of local model with different horizon sizes . The model is trained on FashionMNIST dataset . Method h=1 h=2 h=3 h=4 h=5 ( ID ) Fashion 3.17 2.93 2.89 2.88 2.88 ( OOD ) MNIST 1.54 1.48 1.49 1.50 1.51 ( OOD ) KMNIST 2.54 2.43 2.44 2.46 2.47 These hypotheses indicate two opposing strategies for the considered tasks : OOD detection : distributions are distinguished by dataset-unique features , thus building non-local models , able to discount common local features , improves the detector distinguishability power . 3We down-sample the original CelebA to 32×32×3 , see Appendix A.2 for details . Lossless compression : OOD generalization is possible due to the sharing of local features between distributions . Employing only a local model can encourage OOD generalization , by preventing the model from over-fitting to dataset-unique features , specific to the training distribution . Sections 3 and 4 will further demonstrate how contrasting modeling strategies can benefit these tasks .

Summary: A surprising observation in the OOD literature is that OOD detection based on density estimation fails on some image datasets since the learned density tends to be higher on OOD data. A hypothesis for this phenomenon is that the likelihood is dominated by the local features of the image, resulting in the high-level features being ignored. Motivated by this observation, this paper proposes a model that learns local features of images. This model is used for two different applications: (1) by combining it with a full density estimation model, it gives a better OOD detection method, achieving state-of-the art performance (2) by pairing the model with the artihmetic coding algorithm, it gives a dataset compression method that achieves state-of-the-art compression rates on in- or out-of-distribution datasets.