Improving Compositionality of Neural Networks by Decoding Representations to Inputs
Year: 2022

Abstract:
In traditional software programs, it is easy to trace program logic from variables back to input, apply assertion statements to block erroneous behavior, and compose programs together. Although deep learning programs have demonstrated strong performance on novel applications, they sacrifice many of the functionalities of traditional software programs. With this as motivation, we take a modest first step towards improving deep learning programs by jointly training a generative model to constrain neural network activations to “decode” back to inputs. We call this design a Decodable Neural Network, or DecNN. Doing so enables a form of compositionality in neural networks, where one can recursively compose DecNN with itself to create an ensemble-like model with uncertainty. In our experiments, we demonstrate applications of this uncertainty to out-of-distribution detection, adversarial example detection, and calibration — while matching standard neural networks in accuracy. We further explore this compositionality by combining DecNN with pretrained models, where we show promising results that neural networks can be regularized from using protected features.

heading: 1 Introduction
text: 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 cannot 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.

heading: 2 Background
text: 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].

heading: 3 Neural Networks with Decodable Activations
text: 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.

heading: 4 Composing Decodable Neural Networks
text: 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.

heading: 4.1 Recursive Self-Composition
text: 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.

heading: 4.2 Composing with Pretrained Models
text: Apart from self-composition, we can also compose our neural networks with off-the-shelf pretrained models as a form of regularization or distillation. Suppose we have a pretrained model m that maps an input x to a class prediction. We can modify the ReDecNN objective as follows:
L̃(x; θ, φ) = L̂(x; θ, φ) + D∑ d=1 αd log p(fθ(gφ(h (d−1) ld−1 ))|m(gφ(h(d−1)ld−1 ))) (7)
where L̂ is as defined in Equation 5. The additional term acts as a divergence, bringing the classifer’s predictions close to those of the pretrained model. A similar edit can be made to the DecNN. We revisit Section 4 experiments now using two pretrained models — a linear classifier or a ResNet-18 classifier. For each, we optimize Equation 7 and compute performance on OOD detection.
Table 5 compares the results of composing with either (1) a linear model, (2) a residual network, or (3) no composition (None). Using a linear model, we observe a slight drop in test accuracy but an increase across all OOD experiments and calibration. In fact, these new results rival or surpass MC-Dropout from Table 2. Conversely, with ResNet-18, we observe a 1 point increase in test
accuracy but notable drops in OOD and calibration results. These polarizing outcomes show two use cases of pretrained models: using simpler models allow for regularization that trades off accuracy for robustness while more complex models encourage distillation, making the opposing tradeoff.
Method Eq. 7 Direct Reg.
Accuracy 97.17 96.72 Misclass. 0.904 0.828 Adversarial 0.805 0.707 OOD (Fashion) 0.845 0.699 OOD (CelebA) 0.928 0.782 Corruption (Mean) 0.790 0.667 Calibration (ECE) 0.296 0.280
as well. Intuitively, constraining the model to be similar to what a pretrained model thinks of the reconstruction (which is a dynamic value that varies with model weights) is a stronger learning signal than what the pretrained model thinks of the original input (which is static). In other words, regularizing fθ(gφ(hk)) to be like m(x) at all layers l is too strict of a constraint.

heading: 5 Constraining Decodable Computation
text: A milestone for deep learning would be the ability to specify what information an activation should not capture. This is especially important in cases when we have protected attributes that should not be abused in optimizing an objective. While this milestone remains out of grasp, we take a small step towards it by optimizing a neural network to “be indifferent” with respect to a protected attribute through composition with pretrained models.
Suppose we have access to a pretrained classifierm that predicts the assignment for aK-way protected feature. Then, building on Section 4.2, we can optimize our classifier to “ignore” information about the protected feature through m. To do this, we define the objective:
L̃(x; θ, φ) = L̂(x; θ, φ) + D∑ d=1 αd log p(m(gφ(h (d−1) ld−1 ))| 1 K 1K) (8)
where 1K is a K-dimensional vector of ones. That is, Equation 8 encourages the decoded activation gφ(h (d−1) ld−1
), when passed through the protected classifier m, to predict chance probabilities i.e. have no information about the protected attribute. A high performing classifier trained in this manner must have solved the prediction task without abusing the protected attribute(s).
To test this, we extract two attributes from CelebA: “beardedness” and “baldness”. Suppose we wish to design a neural network to predict beardedness. If we were to naively do so, there could be a discrepancy in performance between groups of individuals — for example, between bald individuals and individuals with hair, the latter group is more common in the CelebA dataset. Indeed, Table 7 shows a 10 point difference in F1 and a 5 point difference in average precision (AP) in classifying beardedness between the two groups. The second row of Table 7 provides a baseline ReDecNN
optimized without knowledge of the protected attribute. Unsurprisingly, we find a similar discrepancy between groups, as with a standard neural network. In the third row, we evaluate a ReDecNN that was optimized to ignore the “baldness” attribute using a pretrained classifier for baldness. Critically, we find much more balanced (and higher) F1, AP, and ECE across the two groups.
Method F1
Standard 0.383 ReDecNN (Eq. 5) 0.371 ReDecNN (Eq. 8) 0.108

heading: 6 Generalizing to other Modalities
text: While we have focused on image classification, the proposed approach is more general and can be extended to other modalities. We apply the same decodable representations to speech, in particular utterance and action recognition. We compare DecNN and ReDecNN to the same baselines, measuring accuracy and uncertainty through similar experiments as we did for images.

heading: 7 Related Work
text: Autoencoders Although DecNNs are reminiscent of autoencoders, there are differences: while autoencoders only reconstruct the final hidden layer, DecNN has a reconstruction for every single hidden layer. This, and because DecNNs have a supervised objective (which autoencoders do not have) critically changes what information is captured in the hidden layers. Whereas the autoencoder seeks perfect reconstruction, the DecNN does not, as evidenced by loss of information for reconstructions later in the network (see Appendix A.3). Finally, we point out that the applications of autoencoders are very different from the DecNN. The former is not used for calibration or composition.
Invertible and Ensemble Networks The DecNN is akin to invertible generative models [6, 36] as DecNN is training a generative model to “invert” the classifier up to a hidden layer. However, unlike invertible flows, DecNN does not impose an explicit prior on the activations (i.e. latent variables), and further, DecNN has a supervised objective. Our proposed recursive network, ReDecNN, has similarities to Bayesian neural networks [11, 4] and ensemble networks [48] — two baselines we compare against for evaluating model uncertainty. However, we find ReDecNN to be easier to train than doing inference over weights, and cheaper in parameter count than naive ensembles.
Probing Neural Networks A related line of work seeks to probe neural network post-training to understand the underlying computation [1]. Most relevant to our work is an approach that inverts CNN activations through reconstruction [28]. Unlike this approach, we are not proposing any additional computation post-hoc. Rather, we are interested in optimizing neural networks such that their representations are more easily mapped back to the input domain.
Robustness and Protected Attributes Out-of-distribution detection [26, 24, 52, 40, 20, 19], selective classification [12, 13], adversarial perturbation detection [49, 10, 30, 34], and neural network calibration [16, 51, 50, 32, 38] each have a rich subfield of algorithms. While task-specific methods likely outperform our results, we are excited about representations that enable uncertainty without task information or additional algorithms. Our approach also shares similarities to [42] where the authors minimize the mutual information between learned representations and a protected attribute. We view Equation 8 as an approximation of this where we treat the pretrained classifier for the protected attribute as a proxy for the mutual information.

heading: 8 Limitations and Future Work
text: Method MNIST Fashion CelebA
Standard 13.6 (0.3) 13.5 (0.2) 82.6 (6.2) DecNN 80.5 (6.8) 99.2 (9.6) 486.2 (12.9) ReDecNN 96.7 (9.7) 78.6 (6.6) 521.6 (9.2) Dropout 13.9 (0.6) 14.0 (0.1) 79.5 (5.7) MC-Dropout 13.9 (0.6) 14.0 (0.1) 79.5 (5.7) Ensemble 109.6 (7.1) 132.6 (8.5) 659.4 (8.2)
mizing recursive networks, although parameter efficient, costs more compute as backpropagation is more expensive. See Table 10 for timings of 1 epoch in seconds. The proposed method does have significant computational overhead, averaging about 6x cost, while baseline methods like Dropout and MC-Dropout impose little to no overhead. On the other hand, ReDecNN has comparable cost to DecNN despite requiring recursive gradients. Also, both DecNN and ReDecNN are both cheaper than naive ensembles, significantly in the CelebA case. Future work could explore weight sharing across layers (not just depth) to reduce compute.
In summary, we explored building neural networks with decodable representations, and leveraged this new property to compose models together. By re-purposing decoded activations as novel inputs, we were able to join neural networks together in useful ways for out-of-distribution detection and calibration, as well as influence what information a network retains in optimization. We are optimistic about promising results and look to future research for more applications of decodability.
Acknowledgments and Disclosure of Funding
We thank the Ermon group for their comments and suggestions. We thank the reviewers for their many iterations of feedback that helped this paper improve significantly. This research was supported in part by NSF (#1651565, #1522054, #1733686), ONR (N000141912145), ONR MURI (N000141612007), AFOSR (FA95501910024), ARO (W911NF-21-1-0125), and Sloan Fellowship. MW is supported by the Stanford Interdisciplinary Graduate Fellowship as the Karr Family Fellow.

Review summaries:
This paper presents an approach to simultaneously train a prediction network along with a generative model to map the neural network activations back to inputs. The use cases of doing so include (1) interpreting the information encoded in an activation, (2) composing the network with itself and other pre-trained networks, which in turn provides a means to get uncertainties and to regularize the network, and (3) to assert that the network satisfies some constraints on what it can and cannot learn in the intermediate layers. The paper is evaluated on several image benchmarks and speech benchmarks and showcases all the above benefits.  

This paper proposes to jointly train a generative model that maps activations back to the input space. The work demonstrates several benefits of this approach.  * Probing: It provides a way to decode the activations which enables the users to make sense out of it since the decoded activations are in the input space.  * Composing networks: The decoded activations allow to compose a DecNN model with another model with the same input type.  * Data Augmentation: The DecNN model can be composed with itself. This can be viewed as data augmentation since the decoded activations are of the same "type" as the input and retain the features that are relevant to the prediction task. * Ensemble network. The DecNN model can be recursively composed with itself to get the “Recursively Decodable Neural Network” (ReDecNN) model. This model can be viewed as an ensemble network. The paper shows the usefulness of ReDecNN for the tasks of measuring uncertainty, OOD detection, and calibration. * Constrain the activations: The decodable activations makes it possible to specify the kind of information an activation is not supposed to capture. This provides assertion-like capability to the neural network. 

This work proposes DecNN, and the main goal is to allow people to interpret the intermediate neural network activations. The base model is an 8-layer MLP. The core idea is to train another generative model that takes the MLP activation as the input, and the generative model aims to decode an image in the model input space. During training, the loss of the generative model aims to reconstruct the input image given the activations of all intermediate layers. They further design ReDecNN, which recursively composes DecNN by taking the output of the generative model as the input of DecNN. They demonstrate that the generative model produces interestingly interpretable outputs to explain why the classifier makes correct or wrong predictions. They also show some promising results for out-of-distribution detection, calibration, and fairness.

This paper argues for inverting intermediate latent representations of convolutional classifiers -- predicting the original input -- as a tool for improving robustness and interpretability. The ability of the model to learn invertible intermediate hidden representations acts as a regularizer during training.  This basic premise is used in several ways. For example, a recursive version of this idea (using the reconstructed input to classify, and the intermediate hiddens of that to re-reconstruct the input, adding that reconstruction loss into the overall loss calculation) leads to an interesting form of drop out: choose which reconstructions (from which layer, and which recursive depth) to use in the final output. That can be used for enhanced forms of Monte Carlo Dropout, improving uncertainty computations over regular MC Dropout.  Another use is in placing constraints on intermediate hidden layers. For a fairness goal, such a constraint might be that, given a good classifier for some protected attribute, that classifier has poor performance on the reconstructed inputs from the original model's hidden layers. In a sense, this constrains the training of the model to lose the information that would allow the external classifier to predict the protected feature.

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.