Bi-level Feature Alignment for Versatile Image Translation and Manipulation

Image translation has achieved remarkable progress in learning the mapping between images of different domains. It could be applied in different tasks such as style transfer  [10, 7, 20], image super-resolution [16, 21, 15, 58], domain adaptation [36, 30, 8, 41, 49], image composition [57, 48, 55, 51, 54] etc. To achieve high-fidelity and flexible translation, existing work uses different conditional inputs such as semantic segmentation [12, 43, 32, 53, 56], scene layouts [38, 60, 19, 52], key points [27, 29, 5, 50], edge maps [12, 6], etc. However, effective style control remains a challenging task in image translation.

Style control has attracted increasing attention in image translation and generation. Earlier works such as [14] regularize the latent feature distribution to control the generation outcome. However, they struggle to capture the complex textures of natural images. Style encoding has been studied to address this issue. For example, [11] and [26] transfer style codes from exemplars to source images via adaptive instance normalization (AdaIN) [10]. [3] employs a style encoder for style consistency between exemplars and translated images. [65] designs semantic region-adaptive normalization (SEAN) to control the style of each semantic region individually. However, encoding style exemplars tends to capture the overall image style and ignores the texture details in local regions. To achieve accurate style guidance for each local region, Zhang et al. [59] build dense semantic correspondences between conditional inputs and exemplars with Cosine similarity to capture accurate exemplar details. To mitigate the quadratic complexity issue and enable high-resolution correspondence building, Zhou et al. [63] introduce the GRU-assisted Patch-Match to efficiently establish the high-resolution correspondence.

The arise of generative adversarial network (GANs) brings revolutionary advances to image editing [64, 9, 31, 2, 33, 45, 44, 46]. As one of the most intuitive representation in image editing, semantic information has been extensively investigated in conditional image synthesis. For example, Park et al. [32] introduce spatially-adaptive normalization (SPADE) to inject guided features in image generation. MaskGAN [17] exploits a dual-editing consistency as auxiliary supervision for robust face image manipulation. Instead of directly learning a label-to-pixel mapping, Hong et al. [9] propose a semantic manipulation framework HIM that generates images guided by a predicted semantic layout. Upon this work, Ntavelis et al. [31] propose SESAME which requires only local semantic maps to achieve image manipulation. However, the aforementioned methods either only learn a global feature without local focus (e.g., MaskGAN [17]) or ignore the features in the editing regions of the original image (e.g., HIM [9], SESAME [31]). To better utilize the fine features in the original image, Zheng et al. [61] adapt exemplar-based image synthesis framework CoCosNet [59] for semantic image manipulation by building a high-resolution correspondence between the original image and the edited semantic map.

The alignment network aims to build the correspondence between conditional inputs and exemplars, and accordingly provide accurate style guidance by warping the exemplars to be aligned with the conditional inputs. As shown in Fig. 2, conditional input and exemplar are fed to feature extractors $F_{X}$ and $F_{Z}$ to extract two sets of feature vectors $X=[x_{1},\cdots,x_{L}]\in\mathbb{R}^{d}$ and $Z=[z_{1},\cdots,z_{L}]\in\mathbb{R}^{d}$, where $L$ and $d$ denote the number and dimension of feature vectors, respectively. Then $X$ and $Z$ can be aligned by building a $L\times L$ dense correspondence matrix where each entry denotes the Cosine similarity between the corresponding feature vectors in $X$ and $Z$.

Semantic Position Encoding. Existing works [59, 63] mainly rely on semantic features to establish the correspondences. However, as textures within a semantic region share the same semantic feature, the pure semantic correspondence fails to preserve the texture structures or patterns within each semantic region. Thus, the position information of features can be facilitated to preserve the texture structures and patterns. A vanilla method to encode the position information is employing a simple coordconv [22] to build a global coordinate for the full image. However, this vanilla position encoding mechanism builds a single coordinate system for the whole image, ignoring region-wise semantic differences. To preserve the fine texture pattern within each semantic region, we design a semantic position encoding (SPE) mechanism that builds a dedicated coordinate for each semantic region as shown in Fig. 3. Specifically, SPE treats the center of each semantic region as the origin of coordinate, and the coordinates within each semantic region are normalized to [-1, 1]. The proposed SPE outperforms the vanilla position encoding significantly as shown in Fig. 6 and to be evaluated in experiments.

Bi-level Feature Alignment. On the other hand, building correspondence has quadratic complexity which incurs large memory and computation costs. Most existing studies thus work with low-resolution exemplar images (e.g. 64 $\times$ 64 in CoCosNet [59]) which often struggle in generating realistic images with fine texture details. In this work, we propose a bi-level alignment strategy via a novel ranking and attention scheme (RAS) that greatly reduces computational costs and allows to build correspondences with high-resolution images as shown in Fig. 6. Instead of building correspondences between features directly, the bi-level alignment strategy builds correspondences at two levels, including the first level that introduces top-$k$ ranking to generate block-wise ranking matrices dynamically and the second level that achieves dense attention between the features within blocks. As Fig. 2 shows, $b$ local features are grouped into a block, thus the features of conditional input and exemplar are partitioned into $N$ blocks ($N=L/b$) as denoted by $X=[X_{1},\cdots,X_{N}]\in\mathbb{R}^{bd}$ and $Z=[Z_{1},\cdots,Z_{N}]\in\mathbb{R}^{bd}$. In the first level of top-$k$ ranking, each block feature of the conditional input serves as a query to retrieve top-$k$ block features from the exemplar according to the Cosine similarity between blocks. In the second level of local attention, the features in each query block further attends to the features in the top-$k$ retrieved blocks to build up local attention matrices within the block features. The correspondence between the exemplar and conditional input can thus be built much more efficiently by combining such inter-block ranking and inner-block attention.

The ranking and attention scheme employs a top-$k$ operation that ranks the correlative blocks. However, the original top-$k$ operation involves index swapping whose gradient cannot be computed and so cannot be integrated into end-to-end network training. Inspired by Xie et al. [47], we tackle this issue by formulating the top-$k$ ranking as a regularized earth mover’s problem which allows gradient computation via implicit differentiation. Earth mover’s problem aims to find a transport plan that minimizes the total cost to transform one distribution to another. Consider two discrete distributions $U=[\mu_{1},\dots,\mu_{N}]^{\top}$ and $V=[\nu_{1},\dots,\nu_{M}]^{\top}$ defined on supports $\mathcal{A}=[a_{1},\cdots,a_{N}]$ and $\mathcal{B}=[b_{1},\cdots,b_{M}]$, with probability (or amount of earth) $\mathbb{P}(a_{i})=\mu_{i}$ and $\mathbb{P}(b_{j})=\nu_{j}$. We define $C\in\mathbb{R}^{N\times M}$ as the cost matrix where $C_{ij}$ denotes the cost of transportation between $a_{i}$ and $b_{j}$, and $T$ as a transport plan where $T_{ij}$ denotes the amount of earth transported between $\mu_{i}$ and $\nu_{j}$. An earth mover’s (EM) problem can be formulated by: ${\rm EM}=\mathop{\min}\limits_{T}\langle C,T\rangle,\quad{\rm s.t.}\ T\vec{1}_{M}=U,\;T^{\top}\vec{1}_{N}=V$ where $\vec{1}$ denotes a vector of ones, $\langle\rangle$ denotes inner product. By treating a correlation scores between a query block and $N$ key blocks as $\mathcal{A}=[a_{1},\cdots,a_{N}],a_{i}\in[-1,1]$ and defining $\mathcal{B}=\{-1,1\}$, $U=[\mu_{1},\cdots,\mu_{N}]$ and $V=[\nu_{1},\nu_{2}]$, it can be proved that solving the Earth Mover’s problem is equivalent to select the largest $K$ elements from $\mathcal{A}=[a_{1},\cdots,a_{N}]$. The detailed proof and optimization of the earth mover’s problem is provided in supplementary material. Fig. 5 illustrates the earth mover’s problem and transport plan $T$ which indicates the top-$k$ elements.

Complexity Analysis. The vanilla dense correspondence has a self-attention memory complexity $\mathcal{O}(L^{2})$ where $L$ is the input sequence length. For our bi-level alignment strategy, the memory complexity of building block ranking matrices and local attention matrices are $\mathcal{O}(N^{2})$ and $\mathcal{O}(b*(kb))$, where $b$, $N$ ($N=L/b$) and $k$ are block size, block number and the number of top-$k$ selection. Thus, the overall memory complexity is $\mathcal{O}(N^{2}+b*(kb))$.

The generation network aims to synthesize images under the semantic guidance of conditional inputs and style guidance of exemplars. The overall architecture of the generation network is similar to SPADE [32]. Please refer to supplementary material for details of the network structure.

State-of-the-art approach [59] simply concatenates the warped exemplar and conditional input to guide the image generation process. However, the warped input image and edited semantic map could be structurally aligned but semantically different especially when they have severe semantic discrepancy. Such unreliably warped exemplars could serve as false guidance and heavily deteriorate the generation performance. Therefore, a mechanism is required to identify the semantic reliability of warped exemplar to provide reliable guidance for the generation network. To this end, we propose a CONfidence Feature Injection (CONFI) module that adaptively weights the features of conditional input and warped exemplar according to the reliability of feature matching.

Confidence Feature Injection. Intuitively, in the case of lower reliability of the feature correspondence, we should assign a relatively lower weight to the warped exemplar which provides unreliable style guidance and a higher weight to the conditional input which consistently provides accurate semantic guidance.

As illustrated in Fig. 5, the proposed CONFI fuses the features of the conditional input and warped exemplar based on a confidence map (CMAP) that captures the reliability of the feature correspondence. To derive the confidence map, we first obtain a block-wise correlation map of size $N\times N$ by computing element-wise Cosine distance between $X=[X_{i},\cdots,X_{N}]$ and $Z=[Z_{i},\cdots,Z_{N}]$. For a block $X_{i}$, the correlation score with $Z$ is denoted by $\mathcal{A}=[a_{1},\cdots,a_{N}]$. As higher correlation scores indicate more reliable feature matching, we treat the peak value of $\mathcal{A}$ as the confidence score of $X_{i}$. Similar for other blocks, we can obtain the confidence map (CMAP) of size $1\times H\times W$ ($N=H*W$) which captures the semantic reliability of all blocks. The features of the conditional input and exemplar (both of size $C\times H\times W$ after passing through convolution layers) can thus be fused via weighted sum based on the confidence map CMAP: $F=X*(1-{\rm CMAP})+(T\cdot Z)*{\rm CMAP}$ where $T$ is the built correspondence matrix. As the confidence map contains only one channel ($1\times H\times W$), the above feature fusion is conducted in $H\times W$ but ignores that in $C$ channel. To achieve thorough feature fusion in all channels, we feed the initial fusion $F$ to convolution layers to generate a multi-channel confidence map (Multi-CMAP) of size $C\times H\times W$. The conditional input and warped exemplar are then thoroughly fused via a full channel-weighted summation according to the Multi-CMAP. The final fused feature is further injected to the generation process via spatial de-normalization [32] to provide accurate semantic and style guidance.

Datasets. We evaluate and benchmark our method over multiple datasets for image translation & manipulation tasks.

$\bullet$ ADE20K [62] is adopted for image translation conditioned on semantic segmentation. For image manipulation, we apply object-level affine transformations on the test set to acquire paired data (150 images) for evaluations as in [61].

$\bullet$ CelebA-HQ [24] is used for two translation tasks by using face semantics and face edges as conditional inputs. We use 2993 face images for translation evaluations as in [59], and manually edit 100 semantic maps which is randomly selected for image manipulation evaluations.

$\bullet$ DeepFashion [23] is used for image translation conditioned key points.

We compare RABIT with several state-of-the-art image translation methods.

Quantitative Results. In quantitative experiments, all methods translate images with the same exemplars except Pix2PixHD [43] which doesn’t support style injection from exemplars. LPIPS is calculated by comparing the generated images with randomly selected exemplars. All compared methods adopt three exemplars for each conditional input and the final LPIPS is obtained by averaging the LPIPS between any two generated images.

Table 1 shows experimental results. It can be seen that RABIT outperforms all compared methods over most metrics and tasks consistently. By building explicit yet accurate correspondences between conditional inputs and exemplars, RABIT enables direct and accurate guidance from the exemplar and achieves better translation quality (in FID and SWD) and diversity (in LPIPS) as compared with the regularization-based methods such as SPADE [32] and SMIS [66], and style-encoding methods such as StarGAN v2 [3] and SEAN [65]. Compared with correspondence-based method CoCosNet [59], the proposed bi-level alignment allows RABIT to build correspondences and warp exemplars at higher resolutions (e.g. $128\times 128$) which offers more detailed guidance in the generation process and helps to achieve better FID and SWD. While compared with CoCosNet v2 [63], the proposed semantic position encoding enables to preserve the texture structures and patterns, thus yielding more accurate warped exemplars as guidance. Besides generation quality, RABIT achieves the best generation diversity in LPIPS except StarGAN v2 [3] which sacrifices the generation quality with much lower FID and SWD.

Qualitative Evaluations. Fig. 7 shows qualitative comparisons on various conditional image translation tasks. It can be seen that RABIT achieves the best visual quality with faithful styles as exemplars. RABIT also demonstrates superior diversity in image translation as illustrated in Fig. 8.

RABIT manipulates images by treating input images as exemplars and edited semantic guidance as conditional inputs. We compare RABIT with several state-of-the-art image manipulation methods including 1) SPADE [32], 2) SEAN [65], 3) MaskGAN [18], 4) Hierarchical Image Manipulation (HIM) [9], 5) SESAME [31], 6) CoCosNet [59].

Quantitative Results: In quantitative experiments, all compared methods manipulate images with the same input image and edited semantic label map. Left side of Table 2 shows experimental results over the synthesized test set of ADE20K [62]. It can be observed that RABIT outperforms state-of-the-art methods over all evaluation metrics consistently. Right side of Table 2 shows experimental results over the CelebA-HQ dataset with manually edited semantic maps. It can be observed that RABIT outperforms the state-of-the-art methods by large margins in all perceptual quality metrics.

Qualitative Evaluation: Fig. 9 shows visual comparisons with state-of-art manipulation methods on ADE20K. Fig. 10 shows the editing capacity of RABIT with various types of manipulation on semantic labels. We also compare RABIT with MaskGAN [17] on CelebA-HQ [18] in Fig. 12.