Generative Hierarchical Features from Synthesizing Images

Introduction

Representation learning plays an essential role in the rise of deep learning. The learned representation is able to express the variation factors of the complex visual world. Accordingly, the performance of a deep learning algorithm highly depends on the features extracted from the input data. As pointed out by Bengio et al. , a good representation is expected to have the following properties. First, it should be able to capture multiple configurations from the input. Second, it should organize the explanatory factors of the input data as a hierarchy, where more abstract concepts are at a higher level. Third, it should have strong transferability, not only from datasets to datasets but also from tasks to tasks.

Deep neural networks supervisedly trained for image classification on large-scale datasets (e.g., ImageNet and Places ) have resulted in expressive and discriminative visual features . However, the developed features are heavily dependent on the training objective. For example, prior work has shown that deep features trained for the object recognition task may mainly focus on the shapes and parts of the objects while remain invariant to rotation , and the deep features from a scene classification model may focus more on detecting the categorical objects (e.g., bed for bedroom and sofa for living room) . Thus the discriminative features learned from solving high-level image classification tasks might not be necessarily good for other mid-level and low-level tasks, limiting their transferability . Besides, it remains unknown how the discriminative features can be used in generative applications like image editing.

Generative Adversarial Network (GAN) has recently made great process in synthesizing photo-realistic images. It considers the image generation task as the training supervision to learn the underlying distribution of real data. Through adversarial training, the generator can capture the multi-level variations underlying the input data to the most extent, otherwise, the discrepancy between the real and synthesized data would be spotted by the discriminator. The recent state-of-the-art StyleGAN has been shown to encode rich hierarchical semantics in its layer-wise representations . However, the generator is primarily designed for image generation and hence lacks the inference ability of taking an image as the input and extracting its visual feature, which greatly limits the applications of GANs to real images.

To solve this problem, a common practice is to introduce an additional encoder into the two-player game described in GANs . Nevertheless, existing encoders typically choose the initial latent space (i.e., the most abstract level feature) as the target representation space, omitting the pre-layer information learned by the generator. On the other hand, the transferability of the representation from GAN models is not fully verified in the literature. Most prior work focuses on learning discriminative features for the high-level image classification task yet put little effort on other mid-level and low-level downstream tasks, such as landmark detection and layout prediction.

In this work, we show that the pre-trained GAN generator can be considered as a learned loss function. Training with it can bring highly competitive hierarchical visual features which are generalizable to various tasks. Based on the StyleGAN model, we tailor a novel hierarchical encoder whose outputs align with the layer-wise representations from the generator. In particular, the generator takes the feature hierarchy produced by the encoder as the per-layer inputs and supervises the encoder via reconstructing the input image. We evaluate such visual features, termed as Generative Hierarchical Features (GH-Feat), on both generative and discriminative tasks, including image editing, image harmonization, image classification, face verification, landmark detection, layout prediction, etc. Extensive experiments validate that the generative feature learned from solving the image synthesis task has compelling hierarchical and transferable properties, facilitating many downstream applications.

Related Work

Visual Features. Visual Feature plays a fundamental role in the computer vision field. Traditional methods used manually designed features for pattern matching and object detection. These features are significantly improved by deep models , which automatically learn the feature extraction from large-scale datasets. However, the features supervisedly learned for a particular task could be biased to the training task and hence become difficult to transfer to other tasks, especially when the target task is too far away from the base task . Unsupervised representation learning is widely explored to learn a more general and transferable feature . However, most of existing unsupervised feature learning methods focus on evaluating their features on the tasks of image recognition, yet seldom evaluate them on other mid-level or low-level tasks, let alone generative tasks. Shocher et al. discover the potential of discriminative features in image generation, but the transferability of these features are still not fully verified.

Generative Adversarial Networks. GANs are able to produce photo-realistic images via learning the underlying data distribution. The recent advance of GANs has significantly improved the synthesis quality. StyleGAN proposes a style-based generator with multi-level style codes and achieves the start-of-the-art generation performance. However, little work explores the representation learned by GANs as well as how to apply such representation for other applications. Some recent work interprets the semantics encoded in the internal representation of GANs and applies them for image editing . But it remains much less explored whether the learned GAN representations are transferable to discriminative tasks.

Adversarial Representation Learning. The main reason of hindering GANs from being applied to discriminative tasks comes from the lack of inference ability. To fill this gap, prior work introduces an additional encoder to the GAN structure . Donahue and Simonyan and Pidhorskyi et al. extend this idea to the state-of-the-art BigGAN and StyleGAN models respectively. In this paper, we also study the representation learning using GANs, with following improvements compared to existing methods. First, we propose to treat the well-trained StyleGAN generator as a learned loss function. Second, instead of mapping the images to the initial GAN latent space, like most algorithms have done, we design a novel encoder to produce hierarchical features that well align with the layer-wise representation learned by StyleGAN. Third, besides the image classification task that is mainly targeted at by prior work , we validate the transferability of our proposed GH-Feat on a range of generative and discriminative tasks, demonstrating its generalization ability.

Methodology

We design a novel encoder to extract hierarchical visual features from the input images. This encoder is trained in an unsupervised learning manner from the image reconstruction loss based on a prepared StyleGAN generator. Sec. 3.1 describes how we abstract the multi-level representation from StyleGAN. Sec. 3.2 presents the structure of the novel hierarchical encoder. Sec. 3.3 introduces the idea of using pre-trained StyleGAN generator as a learned loss function for representation learning.

2 Hierarchical Encoder

We adopt ResNet architecture as the encoder backbone and add an extra residual block to get an additional feature map with lower resolution.In fact, there are totally six stages in our encoder, where the first one is a convolutional layer (followed by a pooling layer) and each of the others consists of several residual blocks. Besides, we introduce a feature pyramid network to learn the features from multiple levels. The output feature maps from the last three stages, $\{R_{4},R_{5},R_{6}\}$, are used to produce GH-Feat. Taking a 14-layer StyleGAN generator as an instance, $R_{4}$ aligns with layer 9-14, $R_{5}$ with 5-8, while $R_{6}$ with 1-4. Here, to bridge the feature map with each style code, we first downsample it to $4\times 4$ resolution and then map it to a vector of the target dimension using a fully-connect (FC) layer. In addition, we introduce a lightweight Spatial Alignment Module (SAM) into the encoder structure to better capture the spatial information from the input image. SAM works in a simple yet efficient way:

where $W_{4}$, $W_{5}$, and $W_{6}$ (all are implemented with an $1\times 1$ convolutional layer) are used to project the feature maps $R_{4}$, $R_{5}$, and $R_{6}$ to have the same number of feature channels respectively. $R_{4}$ and $R_{5}$ are downsampled to the same resolution of $R_{6}$ before fusion. The detailed structure of the encoder can be found in Appendix.

3 StyleGAN Generator as Learned Loss

We consider the pre-trained StyleGAN generator as a leaned loss function. Specifically, we employ a StyleGAN generator to supervise the encoder training with the objective of image reconstruction. We also introduce a discriminator to compete with the encoder, following the formulation of GANs , to ensure the reconstruction quality. To summarize, the encoder $E(\cdot)$ and the discriminator $D(\cdot)$ are jointly trained with

Experiments

We evaluate Generative Hierarchical Features (GH-Feat) on a wide range of downstream applications. Sec. 4.1 introduces the experimental settings, such as implementation details, datasets, and tasks. Sec. 4.2 conducts ablation study on the proposed hierarchical encoder. Sec. 4.3 and Sec. 4.4 evaluate the applicability of GH-Feat on generative and discriminative tasks respectively.

Implementation Details. The loss weights are set as $\lambda_{1}=0.1$, $\lambda_{2}=5e^{-5}$, and $\lambda_{3}=5$. We use Adam optimizer, with $\beta_{1}=0$ and $\beta_{2}=0.99$, to train both the encoder and the discriminator. The learning rate is initially set as $1e^{-4}$ and exponentially decayed with the factor of 0.8.

Datasets and Models. We conduct experiments on four StyleGAN models, pre-trained on MNIST , FF-HQ , LSUN bedrooms , and ImageNet respectively. The MNIST model is with $32\times 32$ resolution and the remaining models are with $256\times 256$ resolution.

2 Ablation Study

We make ablation studies on the training of encoder from two perspectives. (1) We choose the layer-wise style codes ${\rm\bf y}$ over the ${\rm\bf w}$ codes as the representation from StyleGAN. (2) We introduce Spatial Alignment Module (SAM) into the encoder to better handle the spatial information.

Since the encoder is trained with the objective of image reconstruction, we use mean square error (MSE), SSIM , and FID to evaluate the encoder performance. Tab. 1 shows the results where we can tell that our encoder benefits from the effective SAM module and that choosing an adequate representation space (i.e., the comparison between the first row and the last row) results in a better reconstruction. More discussion on the differences between $\mathcal{W}$ space and $\mathcal{Y}$ space can be found in Sec. 4.4.1.

3 Evaluation on Generative Tasks

Thanks to using the StyleGAN as a learned loss function, a huge advantage of GH-Feat over existing unsupervised feature learning approaches , which mainly focus on the image classification task, is its generative capability. In this section, we conduct a number of generative experiments to verify this point.

Image reconstruction is an important evaluation on whether the learned features can best represent the input image. The very recent work ALAE also employs StyleGAN for representation learning. We have following differences from ALAE: (1) We use the $\mathcal{Y}$ space instead of the $\mathcal{W}$ space of StyleGAN as the representation space. (2) We learn hierarchical features that highly align with the per-layer style codes in StyleGAN. (3) Our encoder can be efficiently trained with a well-learned generator by treating StyleGAN as a loss function. Tab. 2 and Fig. 2 show the quantitative and qualitative comparison between GH-Feat and ALAE on FF-HQ faces and LSUN bedrooms . We can tell that GH-Feat better reconstructs the input by preserving more information, resulting a more expressiveness representation.

3.2 Image Editing

In this part, we evaluate GH-Feat on a number of image editing tasks. Different from the features learned from discriminative tasks , our GH-Feat naturally supports sampling and enables creating new data.

Style Mixing. To achieve style mixing, we use the encoder to extract visual features from both the content image and the style image and swap these two features at some particular level. The swapped features are then visualized by the generator, as shown in Fig. 3. We can observe the compelling hierarchical property of the learned GH-Feat. For example, by exchanging low-level features, only the image color tone and the skin color are changed. Meanwhile, mid-level features controls the expression, age, or even hair styles. Finally, high-level features correspond to the face shape and pose information (last two columns).

Global Editing. The style mixing results have suggested the potential of GH-Feat in multi-level image stylization. Sometime, however, we may not have a target style image to use as the reference. Thanks to the design of the latent space in GANs , the generative representation naturally supports sampling, resulting in a strong creativity. In other words, based on GH-Feat, we can arbitrarily sample meaningful visual features and use them for image editing. Fig. 4 presents some high-fidelity editing results at multiple levels. This benefits from the matching between the learned GH-Feat and the internal representation of StyleGAN.

Local Editing. Besides global editing, our GH-Feat also facilitates editing the target image locally by deeply cooperating with the generator. In particular, instead of directly swapping features, we can exchange a certain region of the spatial feature map at some certain level. In this way, only a local patch in the output image will be modified while other parts remain untouched. As shown in Fig. 5, we can successfully manipulate the input face with different eyes, noses, and mouths.

3.3 Image Harmonization

Our hierarchical encoder is robust such that it can extract reasonable visual features even from discontinuous image content. We copy some patches (e.g., bed and window) onto a bedroom image and feed the stitched image into our proposed encoder for feature extraction. The extracted features are then visualized via the generator, as in Fig. 6. We can see that the copied patches well blend into the “background”. We also surprisingly find that when copying a window into the source image, the view from the original window and that from the new window highly align with each other (e.g., vegetation or ocean), benefiting from the robust generative visual features.

4 Evaluation on Discriminative Tasks

In this part, we verify that even the proposed GH-Feat is learned from generative models, it can be applicable to a wide range of discriminative tasks with competitive performances. Here, we do not fine-tune the encoder for any certain task. In particular, we choose multi-level downstream applications, including image classification, face verification, pose estimation, layout prediction, landmark detection, and luminance regression. For each task, we use our encoder to extract visual features from both the training and the test set. A linear regression model (i.e., a fully-connected layer) is learned on the training set with ground-truth and then evaluated on the test set.

Recall that GH-Feat is a multi-scale representation learned by using StyleGAN as a loss function. As a results, it consists of features from multiple levels, each of which correspond to a certain layer in the StyleGAN generator. Here, we would to explore how this feature hierarchy is organized as well as how they can facilitate multi-level discriminative tasks, including face pose estimation, indoor scene layout prediction, and luminanceWe convert images from RGB space to YUV space and use the mean value from Y space as the luminance. regression from face images. In particular, we evaluate GH-Feat on each task level by level. As a comparison, we also train encoders by treating the ${\rm\bf w}$ code, instead of the style code ${\rm\bf y}$, as the representation. From Fig. 7, we have three observations: (1) GH-Feat is discriminative. (2) Features at lower level are more suitable for low-level tasks (e.g., luminance regression) and those at higher level better aid high-level tasks (e.g., pose estimation). (3) $\mathcal{Y}$ space demonstrates a more obvious hierarchical property than $\mathcal{W}$ space.

4.2 Digit Recognition & Face Verification

Image classification is widely used to evaluate the performance of learned representations . In this section, we first compare our proposed GH-Feat with other alternatives on a toy dataset, i.e., MNIST . Then, we use a more challenging task, i.e., face verification, to evaluate the discriminative property of GH-Feat.

MNIST Digit Recognition. We first show a toy example on MNIST following prior work . We make a little modification to ResNet-18 like which is widely used in literatures to handle samples from MNIST in lower resolution. The Top-1 accuracy is reported in Tab. 3 (a). Our GH-feat outperforms ALAE and BiGAN with $1.45\%$ and $1.92\%$, suggesting a stronger discriminative power. Here, ResNet-18 is employed as the backbone structure for both MoCo and GH-Feat.

LFW Face Verification. We directly use the proposed encoder, which is trained on FF-HQ , to extract GH-Feat from face images in LFW and tries three different strategies on exploiting GH-Feat for face verification: (1) using a single level feature; (2) grouping multi-level features (starting from the highest level) together; (3) voting by choosing the largest face similarity across all levels. Fig. 7 (last column) shows the results from the first two strategies. Obviously, GH-Feat from the 5-th to the 9-th levels best preserve the identity information. Tab. 3 (b) compares GH-Feat with other unsupervised feature learning methods, including VAE , MoCo , and ALAE . All these competitors are also trained on FF-HQ dataset with optimally chosen hyper-parameters. ResNet-50 is employed as the backbone for MoCo and GH-Feat. Our method with voting strategy achieves 69.7% accuracy, surpassing other competitors by a large margin. We also visualize some reconstructed LFW faces in Fig. 8, where our GH-Feat well handles the domain gap (e.g., image resolution) and preserves the identity information.

4.3 Large-Scale Image Classification

We further evaluate GH-Feat on the high-level image classification task using ImageNet . Before the training of encoder, we first train a StyleGAN model, with $256\times 256$ resolution, on the ImageNet training collection. After that, we learn the hierarchical encoder by using the pre-trained generator as the supervision. No labels are involved in the above training process.Our encoder can be trained very efficiently, usually $3\times$ faster than the GAN training. For the image classification problem, we train a linear model on top of the features extracted from the training set with the softmax loss. Then, this linear model is evaluated on the validation set.During testing, we adopt the fully convolutional form as in and average the scores at multiple scales. Tab. 4 shows the comparison between GH-Feat and other unsupervised representation learning approaches , where we beat most of the competitors. The state-of-the-art MoCo gives the most compelling performance. But different from the representations learned with contrastive learning, GH-Feat has huge advantages in generative tasks, as already discussed in Sec. 4.3. Among adversarial representation learning approaches, BigBiGAN achieves the best performance, benefiting from the incredible large-scale training. However, GH-Feat presents a stronger generative ability, suggested by the comparison results on image reconstruction shown in Tab. 5. More discussion can be found in Appendix.

4.4 Transfer Learning

In this part, we explore how GH-Feat can be transferred from one dataset to another.

Landmark Detection. We train a linear regression model using GH-Feat on FF-HQ and test it on MAFL , which is a subset of CelebA . This two datasets have a large domain gap, e.g., faces in MAFL have larger poses yet lower image quality. As shown in Fig. 9, GH-Feat shows a strong transferability across these two datasets. We compare our approach with some supervised and unsupervised alternatives . For a fair comparison, we try the multi-scale representations from MoCo (i.e., Res2, Res3, Res4, and Res5 feature maps) and report the best results. Tab. 6 demonstrates the strong generalization ability of GH-Feat. In particular, it achieves on-par or better performance than the methods that are particular designed for this task . Also, it outperforms MoCo on this mid-level discriminative task.

Layout Prediction. We train the layout predictor on LSUN bedrooms and test it on kitchens to validate how GH-Feat can be transferred from one scene category to another. Feature learned by MoCo on the bedroom dataset is used for comparison. We can tell from Fig. 10 that GH-Feat shows better predictions than MoCo, especially on the target set (i.e., kitchens), suggesting a stronger transferability. Like landmark detection, we also conduct experiments with the 4-level representations from MoCo and select the best.

Conclusion

In this work, we consider the well-trained GAN generator as a learned loss function for learning multi-scale features. The resulting Generative Hierarchical Features are shown to be generalizable to a wide range of vision tasks. Acknowledgements: This work is supported in part by the Early Career Scheme (ECS) through the Research Grants Council (RGC) of Hong Kong under Grant No.24206219, CUHK FoE RSFS Grant, SenseTime Collaborative Grant, and Centre for Perceptual and Interactive Intelligence (CPII) Ltd under the Innovation and Technology Fund.

Appendix

Appendix A Encoder Structure

Tab. 7 provides the detailed architecture of our hierarchical encoder by taking a 14-layer StyleGAN generator as an instance. The design of GH-Feat treats the layer-wise style codes used in the StyleGAN model (i.e., the code fed into the AdaIN module ) as generative features. Accordingly, GH-Feat consists of 14 levels that exactly align with the multi-scale style codes yet in a reverse order, as shown in the last two columns of Tab. 7. In particular, these features are projected from the feature maps produced by the last three stages via fully-connected layers.

Appendix B More Details and Results on ImageNet

Training Details. During the training of the StyleGAN model on the ImageNet dataset , we resize all images in the training set such that the short side of each image is 256, and then centrally crop them to $256\times 256$ resolution. All training settings follow the StyleGAN official implementation , including the progressive strategy, optimizer, learning rate, etc. The generator and the discriminator are alternatively optimized until the discriminator have seen $250M$ real images. After that, the generator is fixed and treated as a well-learned loss function to guide the training of the encoder. During the training of the hierarchical encoder, images in the training collection are pre-processed in the same way as mentioned above. After the encoder is ready (usually trained for 25 epochs), we treat it as a feature extractor. We use the output feature map at the “res5” stage (with dimension $2048\times 8^{2}$), apply adaptively average pooling to obtain $2\times 2$ spatial feature, and vectorize it. A linear classifier, i.e., with one fully-connected layer, takes these extracted features as the inputs to learn the image classification task. SGD optimizer, together with batch size 2048, is used. The learning rate is initially set as 1 and decayed to 0.1 and 0.01 at the 60-th and the 80-th epoch respectively. During the training of the final classifier, ResNet-style data augmentation is applied.

Discussion. As shown in the main paper, GH-Feat achieves comparable accuracy to existing alternatives. Especially, among all of methods based on generative modeling, GH-Feat obtains second performance only to BigBiGAN , which requires incredible large-scale training.As reported in , the model train on images of $256\times 256$ resolution requires 256 TPUs running for 48 hours. However, our GH-Feat facilitates a wide rage of tasks besides image classification. Taking image reconstruction as an example, our approach can well recover the input image, significantly outperforming BigBiGAN . As shown in Fig. 11, BigBiGAN can only reconstruct the input image from the category level (i.e., dog or bird). By contrast, GH-Feat is able to recover more details, like shape and texture.

Appendix C Ablation Study

Recall that, during the training of the encoder, we propose to treat the well-trained StyleGAN generator as a learned loss function. In this part, we explore what will happen if we train the generator from scratch together with the encoder. Tab. 8 and Fig. 12 show the quantitative and qualitative results respectively, which demonstrate the strong performance of GH-Feat. It suggests that besides higher efficiency, reusing the knowledge from a well-trained generator can also bring better performance.

Appendix D Style Mixing

In this part, we verify the hierarchical property of GH-Feat on the task of style mixing and further make comparison with ALAE . In particular, we use ALAE and our approach to extract features from same images (including both style images and content images) and then use the extracted features for style mixing at different levels.

Fig. 13 shows the comparison results. Note that all test images are selected following the original paper of ALAE . We can see that when mixing high-level styles, the pose, age, and gender of mixed results are close to those of style images. By comparing with ALAE, results using GH-Feat better preserve the identity information (high-level feature) from style images as well as the color information (low-level feature) from content images. In addition, when mixing low-level styles (bottom two rows), both ALAE and GH-Feat can successfully transfer the color style from style images to content images, but GH-Feat shows much stronger identity preservation.

References