Boundary-aware Decoupled Flow Networks for Realistic Extreme Rescaling

1 Introduction

Image rescaling, involving reconstructing high-resolution (HR) images from their corresponding low-resolution (LR) versions, plays a crucial role in large-size data services such as storage and transmission. Early image rescaling methods [Dong et al., 2014; Lim et al., 2017; Zhang et al., 2018b;

∗This work was done while Jinmin Li was an intern at WeChat Pay Lab33, Tencent.

†Corresponding author: Tao Dai (daitao.edu@gmail.com)

(c) IRN

Figure 1: Visual quality of various IRN-based and GAN-based methods, including (b) GLEAN, (c) IRN, and (d) GRAIN. Existing methods produce (c) over-smoothed results or (b) and (d) fake details. By contrast, our method can generate visually pleasing results with sharper details.

Li et al., 2023; Dai et al., 2023; Guo et al., 2024; Cui et al., 2023] primarily concentrated on non-adjustable downscaling kernels and neglected the compatibility between downscaling operations and reconstruction algorithms, which thus easily lose high-frequency information and thus produce visually unpleasing results.

To further improve the reconstruction performance, recently developed invertible rescaling network (IRN) based and generative adversarial network (GAN) based methods [Xiao et al., 2020; Zhong et al., 2022] have demonstrated impressive performance in image rescaling. Among them, IRN [Xiao et al., 2020] transforms high-frequency components into a latent space, and assumes that the LR components and high-frequency components are independent. To consider the effect of the LR counterparts, HCFlow [Liang et al., 2021] using a hierarchical conditional framework for modeling the LR information. Later, GRAIN [Zhong et al., 2022] introduces a reciprocal invertible image rescaling approach, enabling subtle embedding of HR information into a reversible LR image and generative prior for accurate HR reconstruction. Despite the success of these generative methods in image rescaling, they still suffer from visually unpleas-

(d) GRAIN

ing results, due to the inaccurate modeling way of the highfrequency information. As shown in Fig. 1, existing methods produce over-smoothed results or fake details, thus leading to unrealistic results.

On the other hand, it is observed that the high-frequency part of the input deviates from a standard Gaussian distribution. As shown in Fig. 2(b), we compute the statistics of the high-frequency information ZGT from DIV2K and CelebA, which is obtained by computing the difference between the HR image in Fig. 2(a) and the reconstructed image of LR. We observe that ZGT is a mixture of semantic boundary information and Gaussian distribution with µ = 0 and σ2 = 0.2. This motivates us to decouple the high-frequency information into different parts.

The above observations inspire us to design a more effective modeling way for high-frequency information, in this paper, we propose Decoupled Boundary-aware Decoupled Flow Networks (BDFlow), which decouples the original high-frequency information into approximate semantic Boundary distribution and non-semantic Gaussian distribu-

tion. As shown in Fig. 1(e), we utilize Boundary-aware Mask (BAM) to preserve semantic information, ensuring that the recovered image follows the true distribution. The nonsemantic Gaussian distribution is independent of the lowfrequency distribution and semantic Boundary distribution, and is randomly sampled in the backward process. Furthermore, we impose additional constraints on the model to generate textures consistent with the Ground Truth by utilizing the Boundary-aware Weight (BAW).

In summary, our main contributions are:

• To our best knowledge, the proposed Boundary-aware Decoupled Flow Networks (BDFlow) is the first attempt to decouple high-frequency information into semantic Boundary distribution and non-semantic Gaussian distribution.
• We introduce a general Boundary-aware Mask (BAM) to preserve semantic information, ensuring that the recovered image follows the true distribution. Canny operator is a special case of BAM, where the quantization is binarized, and the magnitude of the gradient is calculated using the 2-Norm. Additionally, our proposed Boundary-aware Weight (BAW) further constrains the model to generate rich texture details.
• Extensive experiments demonstrate that our proposed BDFlow achieves state-of-the-art (SOTA) performance while maintaining a lower computational burden and faster inference time compared to other existing methods. We also explore and analyze image invertibility with respect to different influencing factors, such as scaling factor, JPEG compression, and data domain.

2 RelatedWork

2.1 Image RescalingMethods

Image rescaling aims to downscale a high-resolution (HR) image to a visually pleasing low-resolution (LR) image and later reconstruct the HR image. In recent years, there has been a surge in the use of normalization flow [Xiao et al.,

2020; Liang et al., 2021; Li et al., 2024a] and generative adversarial network (GAN) [Zhong et al., 2022; Chan et al., 2021; Yang et al., 2021] techniques for tackling the issue of image rescaling. IRN [Xiao et al., 2020] transforms highfrequency information into a latent space Z ∼ p(Z) which is a Gaussian distribution. Contrarily, HCFlow[Liang et al., 2021] posits that the high-frequency information is dependent on the LR image. GANs-based GRAIN [Zhong et al., 2022] introduces a reciprocal approach, enabling delicate embedding of high-resolution information into a reversible lowresolution image and generative prior. However, IRN generates overly smooth images due to the substantial loss of high-frequency information. HCFlow is non-decoupled and biased estimation. GRAIN produces fake details and requires extensive computation. To address these issues, we introduce light-weight Boundary-aware Mask (BAM) to constrain the model to produce rich textures.

2.2 Boundary-aware Methods

Boundary-aware techniques are vital for preserving facial structure and appearance in image rescaling, especially for facial images. These methods have proven effective in applications like face recognition and segmentation. Previous Boundary-aware work involves attention mechanisms [Wang et al., 2017; Tang et al., 2024; Zhang et al., 2023], edgeaware filtering [Xu et al., 2011; Chen et al., 2013; Li et al., 2024b], and facial landmark detection [Cao et al., 2014; Kazemi and Sullivan, 2014]. Attention mechanisms focus on important features, improving recognition and segmentation. Edge-aware filtering preserves boundaries while maintaining smoothness elsewhere. Integrating Boundary-aware methods in image rescaling models offers advantages in preserving semantic Boundary distribution and model robustness.

2.3 Invertible Residual Networks

Initially developed for unsupervised learning of probabilistic models [Behrmann et al., 2019; Dinh et al., 2016; Liu et al., 2020; Wang et al., 2023; Gao et al., 2023], Invertible Residual Networks (IRN) facilitate the transformation of one distribution to another through bijective functions, thereby preserving information integrity. This characteristic empowers invertible networks to ascertain the precise density of observations, which can subsequently be employed to generate images with intricate distributions. Owing to these distinctive properties, invertible networks have been effectively utilized in an array of applications, including image rescaling [Xiao et al., 2020; Liang et al., 2021]. In this work, we adopt the Invertible Residual Networks architecture presented in [Dinh et al., 2016], which consists of a series of invertible blocks. For the lth block, the forward process is formulated as:

Here, F(·) and G(·) represent arbitrary transformation functions. The input u is decomposed into ul1 and ul2 . The backward process can be conveniently defined as:

3 Methodology

3.1 Networks Framework

BDFlow can be dissected into three components: the Boundary-aware Mask generation algorithm (shown in Algorithm. 1), the Haar Wavelet transformation [Lienhart and Maydt, 2002], and the Flow Networks (shown in Fig. 3).

Image rescaling focuses on reconstructing a highresolution (HR) image X from a low-resolution (LR) image Y and high-frequency distribution, which are obtained by downscaling X. As the downscaling process is the inverse of upscaling, we employ an invertible neural network to generate the LR image Y and decouple highfrequency distribution to non-semantic Gaussian distribution Z and semantic Boundary distribution B (i.e., [Y,B,Z] = Fθ(X),Z ∼ pZ(Z),B ∼ pB(B)). Inversely, X can be reconstructed through the inverse process from [Y,B,Z]: X = F -1 θ (Y,B,Z). It is important to note that the Boundary distribution preserves semantic information in B, and Z corresponds to the non-semantic high-frequency information, as per the Nyquist-Shannon sampling theorem [Shannon, 1949]. To ensure the model's invertibility, we must verify that Fθ is invertible, which is equivalent to having a non-zero determinant of the Jacobian for each invertible unit F.

Here, F and G represent two transformations of each invertible unit. X corresponds to ul , and Y corresponds to ul+1 in Eq. 1. The value of JF is equal to 1 due to ∂G ∂X1 = ∂G ∂F .

Algorithm 1 Boundary-aware Mask

Input: I, input image

Parameter: T, the threshold to sparsify boundary

Output: Bq, boundary distribution

• 1: Ig ←GaussBlur(I, σ)
• 2: Gx,Gy ←Gradient(Ig)
• 3: M ← getMagnitude(Gx,Gy)
• 4: M ′ ←NonMaximumSuppress(M)
• 5: Bs ← SparsifyBoundary(M ′ , T)
• 6: Bq ←Quantify(Bs)
• 7: return Bq

3.2 Boundary-aware Decoupled FlowNetworks

Boundary-aware Decoupled Flow Networks (BDFlow) decouple high-frequency information into non-semantic Gaussian distribution and semantic Boundary distribution. General Boundary-aware Mask preserves semantic information, ensuring that the recovered image follows the true distribution.

Boundary-aware Mask. The Boundary-aware Mask (BAM) is a crucial component of the BDFlow Networks, and its generation process can be summarized as follows. As shown in Algorithm 1, given an input image, denoted as I, the algorithm begins by applying a Gaussian blur with a standard deviation of σ to produce a blurred image Ig. Next, the gradient components Gx and Gy of the blurred image Ig are computed, which are then used to calculate the gradient magnitude M (1-Norm, 2-Norm or others). Subsequently, nonmaximum suppression is performed on the gradient magnitude M to obtain a thinned boundary map M ′ . The boundary map M ′ is then sparsified by applying a threshold T, resulting in a sparse boundary map Bs. Finally, the sparse boundary map Bs is quantified to create the Boundary-aware Mask Bq .

Canny operator is a special case of BAM, where the quantization is binarized, and the magnitude of the gradient is calculated using the 2-Norm.

Boundary-aware Weight. The proposed Boundary-aware Weight (BAW) further constrains the model to generate rich textures and semantic Boundary distribution, defined as:

where λ2 is a hyper-parameter, α = 0.3 denotes 30 per cent of the training epochs, Ecur is the current training epoch and Emax is the max training epoch. During the pre-training period (α∗Emax), the precision of the BAM is inadequate, and therefore fixed weights are employed. In the later training period, the Bs in the BAM is utilized after normalization, and is used to weight the loss Lback(X,Xback). The purpose of normalisation is to remove the effects of outliers. λBAW penalises the model for errors in the high-frequency textures of Ground Truth, constraining the model to generate a realistic Boundary distribution.

Haar Wavelet. We employ the Haar Wavelet (Other wavelet bases are discussed in the supplementary material.) transformation [Mallat, 1989] to decompose the input X into high and low-frequency information, represented as [A,H, V,D]. Specifically, given an HR image X with shape (H,W,C), Haar Wavelet transformation decomposes X into global frequency features ul1 , ul2 :

where l denotes lth layer of the Flow Networks. u1 and u2 with shape ( 12 H, 12 W,C) and ( 12 H, 12 W, 3C), correspond to low-frequency and high-frequency information, respectively. FlowNetworks. Subsequently, FlowNetworks are employed to transform the frequency information into the semantic Boundary distribution, non-semantic Gaussian distribution, and LR distribution. Specifically, wavelet features ul1 and ul2 are fed into stacked InvBlocks to obtain the LR Y and highfrequency distribution (B,Z). In our framework, (B,Z) is decoupled into semantic Boundary distribution B and nonsemantic Gaussian distribution Z. We adopt the general coupling layer for the invertible architecture [Jacobsen et al., 2018; Behrmann et al., 2019]. The output of each InvBlock can be defined as:

where ul+1 1 and ul+1 2 are the outputs of the current InvBlock and the inputs of the next InvBlock. Notably, ϕ(·), ρ(·), and η(·) are arbitrary transformation functions. We use the residual block [He et al., 2016] to enhance the model's nonlinear expressiveness and information transmission capabilities.

InvBlock extracts image features from ul1 and ul2 to ul+1 1 and ul+1 2 . After the transformation of n InvBlocks, we obtain the downscaled image Y with dimensions ( 1 2 H, 12 W,C), Boundary distribution B with dimensions ( 12 H, 12 W, 1), and Gaussian distribution Z with dimensions ( 12 H, 12 W, 3C - 1).

3.3 Loss Function

Image rescaling aims to accurately reconstruct the HR image while generating visually pleasing LR images. Following IRN [Xiao et al., 2020] and HCFlow [Liang et al., 2021], we train our BDFlow by minimizing the following loss:

where X and Y are the ground-truth HR image and LR image, respectively. B and Z are the ground-truth semantic Boundary distribution and non-semantic Gaussian distribution, respectively.Xback is the reconstructed HR image from the generated LR image Yforw, Boundary distribution Bforw, and sampled Gaussian distribution Z.

Lforw is the l2 pixel loss defined as:

where N is the number of pixels, Yforw = F Y θ (X) is the generated LR image, and λ1 is a hyper-parameter balancing different losses.

Lback is the l1 pixel loss defined as:

where M is the number of pixels, Xback = F -1 θ (Y,B,Z), and λBAW is generated by the generation process of BAM but does not need to be quantified, which is Bs in Algorithm 1 and Eq. 4.

LLPIPS with the weight λ3 is used to enhance the visual effect of the image generated by the model. Similarly, LBAM is the l2 pixel loss defined as:

where N is the number of pixels, λ4 is a hyper-parameter, and Bforw = F B θ (X) is the generated semantic Boundary distribution.

The last term, Llatent(Z), chooses more stable distribution metrics for minimization to ensure that Z follows a nonsemantic Gaussian distribution [Xiao et al., 2020], defined as:

We jointly optimize the invertible architecture F by utilizing both forward and backward losses.

4 Experiments

4.1 Datasets

We utilize the CelebA-HQ dataset [Karras et al., 2020] to train our BDFlow. The dataset comprises 30,000 highresolution (1024 × 1024) human face images. Moreover, we evaluate our models using widely-accepted pixel-wise metrics, including PSNR, SSIM, and LPIPS [Zhang et al., 2018a] (Y channel) on the CelebA-HQ test dataset. We additionally train our BDFlow on the Cat dataset [Zhang et al., 2008] and the LSUN-Church [Yu et al., 2015] dataset to evaluate its generalization ability across different domains.

4.2 Training Details

Our BDFlow consists of n Haar Wavelet Transformation and Flow Networks for 2n×, each containing two Invertible Blocks (InvBlocks). We train these models using the ADAM [Kingma and Ba, 2014] optimizer with β1 = 0.9 and β2 = 0.999. Furthermore, we initialize the learning rate at 2 × 10-4 and apply a cosine annealing schedule, decaying from the initial value to 1 × 10-6 over the total number of iterations. We set λ1, λ2, λ3, λ4, and λ5 to 2, 2, 1, 16, and 4, respectively.

4.3 Comparison with State-of-the-art Methods

We compare our BDFlow with several state-of-the-art methods, including GAN-based inversion methods, CNN-based SR methods, and invertible rescaling methods. Table 1 presents the quantitative evaluations of the various approaches. Evidently, CNN-based methods outperform GANbased methods in terms of PSNR and SSIM, as they tend to produce fuzzy results that follow the overall structure. Our flow-based approach BDFlow achieves superior PSNR and SSIM, with LPIPS results second only to GRAIN [Zhong et al., 2022] on 64×. This demonstrates that our model considers not only pixel-level fidelity but also perceptual visual effects in facilitating the decoupling of high-frequency information into semantic Boundary distribution and nonsemantic Gaussian distribution. Specifically, BDFlow improves the PSNR by 4.4 dB and the SSIM by 0.0976 on average over GRAIN.

Table 2 compares IRN [Xiao et al., 2020], GRAIN [Zhong et al., 2022], and our BDFlow in terms of model parameters, computational effort, and inference time. Our BDFlow surpasses IRN in all dimensions and outperforms GRAIN by using only 74% of the parameters and 20% of the computation. Notably, BDFlow is suitable for real-time applications, as it can achieve 100 fps.

Fig. 4 contrasts our BDFlow with the GAN-based approach. In particular, pSp [Richardson et al., 2021] recovers images of faces with different identities, while the results of GPEN [Yang et al., 2021] and GLEAN [Chan et al., 2021] display slight quality improvements but still exhibit significant flaws. Major differences persist in hair, eyebrows, eyes, and facial expressions. Although GRAIN [Zhong et al., 2022] combines reversible and generative priors, it can only mitigate the limitation of under-informed input. GRAIN achieves good detail generation but produces some unfaithful details, as shown in Fig. 4 (d). In contrast, as shown in Fig. 4 (e), our BDFlow excels in both realism and fidelity, capturing semantic facial features, like vivid facial expressions, flyaway hair, and eye color.

(a) pSp (64X)

(b) GPEN (32X)

(c) GLEAN (64X)

(d) GRAIN (64X)

(e) Ours

(f) GT

Table 4: Ablation study evaluating the impact of quantify level and calculation of magnitude in the CelebA with ×64. T denotes the magnitude threshold.

4.4 Ablation Study

We conducted an ablation study to evaluate the impact of our proposed BAM and BAW. The effects of LPIPS loss, different Norm Loss, different wavelet bases, and hyper-parameters are discussed in the supplementary material.

Boundary Distribution and Semantic High-frequency Textures. Boundary-aware Mask (BAM) and Boundaryaware Weight (BAW) are able to recover more realistic Boundary distribution, such as the hair strands shown in Fig. 5 (c), because they focus on handling image boundaries and regions with high-frequency information more effectively. This is achieved by encouraging the model to decouple the high-frequency information into non-semantic Gaussian distribution and semantic Boundary distribution. However, GAN-based methods, like ESRGAN [Wang et al., 2018] and GRAIN [Zhong et al., 2022], generate fake details without the constraint of semantic Boundary distribution.

Table 3 demonstrates that the BDFlow structure incorporating both BAM and BAW consistently outperforms the basic BDFlow structure in terms of PSNR, SSIM, and LPIPS.

Specifically, Table 4 shows that BAM with a threshold value of T = 50 achieves superior performance due to the smaller values of high-frequency textures, which have a mean of approximately 50. Additionally, using the 2-Norm for magnitude calculation is better than using the 1-Norm because it widens the gap between different magnitudes. It is worth noting that the Canny operator can be regarded as a special case of BAM, where quantization is binarized and the magnitude of the gradient is calculated using the 2-Norm. However, as the level of quantization increases, the model achieves better performance. The Boundary distribution with 1 bit and single channel requires only a 4% ( 18 × 13 ≈ 4%) increase in storage cost, but can be further compressed to 1% due to the sparsity of high frequency.

4.5 Discussions

In our discussions, we conducted various analyses to evaluate the effectiveness of our approach. Firstly, we conducted a comparative analysis against the widely-used JPEG image compression technique to establish its superiority. Next, we evaluated the effect of BDFlow on different domains. More discussion is placed in the supplementary material, including the scale of sampled Z and evaluation of LR Images.

Comparison with JPEG Compression. Table 5 illustrates the superior performance of BDFlow compared to JPEG

compression and GRAIN [Zhong et al., 2022] in terms of PSNR and storage size. BDFlow consistently achieves higher PSNR values while maintaining storage sizes comparable to GRAIN, showcasing its ability to effectively compress images while preserving the quality of the reconstructed images. BDFlow's decoupling of high-frequency information into semantic Boundary distribution and non-semantic Gaussian distribution is at the core of the above gains. In addition, JPEG compression and rescaling are not in conflict, it is possible to continue to compress images using the JPEG algorithm after rescaling.

Evaluation in Different Domains. As shown in Fig. 6 and Table 6, we evaluate our proposed approach in different Domains. These experiments demonstrate the generalisation ability of our proposed BDFlow in reconstructing images from different Domains. This is attributed to the fact that we decouple the true high-frequency distributions into semantic Boundary distribution and non-semantic Gaussian distribution, which helps to achieve the final realistic results.

As shown in Fig. 6, GRAIN has difficulty recovering the hair of the cat and the edges of the church, but our BDFlow is able to do so with high fidelity. As shown in Table 6, our BDFlow achieves much higher PSNR and SSIM than GRAIN at the pixel level.

l1 Loss and l2 Loss. We use l1 loss and l2 loss in some cases in our method for the fairness comparison, since we followed the default setup of the previous works [Xiao et al., 2020; Liang et al., 2021]. Besides, we have conducted experiments to verify the effects of Norm-loss (see in Table 1 of the Supplementary Material). In fact, the forward process utilizes l2 to increase the penalty for outliers in favor of constraining the LR distribution, Gaussian distribution, and Boundary distribution. The backward process utilizes the l1 loss encouraging the model to recover more detail rather than smoothing the image, thus leading to better LPIPS. Thus, (l2, l1) may be a proper choice.

5 Conclusion

In this paper, we proposed the Boundary-aware Decoupled Flow Networks (BDFlow), which decouple high-frequency information into non-semantic Gaussian distribution and semantic Boundary distribution. This approach effectively preserves boundary information and texture details in the image rescaling process. Specifically, by introducing a generated Boundary-aware Mask, our model ensures that the recovered image follows the true high-frequency distribution and generates rich semantic details. Boundary-aware Mask further constrains the generation of the model to obey the true semantic distribution.

Acknowledgements

This work is supported in part by the National Natural Science Foundation of China, under Grant (62302309, 62171248), Shenzhen Science and Technology Program (JCYJ20220818101014030, JCYJ20220818101012025), and the PCNL KEY project (PCL2023AS6-1), and Tencent "Rhinoceros Birds" - Scientific Research Foundation for Young Teachers of Shenzhen University.