AlignMixup: Improving Representations By Interpolating Aligned Features
Abstract
Mixup is a powerful data augmentation method that interpolates between two or more examples in the input or feature space and between the corresponding target labels.
However, how to best interpolate images is not well defined.
Recent mixup methods overlay or cut-and-paste two or more objects into one image, which needs care in selecting regions.
Mixup has also been connected to autoencoders, because often autoencoders generate an image that continuously deforms into another.
However, such images are typically of low quality.
In this work, we revisit mixup from the deformation perspective and introduce
AlignMixup, where we geometrically align two images in the feature space.
The correspondences allow us to interpolate between two sets of features, while keeping the locations of one set.
Interestingly, this retains mostly the geometry or pose of one image and the appearance or texture of the other.
We also show that an autoencoder can still improve representation learning under mixup, without the classifier ever seeing decoded images.
AlignMixup outperforms state-of-the-art mixup methods on five different benchmarks.
Code available at https://github.com/shashankvkt/AlignMixup CVPR22.git
Introduction
Data augmentation [27, 33, 7] is a powerful regularization method that increases the amount and diversity of data, be it labeled or unlabeled [12].
It improves the generalization performance and helps learning invariance [38] at almost no cost, because the same example can be transformed in different ways over epochs.
However, by operating on one image at a time and limiting to labelpreserving transformations, it has limited chances of exploring beyond the image manifold.
Mixup operates on two or more examples at a time, interpolating between them in the input space [51] or feature space [45] while also interpolating between target labels for image classification.
This flattens class representations [45], reduces overly confident incorrect predictions, and smoothens decision boundaries far away from training data.
However, input mixup images are overlays and tend to be unnatural [49].
Interestingly, recent mixup methods focus of combining two [49, 25] or more [24] objects from different images into one in the input space, making efficient use of training pixels.
However, randomness in the patch selection and thereby label mixing may mislead the classifier to learn uninformative features [44], which raises the question: what is a good interpolation of images?
Bengio et al.
[3] show that traversing along the manifold of representations obtained from deeper layers of the network more likely results in finding realistic examples.
This is because the interpolated points smoothly traverse the underlying manifold of the data, capturing salient characteristics of the two images.
Furthermore, [4] show the ability of autoencoders to capture semantic correspondences obtained by decoding mixed latent codes.
This is because the autoencoder may disentangle the underlying factors of variation.
Efforts have followed on mixing latent representations of autoencoders to generate realistic images for data augmentation.
However, these approaches are more expensive, requiring three networks (encoder, decoder, classifier) [4] and more complex, often also requiring an adversarial discriminator [2, 30].
More importantly, they perform poorly compared to standard input mixup on large datasets [30], due to the low quality of generated images.
In this work, we are motivated by the idea of deformation as a natural way of interpolating images, where one image may deform into another, in a continuous way.
Contrary to previous efforts, we do not interpolate directly in the input space, we do not limit to vectors as latent codes and we do not decode.
We rather investigate geometric alignment for mixup, based on explicit semantic correspondences in the feature space.
In particular, we explicitly align the feature tensors of two images, resulting in soft correspondences.
The tensors can be seen as sets of features with coordinates.
Hence, each feature in one set can be interpolated with few features in the other.
By choosing to keep the coordinates of one set or the other, we define an asymmetric operation.
What we obtain is one object continuously morphing, rather than two objects in one image.
Interestingly, observing this asymmetric morphing reveals that we retain the geometry or pose of the image where we keep the coordinates and the appearance or texture of the other. ?? illustrates that our method,
AlignMixup , retains the pose of image 2 and the texture of image 1 , which is different from existing mixup methods.
Note that, as in manifold mixup, we do not decode, hence we are not concerned about the quality of generated images.
AlignMixup 2.1 Preliminaries
Problem formulation Let ( x, y ) be an image x ∈ X with its one-hot encoded class label y ∈ Y , where X is the input image space, Y =
[0 , 1] k and k is the number of classes.
An encoder network
F : X → R c × w × h maps x to feature tensor A = F ( x ) , where c is the number of channels and w × h is the spatial resolution.
A classifier g : R c × w × h → R k then maps A to the vector p =
g ( A ) of probabilities over classes.
Mixup We follow [45] in mixing the representations from different layers of the network, focusing on the deepest layers near the classifier.
We are given two labeled images ( x, y ) , ( x ′ , y ′ ) ∈
X × Y .
We draw an interpolation factor λ ∈ [0 , 1] from Beta( α, α ) [51] and then we interpolate labels y, y ′ linearly by the standard mixup operator
[Equation]
and inputs x, x ′ by the generic formula
[Equation]
where Mix λ is a mixup operator to be defined.
This generic formula allows interpolation of the input or feature as f 2 ◦ f 1 according to
[Equation]
where id is the identity mapping.
For (3), we define Mix λ in (2) as standard mixup mix λ (1), like [51]; while for (4), we define Mix λ as discussed in subsection 2.2.
By default, we train the encoder network and the classifier by using a classification loss L c on the output of the classifier g for mixed examples along with the corresponding mixed labels:
[Equation]
where L c ( p, y ) :
= − (cid:80) ki =
y i log p i is the standard cross-entropy loss.
More options using an autoencoder architecture are investigated in section 3.
2.2 Interpolation of aligned feature tensors
Alignment Alignment refers to finding a geometric correspondence between image elements before interpolation.
The feature tensor is ideal for this purpose, because its spatial resolution is low,
[Figure] 𝐴 [Figure] reducing the optimization cost, and allows for semantic correspondence, because features close to the classifier are small.
Importantly, we are not attempting to combine two or more objects into one image [25], but put two objects in correspondence and then interpolate into one.
We make no assumptions on the structure of input images in terms of objects and we use no ground truth correspondences.
Our feature tensor alignment is based on optimal transport theory [46] and Sinkhorn distance (SD) [8] in particular.
Let A :=
F ( x ) , A ′ :
= F ( x ′ ) be the c × w × h feature tensors of images x, x ′ ∈ X .
We reshape them to c × r matrices A, A ′ by flattening the spatial dimensions, where r := hw .
Then, every column a j , a ′ j ∈ R c of A, A ′ for j = 1 , . . .
, r is a feature vector representing corresponding to a spatial position in the original image x, x ′ .
Let M be the r × r cost matrix with its elements being the pairwise distances of these vectors:
[Equation]
(6) for i, j ∈ { 1 , . . .
, r } .
We are looking for a transport plan , that is, a r × r matrix P ∈
U r , where
[Equation]
and 1 is an all-ones vector in R r .
That is, P is non-negative with row-wise and column-wise sum 1 /r , representing a joint probability over spatial positions of A , A ′ with uniform marginals.
It is chosen to minimize the expected pairwise distance of their features, as expressed by the linear cost function ⟨ P, M ⟩ , under an entropic regularizer:
[Equation]
where H ( P ) :
= − (cid:80) ij p ij log p ij is the entropy of P , ⟨· , ·⟩ is Frobenius inner product and ϵ is a regularization coefficient.
The optimal solution P ∗ is unique and can be found by forming the r × r similarity matrix e − M/ϵ and then applying the Sinkhorn-Knopp algorithm [26], i.e ., iteratively normalizing rows and columns.
A small ϵ leads to sparser P , which improves one-to-one matching but makes the optimization harder [1], while a large ϵ leads to denser P , causing more correspondences and poor matching.
Interpolation The assignment matrix R := rP ∗ is a doubly stochastic r × r matrix whose element r ij expresses the probability that column a i of A corresponds to column a ′ j of A ′ .
Thus, we align A and A ′ as follows:
[Equation]
[Equation]
Here, column (cid:101) a i of c × r matrix (cid:101)
A is a convex combination of columns of A ′ that corresponds to the same column a i of A .
We reshape (cid:101)
A back to c × w × h tensor (cid:101)
A by expanding spatial dimensions and we say that (cid:101)
A represents A aligned to A ′ .
We then interpolate between (cid:101) A and the original feature tensor A :
[Equation]
{+x x ′ 0 .0 .0 .0 .0 .1 .0 .0 .0 .
+}
.λ
mix λ
( A , A ′ ) mix λ
( A , (cid:101) A ) mix λ
( A ′ , (cid:101)
A ′ )
(a) (b) mix λ
( A , A ′ ) mix λ
( A , (cid:101) A ) mix λ
( A ′ , (cid:101)
A ′ ) [Figure] [Figure] As shown in Figure 1 (toy example, top right)
, (cid:101)
A is geometrically close to A .
The correspondence with A ′ and the geometric proximity to A makes (cid:101)
A appropriate for interpolation with A .
Symmetrically, we can also align A ′ to A and interpolate between (cid:101)
A ′ and A ′ : mix λ ( A ′ , (cid:101) A ′ ) .
When mixing feature tensors with alignment (4), we define Mix λ in (2) as the mapping of ( A , A ′ ) to either (11) or (12), chosen at random. 2.3 Visualization and discussion
Decoder We use a decoder to study images generated with or without feature alignment. Let f : R c × w × h → R d be a FC layer mapping tensor A to embedding e = f ( A ) . We use f ◦ F as an encoder and a decoder D : R d → X mapping e back to the image space, reconstructing image ˆ x = D ( e ) .
The autoencoder is trained using only clean images (without mixup) using reconstruction loss L r between x and ˆ x , where L r ( x, x ′ ) :
= ∥ x − x ′ ∥ 2 is the squared Euclidean distance.
We use generated images only for visualization purposes below, but we also use the decoder optionally during AlignMixup training in section 3.
Discussion For different λ ∈ [0 , 1] , we interpolate the feature tensors A , A ′ of x, x ′ without or with alignment, using (11) or (12), and we generate a new image by decoding the resulting embedding through the decoder D .
In Figure 2, we visualize such generated images. Interestingly, by aligning A to A ′ and mixing using (11) with λ = 0 , the generated image retains the pose of x and the texture of x ′ . In Figure 2(a) in particular, when x is ‘penguin’ and x ′ is ‘dog’, the generated image retains the pose of the penguin, while the texture of the dog aligns to the body of the penguin. Similarly, in Figure 2(c), the texture from the goldfish is aligned to that of the stork, while the pose of the stork is retained.
Vice versa, as shown in Figure 2(b,d), by aligning A ′ to A and mixing using (12) with λ = 0 , the generated image retains the pose of x ′ and the texture of x . By contrast, the image generated from unaligned features appears to be an overlay. Randomly sampling several values of λ ∈ [0 , 1] during training generates an abundance of samples, capturing texture from one image and the pose from another. This allows the model to explore beyond the image manifold, thereby improving its generalization and enhancing its performance across multiple benchmarks, as discussed in section 3.
[Table] [Table]
Experiments
Image classification As shown in Table 1(a), AlignMixup is on par or outperforms the SOTA methods by achieving the lowest top-1 error, especially on large datasets.
On CIFAR-10, AlignMixup is on par with Co-Mixup and Puzzlemix with R-18 and WRN16-8.
On CIFAR-100, AlignMixup outperforms Manifold mixup by 1.51% and 0.46% with R-18 and WRN16-8, respectively.
On TI, AlignMixup outperforms Co-Mixup by 2.72% using R-18.
From Table 1(b), AlignMixup outperforms PuzzleMix by 0.56% on ImageNet.
Robustness to FGSM and PGD attacks Following the evaluation protocol of [25] to evaluate
AlignMixup robustness to FGSM and PGD attacks As shown in Table 2, AlignMixup is more robust comparing to SOTA methods.
While AlignMixup is on par with PuzzleMix and Co-Mixup on
CIFAR-10 image classification, it outperforms Co-Mixup and PuzzleMix by 5.36% and 2.28% in terms of robustness to FGSM attacks.
There is also significant gain of robustness to FGSM on
Tiny-ImageNet and to the stronger PGD on CIFAR-100.
Conclusion
We have shown that mixup of a combination of input and latent representations is a simple and very effective pairwise data augmentation method.
The gain is most prominent on large datasets and in combating overconfidence in predictions, as indicated by out-of-distribution detection.
Interpolation of feature tensors boosts performance significantly, but only if they are aligned.
Our work is a compromise between a “good” hand-crafted interpolation in the image space and a fully learned one in the latent space.
A challenge is to make progress in the latter direction without compromising speed and simplicity, which would affect wide applicability.
A Related Work
Mixup [51], concurrently with similar methods [23, 43], introduce mixup , augmenting data by linear interpolation between two examples.
While [51] apply mixup on intermediate representations, it is [45] who make this work, introducing manifold mixup .
Without alignment, the result is an overlay of either images [51] or features [45].
[17] eliminate “manifold intrusion”—mixed data conflicting with true data.
Unlike manifold mixup, AlignMixup interpolates feature tensors from deeper layers after aligning them.
Nonlinear mixing over random image regions is an alternative, e.g .
from masking square regions [10] to cutting a rectangular region from one image and pasting it onto another [49], as well as several variants using arbitrary regions [41, 40, 19].
Instead of choosing regions at random, saliency can be used to locate objects from different images and fit them in one [44, 34, 25, 24].
Exploiting the knowledge of a teacher network to mix images based on saliency has been proposed in [9].
Instead of combining more than one objects in an image, AlignMixup attempts to deform one object into another.
Another alternative is Automix [53], which employs a U-Net rather than an autoencoder, mixing at several layers.
It is limited to small datasets and provides little improvement over manifold mixup [45].
StyleMix and StyleCutMix [21] interpolate content and style between two images, using AdaIN [22], a style transfer autoencoder network.
By contrast, AlignMixup aligns feature tensors and interpolates matching features directly, without using any additional network.
Alignment Local correspondences from intra-class alignment of feature tensors have been used in image registration [6, 31], optical flow [47], semantic alignment [36, 18] and image retrieval [39].
Here, we mostly use inter-class alignment.
In few-shot learning , local correspondences between query and support images are important in finding attention maps, used e.g . by CrossTransformers [11] and DeepEMD [50].
The earth mover’s distance (EMD) [37], or Wasserstein metric , is an instance of optimal transport [46], addressed by linear programming.
To accelerate, [8] computes optimal matching by Sinkhorn distance with entropic regularization .
This distance is widely applied between distributions in generative models [14, 32].
EMD has been used for mixup in the input space, for instance point mixup for 3D point clouds [5] and OptTransMix for images [53], which is the closest to our work.
However, aligning coordinates only applies to images with clean background.
We rather align tensors in the feature space , which is generic.
We do so using the Sinkhorn distance, which is orders of magnitude faster than EMD [8].
B Algorithm
AlignMixup and AlignMixup/AE are summarized in algorithm 1.
By default (AlignMixup), for each mini-batch, we uniformly draw at random one among three choices (line 2) over mixup on input ( x ) or feature tensors ( A , using either (11) or (12) for mixing).
For AlignMixup/AE, there is a fourth choice where we only use reconstruction loss on clean examples (line 7).
For mixup, we use only classification loss (5) (line 24).
Following [45], we form, for each example ( x, y ) in the mini-batch, a paired example ( x ′ , y ′ ) from the same mini-batch regardless of class labels, by randomly permuting the indices (lines 1,10).
Inputs x, x ′ are mixed by (2),(3) (line 12).
Feature tensors A and A ′ are first aligned and then mixed by (2),(11)
( A aligns to A ′ ) or (2),(12) ( A ′ aligns to A ) (lines 14,23).
In computing loss derivatives, we backpropagate through feature tensors A , A ′ but not through the transport plan P ∗ (line 20).
Hence, although the Sinkhorn-Knopp algorithm [26] is differentiable, its iterations take place only in the forward pass. Importantly, AlignMixup is easy to implement and does not require sophisticated optimization like [25, 24]. C Hyperparameter settings CIFAR-10/CIFAR-100 We train AlignMixup using SGD for 2000 epochs with an initial learning rate of 0 . 1 , decayed by a factor 0 .
1 every 500 epochs. We set the momentum as 0 . 9 with a weight decay of 0 . and use a batch size of 128 .
The interpolation factor is drawn from Beta( α, α ) permutation b }
} ▷ i b ▷ current example
▷ mixup
▷
/decode
▷ reconstruction loss else ▷ mixup
▷ interpolation factor
▷ paired example
▷
▷
▷
▷ feature tensors
▷ matrix
▷
▷ tran.
[-By contrast, the image generated from unaligned featuresappears to be an overlay. Randomly sampling several values of λ ∈ [0 , 1] during training generates an abundance of samples,capturing texture from one image and the pose from another. This allows the model to explore x x ′ 0 .0 .0 .0 .0 .1 .0 .0 .0 .0 .0 .
 .λ-]▷ assignments
▷ alignment (9)
[Table]
where α = 2 . 0 .
Using these settings, we reproduce the results of SOTA mixup methods for image classification, robustness to FGSM and PGD attacks, calibration and out-of-distribution detection.
For alignment, we apply the Sinkhorn-Knopp algorithm [26] for 100 iterations with entropic regularization coefficient ϵ = 0 .
TinyImagenet We follow the training protocol of Kim et al .
[25], training R-18 as stage-1 encoder
F using SGD for 1200 epochs.
We set the initial learning rate to 0 .
1 and decay it by 0 .
1 at 600 and 900 epochs.
We set the momentum as 0 .
9 with a weight decay of 0 .
and use a batch size ofon 2 GPUs.
The interpolation factor is drawn from Beta( α, α ) where α = 2 . 0 .
For alignment, we apply the Sinkhorn-Knopp algorithm [26] for 100 iterations with entropic regularization coefficient ϵ = 0 .
ImageNet We follow the training protocol of Kim et al .
[25], where training R-50 as F using SGD for 300 epochs.
The initial learning rate of the classifier and the remaining layers is set to 0 .
1 and 0 .
01 , respectively.
We decay the learning rate by 0 .
1 at 100 and 200 epochs.
We set the momentum as 0 .
9 with a weight decay of 0 .
and use a batch size of 100 on 4 GPUs.
The interpolation factor is drawn from Beta( α, α ) where α = 2 . 0 .
For alignment, we apply the Sinkhorn-Knopp algorithm [26] for 100 iterations with entropic regularization coefficient ϵ = 0 .
We also train R-50 on ImageNet for 100 epochs, following the training protocol described in Kim et al .
CUB200-2011 For weakly-supervised object localization (WSOL), we use VGG-GAP and Rpretrained on ImageNet as F .
The training strategy for WSOL is the same as image classification and the network is trained without bounding box information .
In R-50, following [49], we modify the last residual block ( layer 4 ) to have stride 2 instead of 1 , resulting in a feature map of spatial [Table] [Figure] [Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table]
[Table]
resolution 14 × 14 . The modified architecture of VGG-GAP is the same as described in [52]. The classifier is modified to have 200 classes instead of 1000 . For fair comparisons with [49], during training, we resize the input image to 256 × 256 and randomly crop the resized image to 224 × 224 .
During testing, we directly resize to 224 × 224 .
We train the network for 600 epochs using SGD.
For R-50, the initial learning rate of the classifier and the remaining layers is set to 0 .
01 and 0 .
001 , respectively.
For VGG, the initial learning rate of the classifier and the remaining layers is set to 0 .
001 and 0 . 0001 , respectively.
We decay the learning rate by 0 .
1 every 150 epochs.
The momentum is set to 0 .
9 with weight decay of 0 . 0001 and batch size of 16 .
[Figure]
D Additional experiments
ImageNet classification Following the training protocol of [24], Table 3 reports classification performance when training for 100 epochs on ImageNet.
Using the top-1 error (%) reported for
[Table]
competitors by [24], AlignMixup outperforms all methods, including Co-Mixup [24].
Importantly, while the overall improvement by SOTA methods over Baseline is around 1 .
64% , AlignMixup improves SOTA by another 0 .
Experiments using transformers We apply mixup to LeViT-128S [15] on ImageNet forepochs.
For AlignMixup, we align the feature tensors in the last layer of the convolution stem.
The top-1 accuracy is: baseline 67.4%, input mixup 68.3%, manifold mixup 67.8%, CutMix 68.7%,
AlignMixup 69.9%.
Thus, we outperform input mixup and CutMix by 1.6% and 1.2% respectively, which in turn outperform the baseline by 0.9% and 1.3% respectively.
This means that the improvement brought by mixing is roughly doubled.
Out-of-distribution detection We compare AlignMixup with SOTA methods, training R-18 on
CIFAR-100.
At inference, ID examples are test images from CIFAR-100, while OOD examples are test images from LSUN [48] and Tiny-ImageNet, resizing OOD examples to 32 × 32 to match the resolution of ID images [49].
We also use test images from CIFAR-100 with Uniform and
Gaussian noise as OOD samples.
Uniform is drawn from U (0 , 1) and Gaussian from N ( µ, σ ) with µ = σ = 0 .
All SOTA mixup methods are reproduced using the same experimental settings.
Following [20], we measure detection accuracy (Det Acc) using a threshold of 0 . 5 , area under ROC curve (AuROC) and area under precision-recall curve (AuPR).
As shown in Table 4, AlignMixup outperforms SOTA methods under all metrics by a large margin, indicating that it is better in reducing over-confident predictions.
Calibration We compare AlignMixup with SOTA methods , training R-18 on CIFAR-100.
All
SOTA mixup methods are reproduced using the same experimental settings.
We compare qualitatively by plotting accuracy vs .
confidence.
As shown in Figure 3, while Baseline is clearly overconfident and Input and Manifold mixup are clearly under-confident, AlignMixup results in the best calibration among all competitors.
We also compare quantitatively, measuring the expected calibration error (ECE) [16] and overconfidence error (OE) [42].
As shown in Table 5, AlignMixup outperforms SOTA methods by achieveing lower ECE and OE, indicating that it is better calibrated.
Qualitative results of WSOL Qualitative localization results shown in Figure 4 indicate that
AlignMixup encodes semantically discriminative representations, resulting in better localization performance.
Object detection Following the settings of CutMix [49], we use Resnet-50 pretrained on ImageNet using AlignMixup as the backbone of SSD [29] and Faster R-CNN [35] detectors and finetune it on Pascal VOC07 [13] and MS-COCO [28] respectively.
AlignMixup outperforms CutMix mAP by 0.8% (77.6 → 78.4) on Pascal VOC07 and 0.7% (35.16 → 35.84) on MS-COCO.
I npu t
m i xup [ 51 ]
[Figure]
i xup
[Figure] [Figure]
E Additional ablations
[Figure]
[Table]
in Sinkhorn-Knopp
The default number of iterations for the Sinkhorn-Knopp algorithm in solving (8) is i = 100 .
Here, we investigate more choices, as shown in Table 6.
The case of i = 0 is similar to cross-attention.
In this case, we only normalize either the rows or columns in (7) once, such that P 1 = 1 /r (when A aligned to A ′ ) or P ⊤ 1 = 1 /r (when A ′ aligned to A ).
We observe that while AlignMixup outperforms the best baseline–StyleCutMix (80.66)–in all cases, it performs best for i = 100 iterations.