Cross-Scale MAE: A Tale of Multi-Scale Exploitation in Remote Sensing

Let $p\in\mathbb{R}^{H\times W\times C}$ represent an input image of height $H$ and width $W$ with $C$ channels. First, multi-scale augmentation (Sec. 3.2) is applied to $p$ to generate two images, $p_{h}$ of relatively higher GSD and $p_{l}$ of relatively lower GSD. For each of the augmented images, a sequence of $|S|$ independent patches of height and width of $n$ pixels are created, where each of the patches, $s\in S$, has dimension $s\in\mathbb{R}^{n^{2}C}$. A fraction, $m$, of the patches is removed, while the remaining patches are passed through a projection function (e.g., a linear layer) to project the patches $s$ into $D$ dimensions, $f_{emd}:\mathbb{R}^{n^{2}C}\rightarrow\mathbb{R}^{D}$, to obtain embedded patches $S_{emd}=f_{emd}(S)$. A standard positional encoding vector is added to the embedded patches and fed into the encoder. After the encoder, the removed $m$ patches are placed back into their original location in the sequence of patches where a learned mask token represents the masked patches that were not encoded. Another positional encoding vector is added to all patches, and a sequence of transformer blocks decodes these patches to form the original input image; these steps are illustrated in Fig. 1.

A specialized augmentation method is applied to the input image before patchifying it in the Cross-Scale MAE framework. This augmentation process generates an additional image of lower GSD for the same site, and the two images are then used as inputs through two independent instances of the MAE network. More specifically, given an input image $p$ with a certain GSD, the framework randomly selects two scale ratios, $r_{h}$ and $r_{l}$, from the range [0, 1], where $r_{l}<r_{h}$. The image $p$ is then down-sampled using these scale ratios to generate two images, $p_{h}$ and $p_{l}$. In our experiments, we set $r_{h}=1$, so $p_{h}=p$. Both images, $p_{h}$ and $p_{l}$, are subsequently resized to a fixed size of $128\mathrm{px}^{2}$, and the GSD values of $p_{h}$ and $p_{l}$ are denoted as $g_{h}$ and $g_{l}$, respectively.

Note that this augmentation approach differs from traditional augmentation, where the rescaled images are traditionally only used as separate samples. Instead, the multi-scale inputs are generated dynamically at runtime and used as inputs through multiple network branches, where the branches have shared weights. The constraints enforced upon the learned representations of these network branches are detailed in the following sections.

Typically, the standard MAE trains a network to reconstruct an image by masking out a significant portion of its pixels. However, the standard MAE encoder does not explicitly address the issue of representation consistency for multi-scale input sources. In contrast, the proposed Cross-Scale MAE aims to learn meaningful and consistent scale-invariant representations. To achieve this, Cross-Scale MAE focuses on two fundamental properties in the learned representations: (1) consistency of information across different-scale images and (2) representativeness of the raw image. In the Cross-Scale MAE framework, a contrastive learning strategy is employed within the encoder to ensure consistency among representations derived from different-scale images.

The intuitive idea behind Cross-Scale MAE is that different-scale images captured at the same ground site contain highly relevant information despite potential visual differences. Cross-Scale MAE leverages this notion by maximizing the shared information between representations derived from different-scale images of the same site while minimizing the shared information between images from different sites. This approach aligns with the principles of contrastive learning. In contrastive learning, the InfoNCE loss is commonly used to estimate the lower bound of mutual information between data samples. By minimizing the InfoNCE loss, we effectively maximize the mutual information between representations from different scales [44, 32].

Specifically, the encoder ${E}$, uses the ViT-base [16] as the backbone. Upon feeding an input image $p$ to encoder ${E}$, a high-level representation, ${E}(p)$, can be obtained. Before computing the contrastive loss, a nonlinear projection head ${g}_{f}(\cdot)$ is needed, which has been proven to be effective [8]. So, after the encoder and projection head, we obtain the representation, $z=g_{f}\left(E\left({p}\right)\right)$. The cross-scale consistency loss is thus defined as:

$$\mathcal{L}_{cc}=\frac{1}{2N}\sum_{k=1}^{N}\left(\ell_{info}\left(p_{l}^{k},p_{h}^{k}\right)+\ell_{info}\left(p_{h}^{k},p_{l}^{k}\right)\right)$$

(1)

where $p_{h}^{k}$ and $p_{l}^{k}$ are the different-scale images from the same input image $p_{k}$, $N$ denotes the number of samples in a mini-batch, and the InfoNCE contrastive loss function $\ell_{\text{info}}$ is the same as in [6], which is defined as follows:

$$\ell_{info}\left(p_{l}^{k},p_{h}^{k}\right)=-\log\frac{\exp\left(\operatorname{sim}\left(z_{l}^{k},z_{h}^{k}\right)/\tau\right)}{\sum_{p\in\Lambda}\exp\left(\frac{\operatorname{sim}\left(z_{l}^{k},g_{f}\left(E(p)\right)\right)}{\tau}\right)}$$

(2)

where $z_{l}^{k}=g_{f}\left(E\left({p_{l}^{k}}\right)\right)$, $z_{h}^{k}=g_{f}\left(E\left({p_{h}^{k}}\right)\right)$ and $\Lambda$ is the sample batch. The sim operator is the cosine distance.

The InfoNCE loss in the encoder provides consistency for multi-scale input but cannot guarantee that the learned representation adequately represents the input. In other words, the representation learned by the network from multi-scale images should be not only consistent among different scales, but also representative of the semantic information. According to [3], the self-supervised ViT automatically learns class-specific features leading to unsupervised object segmentation; specifically, the attention block’s last layer includes abundant semantic information. Motivated by [3], the decoder of the Cross-Scale MAE has two tasks: (1) reconstruction of the multi-scale input as close as possible, (2) cross-scale prediction between attentions from different scales, as shown in the decoder part of Fig. 1 and detailed below.

Similar to the standard MAE, Cross-Scale MAE also adopts a light decoder. Decoding follows the standard MAE decoder where, following the encoder, the removed $m$ patches are placed back into their original location in the sequence of patches where a learned mask token represents the masked patches that were not encoded, and a positional encoding is added. Then, a series of transformer layers decode all patches. Given the two representations $f_{e,h}$ and $f_{e,l}$, where $f_{e,h}=E\left({p_{h}}\right)$, $f_{e,l}=E\left({p_{l}}\right)$, after the decoder attention block, we will obtain two new representations, $f_{d,h}$ and $f_{d,l}$. Then, the cross-scale prediction is applied between them to get the cross-scale prediction loss, defined as:

$$\mathcal{L}_{cp}=\frac{1}{N}\sum_{k=1}^{N}\left(\ell_{pred}\left(f_{d,l}^{k},f_{d,h}^{k}\right)\right)$$

(3)

where $N$ is the batch size and the prediction loss $\ell_{pred}$ is defined as:

$$\ell_{pred}\left(f_{d,l}^{k},f_{d,h}^{k}\right)=||f_{d,h}^{k}-g_{p}(f_{d,l}^{k})||_{2}^{2}$$

(4)

It follows the Mean Squared Error (MSE) between the prediction and target representation, where the predictor, $g_{p}(\cdot)$, is a multi-layer perceptron (MLP). Besides the cross-scale prediction loss, the decoder reconstructs images of different scales. The reconstruction loss is thus defined as:

$$\mathcal{L}_{re}=\frac{1}{N}\sum_{k=1}^{N}\left(||p_{l}-\tilde{p}_{l}||_{2}^{2}+||p_{h}-\tilde{p}_{h}||_{2}^{2}\right)$$

(5)

where $\tilde{p}_{l}=D(E(p_{l}))$ and $\tilde{p}_{h}=D(E(p_{h}))$. Via the cross-scale prediction loss and the reconstruction loss, the consistency and effectiveness of the semantic information of the representation can be better preserved.

The total loss of the Cross-Scale MAE is the sum of cross-scale consistency loss at the encoder ($\mathcal{L}_{cc}$), cross-scale prediction loss at the decoder ($\mathcal{L}_{cp}$), and the reconstruction loss ($\mathcal{L}_{re}$), defined as:

$$\mathcal{L}=\mathcal{L}_{cc}+\mathcal{L}_{cp}+\mathcal{L}_{re}$$

(6)

Similar to Scale-MAE, we use both the K-nearest neighbors (KNN) classification accuracy and performance in downstream tasks as metrics to evaluate the quality of the representation learned by Cross-Scale MAE. Specifically, we pre-train the Cross-Scale MAE on the fMoW-RGB dataset and freeze the encoder as a representation generator; then, we test its effectiveness on different datasets via KNN classification.

KNN Performance. Similar to the original MAE configuration, the encoder of Cross-Scale MAE uses a ViT-Large model, and the decoder is a light ViT model with the 8-layer attention blocks. The datasets we use include RESISC45 [10], WHU-RS19 [47], UC Merced [47], EuroSAT [22], as shown in Table  1. To evaluate the representation learning capacity of the proposed method with different GSD images, we evaluate the network performance using images with different scale ratios, $\left\{12.5\%,25\%,50\%,100\%\right\}$, to the raw image to simulate different GSD.

The results are presented in Fig. 2. This figure shows the performance of Cross-Scale MAE compared to both SatMAE and Scale-MAE for different scale ratios. From Fig. 2, we observe that, in general, the Cross-Scale MAE performs better than SatMAE or Scale-MAE in all different scale ratios on all datasets. And, for scale ratio in $\left\{25\%,50\%,100\%\right\}$, the Cross-Scale MAE can obtain more stable performance compared with SatMAE or Scale-MAE. When the scale ratio equals $12.5\%$, all three models perform relatively worse. This is because the scale ratio set during pre-training is from 0.2 to 0.8, and 0.125 is out of this range. Nonetheless, Cross-Scale MAE still presents overwhelmingly better performance than the other two. Table  2 shows the average accuracy at all scale ratios and compares Cross-Scale MAE with SatMAE and Scale-MAE. We can find that Cross-Scale MAE presents overwhelmingly better performance than SatMAE in all datasets and we observe similar trend as in Fig. 2. Cross-Scale MAE generally performs better than Scale-MAE except in the UC Merced dataset. It may be because RESISC45 covers an extensive range of GSD, from $0.2m$ to $30m$, but the other three datasets have fixed GSD. This comparison effectively demonstrates the robustness of representations generated from Cross-Scale MAE, especially on multi-scale datasets.

Downstream Tasks. We further evaluate the performance of Cross-Scale MAE on different downstream tasks to assess the effectiveness of the learned representations in various practical scenarios, including classification and segmentation. We conduct the classification task on the fMoW-RGB dataset. The segmentation downstream tasks are performed on two datasets, Potsdam and Vaihungen. We fine-tune the model for 50 epochs for all downstream tasks, following the same hyper-parameter settings as Scale-MAE [36].

To test the capability of handling multi-scale inputs, we compare the performance of different methods on low-resolution images with different scaling ratios applied on full-resolution images. In Table  3, we report the classification performance regarding Top-1 and Top-5 accuracy. In Table  4, we report the segmentation performance regarding mIoU. The proposed Cross-Scale MAE performs superior to state-of-the-art methods in both classification and semantic segmentation downstream tasks.

The ablation study is four-fold. First, we investigate the critical role of each of the three loss functions to multi-scale inputs. Second, we investigate the effect of negative samples in different levels of the representation. This includes two parts, with the first evaluating the whole Cross-Scale MAE framework and the second under the pure contrastive learning framework (leaving out the masking patch reconstruction part). Third, we compare the effect of multi-scale input and the GSD positional encoding used in Scale-MAE [36]. Although we have shown performance improvement with scale augmentation, we extend the investigation to see whether it has better generalization capacity than the GSD positional encoding. Finally, we explore the effects of some hyper-parameters, including, for example, the number of epochs and the type of backbone used. All experiments in this section are conducted on the RESISC45 dataset [10] and use the KNN classification as a metric.

The weight update of the encoder and decoder networks is mainly controlled by the loss function in Eq. 6 that consists of three modules, the cross consistency loss in the encoder, $\mathcal{L}_{cc}$, the cross prediction loss, $\mathcal{L}_{cp}$, and the reconstruction loss, $\mathcal{L}_{re}$. Table 5 thoroughly compares the classification accuracy using different loss combinations, from which we make some interesting observations. This experiment is conducted on RESISC45 at two scaling ratios, ${0.5,1}$, on the raw image with ViT-Base and pre-trained 300 epochs on fMoW-RGB dataset with input size of $128\times 128$.

The first row is essentially SatMAE, with single scale input, and thus serves as the baseline for the subsequent comparisons. The second row adds the multi-scale input and solely uses reconstruction loss. The improved accuracy displays the effectiveness of multi-scale inputs, especially for lower-scale samples. From the second to the last row, all experiments use multi-scale inputs with different loss components. From the 3rd to the 4th row, where only one additional loss module is applied, we observe that both the consistency in the encoder representation and decoder representation are adequate, with the consistency in the decoder representation playing a more critical role. Finally, in the last row, we observe that the combination of consistency in the encoder and decoder has a performance increment of $3\%-4\%$ on both scaling rates, again demonstrating the effectiveness of the proposed Cross-Scale MAE.

Second, we look into the effectiveness of negative samples in formulating the contrastive loss. Because Cross-Scale MAE leverages contrastive learning at two levels to handle multi-level representations, we must investigate the necessity of involving negative samples at each stage. Specifically, two kinds of losses are evaluated: InfoNCE relating to positive and negative samples and prediction distance relating to only positive samples. First, we assess it under the complete Cross-Scale MAE framework with two scaling ratios, and the results are reported in Table 6. From Table 6, we can find the negative samples in the encoder can improve the classification performance but decrease the performance when involved in the decoder. It may be because the representation from the encoder contains more “structural” information, but the decoder contains more “semantic” information [36]. Hence, the involvement of negative samples in the encoder can help boost the discriminative capacity between multi-scale inputs. On the other hand, with prediction loss between positive samples at the decoder, the consistency in “semantic” information can be better preserved. Having observed this interesting phenomenon, we are intrigued to study if the same pattern would be found in a pure contrastive learning (CL) framework, i.e., Cross-Scale MAE without the reconstruction loss. And we did, the results are shown in Table 7.

Third, we investigate the differences between scale augments and the GSD positional encoding proposed in Scale-MAE [36]. The results are presented in Fig. 3. Two sets of comparisons have been conducted – the left figure shows the performance of Cross-Scale MAE with or without the GSD positional encoding, and the right one shows the SatMAE performance with GSD position or multi-scale augments. The scaling ratios used in the comparison are $\left\{25\%,50\%,75\%,100\%\right\}$. From Fig. 3(a), we find that Cross-Scale MAE adding GSD positional encoding decreases the performance. The reason may be that Cross-Scale MAE has already used multi-scale augmentation, rendering the GSD positional encoding redundant. On the other hand, from Fig. 3(b), we can find both GSD positional encoding and multi-scale augmentation can improve SatMAE’s performance since SatMAE does not exploit cross-scale consistency in the images.

Finally, we show the performance of Cross-Scale MAE with different backbones at different training epochs and the effect of different masking ratios. The results are posted in supplementary materials.

Large-scale models are slow to train and have high memory consumption. Thus, they are usually trained on large clusters of GPUs. We aim to make the model more accessible and optimize the implementation in various ways to minimize the training time and memory footprint for training and inference, such that the model can be feasibly trained end-to-end and used for inference on a single GPU while maintaining competitive results. In this work, we showed using the xFormers library to build an efficient foundation for our models, allowing us to train and conduct experiments more quickly.

Most of the optimizations offered by the xFormers library were made with Tensor-architecture GPUs in mind. We show that we can still gain noticeable improvements in training time even when using an older and less powerful GPU. Thus, for this comparison, we stress-test the xFormers and Timm versions of the baseline implementations by training them on individual Nvidia RTX A6000s. The details of the implementation and ablation studies on attention and loss types have been placed in the supplementary.