A Continuous Mapping For Augmentation Design

1 Introduction . Data augmentation [ 22 ] is one of the advanced techniques or recipes that proliferate the success of deep learning . The crux behind data augmentation is to make a model invariant towards various input transformations that we expect to observe during inference , thus , improves the generalization of the model . Since it is not clear which transformations would benefit most , manual design is difficult . Thus , there has been a growing interest in automatically learning such augmentations , a new thriving subarea of data augmentation known as the Automated Data Augmentation ( ADA ) [ 5 , 16 , 21 , 20 , 19 , 28 ] . ADA is generally formulated as a bilevel optimization problem over a discrete search space , and aims to find an optimal augmentation policy . For example , the search space used in AutoAugment [ 5 ] , the very first promising work in this topic , contains several augmentation candidates , e.g. , rotation/blurring/brightening with different discrete magnitudes , and the augmentation policy is defined as a parameterized discrete distribution over this space . Such approach has been shown to provide ∗Authors contributed equally . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . far more superior performances over several datasets and tasks compared against the traditional handcrafted augmentation counterparts . This , however , comes at the expense of a high computational budget , taking far more time than model training , resulting in an untenable training procedure in the presence of real-world large-scale datasets . While several approaches have been proposed to address the efficiency issue [ 20 , 16 , 21 , 19 ] , the compromise in performance due to such efficiency measures was nontrivial ; some of the current state-of-art remains computationally expensive [ 28 ] . The contradiction between performance and efficiency has raised a key challenge in this field . Another major limitation of such approaches is due to the discretized search space which limits the diversity of the learnable policy . In fact , the discretization results in augmentations that do not always represent the variations in the real world sufficiently . In the real world , the data is collected with a great range of freedom such as varying camera viewpoint and lighting conditions . However , the discrete augmentation would only be able to produce transformed data with specific types and magnitudes . Furthermore , quantization of augmentation search space introduces extra manual design , which is contrary to the aim of ADA . In this paper , we resolve the aforementioned issues of discretization by constructing a continuous and differentiable mapping φ which maps an augmentation vector α ∈ Rd to an image transformation function tα ∈ T that transforms an image I to tα ( I ) . Specifically , we group basic augmentations into three categories : Color Adjustment that adjusts the color values , Image Filtering like blurring or sharpening , and Image Warping which represents the geometric transformation . We then split the augmentation vector α to three sub-vectors ( αCA , αIF and αIW ) , each of which would map to a specific transformation ( tαCA , tαIF or tαIW ) . We finish the construction by regard tα as the composite transformation , i.e. , tα = tαIW ◦ tαIF ◦ tαCA , as visualized in Fig . 1 . Based on this construction , we pose ADA as learning an augmentation distribution in the continuous augmentation space . This gives us the freedom to search over the entire continuous space in order to obtain the optimal augmentations . For optimization and inference , we use the well-known Stochastic Gradient Langevin Dynamics ( SGLD ) and propose a pseudo one-step approximation that allows us to efficiently learn the underlying augmentation distribution and sample augmentations from it . On a wide range of experiments involving state-of-the-art architectures on multiple image classification benchmarks , we show that our approach not only provides results at par with existing ADA methods , it also is orders of magnitude faster than them . Our main contributions can be summarized as : 1 . We propose a continuous augmentation mapping and first pose ADA as a continuous optimization problem , thus , avoiding the pitfalls of the discrete augmentation space . 2 . We develop a gradient-based approach to learn the augmentation distribution in the continuous space that allows for highly efficient sampling of augmentations , offering a favorable speed-performance trade-off . 2 Related Work . Recent automated augmentation methods have been proposed that optimize data augmentation policies on different datasets and tasks [ 5 , 16 , 21 , 20 , 19 , 28 ] . AutoAugment [ 5 ] provides a discrete augmentation search space and an RNN controller to parameterize the augmentation distribution . Reinforcement learning is then employed to search for the augmentation strategy . Subsequent work focuses more on searching efficiently in the same space . To this end , OHL [ 21 ] proposed an online reinforcement learning pipeline to reduce the computation complexity . PBA [ 16 ] proposed a population based search strategy to obtain the augmentation distribution with fitness selection also in an online manner . Fast AutoAugment [ 20 ] proposed a different objective function that applies data augmentation at test time . AWS [ 28 ] investigates the effectiveness of a weight-sharing paradigm , yielding superior performance on various datasets . The most related work , DADA [ 19 ] take a continuous relaxation to the discrete augmentation distribution by Gumbel-softmax to obtain better approximation . Our work eliminated such approximation tricks when computing the gradient thanks to the continuous nature of the proposed search space . Markov Chain Monte Carlo ( MCMC ) is a popular sampling method for approximating the posterior distribution in large-scale Bayesian modeling . MCMC draws samples from high-dimensional probability distributions in the construction of a Markov chain . A broad category of recent MCMC methods , namely Stochastic Gradient MCMC ( SG-MCMC ) [ 3 , 1 , 9 , 30 ] , comes from the traditional stochastic gradient optimization and thus inherits the efficiency of Stochastic Gradient Descent ( SGD ) . To cast a stochastic optimization process to a sampling procedure , artificial noise needs to be added to the gradients . A typical example of gradient MCMC methods is Stochastic Gradient Langevine Dynamics ( SGLD ) [ 30 ] , which injects Gaussian noise into the parameter gradients , pushing the optimization trajectory to asymptotically converge to the actual posterior distribution . In practice , some of these algorithms are also used in combination with several heuristic approximations , such as constant step sizes or diagonal approximations to the Hessian for the sake of simplicity or efficiency [ 25 ] . In Sec . 4.2 of this paper , we also use the simple SGLD with constant noise to optimize our augmentation distribution . 3 A Continuous and Differentiable Mapping . 3.1 Preliminaries . Let the mapping φ : α 7→ tα map a d-dimensional vector α ∈ Rd to an image augmentation transformation tα ∈ T , where the augmentation tα ( · ) maps an image I to its transformed one tα ( I ) . For convenience , we introduce the notation tα to denote the transformation function in T for which tα = φ ( α ) . In other words , α could be regarded as the parameter of the transformation function tα , and φ is the realization function . 3.2 Constructing The Mapping . We start by developing a taxonomy of image augmentation transformations . Popular basic augmentations can be grouped into three categories : Color Adjustment that adjusts the color values pixel by pixel , Image Filtering like blurring or sharpening the whole picture , and Image Warping which applies geometric transformations to the image . Now we discuss how to construct a continuous and differential mapping for each augmentation categories respectively : Color Adjustment is defined as a transformation in the spatial domain that is equally applied to each pixel . Let n be the image size . For every location ( x , y ) with pixel vector Ixy = [ hxy sxy vxy ] T , it would be adjusted to I ′xy = [ h′xy s ′ xy v ′ xy ] T . The augmentation can be described as : I ′ = tαCA ( I ) , where I ′ xy = h′xys′xy v′xy = αh + ( 1 + βh ) ( hxy ) γhαs + ( 1 + βs ) ( sxy ) γs αv + ( 1 + βv ) ( vxy ) γv ( ∀ x , y that 1 ≤ x , y ≤ n ) , ( 1 ) where the nine-dimensional vector αCA = [ αh βh γh αs βs γs αv βv γv ] T maps to a Color Adjustment augmentation tαCA ( · ) . The first three values control the hue , the next three would change the saturation , and the last three control the brightness . This representation is motivated by the Power law ( Gamma ) transformation function [ 24 ] . As shown in Equ . ( 1 ) , this augmentation family is differential w.r.t . all α , β , γ since it can be expressed as power functions of them and both RGB-to-HSV and HSV-to-RGB conversions are known to be continuous and differentiable [ 26 , 29 ] . Image Filtering is a typical image enhancement in the frequency domain . Here we consider using filters to blur or sharpen images . Let a single value s denote the degree of sharpening or smoothing . Let Gk be a sharpness filter obtained by zero-mean normalizing a k × k discrete Gaussian filter , and Ck be a k × k identity filter which has an unique non-zero value 1.0 at its center . The sharpening or smoothing process that transforms an image I to I ′ is : I ′ = tαIF ( I ) = I ∗ ( s ·Gk + Ck ) , ( 2 ) where the asterisk ∗ means convolution and tαIF ( · ) is the Image Filtering function . If s is greater than zero , it represents sharpening . Otherwise , it is equivalent to blurring . Thus , a one-dimensional vector [ αIF ] T could map to an Image Filtering transformation . Image filtering is differentiable w.r.t s by nature . However , the computation grows fast with respect to the filter size k. When k is close to the image size n , a better implementation for computing the convolution would be applying discrete Fast Fourier Transform ( FFT ) and multiplying the frequencies , which achieve a complexity of O ( n2 log n ) . In this work , we take direct convolution since we pick k n for simplicity . Image Warping refers to a geometric distortion on an image where pixels are moved to other locations without changing their pixel values . Perspective transformation ( also known as projective transformation ) is a popular warping method , which can be uniquely represented by a 3× 3 matrix H called homography [ 2 ] . Each point in the original image ( x , y ) would move to ( x′ , y′ ) with the perspective transformation . This Image Warping can be described as : I ′ = tαIW ( I ) ⇐⇒ I ′x′y′ = Ixy , where x′ = H11x+ H12y + H13 H31x+ H32y + H33 , y′ = H21x+ H22y + H23 H31x+ H32y + H33 , ( ∀ x , y that 1 ≤ x , y ≤ n ) , ( 3 ) where tαIW ( · ) denotes the Image Warping augmentation via perspective transformation , and αIW = [ H11 H12 H13 H21 H22 H23 H31 H32 H33 ] T could represent an Image Warping augmentation . Rotations , translations , shearing , scaling , as well as more single geometric transformations and any combinations of them are all special cases of perspective transformation . Perspective transformation is originally non-differentiable to each element in the homography H since index x , y take discrete values . However , we can leverage interpolation between pixels and inverse warping to make the transformation differentiable [ 17 , 11 , 27 ] . In summary , an augmentation vector α = [ αTCA αTIF αTIW ] T lies in continuous real space with d = 9 + 1 + 9 = 19 dimensions could map to a specific augmentation transformation tα = tαIW ◦ tαIF ◦ tαCA . Although composite order matters , for simplicity , leave the matter to future work . One can easily verify that many commonly used augmentation transforms [ 5 , 6 , 20 , 16 , 21 , 19 , 28 ] , such as brightening , blurring , and rotation , are all discrete samples in the augmentation function space T of our mapping . Last but not least , practical constraint functions ( tanh , exp ) and clamp function are used before and after applying augmentations respectively , to ensure all augmentations are valid and all augmented pixels are in [ 0 , 1 ] . The overview of this space is illustrated in Fig . 1 . More details on the augmentation process of this figure are provided in Supplementary A .

Summary:     Authors present a continuous formulation of data augmentation for     images. The formulation is used to define a Bayesian model for augmentation     parameters. The Bayesian model does not rely on a bi-level optimization     problem as commonly used for optimizing data augmentation. Instead, only a     single gradient step approximates a full optimization of a network with a     given set of augmentation parameters. Gradient step is taken starting from     pretrained weights determined through initial training using random     augmentation. Samples from the posterior distribution are gathered using     Langevin dynamics MCMC with random batches. Samples are used to train     the ultimate networks. Experiments show that the proposed model outperforms     existing techniques in terms of accuracy of the final models at a fraction     of computation time.