SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yite Wang $^{1\dagger}$ , Jing Wu $^{1\dagger}$ , Naira Hovakimyan $^{1}$ , Ruoyu Sun $^{2,3*}$
<sup>1</sup>University of Illinois Urbana-Champaign, USA
$^{2}$ School of Data Science, The Chinese University of Hong Kong, Shenzhen, China
<sup>3</sup>Shenzhen International Center for Industrial and Applied Mathematics, Shenzhen Research Institute of Big Data
{yitew2, jingwu6, nhovakim}@illinois.edu, sunruoyu@cuhk.edu.cn
Over the past few years, there has been growing interest in developing larger and deeper neural networks, including deep generative models like generative adversarial networks (GANs). However, GANs typically come with high computational complexity, leading researchers to explore methods for reducing the training and inference costs. One such approach gaining popularity in supervised learning is dynamic sparse training (DST), which maintains good performance while enjoying excellent training efficiency. Despite its potential benefits, applying DST to GANs presents challenges due to the adversarial nature of the training process. In this paper, we propose a novel metric called the balance ratio (BR) to study the balance between the sparse generator and discriminator. We also introduce a new method called balanced dynamic sparse training (ADAPT), which seeks to control the BR during GAN training to achieve a good trade-off between performance and computational cost. Our proposed method shows promising results on multiple datasets, demonstrating its effectiveness. Our code is available at https://github.com/YiteWang/ADAPT.
Generative adversarial networks (GANs) [20, 7, 70, 44] are a type of generative model that has gained significant attention in recent years due to their impressive performance in image-generation tasks. However, the mainstream models in GANs are known to be computationally intensive, making them challenging to train in resource-constrained settings. Therefore, it is crucial to develop methods that can effectively reduce the computational cost of training GANs while maintaining their performance, making GANs more practical and applicable in real-world scenarios.
Neural network pruning has recently emerged as a powerful tool to reduce the training and inference costs of DNNs for supervised learning. There are mainly three genres of pruning methods, namely pruning-at-initialization, pruning-during-training, and post-hoc pruning methods. Post-hoc pruning [32, 42, 25] can date back to the 1980s, which was first introduced for reducing inference time and memory requirements for efficient deployment; hence does not align with our purpose of efficient training. Later, pruning-at-initialization [46, 78, 75] and pruning-during-training methods [84] were introduced to circumvent the need to fully train the dense networks. However, early pruning-during-training algorithms [58] do not bring much training efficiency compared to post-hoc pruning, while pruning-at-initialization methods usually suffer from significant performance drop [18]. Recently, advances in dynamic sparse training (DST) [62, 16, 51, 52, 54] for the first time show that pruning-during-training methods can have comparable training FLOPs as pruning-at-initialization methods
Equal contribution
while having competing performance to post-hoc pruning. Therefore, applying DST on GANs seems to be a promising choice.
Although DST has attained remarkable achievements in supervised learning, the application of DST on GANs is not successful due to newly emerging challenges. One challenge is keeping the generator and the discriminator balanced. In particular, using overly strong discriminators can lead to overfitting, while weaker discriminators may fail to effectively prevent mode collapse [3, 4]. Hence, balancing the sparse generator and the (possibly) sparse discriminator throughout training is even more difficult. To mitigate the unbalance issue, a recent work STU-GAN [53] proposes to apply DST directly to the generator. However, we find empirically that such an algorithm is likely to fail when the generator is already more powerful than the discriminator. Consequently, it remains unclear how to conduct balanced dynamic sparse training for GANs.
To this end, we propose a metric called balance ratio (BR), which measures the degree of balance of the two components, to study sparse GAN training. We find that BR is useful in (1) understanding the interaction between the discriminator and the generator, (2) identifying the cause of a certain training failure/collapse [7, 8], and (3) helping stabilize sparse GAN training as an indicator. To our best knowledge, this is the first study to quantify the unbalance of sparse GANs and may even provide new insights into dense GAN training.
Furthermore, using BR as an indicator, we propose bAlancedDynAmic sParse Training (ADAPT) to adjust the density and the connections of the discriminator automatically during training.
Our main contributions are summarized below:
- We introduce a novel quantity named balance ratio to study the degree of balance in sparse GAN training.
- We find empirically that the balance ratio is problematic in certain practical training scenarios and that existing methods are inadequate for resolving this issue.
- We propose ADAPT, which makes real-time monitoring of the balance ratio. By dynamically adjusting the discriminator, ADAPT enables effective control of the balance ratio throughout training. Empirically, ADAPT achieves a good trade-off between performance and computational cost on several datasets.
2.1 Neural network pruning
In deep learning, efficiency is achieved through several methods. This paper primarily focuses on model training and inference efficiency, which is different from techniques for data efficiency [83, 85, 86]. These include neural architecture search (NAS) [80, 49] to discover optimal network structures, quantization [31, 67] for computational efficacy, knowledge distillation [28] to leverage the knowledge of larger models for smaller counterparts, and neural network pruning to remove unnecessary connections. Among these, neural network pruning is the focal point of our research. More specifically, we narrow our focus on unstructured pruning [25, 17], where individual weight is the finest resolution. This contracts with structured pruning [56, 59, 57, 30] where entire neurons or channels are pruned.
Post-hoc pruning. Post-hoc pruning method prunes weights of a fully-trained neural network. It usually requires high computational costs due to the multiple rounds of the train-prune-retrain procedure [25, 68]. Some use specific criteria [42, 26, 25, 23, 15, 12, 91, 63] to remove weights, while others perform extra optimization iterations [77]. Post-hoc pruning was initially proposed to reduce the inference time, while lottery ticket works [17, 68] aimed to mine trainable sub-networks.
Pruning-at-initialization methods. SNIP [46] is one of the pioneering works that aim to find trainable sub-networks without any training. Some follow-up works [78, 75, 13, 65, 1] aim to propose different metrics to prune networks at initialization. Among them, Synflow [75], SPP [45], and FORCE [13] try to address the problem of layer collapse during pruning. NTT [55], PHEW [65], and NTK-SAP [82] draw inspiration from neural tangent kernel theory.
Pruning-during-training methods. Another genre of pruning algorithms prunes or adjusts DNNs throughout training. Early works add explicit  $\ell_0$  [58] or  $\ell_1$  [84] regularization terms to encourage a
sparse solution, hence mitigating performance drop incurred by post-hoc pruning. Later works learn the subnetworks structures through projected gradient descent [94] or trainable masks [73, 88, 35, 41, 50, 71]. However, these pruning-during-training methods often do not introduce memory sparsity during training. As a remedy, DST methods [5, 62, 64, 14, 16, 51, 52, 54, 21] were introduced to train the neural networks under a given parameter budget while allowing mask change during training.
2.2 Generative adversarial networks
Generative adversarial networks (GANs). GANs [20] have drawn considerable attention and have been widely investigated for years. Deep convolutional GANs [66] replace fully-connected layers in the generator and the discriminator. Follow-up works [22, 36, 7, 92] employed more advanced methods to improve the fidelity of generated samples. After that, several novel loss functions [69, 60, 2, 22, 74], normalization and regularization methods [61, 76, 87] were proposed to stabilize the adversarial training. Besides the efforts devoted to training GANs, image-to-image translation is also extensively explored [96, 95, 10, 38, 43, 81].
GAN balance. Addressing the balance between the generator and discriminator in GAN training has been the focus of various works. However, directly applying existing methods to sparse GAN training poses challenges. For instance, [3, 4] offer theoretical analyses on the issue of imbalance but may have limited practical benefits, e.g., they require training multiple generators and discriminators. Empirically, BEGAN [6] proposes to use proportional control theory to maintain a hyper-parameter  $\frac{\mathbb{E}[|G(z) - D(G(z))|^{\eta}]}{\mathbb{E}[|x - D(x)|^{\eta}]}$ , but it is only applicable when the discriminator is an auto-encoder. Unbalanced GAN [24] pretrains a VAE to initialize the generator, which may only address the unbalance near initialization. GCC [48] considers the balance during GAN compression, while its criterion requires a trained (dense) GAN, which is not given in the DST setting. Finally, STU-GAN [53] proposes to use DST to address the unbalance issues but may fail under certain conditions, as demonstrated in our experiments.
GAN compression and pruning. One of the promising ways is based on neural architecture search and distillation algorithm [47, 19, 29]. Another part of the work applied pruning-based methods for generator compression [72, 90, 33]. Later, works by [79] presented a unified framework by combing the methods mentioned above. Nevertheless, they only focus on the pruning of generators, thus potentially posing a negative influence on Nash Equilibrium between generators and discriminators. GCC [48] compresses both components of GANs by letting the student GANs also learn the losses. Another line of work [34, 9, 8] tries to test the existence of lottery tickets in GANs. To the best of our knowledge, STU-GAN [53] is the only work that tries to directly train sparse GANs from scratch.
3 Preliminary and setups
Generative adversarial networks (GANs) have two fundamental components, a generator  $G(\cdot; \pmb{\theta}_G)$  and a discriminator  $D(\cdot; \pmb{\theta}_D)$ . Specifically, the generator maps a sampled noise  $z$  from a multivariate normal distribution  $p(z)$  into a fake image to cheat the discriminator. In contrast, the discriminator distinguishes the generator's output and the real images  $x_r$  from the distribution  $p_{\mathrm{data}}(x_r)$ .
Formally, the optimization objective of the two-player game defined in GANs can be generally defined as follows:
$$
\mathcal {L} _ {D} \left(\boldsymbol {\theta} _ {D}, \boldsymbol {\theta} _ {G}\right) = \mathbb {E} _ {\boldsymbol {x} _ {r} \sim p _ {\mathrm {d a t a}}} \left[ f _ {1} \left(D \left(\boldsymbol {x} _ {r}; \boldsymbol {\theta} _ {D}\right)\right) \right] + \mathbb {E} _ {\boldsymbol {z} \sim p} \left[ f _ {2} \left(D \left(G \left(\boldsymbol {z}; \boldsymbol {\theta} _ {G}\right)\right)\right) \right] \tag {1}
$$
$$
\mathcal {L} _ {G} \left(\boldsymbol {\theta} _ {G}\right) = \mathbb {E} _ {\boldsymbol {z} \sim p} \left[ g _ {1} \left(D \left(G \left(\boldsymbol {z}; \boldsymbol {\theta} _ {G}\right)\right)\right) \right]. \tag {2}
$$
To be more specific, different losses can be used, including the loss in the original JS-GAN [20] where  $f_{1}(x) = -\log (x)$ ,  $f_{2}(x) = -g_{1}(x) = -\log (1 - x)$ ; Wasserstein loss [22] where  $f_{1}(x) = -f_{2}(x) = g_{1}(x) = -x$ ; and hinge loss [61] where  $f_{1}(x) = \max (0,1 - x)$ ,  $f_{2}(x) = \max (0,1 + x)$ , and  $g_{1}(x) = -x$ . The two components are optimized alternately to achieve the Nash equilibrium.
Figure 1: FID  $(\downarrow)$  comparison of SDST against STATIC sparse training for SNGAN on CIFAR-10 with different sparsity ratio combinations. The shaded areas denote the standard deviation.
GAN sparse training. In this work, we are interested in sparse training for GANs. In particular, the objective of sparse GAN training can be formulated as follows:
$$
\boldsymbol {\theta} _ {D} ^ {*} = \min  _ {\boldsymbol {\theta} _ {D}} \mathcal {L} _ {D} \left(\boldsymbol {\theta} _ {D} \odot \boldsymbol {m} _ {D}, \boldsymbol {\theta} _ {G} \odot \boldsymbol {m} _ {G}\right)
$$
$$
\boldsymbol {\theta} _ {G} ^ {*} = \min  _ {\boldsymbol {\theta} _ {G}} \mathcal {L} _ {G} (\boldsymbol {\theta} _ {G} \odot \boldsymbol {m} _ {G})
$$
s.t.  $\pmb{m}_D\in \{0,1\}^{|\pmb{\theta}_D|},\quad \pmb{m}_G\in \{0,1\}^{|\pmb{\theta}_G|},\quad \| \pmb{m}_D\| _0 / |\pmb{\theta}_D|\leq d_D,\quad \| \pmb{m}_G\| _0 / |\pmb{\theta}_G|\leq d_G,$
where  $\odot$  is the Hadamard product;  $\pmb{\theta}_{D}^{*}$ ,  $\pmb{m}_{D}$ ,  $|\pmb{\theta}_{D}|$ ,  $d_{D}$  are the sparse solution, mask, number of parameters, and target density for the discriminator, respectively. The corresponding variables for the generator are denoted with subscript  $G$ . For pruning-at-initialization methods, masks  $\pmb{m}$  are determined before training, whereas  $\pmb{m}$  are dynamically adjusted for dynamic sparse training (DST) methods.
Dynamic sparse training (DST). DST methods [62, 16] usually start with a sparse network parameterized by  $\theta \odot m$  with randomly initialized mask  $m$ . After a constant time interval  $\Delta T$ , it updates mask  $m$  by removing a fraction of connections and activating new ones with a certain criterion. The total number of active parameters  $\| m \|_0$  is hence kept under a certain threshold  $d|\theta|$ . Please see Appendix B for more details.
4 Motivating observations: The unbalance in sparse GAN training
As discussed in section 1, it is essential to maintain the balance of generator and discriminator during GAN training. As strong discriminators may lead to over-fitting, whereas weak discriminators may be unable to detect mode collapse. When it comes to sparse GAN training, the consequences caused by the unbalance can be further amplified. For example, sparsifying a weak generator while keeping the discriminator unmodified may lead to an even more unbalanced worst-case scenario.
To support our claim, we conduct experiments with SNGAN [61] on the CIFAR-10 dataset. We consider the following sparse training algorithms:
$\bullet$  Static sparse training (STATIC). For STATIC, layer-wise sparsity ratio and masks  $m_G, m_D$  are fixed throughout the training.
$\Theta$  Single dynamic sparse training (SDST). SDST is a direct application of the DST method on GANs where only the generator dynamically adjusts masks during the training. We name such method SDST as only one component of the GAN, i.e., the generator, is dynamic. Furthermore, we call the variant which grows connections based on gradients magnitude as  $\clubsuit$  SDST-RigL [16], and randomly as  $\clubsuit$  SDST-SET [62]. Note that STU-GAN [53] is almost identical to  $\clubsuit$  SDST-RigL with EMA [89] tailored for DST. We do not consider naively applying DST on both generators and discriminators, as in STU-GAN, it is empirically shown that simply adjusting both components generates worse performance with more severe training instability. $^{2}$
We test the considered algorithms with  $d_{G} \in \{10\%, 20\%, 30\%, 50\% \}$  and  $d_{D} \in \{10\%, 20\%, 30\%, 50\%, 100\% \}$ . More experiment details can be found at Appendix A and Appendix B.
2We also perform a small experiment in subsection E.2 to validate their findings.
Figure 2: BR comparison of SDST against STATIC sparse training for SNGAN on CIFAR-10 with different sparsity ratio combinations.
4.1 Key observations
We report the results in Figure 1 and summarize our critical findings as follows:
$\bullet$  Observation 1: Neither strong nor weak sparse discriminators can provide satisfactory results. The phenomenon is most noticeable when  $d_{G} = 10\%$ , where the FID initially decreases but then increases. The reasons may be as follows: (1) Overly weak discriminators may cause training collapse as they cannot provide useful information to the generator, resulting in a sudden increase in FID at the early stage of sparse GAN training. (2) Overly strong discriminators may not yield good FID results because they learn too quickly, not allowing the generator to keep up. Hence, to ensure a balanced training of GAN for sparse training methods, it is crucial to find an appropriate sparsity ratio for the discriminator.
$\Theta$  Observation 2: SDST is unable to give stable performance boost compared to the STATIC baseline. Another critical observation is that SDST is better than STATIC only when the discriminator is strong enough. More specifically, for all selected discriminator density ratios, SDST method is not better than STATIC when using a small discriminator density ( $d_{D} = 10\%$ ). On the contrary, for the cases where  $d_{D} > d_{G}$ , we generally see a significant performance boost brought by SDST.
5 Balance ratio: Towards quantifying the unbalance in sparse GAN training
5.1 Formulation of the balance ratio
To gain a deeper understanding of the phenomenon observed in the previous section, and to better monitor and control the degree of unbalance in sparse GANs, we introduce a novel quantity called the balance ratio (BR). This quantity is defined as follows.
At each training iteration, we draw random noise  $z$  from a multivariate normal distribution and real images  $\boldsymbol{x}_r$  from the training set. We denote the discriminator after gradient descent update as  $D(\cdot; \boldsymbol{\theta}_D)$ . We denote generator before and after gradient descent training as  $G^{\mathrm{pre}}(\cdot; \boldsymbol{\theta}_G)$  and
Figure 4: FID  $(\downarrow)$  of STATIC sparsely trained SNGAN with and without DDA on CIFAR-10 with different sparsity ratio combinations. The result of DDA is independent of  $d_{D}$  as it is determined automatically. The shaded areas denote the standard deviation.
$G^{\mathrm{post}}(\cdot ;\pmb{\theta}_G^{\prime})$  , respectively. Then the balance ratio is defined as:
$$
\mathrm {B R} = \frac {\mathbb {E} _ {\boldsymbol {z} \sim p} \left[ D \left(G ^ {\text {p o s t}} (\boldsymbol {z})\right) - D \left(G ^ {\text {p r e}} (\boldsymbol {z})\right) \right]}{\mathbb {E} _ {\boldsymbol {x} _ {r} \sim p _ {\text {d a t a}}} \left[ D \left(\boldsymbol {x} _ {r}\right) \right] - \mathbb {E} _ {\boldsymbol {z} \sim p} \left[ D \left(G ^ {\text {p r e}} (\boldsymbol {z})\right) \right]} = \frac {\alpha}{\beta}. \tag {3}
$$
Precisely, BR measures how much improvement the generator can achieve in the scale measured by the discriminator for a specific random noise  $z$ . When BR is too small (e.g.,  $\mathrm{BR} < 30\%$ ), the updated generator is too weak to trick the discriminator, as the generated images are still considered fake. Similarly, for the case where BR is too large (e.g.,  $\mathrm{BR} > 80\%$ ), the discriminator is considered too weak hence it may not provide useful information to the generator. We also illustrate BR in Figure 3.
Figure 3: Illustration of balance ratio.
5.2 Understanding observation 1: Analysing GAN balance with the balance ratio
We visualize the BR evolution throughout the training for the experiments in section 4 to show the effectiveness of BR in quantifying the balance of sparse GANs. We show the results in Figure 2.
It illustrates that BR can distinguish the density difference (hence the representation power difference) of the discriminator. Specifically, we can see that for larger discriminator density  $d_{D}$ , the BR is much lower throughout the training, indicating strong discriminators. On the contrary, for the cases where the discriminators are too weak compared to the generators, e.g., all cases where  $d_{D} = 10\%$ , we can observe BR first increases and then oscillates wildly. We believe this oscillatory behavior is related to the training collapse. Empirical results also show that the FID metric experiences a sudden increase after this turning point.
5.3 Dynamic density adjust: A first attempt to utilize the balance ratio
As demonstrated in the previous section, the balance ratio (BR) effectively captures the degree of balance between the generators and discriminators in sparse GANs. Hence, it is natural to leverage BR to dynamically adjust the density of discriminators during sparse GAN training so that a reasonable discriminator density can be found.
To demonstrate the value of BR, we propose a simple yet powerful modification to the STATIC baseline. This method, which we call dynamic density adjust (DDA), is explained below. Specifically, we initialize the initial density of the discriminator  $d_{D}^{\mathrm{init}} = d_{G}$ . After a specific training iteration interval, we adjust the density of the discriminator based on the BR over the last few iterations with a pre-defined density increment  $\Delta d$ . With a pre-defined BR bounds  $[B_{-}, B_{+}]$ , we decrease  $d_{D}$  by  $\Delta d$  when BR is smaller than  $B_{-}$ , and vise versa. We show the algorithm in Appendix C Algorithm 1.
Comparison to ADA [37]. In this paragraph, we compare ADA and DDA. (1) Notice that DDA algorithm is orthogonal to ADA in a sense that StyleGAN2-ADA adjusts the data augmentation
probability while DDA adjusts the discriminator density. (2) Moreover, the criterion used in DDA, i.e. BR, is very different from the criterion proposed in StyleGAN2-ADA, i.e.  $r_v = \frac{\mathbb{E}[D(x_{\mathrm{train}})] - \mathbb{E}[D(x_{\mathrm{val}})]}{\mathbb{E}[D(x_{\mathrm{train}})] - \mathbb{E}[D(G(z))]}$  and  $r_t = \mathbb{E}[\mathrm{sign}(D(x_{\mathrm{train}}))]$ . In particular,  $r_v$  requires a separate validation set, while  $r_t$  only quantifies the overfitting of the discriminator to the training set. (3) Another note is that DDA is a flexible framework, where its criterion, i.e. BR, can be potentially replaced by  $r_v$ ,  $r_t$ , and such.
Experiment results. We test DDA with target BR interval  $[B_{-}, B_{+}] = [0.5, 0.65]$ . Precisely, DDA tends to find a suitable discriminator where the generator can just trick the discriminator throughout the training. We show the results in Figure 4 with red lines. The experiments show that DDA can identify reasonable discriminator densities to enable balanced training for sparse GANs.
5.4 Understanding observation 2: Analysing the failure of SDST with the balance ratio
By leveraging BR, we can also gain further insights into why some configurations do not benefit from SDST as compared to STATIC.
Regarding SDST as a way of increasing the generator capacity. Our findings suggest that SDST can possibly enhance the generator's representation power, as demonstrated by the higher BR values compared to STATIC observed in Figure 2. We attribute this effect to the in-time over-parameterization (ITOP) [54] induced by dynamic sparse training.
SDST does not address training collapse. The increase in generator's representation power resulting from SDST is only beneficial when the discriminator has matching or superior representation power. Therefore, if the training has already collapsed for the static baseline methods (STATIC), meaning that the generator is already stronger than the discriminator, SDST may not be effective in stabilizing sparse GAN training. This is evident from the results shown in Figure 2 first-row column 1, second-row column 1, third-row columns 1-2, and fourth-row columns 1-3.
Despite the superior performance of STU-GAN (or SDST in general) at higher discriminator density ratios  $d_{D}$ , there exist some limitations for SDST, which we summarize below: