AFN: Adaptive Fusion Normalisation via an Encoder-Decoder Framework

Consider an input feature map with shape $(N,C,H,D)$, where $N$ is the batch size, $C$ is the number of channels and $H,D$ is the width and height of the feature map, respectively. First, we reshape the feature map to $(N\cdot H\cdot D,C)$, just as in the BN using statistics from the dimension $(N,H,D)$ instead of merely from $(N)$ to be more efficient without decreasing the performance. Then, we compute the mean and variance of this input batch by using each sample $x_{i}$ from $(N\cdot H\cdot D,C)$ dimension as:

Subsequently, we follow the setting of ASRNorm, which uses an encoder-decoder structure, whose bottleneck ratio is set to 16, to force our neural network to provide more suitable statistics for this mini-batch. The encoder extracts global information by interacting with the information and the decoder learns to decompose the information. For efficiency, both the encoder and decoder consist of one fully connected layer. ReLU [14] is used to render the feature non-linear and ensure that $\sigma_{stan}$ is non-negative:

The encoders project the input onto the hidden space $\mathbb{R}^{C_{stan}}$, while the decoders project it back onto the space $\mathbb{R}^{C}$, where $C_{stan}<C$, and $f_{enc}=g_{enc}$.

Then, we linearly combine the original statistics and statistics from our neural network, using the residual learning framework to make our training process stable. In the early training stages, the scale of $\hat{x}$ can be large due to the scale of input $x$, potentially leading to gradient explosion. To balance the scale between $\mu_{stan}$ and $\sigma_{stan}$, we introduce a residual term for regularisation.

where $\lambda_{\mu}$ and $\lambda_{\sigma}$ are learnable parameters ranging from 0 to 1 (bounded by sigmoid). Because the neural network cannot give a good statistic at the beginning of training, we initialise $\lambda_{\mu}$ and $\lambda_{\sigma}$ into $sigmoid(-3)$ to make $\lambda_{\mu}$ and $\lambda_{\sigma}$ approximate 0. Smaller than $sigmoid(-3)$ values will lead to a gradient vanishing problem, which is not desirable at training time.

Finally, we will get our normalised features by:

BN utilises two additional parameters to rescale the normalised features. Just like ASRNorm, we also use two additional neural networks to compute the weights and bias for the rescaling process.

where $\beta_{bias}$ and $\gamma_{bias}$ are two learnable parameters identical to those of BN, and $\lambda_{\beta},\lambda_{\gamma}$ are learnable parameters, which are initialised with a small value $sigmoid(-5)$, to smoothen the learning process. We resume the $\beta_{bias}$ and $\gamma_{bias}$ from the BN layers in the pre-trained model, so each stage of the neural network acts like an identity function at the beginning of the training. Also, encoders and decoders are fully connected layers, and sigmoid and tanh functions are used to ensure the rescaling statistics are bound. The encoders project the inputs onto the hidden space $\mathbb{R}^{C_{rescale}}$ with $C_{rescale}<C$. The decoders project the encoded feature back onto the space $\mathbb{R}^{C}$.

Next, we rescale the normalised feature just like in any other normalisation, and send it to the next module of the neural network:

As shown in Fig 3, we take VGG19_BN architecture as an example to illustrate ASRNorm has gradient explosion/vanishing problems. With the increase of iteration, $\bar{x}$ from Eqn 4 is gradually approaching 0, leading to gradient vanishing. On the contrary, in AFN, $\bar{x}$ tends towards $-\infty$, leading to gradient explosion. To solve this problem, we bound the gradient to successfully solve this problem. More details please see in Section 4.2.

From our point of view, IN may struggle to capture the inter-instance information within a batch, which can hinder the learning of relationships between different images. Motivated by ASRNorm, we introduce additional parameters to BN, enhancing the generalisation capability of neural networks. This modification results in a more stable training process compared to ASRNorm which will be studied in the experiments presented in the next section.

Datasets. We conduct experiments on four standard benchmarks for image classification [15, 16, 17], including CIFAR-10, CIFAR-100, SVHN, MNIST-M, and three standard benchmarks for domain generalisation [18, 19, 20, 21], including Digits, CIFAR-10-C and PACS.

(1) CIFAR-10, CIFAR-100: This benchmark consists of 60,000 32x32 RGB images in 10 different classes, with 6,000 images per class. CIFAR-100 dataset is an extension of CIFAR-10 and contains 100 classes, with 600 images per class. Each image is 32x32 color image.

(2) Digits: This benchmark consists of four digits datasets: MNIST [22], SVHN [17], MNIST-M [16], USPS [23]. Images in MNIST and USPS are grey images, but in SVHN and MNIST-M are color images. In the domain generalisation task, to make four datasets have compatible shapes, all images are resized to 32 x 32 pixels, and the channels of the datasets are verified to be the same. We use one dataset as a training dataset, and the rest of them as a testing dataset.

(3) CIFAR-10-C [24]: This benchmark is proposed to evaluate the robustness of 109 types of corruptions with 5 levels of intensities. The original CIFAR-10 is used for training and corruptions are only applied to the testing images. The corruption intensity can measure the level of domain discrepancy.

(4) PACS [25]: Figure 4 shows the benchmark for domain generalisation containing four domains: art paint, cartoon, sketch, and photo, which share the same categories that include dog, elephant, giraffe, guitar, house, horse and person. To align the results in the previous work, for this dataset, we consider the same settings: 1) training a model with one domain data and testing with the rest three; 2) training a model on three domains and testing on the remaining one. In both settings, we remove the domain label and mix the data from different domains.

Experiments Details. It is noted that $C_{stan},C_{rescale}$ are set to $C/2,C/16$, respectively. The initialisation of hyperparameters follows the settings of ASRNorm [8].

(1) For Digits. In the domain generalisation task, we use the ConvNet architecture [22] (conv-pool-conv-pool-fc-fc-softmax) with ReLU following each convolution. Since this model has no normalisation layer, AFN is inserted after each convolution layer before ReLU. We follow this experimental setting: using Adam [26] optimiser with a learning rate $10^{-3}$ and training batch size 128, testing batch size 256. In the image classification task, we conduct experiments on various backbones, by using SGD as the optimiser with the learning rate initilised with $0.1$.

(2) For CIFAR-10-C, we use WRN-40-4 [27], SGD with Nesterov momentum(0.9), and a batch size of 128. The initial learning rate is 0.1, annealing with ALRS scheduler [28].

(3) For PACS, we use ResNet-18 pretrained on ImageNet, the Adam optimiser with an initial learning rate $10^{-4}$, ALRS as a learning rate scheduler with training epochs 50, and mini-batch size is 32. The model learned AFN with the RSC [29] procedure, which is an algorithm to virtually augment challenging data by shutting down the dominant neurons that have the largest gradients during training.

(4) For CIFAR-10, CIFAR-100 datasets, we follow the same experimental setting: Adam optimiser with an initial learning rate $10^{-3}$, scheduled by ARLS. The model is trained with a mini-batch size 128 for about 200 epochs.

Our approach outperforms the previous SOTA normalisation method (ASRNorm) on single domain generalisation tasks by $0.9\%,0.6\%,1.3\%,1.6\%$ on the Digits(two experiments), CIFAR-10-C, and PACS benchmarks, respectively.

Results on Digits: Table 1 and 2 show the results on the Digits benchmark. The proposed AFN is compared with the previous SOTA normalisation method ASRNorm. Our method outperforms both the baseline and the SOTA methods on average.

Results on CIFAR-10-C: We show average accuracies of the 19 corruption types for each level of intensity of CIFAR-10-C in Table 3. Similar to the results on the Digits benchmark, AFN achieves larger improvements on the more challenging domains, because the early stage of training, behaving like BN, helps collect the information from different samples. From Figure 1, it is obvious that our method has better generalization ability with respect to other methods, with the increase of corruption levels.

Results on PACS: In table 4, we show the results on PACS where we use one domain for training and the remaining three for testing. Our method improves the performance of RSC in the Art painting and Photo domains, and reaches better average accuracy on this task. Table 5 shows the results on PACS for the multi-source domain setting. Similarly, AFN slightly outperforms the previous SOTA method.

Results on image classification. Table 6 and 7 show the results on SVHN and MNIST-M. Our method surpasses BN and ASRNorm, also we see there exist nan values in the tables, which mean the inability to keep the gradients steady during training, leading to training failure. However, our method is capable of keeping the gradient steady to finish our training.

Table 8 and 9 shows the results on CIFAR-10, CIFAR-100. Our method can be successfully applied to image classification tasks and outperform the original BN, which means our method can replace BN in the network to improve performance. It is noted that the previous SOTA method cannot be applied to this task, and may even deteriorate the network performance. ASRNorm is not suitable for image classification tasks due to its inability to consider mini-batch information and address covariance shift effectively, in the early training stage. In contrast, our method leverages Batch Normalisation statistics to capture mini-batch information, making it more suitable. Additionally, our method demonstrates superior stability during training, reducing the risk of gradient explosion associated with ASRNorm.