Masked Angle-Aware Autoencoder for Remote Sensing Images

MAE [25] employs an asymmetric encoder-decoder architecture for efficient masked image modeling. An input image, $x\in{\mathbb{R}}^{H\times W\times C}$, is first reshaped into a series of non-overlapping image patches of size $p\times p$, denoted as ${x^{p}}\in{\mathbb{R}}^{N\times{p^{2}}C}$, where $N=HW/{p^{2}}$ is the number of patches. Then, ${x^{p}}=\{x_{i}^{p}|i=1,2,...,N\}$ is linearly mapped into patch embeddings. MAE adds positional encoding information to these embeddings and randomly masks them at a certain ratio, such as 75%. The masked patches, ${x^{m}}=\{x_{i}^{m}|i=1,2,...,N^{m}\}$, are discarded, and only the remaining visible patches, ${x^{v}}=\{x_{i}^{v}|i=1,2,...,N^{v}\}$, are fed into the encoder to extract latent features. These latent features, along with the shared and learnable mask tokens representing the substituted masked patches, constitute the input to the decoder, with positional embeddings also added. After obtaining the decoder output $\hat{x}=\{{\hat{x}_{i}}|i=1,2,...,N\}$, MAE predicts only the pixel values of the masked patches, using the original image patches as the reconstruction targets. This is achieved by computing the mean squared error (MSE) loss:

where ${\hat{x}^{m}}\in{\mathbb{R}}^{{N^{m}}\times{p^{2}}C}$ denotes the output of the decoder for masked patches. The decoder is adopts for pre-training only, while the encoder is further fine-tuned for downstream tasks. This often results in the common practice of using a lightweight decoder and a complete transformer encoder. The proposed MA3E shares the similar principle with MAE, described next.

MA3E aims to be aware of diverse angles and learn rotation-invariant visual representations. Fig. 2(a) illustrates our pipeline. MA3E constructs rotated crops at arbitrary positions on the RS images by deploying the designed scaling center crop operation to introduce explicit angle variations. These rotated crops have random angles and replace the scenes at their original locations. We add an angle embedding in each rotated crop, and mask the rotated crop and remaining background separately. For the reconstruction of the rotated crop, MA3E automatically allocates similar image patches as reconstruction targets for each rotated crop patch based on the transportation plan solved by the Sinkhorn-Knopp [14] algorithm. This avoids biases introduced by the crop operation.

Rotated crop. Using a simple random rotation operation to construct the rotated crop of side length $a$ for each RS image would lead to adverse results as shown in Fig. 3(b). The model struggles to learn high-quality representations from these regions, resulting in wasted computational resources. Therefore, we propose a scaling center crop operation to create diverse rotated crops that preserve scenes to a substantial extent, as shown in Fig. 3(a). For a square region (blue) of side length $h$ at arbitrary positions on an image, rotating this region at random angle would lead to loss of edge scenes (gray). However, the scenes within its largest inscribed circle (red) are fully preserved. Hence, we perform center cropping to extract the largest inscribed square region with side length $a=\frac{\sqrt{2}}{2}h$ from the red circle as the rotated crop. This region holds an arbitrary orientation, replacing the original scene and introducing the explicit angle variation to the composite image. To ensure that each rotated crop can be entirely patchified, the side length $a$ needs to be divisible by the patch size $p$, with the starting position being a multiple of $p$.

Angle embedding. In addition to adding positional embeddings for the divided patches of composite images, MA3E also includes the angle embeddings for the rotated crops. Each angle embedding is a learnable vector shared across each patch within a rotated crop. They serve as implicit cues for the model to perceive the angle variation in the rotated crop, while also distinguishing them from the remaining background.

Random masking. Given $N_{r}={a^{2}}/{p^{2}}$ patches from the rotated crop (denoted as $r=\{{r_{i}}|i=1,2,...,{N_{r}}\}$) and $N_{b}=N-{N_{r}}$ patches from the background (denoted as $b=\{{b_{i}}|i=1,2,...,{N_{b}}\}$), to avoid the random masking strategy that removes too much or all patches from the rotated crop, we separately mask $r$ and $b$ at a certain ratio, e.g., 75%. Thus, visible patches and masked patches from the background are denoted as ${b^{v}}=\{b_{i}^{v}|i=1,2,...,N_{b}^{v}\}$ and ${b^{m}}=\{b_{i}^{m}|i=1,2,...,N_{b}^{m}\}$ respectively. These definition holds similarly for the rotated crop patches $r$.

Reconstruction. MA3E uses MSE loss to predict the pixel values of masked patches in the background. For the rotated crop, an OT loss ${\cal L}_{OT}$ is proposed to minimize the distance in pixel space between each patch and its matched image patches. The overall loss can be written as follows:

where ${\hat{b}^{m}}\in{\mathbb{R}}^{N_{b}^{m}\times{p^{2}}C}$ means predictions for the masked patches of the background, and $\hat{r}\in{\mathbb{R}}^{{N_{r}}\times{p^{2}}C}$ denotes all predicted patches of the rotated crop.

After the scaling center crop operation, the scenes in all rotated crop patches are offset compared to the original image patches at the same positions. Directly calculating MSE between the masked patches and image patches not only results in biases but also overlooks the changes in angles and scenes on the visible patches. Inspired by [20], this paper treats the reconstruction for rotated crops as an OT problem, allowing each predicted patch to automatically match similar image patches for reconstruction.

Optimal transport. Supposing there are $M$ suppliers and $N$ demanders, where the $i$-th supplier holds $u_{i}$ units of goods, and the $j$-th demander needs $v_{j}$ units of goods. The transportation cost for transporting one unit of goods from the $i$-th supplier to the $j$-th demander is denoted as $c_{ij}$. OT aims to find a transportation plan, denoted as $\Omega=\{\omega_{i,j}|i=1,2,...,M,j=1,2,...,N\}$, that minimizes the total transportation costs, ensuring that all goods are transported from suppliers to demanders:

OT for reconstruction. The context of reconstructing the rotated crop is described in Fig. 2(b), considering $N_{r}$ original image patches and $N_{r}$ predicted patches of the rotated crop, each image patch is treated as a supplier, holding ${p^{2}}C$ units of pixel values (i.e., $u_{i}={p^{2}}C,i=1,2,...,{N_{r}}$), and each predicted patch as a demander, with ${p^{2}}C$ units of channels (i.e., $v_{j}={p^{2}}C,j=1,2,...,{N_{r}}$) needing ${p^{2}}C$ units of pixel values for reconstruction. The similarity between the each unit of pixel value in the image patch and the any unit of channel in the predicted patch represents the transportation cost $c_{ij}$. This is extended to matrix-wise MSE computation for GPU acceleration. Thus, the transportation cost from the $i$-th image patch to the $j$-th predicted patch is given by:

where image patches closer in L2 distance to the predicted patch have higher similarity, tending to lower costs. A fast iterative algorithm, Sinkhorn-Knopp [14], is employed to calculate the transportation plan $\Omega$ in Eq. 3. According to the solved $\Omega=\{\omega_{i,j}\}$, as shown in Fig. 2(b), the OT loss automatically allocates similar multiple image patches for $j$-th predicted patch as the reconstruction targets, defined as follows:

The proposed ${\cal L}_{OT}$, during the pixel reconstruction and angle restoration, guides our model to perceive the angle variations of rotated crops. As the iterations progress, MA3E can effectively learn rotation-invariant visual representations. Supplementary material provides more details about ${\cal L}_{OT}$.

Unless otherwise stated, all experiments are implemented using PyTorch and conducted on a machine equipped with eight 24GB RTX 3090 GPUs. More experimental setups and datasets are detailed in the supplementary material.

Pre-training details. The testing set of MillionAID [38], consisting of 990,848 RS images, is used for pre-training. Each image are resized to $224\times 224$ pixels. The patch size $p$ is 16. For each rotated crop, the side length $a$ is set to 96, and the rotation range is $[-45^{\circ},+45^{\circ}]$. We randomly mask the rotated crop and background patches with a ratio of 75% respectively. MA3E employs a plain ViT-B [16] as the encoder and 8 ViT blocks with 512-D as the decoder. Except for the batchsize of 1024, all other pre-training configurations follow [25].

Scene classification. All fine-tuning and linear probing experiments are conducted on NWPU-RESISC45 [7] (NU45), AID [60], and UC Merced [65] (UCM) datasets. For NU45, 20% of images from each class are randomly sampled as the training set, and the remaining 80% are used for testing. For AID and UCM, these two ratios are both 50%. Fine-tuning is performed with a batchsize of 512 for 200 epochs, and linear probing is trained with a batchsize of 2048 for 100 epochs. We follow the other default fine-tuning and linear probing settings outlined in [25] and report the Top-1 accuracy on each testing set.

Rotated object detection. Experiments for detection are conducted on DOTA1.0 [59] and DIOR-R [8] using the Oriented R-CNN [63] detector. MA3E pre-trained models serves as the backbone of the detector and undergoes end-to-end fine-tuning. The detector uses the batchsize of 2 for DOTA1.0 and 4 for DIOR-R. We train for 12 epochs, with other hyper-parameters following default settings of the detector. Mean Average Precision (mAP) on each testing set is reported, where results on DOTA1.0 are obtained from the official evaluation server. Due to the limited GPU memory, above experiments are deployed on two GPUs and implemented by the OBBDetection and the ViTDet [32] codebases.

Semantic segmentation. Similarly, segmentation experiments are conducted using the UperNet [61] framework for end-to-end supervised fine-tuning on iSAID [57] and Potsdam. The UperNet is trained for $160k$ iterations with a batchsize of 4, while other hyper-parameters remain at the default settings. Mean Intersection over Union (mIoU) on the iSAID validation set and mean F1 score (mF1) on the Potsdam testing set are reported. These experiments are implemented using the mmsegmentation [12] library and also run on two GPUs.

MA3E is compared with eight state-of-the-art pre-training methods for RS images, including seasonal-contrasted SeCo [42], change-aware contrasted CACo [41], GFM [43] with continual pre-training, RingMo [50] adopting incomplete masking, CMID [44] combining contrastive learning and MIM, [54] using rotated varied-size window attention (RVSA) to replace the original global attention during downstream fine-tuning, sptaio-temporal encoded SatMAE [11], and scale-aware ScaleMAE [46]. However, these methods adopt different datasets and fine-tuning settings for downstream tasks. For fairness, we normalize the experimental setups and further compared with six pre-training methods for natural images, including the popular MoCo v3 [6], DINO [2], MAE [25], SimMIM [64], region-reconstructed LoMaR [3], and input-mixed MixMAE [34].

Scene classification. The fine-tuning and linear probing results on three datasets are shown in Table 1 and Table 2, repectively. MA3E pre-trained for 300 epochs achieves competitive results. Although the fine-tuning accuracy on UCM is lower than ViT+RVSA [54], MA3E requires only 52% of the latter’s FLOPs. The fine-tuning and linear probing results on three datasets continually improve as training progresses. MA3E pre-trained for 1600 epochs comprehensively leads, demonstrating that MA3E effectively learns the discriminative rotation-invariant representations of RS objects. In addition, the training time per epoch on a single GPU is calculated. Compared to MAE [25], we achieve a significant improvement in accuracy with only about 0.2 hours of extra training time.

Rotated object detection. Table 3 presents the fine-tuning results on DOTA1.0 and DIOR-R for different methods, where MA3E, using a simple backbone, obtains superior detection performance. The version pre-trained for 300 epochs outperforms other methods with a similar number of epochs. MA3E pre-trained for 1600 epochs surpasses all methods. Compared to [54] with ViTAE+RVSA, the mAP on DOTA1.0 and DIOR-R has increased by 0.51 and 0.87, respectively. This significant improvement in detection performance demonstrates the effectiveness of MA3E in angle perception during pre-training.

Semantic segmentation. The fine-tunning results on iSAID and Potsdam are also reported in Table 3. With fewer pre-training epochs, MA3E demonstrates competitive performance, achieving mF1 only 0.04 lower than CMID [44] on Potsdam. When pre-trained for 1600 epochs, MA3E again achieves the best results, outperforming the second-best [54] by 0.3 mIoU and 0.28 mF1 on iSAID and Potsdam, respectively. It is indicated that the significance of the rotation-invariant representations learned by our model in semantic segmentation.

In this section, a series of ablation studies are conducted to analyze how each key design is available and demonstrate the effectiveness of each component in MA3E. By default, MA3E is pre-trained for 300 epochs. We report the Top-1 accuracy on NU45 [7] after fine-tuning, mAP on DOTA1.0 [59], and mIoU on iSAID [57]. The results of MA3E with default settings are marked in green. Our supplementary material provides additional ablation results.

Each component. The ablation results on different components are shown in Table 5. Each key design of MA3E improves the performance of the baseline MAE across the three RS tasks, and the combination of all components in MA3E yields the best results. This demonstrates the effectiveness of our proposal.

Side length $a$ and crop numbers. We study the effects of different side lengths $a$ for the rotated crop and the number rotated crops in each image, as shown in Table 5. As $a$ increases, MA3E’s performance gradually improves until $a=96$. However, for detection and segmentation tasks, there is a sudden decrease in the performance when $a=128$. Furthermore, an increase in the number of rotated crops also negatively impacts performance. Excessive large $a$ and a more number of crops make angle restoration challenging, and this decline is more obvious in detection where angles are crucial.

Position of rotated crops. We ablate three schemes for selecting the position of the rotated crop: fixed positions at the center of the images, random positions, and positions determined by the selective search [53]. The corresponding results are shown in Table 7. The selective search algorithm generates candidate bounding boxes for potential objects in an unsupervised manner. Although determining the position of the rotated crop through this method results in an improvement of 0.19 in classification accuracy and 0.15 mAP in detection compared to randomly choosing positions, the limited performance gain comes at the cost of at least 10% extra training time per epoch. Therefore, it is not employed in the default settings.

Method to create rotated crops. We compare MA3E with a model using the simple random rotation operation to generate rotated crops with $a=96$. The results in Table 7 demonstrate that rotated crops constructed using the proposed scaling center crop operation significantly enhance the performance in all tasks. Note that our method may still incur minor losses in edge scenes. But this corresponds to deliberately increasing the difficulty of reconstructing the complete original image, thereby improving the quality of learned rotation-invariant representations.

Rotation ranges. Table 9 shows the results for different rotation ranges of rotated crops. An reasonable range maximizes the ability of ${\cal L}_{OT}$ to restore the angle variations introduced on rotated crops, promoting MA3E to learn rotation-invariant visual representations and achieving good performance.

Masking strategy. Table 9 report the effect of different masking strategies on MA3E. The random masking strategy [25] performs better. Block-wise masking [1] increases the reconstruction difficulty. Meanwhile, uniform sampling [31] with masking in adjacent four patches makes the reconstruction of rotated crops easier but leads to lower quality representations learned by the model.

Fig. 4 quantitatively visualizes the reconstruction performance of MA3E on RS images. Example images are randomly sampled from the training set of MillionAID [38], which contains $10k$ images. We resize each image to $224\times 224$ (196 patches, $p=16$) and set the rotated crop of $96\times 96$ (36 patches) for reconstruction. The rotated crop and background are masked at a 75% ratio, corresponding to 9 and 40 visible patches, respectively. It can be seen that MA3E effectively restores the preset angle variations of the rotated crops during reconstructing original pixels. In addition, some uneven color patches may appear in the reconstructed rotated crop, such as the water surface and playground in Fig. 4. This phenomenon is attributed to the proposed OT loss, which calculates the mean square error between each predicted patch and multiple target patches. The supplementary material shows more visualizations.