SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Liang Hou $^{1,3,4}$ , Qi Cao $^{1}$ , Yige Yuan $^{1,3}$ , Songtao Zhao $^{4}$ , Chongyang Ma $^{4}$ , Siyuan Pan $^{4}$ , Pengfei Wan $^{4}$ , Zhongyuan Wang $^{4}$ , Huawei Shen $^{1,3}$ , Xueqi Cheng $^{2,3}$
$^{1}$ CAS Key Laboratory of AI Safety and Security, Institute of Computing Technology, Chinese Academy of Sciences  $^{2}$ CAS Key Laboratory of Network Data Science and Technology, Institute of Computing Technology, Chinese Academy of Sciences
<sup>3</sup>University of Chinese Academy of Sciences <sup>4</sup>Kuaishou Technology lianghou96@gmail.com
Training generative adversarial networks (GANs) with limited data is challenging because the discriminator is prone to overfitting. Previously proposed differentiable augmentation demonstrates improved data efficiency of training GANs. However, the augmentation implicitly introduces undesired invariance to augmentation for the discriminator since it ignores the change of semantics in the label space caused by data transformation, which may limit the representation learning ability of the discriminator and ultimately affect the generative modeling performance of the generator. To mitigate the negative impact of invariance while inheriting the benefits of data augmentation, we propose a novel augmentation-aware self-supervised discriminator that predicts the augmentation parameter of the augmented data. Particularly, the prediction targets of real data and generated data are required to be distinguished since they are different during training. We further encourage the generator to adversariably learn from the self-supervised discriminator by generating augmentation-predictable real and not fake data. This formulation connects the learning objective of the generator and the arithmetic - harmonic mean divergence under certain assumptions. We compare our method with state-of-the-art (SOTA) methods using the class-conditional BigGAN and unconditional StyleGAN2 architectures on data-limited CIFAR-10, CIFAR-100, FFHQ, LSUN-Cat, and five low-shot datasets. Experimental results demonstrate significant improvements of our method over SOTA methods in training data-efficient GANs.<sup>1</sup>
Generative adversarial networks (GANs) [10] have achieved great progress in synthesizing diverse and high-quality images in recent years [2, 17, 19, 20, 13]. However, the generation quality of GANs depends heavily on the amount of training data [47, 18]. In general, the decrease of training samples usually yields a sharp decline in both fidelity and diversity of the generated images [39, 48]. This issue hinders the wide application of GANs due to the fact of insufficient data in real-world applications. For instance, it is valuable to imitate the style of an artist whose paintings are limited. GANs typically consist of a generator that is designed to generate new data and a discriminator that guides the generator to recover the real data distribution. The major challenge of training GANs under limited data is that the discriminator is prone to overfitting [47, 18], and therefore lacks generalization to teach the generator to learn the underlying real data distribution.
1Our code is available at https://github.com/liang-hou/augself-gan.
In order to alleviate the overfitting issue, recent researches have suggested a variety of approaches, mainly from the perspectives of training data [18], loss functions [40], and network architectures [28]. Among them, data augmentation-based methods have gained widespread attention due to its simplicity and extensibility. Specifically, DiffAugment [47] introduced differentiable augmentation techniques for GANs, in which both real and generated data are augmented to supplement the training set of the discriminator. However, this straightforward augmentation method overlooks augmentation-related semantic information, as it solely augments the domain of the discriminator while neglecting the range. Such a practice might introduces an inductive bias that potentially forces the discriminator to remain invariant to different augmentations [24], which could limit the representation learning of the discriminator and subsequently affect the generation performance of the generator [12].
In this paper, we propose a novel augmentation-aware self-supervised discriminator that predicts the augmentation parameter of augmented data with the original data as reference to address the above problem. Meanwhile, the self-supervised discriminator is required to be distinguished between the real data and the generated data since their distributions are different during training, especially in the early stage. The proposed discriminator can benefit the generator in two ways, implicitly and explicitly. On one hand, the self-supervised discriminator can transfer the learned augmentation-aware knowledge to the original discriminator through parameter sharing. On the other hand, we allow the generator to learn adversarially from the self-supervised discriminator by generating augmentation-predictable real and not fake data (Equation (6)). We also theoretically analyzed the connection between this objective function and the minimization of a robust  $f$ -divergence divergence (the arithmetic - harmonic mean divergence [37]). In experiments, we show that the proposed method compares favorably to the data augmentation counterparts and other state-of-the-art (SOTA) methods on common data-limited benchmarks (CIFAR-10 [21], CIFAR-100 [21], FFHQ [17], LSUN-Cat [45], and five low-shot image generation datasets [36]) based on the class-conditional BigGAN [2] and unconditional StyleGAN2 [19] architectures. In addition, we carried out extensive experiments to demonstrate the effectiveness of the objective function design, the adaptability to stronger data augmentations, and the robustness of hyper-parameter selection in our method.
In this section, we provide an overview of existing work related to training GANs in data-limited scenarios. We also discuss methodologies incorporating self-supervised learning techniques.
2.1 GANs under Limited Training Data
Recently, researchers have become interested in freeing training GANs from the need to collect large amounts of data for adaptability in real-world scenarios. Previous studies typically fall into two main categories. The first one involves adopting a pre-trained GAN model to the target domain by fine-tuning partial parameters [41, 31, 42, 30]. However, it requires external training data, and the adoption performance depends heavily on the correlation between the source and target domains.
The other one focuses on training GANs from scratch with elaborated data-efficient training strategies. DiffAugment [47] utilized differentiable augmentation to supplement the training set to prevent discriminator from overfitting in limited data regimes. Concurrently, ADA [18] introduced adaptive data augmentation with a richer set of augmentation categories. APA [16] adaptively augmented the real data with the most plausible generated data. LeCam-GAN [40] proposed adaptive regularization for the discriminator and showed a connection to the Le Cam divergence [23]. [3] discovered that sparse sub-network (lottery tickets) [5] and feature-level adversarial augmentation could offer orthogonal gains to data augmentation methods. InsGen [43] improved the data efficiency of training GANs by incorporating instance discrimination tasks to the discriminator. MaskedGAN [14] employed masking in the spatial and spectral domains to alleviate the discriminator overfitting issue. GenCo [7] discriminated samples from multiple views with weight-discrepancy and data-discrepancy mechanisms. FreGAN [44] focused on discriminating between real and fake samples in the high-frequency domain. DigGAN [8] constrains the discriminator gradient gap between real and generated data. FastGAN [28] designed a lightweight generator architecture and observed that a self-supervised discriminator could enhance low-shot generation performance. Our method falls into the second category, supplementing data augmentation-based GANs and can be also applied to other methods.
$\omega_{\mathrm{color}} = (0.2, 0.3, 0.5)$
$\omega_{\mathrm{translation}} = (0.1, 0.8)$
Figure 1: Examples of images with different kinds of differentiable augmentation (including the original unaugmented one) and their re-scaled corresponding augmentation parameters  $\omega \in [0,1]^d$ .
$\omega_{\mathrm{cutout}} = (0.4, 0.6)$
2.2 GANs with Self-Supervised Learning
Self-supervised learning techniques excel at learning meaningful representations without human-annotated labels by solving pretext tasks. Transformation-based self-supervised learning methods such as rotation recognition [9] have been incorporated into GANs to address catastrophic forgetting in discriminators [4, 38, 12]. Various other self-supervised tasks have also been explored, including jigsaw puzzle solving [1], latent transformation detection [33], and mutual information maximization [25]. Moreover, ContraD [15] decouples the representation learning and discrimination of the discriminator, utilizing contrastive learning for representation learning and a discriminator head for distinguishing real from fake upon the contrastive representations. In contrast to ours, CR-GAN [46] and ICR-GAN proposed consistency regularization for the discriminator, which corresponds to an explicit augmentation-invariant of the discriminator. both our proposed method and SSGAN-LA [12] belong to adversarial self-supervised learning, they differ in the type of self-supervised signals and model inputs. SSGAN-LA is limited to categorical self-supervision [9], which is incompatible with popular augmentation-based GANs like DiffAugment [47]. Our method is applicable for continuous self-supervision and integrates seamlessly with DiffAugment. Furthermore, continuous self-supervision have a magnitude relationship and thus can provide more refined gradient feedback for the model to overcome overfitting in data-limited scenarios. Additionally, unlike SSGAN-LA, our method does not constrain the invertibility of data transformations (Theorem 1) because it additionally take the original sample as input for the self-supervised discriminator (Equation (5)).
In this section, we introduce the necessary concepts and preliminaries for completeness of the paper.
3.1 Generative Adversarial Networks
Generative adversarial networks (GANs) [10] typically contain a generator  $G: \mathcal{Z} \to \mathcal{X}$  that maps a low-dimensional latent code  $\mathbf{z} \in \mathcal{Z}$  endowed with a tractable prior  $p(\mathbf{z})$ , e.g., multivariate normal distribution  $\mathcal{N}(\mathbf{0}, \mathbf{I})$ , to a high-dimensional data point  $\mathbf{x} \in \mathcal{X}$ , which induces a generated data distribution (density)  $p_{G}(\mathbf{x}) = \int_{\mathcal{Z}} p(\mathbf{z}) \delta(\mathbf{x} - G(\mathbf{z})) \mathrm{d}\mathbf{z}$  with the Dirac delta distribution  $\delta(\cdot)$ , and also contain a discriminator  $D: \mathcal{X} \to \mathbb{R}$  that is required to distinguish between the real data sampled from the underlying data distribution (density)  $p_{\mathrm{data}}(\mathbf{x})$  and the generated ones. The generator attempts to fool the discriminator to eventually recover the real data distribution, i.e.,  $p_{G}(\mathbf{x}) = p_{\mathrm{data}}(\mathbf{x})$ . Formally, the loss functions for the discriminator and the generator can be formulated as follows:
$$
\mathcal {L} _ {D} = \mathbb {E} _ {\mathbf {x} \sim p _ {\mathrm {d a t a}} (\mathbf {x})} [ f (D (\mathbf {x})) ] + \mathbb {E} _ {\mathbf {z} \sim p (\mathbf {z})} [ h (D (G (\mathbf {z}))) ], \tag {1}
$$
$$
\mathcal {L} _ {G} = \mathbb {E} _ {\mathbf {z} \sim p (\mathbf {z})} [ g (D (G (\mathbf {z}))) ]. \tag {2}
$$
Different real-valued functions  $f, h$ , and  $g$  correspond to different variants of GANs [32]. For example, the minimax GAN [10] can be constructed by setting  $f(x) = -\log(\sigma(x))$  and  $h(x) = -g(x) = -\log(1 - \sigma(x))$  with the sigmoid function  $\sigma(x) = 1 / (1 + \exp(-x))$ . In this study, we follow the practices of DiffAugment [47] to adopt the hinge loss [26], i.e.,  $f(x) = h(-x) = \max(0, 1 - x)$  and  $g(x) = -x$ , for experiments based on BigGAN [2] and the log loss [10], i.e.,  $f(x) = g(x) = -\log(\sigma(x))$  and  $h(x) = -\log(1 - \sigma(x))$ , for experiments based on StyleGAN2 [19].
Figure 2: Comparison of representation learning ability of discriminator between BigGAN + DiffAugment and our AugSelf-BigGAN on CIFAR-10 and CIFAR-100 using linear logistic regression.
3.2 Differentiable Augmentation for GANs
DiffAugment [47] introduces differentiable augmentation  $T: \mathcal{X} \times \Omega \to \hat{\mathcal{X}}$  parameterized by a randomly-sampled parameter  $\omega \in \Omega$  with a prior  $p(\omega)$  for data-efficient GAN training. The parameter  $\omega$  determines exactly how to transfer a sample  $\mathbf{x}$  to an augmented one  $\hat{\mathbf{x}} \in \hat{\mathcal{X}}$  for the discriminator. After manually re-scaling (for  $\omega \in [0,1]^d$ ), the parameters of all three kinds of differentiable augmentation used in DiffAugment for 2D images can be expressed as follows:
color:  $\omega_{\mathrm{color}} = (\lambda_{\mathrm{brightness}},\lambda_{\mathrm{saturation}},\lambda_{\mathrm{contrast}})\in [0,1]^3$
- translation:  $\omega_{\text{translation}} = (x_{\text{translation}}, y_{\text{translation}}) \in [0,1]^2$ ;
- cutout:  $\omega_{\mathrm{cutout}} = (x_{\mathrm{offset}},y_{\mathrm{offset}})\in [0,1]^2$
Figure 1 illustrates the augmentation operations and their parameters. Formally, the loss functions for the discriminator and the generator of GANs with DiffAugment are defined as follows:
$$
\mathcal {L} _ {D} ^ {\mathrm {d a}} = \mathbb {E} _ {\mathbf {x} \sim p _ {\text {d a t a}} (\mathbf {x}), \omega \sim p (\omega)} [ f (D (T (\mathbf {x}; \omega))) ] + \mathbb {E} _ {\mathbf {z} \sim p (\mathbf {z}), \omega \sim p (\omega)} [ h (D (T (G (\mathbf {z}); \omega))) ], \tag {3}
$$
$$
\mathcal {L} _ {G} ^ {\mathrm {d a}} = \mathbb {E} _ {\mathbf {z} \sim p (\mathbf {z}), \omega \sim p (\omega)} [ g (D (T (G (\mathbf {z}); \omega))) ], \tag {4}
$$
where  $\omega$  can represent any combination of these parameters. We choose all augmentations by default, which means augmentation color, translation, and cutout are adopted for each image sequentially.
Data augmentation for GANs allows the discriminator to distinguish a single sample from multiple perspectives by transforming it into various augmented samples according to different augmentation parameters. However, it overlooks the differences in augmentation intensity, such as color contrast and translation magnitude, leading the discriminator to implicitly maintain invariance to these varying intensities. The invariance may limit the representation learning ability of the discriminator because it loses augmentation-related information (e.g., color and position) [24]. Figure 2 confirms the impact of this point on the discriminator representation learning task [4]. We argue that a discriminator that captures comprehensive representations contributes to better convergence of the generator [35, 22]. Moreover, data augmentation may lead to augmentation leaking in generated data, when using specific data augmentations such as random 90-degree rotations [18, 12]. Therefore, our goal is to eliminate the unnecessary potential inductive bias (invariance to augmentations) for the discriminator while preserving the benefits of data augmentation for training data-efficient GANs.
To achieve this goal, we propose a novel augmentation-aware self-supervised discriminator  $\hat{D}:\hat{\mathcal{X}}\times \mathcal{X}\to \Omega^{+}\cup \Omega^{-}$  that predicts the augmentation parameter and authenticity of the augmented data given the original data as reference. Distinguishing between the real data and the generated data with different self-supervision is because they are different during training, especially in the early stage. Specifically, the predictive targets of real data and generated data are represented as  $\omega^{+}\in \Omega^{+}$  and  $\omega^{-}\in \Omega^{-}$ , respectively. They are constructed from the augmentation parameter  $\omega$  with different transformations, i.e.,  $\omega^{+} = -\omega^{-} = \omega$ . Since the augmentation parameter is a continuous
Figure 3: Diagram of AugSelf-GAN. The original augmentation-based discriminator is  $D(T(\cdot)) = \psi(\phi(T(\cdot)))$ . The augmentation-aware self-supervised discriminator is  $\hat{D}(T(\cdot), \cdot) = \varphi(\phi(T(\cdot)) - \phi(\cdot))$ , where  $\varphi$  is our newly introduced linear layer with negligible additional parameters.
vector, we use mean squared error loss to regress it. The proposed method combines continuous self-supervised signals with real-vs-fake discrimination signals, thus can be considered as soft-label augmentation [12]. Comparison with self-supervision that does not distinguish between real and fake is referred to Table 6 in Appendix C. Notice that the predictive targets (augmentations) can be a subset of performed augmentations (see Table 7 in Appendix C for comparison). Mathematically, the loss function for the augmentation-aware self-supervised discriminator is formulated as the following:
$$
\mathcal {L} _ {\hat {D}} ^ {\mathrm {s s}} = \mathbb {E} _ {\mathbf {x}, \boldsymbol {\omega}} \left[ \| \hat {D} (T (\mathbf {x}; \boldsymbol {\omega}), \mathbf {x}) - \boldsymbol {\omega} ^ {+} \| _ {2} ^ {2} \right] + \mathbb {E} _ {\mathbf {z}, \boldsymbol {\omega}} \left[ \| \hat {D} (T (G (\mathbf {z}); \boldsymbol {\omega}), G (\mathbf {z})) - \boldsymbol {\omega} ^ {-} \| _ {2} ^ {2} \right]. \tag {5}
$$
In our implementations, the proposed self-supervised discriminator  $\hat{D} = \varphi \circ \phi$  shares the backbone  $\phi : \mathcal{X} \to \mathbb{R}^d$  with the original discriminator  $D = \psi \circ \phi$  except the output linear layer  $\varphi : \mathbb{R}^d \to \Omega^+ \cup \Omega^-$ . This parameter-sharing design not only improves the representation learning ability of the original discriminator but also saves the number of parameters in our model compared to the base model, e.g.,  $0.04\%$  more parameters in BigGAN and  $0.01\%$  in StyleGAN2. More specifically, the self-supervised discriminator predicts the target based on the difference between learned representations of the augmented data and the original data, i.e.,  $\hat{D}(T(\mathbf{x};\omega),\mathbf{x}) = \varphi(\phi(T(\mathbf{x};\omega)) - \phi(\mathbf{x}))$  (see Table 8 in Appendix C for comparison with other architectures). The philosophy behind our design is that the backbone  $\phi$  should capture rich (which necessitates the design of a simple head  $\varphi$ ) and linear (inspiring us to perform subtraction on the features) representations.
In order for the generator to directly benefit from the self-supervision of data augmentation, we establish a novel adversarial game between the augmentation-aware self-supervised discriminator and the generator with the objective function for the generator defined as follows:
$$
\mathcal {L} _ {G} ^ {\mathrm {s s}} = \mathbb {E} _ {\mathbf {z}, \omega} \left[ \| \hat {D} (T (G (\mathbf {z}); \omega), G (\mathbf {z})) - \omega^ {+} \| _ {2} ^ {2} \right] - \mathbb {E} _ {\mathbf {z}, \omega} \left[ \| \hat {D} (T (G (\mathbf {z}); \omega), G (\mathbf {z})) - \omega^ {-} \| _ {2} ^ {2} \right]. \tag {6}
$$
The objective function is actually the combination of the non-saturating loss (regarding the generated data as real,  $\min_G\mathbb{E}_{\mathbf{z},\omega}[\| \hat{D} (T(G(\mathbf{z});\omega),G(\mathbf{z})) - \omega^+ \| _2^2 ])$  and the saturating loss (reversely optimizing the objective function of the discriminator,  $\max_G\mathbb{E}_{\mathbf{z},\omega}[\| \hat{D} (T(G(\mathbf{z});\omega),G(\mathbf{z})) - \omega^{-}\| _2^2 ])$  (see Table 9 in Appendix C for ablation). Intuitively, the non-saturating loss encourages the generator to produce augmentation-predictable data, facilitating fidelity but reducing diversity. Conversely, the saturating loss strives for the generator to avoid generating augmentation-predictable data, promoting diversity at the cost of fidelity. We will elucidate in Section 5 how this formalization assists the generator in matching the fidelity and diversity of real data, ultimately leading to an accurate approximation of the target data distribution.
The total objective functions for the original discriminator, the augmentation-aware self-supervised discriminator, and the generator of our proposed method, named AugSelf-GAN, are given by:
$$
\min  _ {D, \hat {D}} \mathcal {L} _ {D} ^ {\mathrm {d a}} + \lambda_ {d} \cdot \mathcal {L} _ {\hat {D}} ^ {\mathrm {s s}}, \tag {7}
$$
$$
\min  _ {G} \mathcal {L} _ {G} ^ {\mathrm {d a}} + \lambda_ {g} \cdot \mathcal {L} _ {G} ^ {\mathrm {s s}}, \tag {8}
$$
Figure 4: Comparison of the function  $f$  in different  $f$ -divergences. The  $f$ -divergence between two probability distributions  $p(\mathbf{x})$  and  $q(\mathbf{x})$  is defined as  $D_{f}(p(\mathbf{x})\| q(\mathbf{x})) = \int_{\mathcal{X}}q(\mathbf{x})f(p(\mathbf{x}) / q(\mathbf{x}))\mathrm{d}\mathbf{x}$  with a convex function  $f:\mathbb{R}_{\geq 0}\to \mathbb{R}$  satisfying  $f(1) = 0$ . The x- and y-axis denote the input and the value of the function  $f$  in the  $f$ -divergence. The function  $f$  of the AHM divergence yields the most robust value for large inputs.
where the hyper-parameters are set as  $\lambda_d = \lambda_g = 1$  in experiments by default unless otherwise specified (see Figure 6 for empirical studies). Details of objective functions are referred to Appendix B.
In this section, we analyze the connection between the theoretical learning objective of AugSelf-GAN and the arithmetic - harmonic mean (AHM) divergence [37] under certain assumptions.
Proposition 1. For any generator  $G$  and given unlimited capacity in the function space, the optimal augmentation-aware self-supervised discriminator  $\hat{D}^*$  has the form of:
$$
\hat {D} ^ {*} (\hat {\mathbf {x}}, \mathbf {x}) = \frac {\int p _ {\mathrm {d a t a}} (\mathbf {x} , \boldsymbol {\omega} , \hat {\mathbf {x}}) \boldsymbol {\omega} ^ {+} \mathrm {d} \boldsymbol {\omega} + \int p _ {G} (\mathbf {x} , \boldsymbol {\omega} , \hat {\mathbf {x}}) \boldsymbol {\omega} ^ {-} \mathrm {d} \boldsymbol {\omega}}{p _ {\mathrm {d a t a}} (\mathbf {x} , \hat {\mathbf {x}}) + p _ {G} (\mathbf {x} , \hat {\mathbf {x}})} \tag {9}
$$
The proofs of all theoretical results (including the following ones) are deferred in Appendix A.
Theorem 1. Assume that  $\omega^{+} = -\omega^{-} = c$  is constant, under the optimal self-supervised discriminator  $\hat{D}^*$ , optimizing the self-supervised task for the generator  $G$  is equivalent to:
$$
\min  _ {G} 4 c \cdot M _ {\mathrm {A H}} \left(p _ {\text {d a t a}} (\mathbf {x}, \hat {\mathbf {x}}) \| p _ {G} (\mathbf {x}, \hat {\mathbf {x}})\right), \tag {10}
$$
where  $c = \|c\|_2^2$  is constant and  $M_{\mathrm{AH}}$  is the arithmetic - harmonic mean divergence [37], of which the minimum is achieved if and only if  $p_G(\mathbf{x}, \hat{\mathbf{x}}) = p_{\mathrm{data}}(\mathbf{x}, \hat{\mathbf{x}}) \Rightarrow p_G(\mathbf{x}) = p_{\mathrm{data}}(\mathbf{x})$ .
Theorem 1 reveals that the generator of AugSelf-GAN theoretically still satisfies generative modeling, i.e., accurately learning the target data distribution, under certain assumptions. Although AugSelf-GAN does not obey the strict assumption, we note that this is not rare in the literature.2 Under this assumption, AugSelf-GAN can be regarded as a multi-dimensional extension of LS-GAN [29] in terms of the loss function, while excluding that of the generator. Additionally, our analysis offers an alternative theoretically grounded generator loss function for the LS-GAN family.3
Corollary 1. The following equality and inequality hold for the AHM divergence: