Mini-batch graphs for robust image classification

Our network has two components: a feature encoder and a mini-batch graph neural network (MBGNN), as illustrated in Fig. 1. We obtain the encoder $f_{\theta}(\cdot)$ by removing the final layer of a standard residual network, such as Resnet-50. Consider a typical training setup which takes a mini-batch of samples as input for a downstream classification task. We denote the input samples in a mini-batch as $\mathcal{X}=\{x_{1},x_{2},...,x_{B}\}$, where $B$ is the batch size. The encoder provides representations for each sample, $h_{i}^{(0)}=f_{\theta}(x_{i}),\forall i\in B$. The MBGNN induces a graph on the set of encoded representations and processes them using graph convolution operations. We denote the representations in the $l^{\text{th}}$ layer by the set $\mathcal{H}^{(l)}=\{h_{1}^{(l)},h_{2}^{(l)},...,h_{B}^{(l)}\}$. To dynamically induce a graph on $\mathcal{H}^{(l)}$, we obtain the adjacency matrix $A^{(l)}$ by computing pairwise cosine similarity between representations and consider the top $k$ similar representations for each sample as its neighbours, removing self connections. The extent of the neighbourhood for each node can be controlled with the parameter $k$. The layer-wise propagation rule of the MBGNN in vector form, is given by,

where $h^{(l)}_{i}$ is stacked row-wise to form $H^{(l)}$, $W^{(l)}$ is the weight matrix, $b^{(l)}$ is the bias, and $\sigma(\cdot)$ is a non-linear function, usually ReLU. In equation (1), $f_{\textrm{self}}^{(l)}=\bar{H}^{(l)}$ contains the ‘self’ information for each node, and $f_{\textrm{neigh}}^{(l)}=(1/k)A^{(l)}\bar{H}^{(l)}$ contains the ‘neighbour’ information, since it is based on the average of the encoded representations of the neighbors of each node, as reflected by the adjacency matrix $A^{(l)}$.

There are several methods to combine the self and neighbour information in the $\text{combine}(\cdot,\cdot)$ function, which are explained below.

Concatenation: Here the self features, $f_{\textrm{self}}^{(l)}$ and neighbour features $f_{\textrm{neigh}}^{(l)}$ are concatenated in the representation dimension. This allows the network the flexibility to learn separate weights for both.

Weighted Addition: This is a convex combination of self and neighbour features, which forces the network to use neighbour information. For $\alpha\in[0,1]$, we have,

This reduces to a standard GCN formulation if we set $\alpha=1/{k+1}$.

Drop Features: We propose a different method in to mitigate the effect of co-adaptation and to make the neighbours contribute meaningful information. During training, we drop the self features with probability $p$ and the neighbour features with probability $1-p$. We then make the testing phase deterministic by taking the expected output feature, which leads to a similar expression as the sum combination:

We explore these different methods in our experiments.

We now show how our MBGNN provides robustness in the face of perturbations in the input, which we have argued is a desirable property for downstream computer vision tasks.

The above proposition effectively states that for any perturbation in the input, the corresponding perturbation in the output is inversely proportional to the number of neighbours, for each node in the mini-batch graph, when using MBGNNs as opposed to standard networks. Similarity based aggregation aids in transductive inference, where a prediction is made for a single corrupted test image within a mini-batch of randomly sampled uncorrupted training set images. In Section 4.4 we carry out experiments to verify this property of robustness to image perturbations.

We also test the robustness of the model against black-box adversarial attacks. These adversaries craft perturbations which cause the model to classify legitimate looking input images incorrectly. Black-box adversaries do not have access to the model parameters and the gradients, and must query the model to observe the output class probabilities. They query the model repeatedly with a chosen image, but perturbing it with each iteration, based on the results of the previous query. A model which has a lower attack success rate and/or requires a higher number of queries on average against an adversary before a successful attack, is considered more robust. We test the MBGNN model against two recently proposed and popular black-box adversarial attack methods, simBA [9] and Bandits-TD [12]. We choose these two methods since they use different methodologies - simBA uses local search in order to craft adversarial perturbations whereas Bandits-TD estimates the gradient by repeatedly querying the model to create the adversarial input. The results of these experiments are also provided in Section 4.4.

Generative Adversarial Networks (GANs), first introduced in [7], are a family of generative models that are used in several computer vision tasks including high resolution image generation[28], image super-resolution[18], domain adaptation[14, 34] and image compression [1]. GANs suffer from the problem of mode collapse, where the generated samples belong to a few modes in the dataset while still being successful at fooling the discriminator. This leads to a lack of diversity in the generated samples. One way to mitigate mode collapse in GANs is to use a technique known as mini-batch discrimination [23]. Here, instead of the discriminator being required to label individual samples as fake or real, it discriminates between an entire mini-batch of generated or real samples.

As it turns out, our proposed Attn-MBGNN can be interpreted as an extension of mini-batch discrimination to the classification task. To establish this connection, we present a variation of our MBGNN, which we refer to as a Mode Collapse MBGNN (MC-MBGNN). Rather than aggregating node features weighted by attention, as in the case of Attn-MBGNN layers, we aggregate the edge features in this model. Note that these edge features capture similarity and can be interpreted as the unnormalized attention values.

In MC-MBGNN, we modify the discriminator network by extracting features from real/fake images using the encoder, $f_{\theta}(\cdot)$, and denote them by $\{h_{1},h_{2},...,h_{B}\}$ where $h_{i}=f_{\theta}(x_{i})$ and $B$ is the batch size. We induce complete graphs for both the batch of generated and real samples, and process them separately. The output for the $n$-th aggregated edge feature for the $i$-th sample in the mini-batch, which is computed in a manner similar to aggregating unnormalized attention weights for the $n$-th head in Attn-MBGNN, is given by:

where $\phi(\cdot)$ is a neural network, $W_{n}$ is a trainable matrix, and $\psi$ can either be concatenation or absolute difference. We can then concatenate the aggregated edge features with the independent node features and use a final layer to get the $i$-th scalar output: $o_{i}=\sigma(W_{final}(h_{i}\,\|\,\bar{h}_{i}^{1}\,\|\,....\,\|\,\bar{h}_{i}^{N})).$ where $\sigma(\cdot)$ is the sigmoid function. If we restrict $\psi(\cdot)$ to absolute difference, $\phi(\cdot)$ to a fixed $\exp{}(-x)$ function and take all the $\bar{h}_{i}^{n}$ as scalars then the model reduces to standard mini-batch discrimination [23]. This shows how MBGNNs are connected to mini-batch discrimination. Our MC-MBGNN model is more expressive than mini-batch discrimination, and can significantly mitigate mode collapse in GANs, as demonstrated by the experiments in Section 4.5.

We perform experiments on two standard computer vision object classification datasets, CIFAR-10 and CIFAR-100. In order to expand the scope of the experiments to include scene classification, which is more complex, we also test our models on the MIT 67 scene dataset. For our GAN experiments we use the CIFAR-10 and the CelebA datasets.
CIFAR-10 consists of 60,000 colour images in 10 classes, with 6,000 images per class. We use the standard split with 50,000 training images and 10,000 test images. Each image is 32 $\times$ 32.
CIFAR-100 is similar to CIFAR-10, except that it has 100 classes, which significantly increases the difficulty of the classification task. Each class contains 600 images with 500 training images and 100 testing images per class. Each image is 32 $\times$ 32.
MIT 67 contains indoor scene images belonging to 67 categories, with a total of 15620 images. The number of images varies across categories, but there are at least 100 images per category. For our experiments we reduce the image size by approximately a factor of 10 in each dimension to 64 $\times$ 64, to ease the computational resources. As a consequence the scene classification task becomes harder.
CelebA is a large-scale face attributes dataset containing 10,177 identities with 202,599 colour face images in total. Each image is 64$\times$64.

Baselines. We use a ResNet-50 model trained using cross-entropy loss as the first baseline. We also reproduce the results of [15] with a batch size of 256 and reproduce results of [27] with our ResNet-50 baseline. We focus on relative improvement between our proposed model and these baselines.

For our proposed model, we use a ResNet-50 network without the final layer as an encoder and add an MBGNN with different combine methods. The entire model is trained end-to-end using cross-entropy loss. We provide the results for both MBGNN and Attn-MBGNN (with 95% confidence intervals over 5 runs) with all the combine options in Table 1, under both the inductive and the transductive settings. Our models perform better than all three baselines across the datasets considered. We observe that there is no significant difference in test accuracy between the inductive and transductive settings. Among the different combine methods, the drop feature seems to perform best in general; however, the difference in performance between these variations is small. Also, Attn-MBGNN is slightly better, owing to the model’s higher expressivity due to learnable attention weights. Figure 2 compares the training and test accuracy of all the models, with the standard supervised baseline versus the number of epochs on the CIFAR-10 dataset. The MBGNN models train faster than the standard network, giving a significant performance difference during early parts of the training process.

The most important parameters for graph-based learning are the size of the graph and the degree of each node, i.e., the neighborhood size $k$. Here we use the ResNet-50 encoder and the CIFAR-10 dataset for all our experiments. We expect the best value for $k$ to be close to the batch size divided by the number of classes, since this is the expected number of samples in a mini-batch having the same class label. For CIFAR-10 with a batch size of $256$, this value is $\nicefrac{{256}}{{10}}\approx 25$. We also expect performance to improve with larger batch sizes, since this translates to larger graphs, and hence more samples per class. We restrict ourselves to inductive predictions and provide the results for the drop feature model in Tables 2 and 3. The results match our prediction. First, for CIFAR-10 the best performance is indeed when $k$ is close to 25. Second, classification accuracy improves with larger batch sizes.

In order to test the robustness of MBGNNs to image perturbations, we first consider random (pixel-wise) Gaussian noise and local Gaussian blurring on input images, with varying levels of corruption severity. The evaluation of test accuracy is done via transductive testing, where a mini-batch consists of a single corrupted image along with a training set of uncorrupted images. The class label prediction made by the model for the corrupted image is compared against the true label. Figure 3 shows plots of test accuracy, measured in the manner described above, for the best performing variations of MBGNN (sum) and Attn-MBGNN (sum) on the CIFAR-10 dataset for different levels of corruption severity. We observe that models using MBGNNs are far better at accommodating the effects of local perturbations to the images (Fig. 3 first row) and that although the effects of local Gaussian blur are less harmful, models based on MBGNNs are still better by about 2% over the range of corruption severities we have considered (Fig. 3 second row).

We also test the robustness of the model to two recently proposed and popular black-box adversarial attack methods, simBA [9], and Bandits-TD [12]. Table 4 shows the mean and median number of queries before a successful attack, and the attack success rate, for different MBGNN models for the two attack methods. The MBGNN models have slightly lower attack success rates and higher mean/median queries before a successful attack when compared to the baseline ResNet model. The large increase in mean queries can be attributed to a heavier tail in the distribution of queries, as can be seen from the histogram plots in Figure 4, which show results for the best performing variations of MBGNN (dropfeat) and Attn-MBGNN (dropfeat). Plots for other variations are provided in the supplementary material. The Attention MBGNN models outperform the baseline model but do not perform as well as the MBGNN models. One reason for this might be that any perturbation in the input samples has a compounding effect on calculating the attention weights and, therefore, the aggregation.

Lastly, we provide results for training GANs using MC-MBGNN, and compare this strategy with both a vanilla GAN and minibatch discrimination [23]. We test the diversity of our generated samples using the technique of Number of statistically Different Bins (NDB), which was proposed in [21] as a metric to quantify mode collapse in GANs. To compute this metric, we first cluster the training dataset in $K$ different bins and then allocate the generated images in these bins based on their proximity to the centroid of each bin. Then, we measure the statistical similarity between the real and fake images in each of the bins and compute the fraction of statistically different bins that give us the NDB score. In the case of mode collapse, the number of statically different bins is close to $K$, and the NDB score is close to $1$. Using the NDB score as the metric, in Figure 5, we show that our proposed MC-MBGNN helps the generator learn faster and generate more diverse samples, which is substantiated by lower NDB scores. We provide the details of the architecture, experiment setup and generated images in the Supplementary material.

Table 5 provides results for our MBGNN with all the combine options, using Wide ResNet-28-10 as the encoder network. We also use this network for all the baselines in this table. We observe a consistent improvement across all datasets, which is consistent with the results of our experiments using ResNet-50 in the main paper.

We use standard generator and discriminator architectures for our GAN model. Let us denote the following operations,
Basic convolution: $Conv(in\_channels,out\_channels,filter\_size,stride,padding)$,
Linear layer: $Linear(in\_dim,out\_dim)$,
Deconvolution: $Conv\_trans(in\_channels,out\_channels,filter\_size,stride,padding,output\_padding)$
The generator architecture using the above notation, is given by
$Linear(128,16384)-batch\_norm-ReLU\\ Conv\_trans(1024,512,5,2,2,1)-batch\_norm-ReLU\\ Conv\_trans(512,256,5,2,2,1)-batch\_norm-ReLU\\ Conv\_trans(256,128,5,2,2,1)-batch\_norm-ReLU\\ Conv\_trans(128,64,5,2,2,1)-batch\_norm-ReLU\\ Conv\_trans(64,3,5,1,2,0)-tanh$
The discriminator architecture is given by
$Conv(3,64,4,2,1)-batch\_norm-LeakyReLU\\ Conv(64,128,4,2,1)-batch\_norm-LeakyReLU\\ Conv(128,256,4,2,1)-batch\_norm-LeakyReLU\\ Conv(256,512,4,2,1)-batch\_norm-LeakyReLU\\ Linear(8192,1)-sigmoid$
This defines the GAN architecture for the CelebA dataset. For CIFAR-10, we slightly modify the $Linear(\cdot)$ layers to adjust to the reduced image sizes. We add both the minibatch discrimination and MC-MBGNN layers before the final $Linear(\cdot)$ layer of the discriminator. For the sake of our experiments, we use the Adam optimizer with a learning rate of $10^{-4}$ for the discriminator and $2\times 10^{-4}$ for the generator, with a batch size of $128$. We provide the generated images using different discriminator architectures in Figure 12 and Figure 13. We train the model on CIFAR-10 for $10000$ iterations and on CelebA for $7500$ iterations.