Improving Axial-Attention Network Classification
via Cross-Channel Weight Sharing^††thanks: Partially supported by an Alfred and Helen M. Lamson endowed chair award.

A useful introduction to quaternion algebra and neural networks appears in [9].

The paper [12] extended the principles of convolution, batch normalization, and weight initialization from real-valued networks to complex-valued networks. [2], in turn, extended these principles to quaternion-valued convolutional networks (QCNNs). For implementation purposes, quaternion calculations can be decomposed into operations on 4-tuples of real numbers and quaternion operations can be implemented on these 4-tuples. The take home message from this is that a quaternion convolution module must accept four channels of input. A quaternion layer can accept more than four input channels, say $m$, as long as $m$ a multiple of four. In this case, the layer must hold $m/4$ separate quaternion convolution modules, each with their own weight sets.

The paper [8] analyzed the Hamilton product, which underlies the quaternion convolution, to identify the mechanism for improved performance in quaternion models. Let us denote the Hamilton product of I and Q as

This section explains 4D generalized hypercomplex networks for fully-connected layer, using parameterized hypercomplex multiplication (PHM) [14]. The 4D (quaternion) PHM (QPHM) layer works in a similar way of the quaternion networks. The Kronecker product is used to construct the parameter matrix. The PHM based fully connected hypercomplex transformation which transforms the input $\mathbf{x}\in\mathbb{R}^{d}$ into an output $\mathbf{y}\in\mathbb{R}^{k}$, is defined as $y=Hx+b$ where, $H\in\mathbb{R}^{k\times d}$ represents the PHM layer. H is calculated as $H=\sum_{i=1}^{n}\mathbf{I}_{i}\otimes\mathbf{A}_{i}$ where $\mathbf{I}_{i}\in\mathbb{R}^{n\times n}$ and $\mathbf{A}_{i}\in\mathbb{R}^{k/n\times d/n}$ are parameter matrices and $i=1\ldots n$ ($n=4$). Parameter reduction comes from reusing matrices $\mathbf{I}$ and $\mathbf{A}$ in the PHM layer. The $\otimes$ is the Kronecker product. The inputs are split as $Q_{in}=Q_{r}+Q_{x}+Q_{y}+Q_{z}$ and the outputs are merged into $Q_{out}$ as $Q_{out}=Q_{ro}+Q_{xo}+Q_{yo}+Q_{zo}$ for 4D hypercomplex which are graphically explained in Figure 2. The 4D hypercomplex parameter matrix is discussed in [14] which expresses the Hamiltonian product for 4D by preserving all PHM layer properties.

As stated in the introduction, standalone attention networks consisting of stacked attention modules can learn to outperform deep CNNs and hybrid CNN/attention models for image classification. Their advantage stems from their ability to detect similarities, or affinities, between pixels that have large spatial separation in the image. The drawback of directly using self-attention is that it is impractically computationally expensive, consuming $\mathcal{O}(N^{2})$ resources for an image of length $N$ (using a flattened pixel set) and using a global window to compare any pair of pixels in the image. For a 2D image of height, $h$, and width, $w$, where $N=hw$, so the cost is $\mathcal{O}((hw)^{2})=\mathcal{O}(h^{2}w^{2})$ to detect similarities for any pair of pixels in the image [6, 13].

Axial-attention networks reduce the cost of computing attention by decomposing the problem into consecutive 1D operations. They were introduced in [6] for generative modeling in auto-regressive models and use the assumption that images are approximately square so that $h$ and $w$ are both much less than the total pixel count, $hw$. For simplicity, assume a square 2D image where $h=w$, so $w^{2}=N$. Axial attention only operates on one dimension at a time. It is first applied to, say, the $h$ axis and then to the $w$ axis. When applied to the $h$ axis, then $h=w$ attention calculations are applied to a 1D region of length $h$. Axial attention applied to the columns, $h$, performs $w$ self-attention operations on each column whose total cost is $\mathcal{O}(h\cdot h^{2})=\mathcal{O}(\sqrt{N}\cdot N)$. This bound is the same for using an axial column module followed by an axial row module.

The paper [13] incorporated axial attention into a ResNet architecture [5] to develop a model called Axial-ResNet, which reduced the computational requirements (described above) of the original standalone attention networks [10, 7] used for image classification. The conversion from ResNet to Axial-ResNet was based entirely on modifying the bottleneck blocks within ResNet. Figure 1(a) shows the bottleneck block used in the original convolution-based ResNet. Figure 1(b) shows the axial-bottleneck block used in Axial-ResNet. The 3x3 2D convolution used in the original ResNet is replaced by an axial attention module (consisting of two 1D layers) that processes the $h$ axis followed by the $w$ axis.

The models we compare are: the standard convolution-based ResNet [5], the QCNN [2], the axial-ResNet [13], and our novel quaternion-enhanced axial-ResNet. The main objective is to see if the representations generated by the quaternion front end improve classification performance of the axial attention model. The axial-ResNet and QuatE-1 axial-ResNet models are depicted in Table 1. We used a 26-layer version with the block multipliers “[1, 2, 4, 1]” and a 50-layer version with the block multipliers “[3, 4, 6, 3]” of the models. Table 1 shows the 50-layer versions. The bracketed expressions show the bottleneck blocks with the operations used and the number of output channels for each stage. If the quaternion modules are counted then the layer counts for our model are 33 layers and 66 layers, respectively. We count the two-1D-layer axial-attention module as one layer because two 1D layers are equivalent to one 2D layer.

Because of hardware limitations, our experiment conducts image classification on a subset of the ImageNet dataset [11]. Specifically, it took over two hours to train one of the axial-ResNet models for one epoch. This smaller dataset which requires lower computing resources, called ImageNet300k, uses the same 1,000 image categories as the original ImageNet [1]. The full ImageNet dataset has 1.28 million training images and 50,000 validation images. Our smaller dataset, which we call ImageNet300k, is sampled from the full ImageNet by taking the first 300 images for each category contained in the original dataset, yielding 300,000 training images. There are also 50,000 validation images with 50 images per category. These are from the original dataset. Although the training dataset is smaller, it still allowed us to train our 50-layer networks without overfitting. Overfitting was assessed by examining performance on the validation dataset in comparison to the training set. The validation dataset was the same as that used in the original ImageNet dataset.

All network models (Experiment-1, 2, and 3) were trained using the same stochastic gradient descent optimizer and hyperparameters. All networks were trained for 150 epochs using stochastic gradient descent optimization with momentum set to 0.9 and the learning rate which is warmed up linearly from $\epsilon$ near zero to 0.1 for 10 epochs [13]. The learning rate was then cut by a factor of 10 at epochs 20, 40, and 70. We adopt the same training protocol as [13] except for batch size. Our batch size is limited to ten because of memory limitations. Also, we used batch normalization with 0.00009 weight decay. For attention models, the number of attention heads is fixed to eight in all attention layers [13].

The main results are seen in Table 2. The top half of the table shows the results for the 26-layer models (33-layers for the QuatE-1 model). The bottom half shows the results for the 50-layer models (66-layers for the QuatE-1 model). These are the parameter count, the inference time, and the percent accuracy results for a single simulation run for each model. Both of the axial-ResNets use far fewer parameters than the convolution-based ResNets.

The most important comparison is between the axial-ResNet and the QuatE-1 axial-ResNet because this directly shows the effect of the quaternion-generated interwoven/interlinked representations. There are two such comparisons, one for the 26/33 layer models and the other for the 50/66 layer models. In both cases, the quaternion enhanced versions of axial-ResNet produced higher classification accuracy performance. This was true for both the training and validation data. This supports our main hypothesis that quaternion modules can produce more usable interlinked/interwoven representations. Of course, there is the alternative explanation that the quaternion enhanced axial-ResNet had more layers and this is the reason for the better performance. This is addressed in the next section.

Another result is that both of the axial-ResNet architectures provide higher performance for the validation set than either of the convolution-based ResNets. This is true for both the 26-layer (33-layer) and 50-layer (66-layer) versions.

Finally, it is surprising that the 33-layer quaternion-enhanced axial-ResNet gives better classification accuracy performance than the 66-layer version. This is true for both the training and validation data. We have no explanation for this. If this were only true for the validation data, then overfitting would be a possible explanation. This is addressed in the Experiment-3 section.

This experiment increases the layer count of the baseline networks from 26 to 35 layers. This was done by using a block multiplier of [2, 3, 4, 2] for these models. Although this did not give us exactly 33 layers to exactly match the QuatE-1 axial-attention model, it preserved the bottleneck structure of the original design and enhances comparability.

Table 3 shows the parameter count and accuracy for these 35-layer models. The most important comparison is the 35-layer Axial-Attention Network in Table 3 with the 33-layer QuatE-1 Network in Table 2. Now the layer count for the non-quaternion version is slightly larger than that for the quaternion version but the quaternion version still gives better performance on both the training and validation accuracy. The increased layer count did improve the performance of the 35-layer Axial Attention Network but not enough to surpass the quaternion version. Thus, we conclude that the quaternion frontend has more impact than the layer count.

This experiment used the RepAA model on ImageNet300k dataset. The performances of QuatE axial-ResNets with QPHM in the back-end, and RepAA networks are compared in Tables 2, 3, and 4. The architectural details of RepAA network (50-layer version) are shown in Table 1. In this experiment, we ran our models for 26-layer, 35-layer, and 50-layer architectures with the same multipliers and we trained RepAA networks using the same hyperparameters as like in Experiment-1 and 2.

We next investigate the use of our final proposed method, the representational axial-attention network, to see how it performs on the task of ImageNet300k dataset. Table 4 shows our overall experiment results for the 26, 35, and 50-layer architectures. Among them, our final proposed method performs better in accuracy and the trainable parameters.

Figure 3 assesses the overfitting problem by examining the performance on the validation dataset in comparison to the training dataset and yields no overfitting for ImageNet300k dataset. The most important comparisons are between QuatE axial-ResNet & our proposed RepAA model, and axial-ResNet & RepAA model as these directly show the effect of removing extra $1\times 1$ quaternion layer from QuatE axial-ResNet and applying representational effect throughout the attention network. In both cases, the RepAA networks produced higher classification accuracy with far fewer parameters. In general, our proposed method outperforms then the other networks with 1.5 times, 1.8 times and 1.7 times fewer parameters in compared to original axial ResNet for 26, 35, and 50-layer networks, respectively.

Although, we ran our model for only 90 epochs, which is less than the other models, but it’s performance was unchanged for 150 epochs. Unlike QuatE axial-ResNets, the validation performance for 35-layers is higher than the 26-layers, and 50-layers is outperformed than the others. This supports our main hypothesis, quaternion modules can produce more usable interlinked/interwoven representations.