Improving Compositionality of Neural Networks by Decoding Representations to Inputs

1 Introduction . Traditional hand-written computer programs are comprised of a computational graph of typed variables with associated semantic meaning . This structure enables practitioners to interact with programs in powerful ways ( even if they are not the author ) — such as debug code by tracing variables back to inputs , apply assertions to block errors , and compose programs together for more complex functionality . However , traditional software has its limitations : it is difficult to hand-write programs to classify images or extract sentiment from natural language . For these functionalities , deep learning and neural networks [ 41 ] have become the dominant approach [ 23 , 2 , 46 ] . While neural networks have made impressive progress on complex tasks , they come at a sacrifice of many of the desirable properties of traditional software . Specifically , the closest approximation to a “ variable ” in a neural network is an activation . Yet it is difficult to understand a neural network ’ s computation from an activation value , and there is little to associate an activation with semantic meaning . A practitioner can not write assertion statements to constrain valid neural network logic : checking the values that an activation takes is usually not enough to gauge correctness nor meaning . Moreover , given multiple neural networks , composing them together requires retraining from scratch . In this paper , we take a modest first step towards bridging the expressivity of deep learning with the engineering practicality of traditional software . We aim to uncover new ways for practitioners to build and use neural networks by leveraging compositionality . Specifically , we propose to train a neural network classifier jointly with a generative model whose role is to map the classifier ’ s activations back to the input , approximating invertibility . For any input , a neural network ’ s computation can 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . be represented by a sequence of activations ( from each layer ) , each of which can now be mapped back to the input space . It is this insight that enables a special form of compositionality for neural networks , as inputs derived from activations can be fed back into other models . In our experiments , we study this compositionality by ( 1 ) recursively composing a neural network with itself , creating an ensemble-like model with a measure of uncertainty that is useful for out-ofdistribution detection , adversarial example detection , and model calibration ; ( 2 ) composing neural networks with pretrained models as forms of regularization and distillation ; and ( 3 ) using decodable activations to discourage a neural network from using protected attributes . Throughout , we show that decodability comes at low cost as we find equivalent accuracy to a standard neural network . 2 Background . Neural Networks We will focus on supervised neural networks for classification , although it is simple to extend to regression tasks . We denote a neural network by fθ where θ represents its trainable parameters . A neural network fθ maps an example x to predicted probabilities over K classes . Typically , a neural network is composed of L blocks . For example , if fθ is a multi-layer perception , each block is a linear layer followed by a non-linearity . If fθ is a ResNet [ 17 ] , each layer is a residual block . In computing its prediction , a neural network fθ produces activations { h1 , . . . , hL } where hl is the output of the l-th block . To train a neural network , we solve an optimization function of the form Lnn ( x ; θ ) = log p ( y|fθ ( x ) ) + βΩ ( x ; θ ) ( 1 ) using stochastic gradient descent . In Equation 1 , the notation x denotes an input example and y its label . The function Ω represents an auxiliary objective such as a regularizer . The hyperparameter β > 0 is used to scale the auxiliary loss . For classification , log p ( y|· ) is cross entropy . Generative Models Broadly , we are given a latent variable H and an observed variable X . A generative model gφ captures the distribution p ( x|h ) , mapping samples of the latent code h to samples no hat Reconstruction loss per layer ( total of 4 here ) has hat Prediction Input Decoded activations through reconstruction . Information is discarded ( removal of hat ) as activations interpolate between the input and ( incorrect ) prediction . Hidden layers of the neural network , True label Cross entropy loss h2 h3 fθ gφ h4h1 ( a ) Decodable Neural Networks gφ ( b ) Ensembling through Self-Composition Continue to recurse Parameter Sharing Parameter SharingTreat decoded activation as input Has hat Uncertainty No hat has hat has hat no hat fθ fθ fθ of the observed variable x . For a deep generative model , gφ is parameterized by a neural network . In the simplest case where gφ is an autoencoder , we are also given an encoder q that maps observed samples x to latent codes h. Then , we optimize the objective Lae ( x ; φ ) = log p ( x|gφ ( q ( x ) ) ) ( 2 ) where for continuous x , the expression log p ( x|h ) is the log Gaussian probability with mean h and fixed variance ( or equivalently , mean squared error ) . More complex generative models might impose a prior distribution on h [ 22 ] , an adversarial discriminator [ 29 , 14 ] , or invertibility of gφ [ 35 ] . 3 Neural Networks with Decodable Activations . In Section 1 , we motivated our desiderata for activations that can be mapped back to the input space . We now propose to do this by “ inverting ” a neural network at a particular activation . Concretely , given an input x , a neural network fθ , we can compute an activation h by computing the forward pass fθ ( x ) . Then , suppose we could construct the set : S ( h ) = { x̂ : x̂ ∼ pD , fθ ( x̂ ) = h } ( 3 ) where pD represents the empirical distribution that x is sampled from . Now , any input x̂ ∈ S ( h ) is a valid mapping of h back to the input space . ( Trivially , x ∈ S ( h ) . ) We call this x̂ a decoding of the activation h. Unfortunately , in practice we do not know how to construct such a set S ( h ) . But , we do know how to ( approximately ) sample from S ( h ) . We can train a generative model gφ to reconstruct the original input x from neural network activations h ( recall x ∈ S ( h ) ) . Then , we can sample x ∼ gφ ( h ) , which we know how to do , to sample from S ( h ) . If we jointly train this generative model gφ with the neural network fθ , we can optimize for activations that both solve the classification task and decode back to inputs at once . So far , we have only considered a single activation despite our neural network having L such activations . Although one could train a separate generative model per layer , we found it efficient and performant to share weights φ across all layers . Our joint objective is : L ( x ; θ , φ ) = log p ( y|fθ ( x ) ) + β 1 L L∑ l=1 log p ( x|gφ ( hl ) ) , h1 , ... , hL come from fθ ( x ) ( 4 ) where the first term is cross entropy as in Equation 1 and the second term is a reconstruction error ( Equation 2 ) averaged over layers . Comparing Equation 4 to Equation 1 , we can interpret the generative loss as a regularization term Ω . Intuitively , in a standard neural network , the activations h1 , . . . , hL are unconstrained , chosen only to minimize the loss . However , in Equation 4 , the network fθ is constrained to produce activations that the generative model gφ can invert . In practice , gφ is a ResNet-18 decoder with upsampling layers ( see Appendix A.1 ) . We call this setup a Decodable Neural Network , or DecNN . Figure 1 depicts an illustration of the DecNN architecture . We highlight that although Equation 4 requires good reconstruction , it does not enforce that reconstructed inputs f ( gφ ( h ) ) must map to the same label y as f ( x ) . We show a few examples of this by visualizing decoded activations throughout a network in Appendix A.3 . 4 Composing Decodable Neural Networks . In traditional software , we can easily compose two compatible functions f1 and f2 together by f2 ( f1 ( x ) ) for some input x . With this as loose inspiration , we wish to similarly compose neural networks together . As motivation , we may want to do this to faciliate interaction between networks , such as for regularization or distillation . But given two image classifiers , it is not clear how to tie outputs from one neural network to inputs for the other ? Both networks expect images as input . The workaround comes from decodable activations . Since they are images of the same shape as inputs , we are free to treat them as inputs for another model . So , we can compose neural networks together using decoded activations as an intermediary : given two models f1 and f2 , and an input x , we can take an activation h coming from f1 ( x ) , decode it , then provide it to the other model , f2 ( gφ ( h ) ) . In the following subsections , we explore two instances of this compositionality : one by composing a model with itself ( i.e. , f1 = f2 ) , and the other by composing a model f1 with a pretrained one , f2 . See Figure 1b , c for a graphical depiction of both compositions . 4.1 Recursive Self-Composition . Our first composition will be to compose a DecNN model with itself . But if we can do that , we can also recursively compose a DecNN model with itself infinitely : Given a decoded activation gφ ( h ( 0 ) l ) for any layer l , we can feed it as input back into the classifier , fθ ( gφ ( h ( 0 ) l ) ) and maximize agreement with the true label y . Intuitively , if gφ ( hl ) is a good reconstruction of the input x ( i.e . gφ ( hl ) and x belong to the same set S ( hl ) ) , then a good classifier fθ should predict the correct label with it as input . Now , we also observe that the computation fθ ( gφ ( h ( 0 ) l ) ) produces L additional activations , h ( 1 ) 1 , . . . , h ( 1 ) L where we use the superscript to represent the number of compositions . We can repeat this exact process with the same intuition , decoding h ( 1 ) l and passing it to fθ to produce h ( 2 ) l . Computational limits aside , this can be repeated ad infinitum . One interpretation of this procedure is as data augmentation : we saw in Section A.3 that decoded activations share features relevant to the prediction task ( e.g . generated celebrities wear hats ) but vary other irrelevant features ( e.g . face shape , hair color , etc. ) . Notably , these generated images are not in the original dataset . Thus , training a model using decoded activations as inputs is similar to augmenting the initial dataset . The only difference is that gradients are propagated back through recursion to the inputs , which is also a form of regularization . Admittedly , what we have described so far is intractable . Besides recursing infinitely , each activation produces L more activations , creating a tree ( with root node x ) with a large branching factor L. In practice , we make two adjustments — First , we limit the recursion to a max depth D. Second , instead of constructing the full graph of activations , which would have O ( LD ) nodes , we randomly sample “ paths ” in the graph . Specifically , starting at the root , we compute { h ( 0 ) 1 , . . . , h ( 0 ) L } through fθ ( x ) and randomly choose an integer l0 ∈ [ 1 , L ] . Then , we compute { h ( 1 ) 1 , . . . , h ( 1 ) L } through fθ ( gφ ( h ( 0 ) l0 ) ) i.e . we classify the l0-th decoded activation from depth 0 . Again , we randomly pick l1 to compute { h ( 2 ) 1 , . . . , h ( 2 ) L } through fθ ( gφ ( h ( 1 ) l1 ) ) and repeat until depth D. Building on Equation 4 , we can write this “ recursive ” objective as : L̂ ( x ; θ , φ ) = log p ( y|fθ ( x ) ) + D∑ d=1 αd log p ( y|fθ ( gφ ( h ( d−1 ) ld−1 ) ) ) +β D∑ d=1 αd log p ( x|gφ ( h ( d−1 ) ld−1 ) ) ( 5 ) where the sequence of integers l0 , . . . lD−1 are each sampled uniformly from [ 1 , L ] . We introduce a new hyperparameter α ∈ [ 0 , 1 ] that geometrically downweights later depths ( if α = 0 , Equation 5 reduces to maximum likelihood ) . The first term in Equation 5 is the usual classification objective . The second term is a sum of cross entropy terms that encourage the decoded activation at depth d to predict the correct label through fθ . The third and final term is a sum of reconstruction losses that encourage the decoded activation at depth d to approximate the original input x . We call this model a “ Recursively Decodable Neural Network ” , abbreviated ReDecNN . Note that ReDecNN is a special case of an ensemble network . That is , every path sampled in the activation tree by picking l0 , . . . lD−1 builds a legitimate classifier . There are O ( LD ) such classifiers implicit in the ReDecNN . We can interpret optimizing Equation 5 as training individual classifiers that are Monte Carlo sampled from the tree every gradient step . See Figure 1b for an illustration . Our goal in the next few paragraphs is to utilize this ensemble to measure uncertainty . Method MNIST Fashion CelebA Standard 97.1 ( 0.11 ) 87.7 ( 0.19 ) 90.8 ( < 0.1 ) DecNN 97.8 ( 0.17 ) 88.8 ( 0.19 ) 90.8 ( < 0.1 ) ReDecNN 97.5 ( 0.20 ) 88.1 ( 0.16 ) 90.8 ( 0.16 ) Dropout 95.9 ( 0.13 ) 80.7 ( 0.24 ) 88.5 ( 0.12 ) MC-Dropout 91.3 ( 0.14 ) 78.6 ( 0.23 ) 87.8 ( 0.19 ) BayesNN 96.3 ( 0.51 ) 87.0 ( 0.36 ) 88.2 ( 0.22 ) Ensemble 97.6 ( 0.11 ) 88.4 ( 0.12 ) 90.9 ( 0.11 ) ( a ) Classification performance on test set . Method MNIST Fashion CelebA ReDecNN 0.869 0.795 0.692 MC-Dropout 0.878 0.818 0.657 BayesNN 0.537 0.574 0.587 Ensemble 0.531 0.559 0.605 ( b ) ROC-AUC of separating correctly classified and misclassified examples using Eq . 6 as a score . The stdev . for all entries over 3 runs are < 0.01 . Measuring Uncertainty As a first step , we compare ReDecNN to several baselines : a standard neural network , a neural network with dropout , two kinds of Bayesian neural networks — Monte Carlo dropout [ 11 ] , abbreviated MC-Dropout , and weight uncertainty [ 4 ] , abbreviated BayesNN — and a naive ensemble network where we train D copies of the standard neural network with different initializations . For all experiments , see Appendix A.1 and A.2 for training and task details . Table 1a shows accuracies over a held-out test set , averaged over three runs . While we observe lower accuracy from Bayesian NNs , we observe ( perhaps surprisingly ) equal performance of DecNN and ReDecNN to a standard neural network . Typically , adding uncertainty to neural networks comes at a cost to performance ( as is true for other baselines ) , but this seems to not be the case here . The next step is to compare the quality of model uncertainty . For the ReDecNN , for a given input x , we measure uncertainty as follows : sample N different classifiers from the activation tree , and make predictions with each classifier on x . The uncertainty is the entropy of the predictions : Uncertainty ( x ) = − K∑ c=1 p̂ ( y = c ) log p̂ ( y = c ) ( 6 ) where p̂ ( y = c ) is the empirical probability ( e.g . normalized count ) of predicting class c out of N classifiers . A higher uncertainty metric would represent more disagreement between classifiers . We can compute the same metric for a naive ensemble , as well as for MC-Dropout ( where repeated calls drop different nodes ) and BayesNN ( by sampling weights from the learned posterior ) . We want to test if the model ’ s uncertainty is useful . One way is to correlate the uncertainty with when the model makes prediction errors ( we call this “ misclassification ” ) . It would be useful if the model was less uncertain on correct predictions and more uncertain on incorrect ones . A practitioner could then use uncertainty to anticipate when a model might make a mistake . Table 1b reports the area under the ROC curve , or ROC-AUC of separating correctly and incorrectly classified examples in the test set ( using uncertainty as a score ) . A higher ROC-AUC , closer to 1 , represents a better separation . We find that ReDecNN is competitive with MC-Dropout , even out performing it on CelebA . In Section A.5 , we include examples of images with low and high uncertainty . Out-of-Distribution Detection A second application of uncertainty is detecting out-of-distribution ( OOD ) inputs . Given an input that is far from the training distribution , a model is likely to make a mistake . Instead , the model should recognize its uncertainty , and refuse to make a prediction . There is a rich literature on detecting OOD inputs using downstream computation on a trained model [ 18 , 25 , 24 , 52 , 19 ] . Unlike those works , we study OOD detection using uncertainty alone . We explore three kinds of OOD examples : ( 1 ) adversarial examples crafted for a single classifier using FGSM [ 15 ] , ( 2 ) examples from a different dataset ( e.g . train on MNIST and use FashionMNIST as OOD ) as done in [ 25 ] , and ( 3 ) corrupted examples by 14 image transformations ( e.g . adding pixel noise ) , borrowed from [ 31 ] . See Appendix A.4 for expanded corruption results . Table 2 reports the ROC-AUC of separating inlier examples taken from the test split of the dataset the models were trained on , and outlier examples . From Table 2 , we observe ReDecNN is just under MC-Dropout , but while MC-Dropout achieves these results at the cost of classification accuracy , ReDecNN does not ( Table 1a ) . Furthermore , we find ReDecNN generalizes better to CelebA , a more complex image dataset , where it outperforms MC-Dropout ( and other baselines ) . Focusing on CelebA , we study a domain-specific OOD challenge by holding out all images with positive annotations for a single attribute as out-of-distribution . Here , inlier and outlier examples are very similar in distribution . We report ROC-AUCs in Table 3 for three held-out attributes ( randomly chosen ) : is the celebrity in the image “ wearing a hat ” , has “ blonde hair ” , or “ bald ” ? As this is a more challenging problem ( since inlier and outlier examples are from the same dataset versus different datasets ) , the performance is lower than in Table 2 . But while MC-Dropout , BayesNN , and a naive ensemble all perform near chance , ReDecNN is consistently near 0.6 . Calibration A third application of uncertainty we explore is calibration of predicted probabilities . Standard neural networks are notoriously over-confident and are incentived in training to predict with extreme probabilities ( e.g . 0.99 ) . A model with useful uncertainty would make predictions with proper probabilities . Although many calibration algorithms exist [ 33 , 16 ] , we want to measure how calibrated each model is out-of-the-box to compare the quality of model uncertainties . Table 4a reports the expected calibration error , or ECE [ 32 ] , which approximates the difference in expectation between confidence and accuracy using discrete bins . A lower number ( closer to 0 ) represents a more calibrated model . We observe that while a standard neural network has a relatively high ECE , many of the approaches ( dropout , MC-Dropout , BayesNN , and ReDecNN ) reduce ECE . On the other hand , naive ensembles increase the calibration error , as they are more prone to overfitting . In two of three datasets , ReDecNN achieves the lowest ECE whereas MC-Dropout achieves the lowest in the third . These results should be viewed in parallel to Table 1a , which measures “ sharpness ” . Otherwise , we can trivially reduce ECE to zero by ignoring the input . Method MNIST Fashion CelebA Standard 1.18 ( 0.10 ) 0.56 ( 0.14 ) 2.75 ( 0.13 ) ReDecNN 0.33 ( 0.05 ) 0.21 ( 0.03 ) 1.92 ( 0.13 ) Dropout 0.63 ( 0.09 ) 0.48 ( 0.09 ) 2.51 ( 0.13 ) MC-Dropout 0.41 ( 0.06 ) 0.13 ( 0.03 ) 2.99 ( 0.15 ) BayesNN 0.91 ( 0.20 ) 0.18 ( 0.03 ) 2.41 ( 0.19 ) Ensemble 1.80 ( 0.12 ) 0.89 ( 0.07 ) 3.12 ( 0.12 ) ( a ) We compare the expected calibration error ( ECE ) of ReDecNN to a standard neural network with various regularization and ensembling . Depth Acc . OOD ( Fashion ) OOD ( CelebA ) 2 97.7 0.665 0.752 4 97.6 0.788 0.804 6 97.7 0.805 0.868 8 97.5 0.812 0.893 ( b ) Effect of recursion depth D on classification accuracy and OOD detection for a ReDecNN model ( 8-layer MLP ) trained on MNIST . Effect of Depth Finally , we study the effect of depth on the efficacy of ReDecNN . For the experiments above , we chose the recursive depth to be equal to the number of layers in the MLP ( D = L = 8 ) . In Table 4b , we vary the depth from 2 to 8 ( while keeping the number of layers fixed to 8 ) and study the effect on classification accuracy and OOD detection . We find increasing OOD performance as D increases , although with marginally decreasing gains . Moreover , we observe constant accuracy , matching standard neural networks regardless of the choice of D .

Summary: This paper explores jointly training generative model probes alongside a neural architecture to imbue the model with capabilities analogous to "debugging" in traditional software. By training the generative model probes to reproduce the model inputs, every intermediate activation can be decoded to input space (images). The authors use this (1) to interpret the partial progress of the model made at intermediate layers, (2) to restrict the model from using sensitive information from the inputs, and (3) as an interface to compose the model with itself and with other pretrained image models. The authors further use self-composition to produce an ensemble from a tree of connected classifiers, which in turn gives uncertainty measurements that enable (4) out-of-distribution detection and (5) calibration. Experimental evidence backs each of these contributions.