GLL: A Differentiable Graph Learning Layer for Neural Networks
Authors: Jason Brown, Bohan Chen, Harris Hardiman-Mostow, Jeff Calder, Andrea L. Bertozzi
Organizations: Department of Mathematics University of California, Los Angeles Los Angeles, CA 90095, USA · Computing + Mathematical Sciences (CMS) Department California Institute of Technology Pasadena, CA 91125, USA · School of Mathematics University of Minnesota Minneapolis, MN 55455, USA
Standard deep learning architectures used for classification generate label predictions with a projection head and softmax activation function. Although successful, these methods fail to leverage the relational information between samples for generating label predictions. In recent works, graph-based learning techniques, namely Laplace learning, have been heuristically combined with neural networks for both supervised and semi-supervised learning (SSL) tasks. However, prior works approximate the gradient of the loss function with respect to the graph learning algorithm or decouple the processes; end-to-end integration with neural networks is not achieved. In this work, we derive backpropagation equations, via the adjoint method, for inclusion of a general family of graph learning layers into a neural network. The resulting method, distinct from graph neural networks, allows us to precisely integrate similarity graph construction and graph Laplacian-based label propagation into a neural network layer, replacing a projection head and softmax activation function for general classification task. Our experimental results demonstrate smooth label transitions across data, improved generalization and robustness to adversarial attacks, and improved training dynamics compared to a standard softmax-based approach.
Figures & tables
Figure 1: Visualization of the graph learning layer (GLL) within a neural network pipeline compared to a standard multilayer perceptron (MLP) projection head. For the GLL, a combination of labeled and unlabeled input images are batched and normalized into data for the feature extractor. The feature network encodes the images in feature space. The GLL is a combination of two steps; it generates a similarity graph from the encoded data and then propagates the labels across the graph, giving predictions on the unlabeled data. These predictions are used in the desired loss function and gradients can flow through the GLL to update the feature extractor. By contrast, the MLP projection head pipeline does not involve any initially labeled nodes and the classification method ignores relational information in the encodings. The + indicates learnable parameters in the network; note GLL has none.
Figure 2: Comparison between the classic structure and our proposed graph learning layer (GLL) neural network structures for image classification. Both structures consist of a feature encoder and a classifier, but we replace the multilayer perceptron (MLP) classifier - consisting of one or more linear and activation layers and a final softmax classifier - with our data-dependent GLL with no trainable parameters. We reserve a subsample of the training data - denoted base samples - to use as the labeled nodes for label propagation. The samples that are not from the base set are use to compute the loss.
Figure 3: The initial embedding of the two moons dataset from the random seed. The blue and red correspond to the different classes, with the starred samples corresponding to the base samples for the respective classes.
Figure 4: Visualization of effect of τ on spatial embeddings. Each row corresponds to a graph learning model with a specified value of τ or a softmax classification head. The columns correspond to the epoch of the training. With low values of τ , the model separates the data by class and resembles the original two moons shape. Increasing τ forces tighter clustering and localization. The softmax layer simply learns a linear separation with no regard to relational information. Further training with the softmax layer increases the scale, which is not seen without the axes.
Structure
Layer
Details
Feature Size
Input Image
-
From FashionMNIST Dataset
28 × 28 × 1
Feature Encoder
Conv1
1 → 64, kernel size: 3, padding: 1
28 × 28 × 64
Conv2
64 → 128, kernel size: 3, padding: 1
28 × 28 × 128
MaxPool
kernel size: 2, stride: 2
14 × 14 × 128
Conv3
128 → 256, kernel size: 3, padding: 1
14 × 14 × 256
MaxPool
kernel size: 2, stride: 2
7 × 7 × 256
Table 1: Network structure used in the FashionMNIST experiments (Section 4.3 ). The feature encoder is deliberately over-parameterized to create a challenging training scenario, and the model contains 1,975,424 parameters in the feature encoder and 142,346 parameters in the MLP classifier.
Dataset
Trial
Epochs
Test error (%)
FashionMNIST
MLP only (MLP)
100 MLP
45.42
FashionMNIST
MLP only (GLL)
100 MLP
14.66
FashionMNIST
GLL only (GLL-0)
50 GLL
8.90
FashionMNIST
GLL-50
50 MLP + 50 GLL
8.68
FashionMNIST
GLL-75
75 MLP + 25 GLL
8.91
Table 2: Test classification error rates (%, lower is better) on FashionMNIST for the customized network in Table 1 under different training strategies. For the “MLP only” strategy we report both the error of the MLP classifier and that of a GLL classifier evaluated on the same encoder (GLL). For the remaining strategies, GLL is used as the classifier during training and evaluation.
Figure 5: Training loss and test accuracy of four training strategies on the FashionMNIST dataset with the customized feature encoder in Table 1 . The strategies are: (1) MLP trained from scratch, (2) GLL trained from scratch (GLL-0), (3) GLL-50 (GLL after 50 epochs of MLP training), and (4) GLL-75 (GLL after 75 epochs of MLP training). For the MLP-only strategy, both the MLP and GLL losses and accuracies are recorded during training; for the remaining strategies, only the GLL classifier is used.
Network
Params
Encoder
MLP head
Total
VGG11
9,223,232
332,778
9,556,010
VGG13
9,407,936
332,778
9,740,714
ResNet20
271,824
16,938
288,762
ResNet32
466,256
16,938
483,194
ResNet44
660,688
16,938
677,626
Table 3: Parameter counts by network
Train Num
Whole (50000)
20% (10000)
2% (1000)
Network
MLP
WNLL
GLL
MLP
WNLL
GLL
MLP
WNLL
GLL
VGG11
7.77 ± 0.22
7.71 ± 0.10
7.48 ± 0.19
15.16 ± 0.58
14.74 ± 0.27
13.14 ± 0.09
34.76 ± 1.38
34.97 ± 1.57
29.59 ± 0.12
VGG13
6.16 ± 0.31
6.20 ± 0.23
5.80 ± 0.03
12.68 ± 0.48
12.27 ± 0.19
11.03 ± 0.04
31.61 ± 1.17
31.31 ± 0.49
26.90 ± 0.18
ResNet20
9.25 ± 0.96
7.99 ± 0.33
7.49 ± 0.24
15.60 ± 0.74
14.76 ± 0.68
14.00 ± 0.31
33.35 ± 1.81
32.04 ± 0.58
31.25 ± 0.19
ResNet32
8.45 ± 0.96
7.26 ± 0.29
6.61 ± 0.15
14.78 ± 0.58
13.63 ± 0.24
13.04 ± 0.13
33.93 ± 1.76
30.82 ± 0.79
32.96 ± 1.44
ResNet44
8.07 ± 0.78
6.57 ± 0.11
6.12 ± 0.12
14.38 ± 0.82
13.23 ± 0.33
12.51 ± 0.34
33.76 ± 2.07
31.85 ± 1.46
31.69 ± 0.70
Table 4: CIFAR-10 test error rates (%, lower is better) comparing MLP, WNLL, and GLL output heads across architectures, evaluated under three training-set sizes: “Whole” (50000 images), “20%” (10000), and “2%” (1000). Entries report mean ± one standard deviation of the test error. Best results for each backbone and label rate are bolded and results within a standard deviation of the best result are underlined . Network parameter counts follow Table 3 .
Train Num
Whole (112800)
20% (22560)
2% (2256)
Network
MLP
WNLL
GLL
MLP
WNLL
GLL
MLP
WNLL
GLL
VGG11
10.72 ± 0.14
12.18 ± 0.79
10.76 ± 0.09
12.06 ± 0.13
13.46 ± 0.69
11.96 ± 0.08
15.86 ± 0.35
16.64 ± 0.78
14.98 ± 0.13
VGG13
10.50 ± 0.06
10.73 ± 0.15
10.60 ± 0.08
11.85 ± 0.08
15.43 ± 1.15
11.94 ± 0.10
15.91 ± 0.38
16.48 ± 0.82
15.08 ± 0.20
ResNet20
9.72 ± 0.07
10.77 ± 0.29
10.34 ± 0.05
11.45 ± 0.06
12.33 ± 0.21
11.60 ± 0.07
15.79 ± 0.33
17.10 ± 1.01
14.63 ± 0.12
ResNet32
9.73 ± 0.06
10.78 ± 0.37
10.15 ± 0.14
11.41 ± 0.05
12.26 ± 0.38
11.54 ± 0.11
15.85 ± 0.46
17.52 ± 0.90
14.30 ± 0.08
ResNet44
9.92 ± 0.06
11.30 ± 0.36
10.10 ± 0.09
11.55 ± 0.16
12.04 ± 0.28
11.43 ± 0.14
15.87 ± 0.35
17.81 ± 0.98
14.98 ± 0.18
Table 5: EMNIST test error rates (%, lower is better) comparing MLP, WNLL, and GLL output heads across architectures, evaluated under three training-set sizes: “Whole” (112800 images), “20%” (22560), and “2%” (2256). “Whole” uses the full standard training set. Best results for each backbone and label rate are bolded and results within a standard deviation of the best result are underlined . Entries report mean ± one standard deviation of the test error. Network parameter counts follow Table 3 .
Train Num
Whole (112800)
20% (22560)
2% (2256)
Network
MLP
WNLL
GLL
MLP
WNLL
GLL
MLP
WNLL
GLL
VGG11
35.88 ± 0.48
36.92 ± 1.15
35.91 ± 0.28
38.51 ± 0.42
36.57 ± 0.55
38.10 ± 0.18
42.91 ± 0.76
45.05 ± 0.33
42.14 ± 0.33
VGG13
34.88 ± 0.32
36.05 ± 0.37
35.38 ± 0.34
37.90 ± 0.23
38.68 ± 0.74
38.11 ± 0.33
44.51 ± 1.48
45.28 ± 0.53
42.29 ± 0.36
ResNet20
31.66 ± 0.21
35.76 ± 0.59
33.56 ± 0.38
35.80 ± 0.59
36.00 ± 0.36
35.53 ± 0.23
43.17 ± 1.29
46.61 ± 0.96
41.45 ± 0.41
ResNet32
31.69 ± 0.29
34.47 ± 0.81
32.87 ± 0.49
36.09 ± 0.32
36.72 ± 0.60
36.09 ± 0.42
43.66 ± 1.35
45.80 ± 0.62
39.94 ± 0.33
ResNet44
31.84 ± 0.27
35.61 ± 0.56
32.83 ± 0.25
37.01 ± 0.56
36.47 ± 0.44
35.98 ± 0.57
43.70 ± 1.44
46.68 ± 0.39
41.51 ± 0.51
Table 6: EMNIST test error rates (%, lower is better) for MLP, WNLL, and GLL output heads across architectures, evaluated under three training-set sizes: “Whole” (112800 images), “20%” (22560), and “2%” (2256). “Whole” uses the full standard training set. Error rates are computed on a class subindex of the EMNIST (balanced) label space: although the dataset contains 47 classes in total, we evaluate only the 10 hardest classes: 0, 1, 9, F, I, L, O, f, g, and q. Entries report mean ± one standard deviation of the test error. Best results for each backbone and label rate are bolded and results within a standard deviation of the best result are underlined . Network parameter counts follow Table 3 .
Batch Size
Metric
GLL
MLP
128
Training time / batch (s)
0.0457±0.0093
0.0044±0.0006
Testing time / batch (s)
0.0308±0.0085
0.0024±0.0027
Testing error (%)
0.76
2.11
256
Training time / batch (s)
0.0685±0.0146
0.0061±0.0005
Testing time / batch (s)
0.0457±0.0108
0.0039±0.0005
Testing error (%)
1.33
1.69
Table 7: Comparison of GLL and MLP performance across different batch sizes on MNIST. The architecture backbone is a four-layer CNN with two linear layers. We use identical training settings for both models. For time metrics, each entry is reported as mean ± standard deviation over five epochs of training. Testing error is reported in percentage points. The experiments were run on a T4 GPU.
Batch Size
Metric
GLL
MLP
128
Training time / batch (s)
0.1537±0.0121
0.0597±0.0083
Testing time / batch (s)
0.0659±0.0068
0.0220±0.0026
Testing error (%)
0.84
0.86
256
Training time / batch (s)
0.2295±0.0193
0.1157±0.0168
Testing time / batch (s)
0.1045±0.0162
0.0426±0.0060
Testing error (%)
0.94
1.04
Table 8: Comparison of GLL and MLP performance across different batch sizes on MNIST. The architecture backbone is ResNet-18. We use identical training settings for both models. For time metrics, each entry is reported as mean ± standard deviation over five epochs of training. Testing error is reported in percentage points. The experiments were run on a T4 GPU.
Batch Size
Metric
GLL
MLP
128
Training time / batch (s)
0.1257±0.0044
0.0471±0.0024
Testing time / batch (s)
0.0629±0.0035
0.0224±0.0037
Testing error (%)
1.65
2.39
256
Training time / batch (s)
0.1772±0.0065
0.0725±0.0028
Testing time / batch (s)
0.0950±0.0103
0.0404±0.0043
Testing error (%)
2.10
1.18
Table 9: Comparison of GLL and MLP performance across different batch sizes on MNIST. The architecture backbone is ResNet-110. We use identical training settings for both models. For time metrics, each entry is reported as mean ± standard deviation over five epochs of training. Testing error is reported in percentage points. The experiments were run on an A100 GPU.
Dataset
Model
Natural
FGSM
IFGSM
CW
MNIST
GL Robust
0.55
2.28
2.30
1.16
WNLL Robust
0.58
2.48
2.97
2.72
MLP Robust
0.80
4.83
6.43
2.59
GL Natural
0.47
3.49
3.56
1.29
WNLL Natural
0.50
3.54
4.64
3.19
MLP Natural
0.75
5.74
8.48
2.76
Table 10: Summary of adversarial results of models trained with our method compared to WNLL and the standard MLP classifier, on both robustly and naturally trained models. We report natural error rate (i.e. no attacks), and adversarial error rate on FGSM, IFGSM, and CW attacks. Best error rates are in bold. For MNIST, we report ϵ=0.3 for the FGSM and IFGSM attacks, and ϵ=0.05 for the other datasets. We report CW results for c=20 . We report more complete results across a range of values for ϵ and c in Figures 6 , 7 , and 8 . We include the natural error rates to show that the GLL does not sacrifice performance in this regime, while also being significantly more robust to adversaries across datasets. One mild exception is WNLL slightly outperforms GLL for the fast gradient attacks on FashionMNIST, however Figure 7 shows that in the case of IFGSM, GLL is much better against stronger attacks. The slight increase in natural test error on robustly trained models compared to naturally trained models has been observed previously in madry2017towards ; however, as expected, the robust models perform significantly better on adversarially-perturbed data than their naturally-trained counterparts.
Figure 6: Adversarial robustness to FGSM, IFGSM, and CW attacks on the MNIST dataset. The backbone network is the same in all cases; the only difference is the classifier (either MLP with softmax, WNLL, or our novel GLL). We report results for both robustly and naturally trained models.
Figure 7: Adversarial robustness to FGSM, IFGSM, and CW attacks on the FashionMNIST dataset. The backbone network is the same in all cases (ResNet-18); the only difference is the classifier (either MLP with softmax, WNLL, or our novel GLL). We report results for both robustly and naturally trained models.
Figure 8: Adversarial robustness to FGSM, IFGSM, and CW attacks on the CIFAR-10 dataset. The backbone network is the same in all cases (PreActResNet-18); the only difference is the classifier (either MLP with softmax, WNLL, or our novel GLL). We report results for both robustly and naturally trained models.
Figure 9: Comparison of accuracy vs. average distance of attack for CW attack on all three datasets. Each point represents a different value of c . We see that GLL resists the adversary from moving the target images significantly more than the standard MLP or WNLL classifier, suggesting that graph-based classification and precise tracking of gradients is key for adversarial robustness.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 10: Visualization of FGSM, IFGSM, and CW attacks. The left, middle, and right columns show the original image, the attack, and the adversarial image, respectively. The title in the rightmost column indicates the original classification and the classification after the attack. These images were chosen to demonstrate attacks where the model was fooled by the advesary.
Figure 11: FGSM results on MNIST. The backbone network was trained in the standard way (with a softmax classifier). In the legend, MLP indicates that the FGSM attacks used the gradients from the MLP classifier (ie a linear layer and a softmax), whereas GLL used our graph learning layer. In all cases, the final classifier on the attacked images was GLL. Hence, the only difference between the results is which gradient was used for the attack. We see that the models that utilized the GLL gradients were significantly more robust to the FGSM attacks than using the standard MLP gradients.
Figure 12: Comparison of accuracy vs. average distance of attack for CW attack on both datasets. Each point represents a different value of c , where we plot additional values of c for the GLL-based models compared to Figure 9 . We see that GLL is still superior at resisting perturbations by CW attack, even when c is orders of magnitude larger than for WNLL or MLP, and even when controlling for distance on CIFAR-10.
TMLR Group, Department of Computer Science, Hong Kong Baptist University, Kowloon Tong · Cooperative Medianet Innovation Center, Shanghai Jiao Tong University, Shanghai, China · Department of Computer Science, Stanford University, Stanford, CA 94305, USA +1