Every Pixel Has its Moments: Ultra-High-Resolution Unpaired Image-to-Image Translation via Dense Normalization

Unpaired I2I translation aims to train a generator $\mathcal{G}$ to translate an image $\boldsymbol{X}$ in domain $\mathcal{X}$ to another domain $\mathcal{Y}$ even when there are no corresponding paired images. The output of $\mathcal{G}$ is denoted as $\hat{\boldsymbol{Y}}$ and is expected to be in domain $\mathcal{Y}$. In the context of UHR unpaired I2I translation, all images in domain $\mathcal{X}$ have a high resolution of $H\times W$. To enable $\mathcal{G}$ to handle images with infinite resolution, the model must be trained and executed in a patch-wise manner to reduce the GPU space complexity to a constant.

After training an I2I generator $\mathcal{G}$ with patches, the IN layers are replaced with DN layers. During the inference process, an UHR image $\boldsymbol{X}\in\mathbb{R}^{H\times W}$ is divided into patches $\boldsymbol{x}^{\text{patch}}_{c,r}\in\mathbb{R}^{N\times N}$, each with a size of $N\times N$, and their coordinates ($c$, $r$) relative to the original image are recorded. Here, $c\in\{0,1,...,\lceil\frac{H}{N}\rceil-1\}$ and $r\in\{0,1,...,\lceil\frac{W}{N}\rceil-1\}$. Then, a dispatcher sequentially inputs these patches and coordinates into the generator.

Fig. 3 illustrates the overall pipeline of our framework. Specifically, during each dispatch, two patches along with their coordinates are sent: one to the prefetching branch and the other to the inference branch. Both patches simultaneously go through all layers except the DN layer. In the DN layer, the patch in the prefetching branch undergoes a standard instance normalization [17], storing the resultant statistics in the cache table ($T_{\mu},T_{\sigma}$) with the coordinates as keys. The patch in the inference branch uses its coordinates to query the cache table for its own and its eight neighbors’ statistics, forming two $3\times 3$ matrices of coarse-level (patch-level) statistical moments ($\tilde{\boldsymbol{\mu}}^{\text{patch}}_{c,r}$, $\tilde{\boldsymbol{\sigma}}^{\text{patch}}_{c,r}\in\mathbb{R}^{3\times 3}$). Fast interpolation is then applied to these matrices to estimate fine-level (pixel-level) statistical measures ($\hat{\boldsymbol{\mu}}^{\text{pixel}}_{c,r}$, $\hat{\boldsymbol{\sigma^{*}}}^{\text{pixel}}_{c,r}\in\mathbb{R}^{N\times N}$), which are subsequently used for dense normalization. The translated patches are then reassembled into an UHR image, with the DN operation effectively reducing tiling artifacts.

Here, $\mathbb{E}[\boldsymbol{x}^{\text{patch}}_{c,r}]$ and $\sqrt{Var[\boldsymbol{x}^{\text{patch}}_{c,r}]}$ are denoted as ${\mu}^{\text{patch}}_{c,r}$ and ${\sigma}^{\text{patch}}_{c,r}$, respectively. These represent the mean and standard deviation of $\boldsymbol{x}^{\text{patch}}_{c,r}$. Subsequently, these two statistics are stored in the cache table using their coordinates as keys; specifically, $T_{\mu}[c][r]:={\mu}^{\text{patch}}_{c,r}$ and $T_{\sigma}[c][r]:={\sigma}^{\text{patch}}_{c,r}$.

We assume that the patch-level statistics represent the statistics for the central pixel of the patch; for example, $\hat{\boldsymbol{\mu}}^{\text{pixel}}_{c,r}\left[\frac{N}{2}\right]\left[\frac{N}{2}\right]={\mu}^{\text{patch}}_{c,r}$. Hence, we can utilize the process below to derive $\hat{\boldsymbol{\mu}}^{\text{pixel}}_{c,r}$ from $\tilde{\boldsymbol{\mu}}^{\text{patch}}_{c,r}$, with Fig. 4(a) providing a visual representation of the entire interpolation process.

1. 1.

   Perform fast interpolation on each corner of the $3\times 3$ matrix $\tilde{\boldsymbol{\mu}}^{\text{patch}}_{c,r}$. For instance, for the top-left corner $\tilde{\boldsymbol{\mu}}^{\text{patch}}_{c,r}[0:1,0:1]$ (a $2\times 2$ submatrix), apply fast interpolation.

2. 2.

   This interpolation is performed on a $2\times 2$ submatrix to expand it to an $N\times N$ matrix.

3. 3.

   By interpolating each $2\times 2$ submatrix into an $N\times N$ matrix, a larger $2N\times 2N$ matrix is constructed.

4. 4.

   The central $N\times N$ submatrix is then extracted from this $2N\times 2N$ matrix, serving as the pixel-level statistical estimation $\hat{\boldsymbol{\mu}}^{\text{pixel}}_{c,r}$ for the patch $\boldsymbol{x}^{\text{patch}}_{c,r}$.

For $\tilde{\boldsymbol{\sigma}}^{\text{patch}}_{c,r}$, we first calculate the inverse of each element to form $\tilde{\boldsymbol{\sigma}^{*}}^{\text{patch}}_{c,r}$. Then, the same interpolation and cropping processes are conducted to obtain the pixel-level statistical estimation $\hat{\boldsymbol{\sigma^{*}}}^{\text{pixel}}_{c,r}\in\mathbb{R}^{N\times N}$ for the patch $\boldsymbol{x}^{\text{patch}}_{c,r}$.

Now, we can utilize the pixel-level statistical moments to perform dense normalization, denoted as $DN(\cdot)$.

Given a matrix $\boldsymbol{Q}\in\mathbb{R}^{2\times 2}$, if we wish to interpolate it into $\boldsymbol{Q}^{\prime}\in\mathbb{R}^{N\times N}$, using standard bilinear interpolation can be formulated as follows:

This can be further reformulated as:

where $\left\langle,\right\rangle$ denotes the Frobenius inner product. It can be noted that for a given coordinate $(i,j)$, the interpolated value $\boldsymbol{Q}^{\prime}[i][j]$ is a weighted sum of elements in $\boldsymbol{Q}$ with a fixed set of weights since $N$ is a constant.

Thus, the final interpolated result $\boldsymbol{Q}^{\prime}$ can be written as:

This equation represents the fast interpolation process (see Fig. 4 (b)), where the elements in each matrix $\boldsymbol{M}_{k,l}\in\mathbb{R}^{N\times N}$ are defined as follows:

This reformulation highlights fast interpolation’s desirable features (see Table 4). First, it consists solely of matrix multiplication, which can be accelerated by a GPU. Second, all matrices $\boldsymbol{M}_{k,l}$ are consistent across all interpolation operations in our dense normalization, allowing for precomputation and caching.

In our experiments with the aforementioned datasets, we cropped the UHR images into 512$\times$512 patches and trained CycleGAN, CUT, and L-LSeSim frameworks for 100 epochs with default hyperparameters. We replaced the IN layers with our DN layers for the inference process and compared the results with patch-wise IN, TIN, and KIN methods. The experiments were conducted on an NVIDIA RTX 3090 GPU. However, due to the GPU memory limitation, we were unable to directly use the IN layer with UHR images as input without cropping. We presented results obtained using the CUT model, with further findings available in the supplementary material.

We compared the results from our DN, patch-wise IN, TIN, and KIN methods using qualitative and quantitative evaluation techniques. The translated UHR images were assessed using the standard Fréchet inception distance (FID) [35] metric. Additionally, we conducted a downstream task to differentiate between patches from the source and target domains. To explicitly showcase the adverse effects of tiling artifacts or over/under-colorization, we intentionally cropped new patches across the raw translated patches.

However, since these metrics are not designed to evaluate UHR images and may not accurately identify tiling artifacts or over/under-colorizing, we conducted a detailed human evaluation consisting of three quality challenges. Participants were shown the source and translated images generated by DN, patch-wise IN, TIN, and KIN and asked to identify the image with the best quality, the fewest tiling artifacts, and the best color and hue. We recruited computer vision and pathology specialists in addition to the general population for the evaluation. For the translated WSIs, we conducted a fidelity challenge, following the AMT perceptual studies protocol [36].

Fig. 6, Fig. S4, and Fig. S5 in the supplementary material show UHR images translated from natural and pathological WSI datasets. These images reveal that patch-wise IN generates a significant amount of gap-type tiling artifacts, while KIN mitigates some of these artifacts but at the expense of details, hue, and color. TIN reduces gap-type tiling artifacts but results in over/under-colorizing, loss of local hue details, and creates jitter-type tiling artifacts. Conversely, our DN approach is the only method that successfully diminishes tiling artifacts while maintaining local hue and color details, producing the best results.

Table 1 (left part) presents the standard FID evaluation results on various datasets. Since the backbone model is optimized using patch-wise IN, this method can be considered as having the optimal FID values, thereby setting an intuitive lower bound, while other methods aim to minimally disturb the translation process to remove tiling artifacts. Overall, our hyperparameter-free DN ($\mathcal{P}_{4}$) outperforms TIN and KIN, indicating that it introduces the smallest adjustment necessary to achieve the goal, and local color and hue are preserved ($\mathcal{P}_{1}$). Remarkably, in some cases, it even surpasses the intuitive lower bound. On the other hand, KIN secures second place due to its balance between patch-wise IN and TIN. TIN inevitably yields the worst results because the use of global statistics introduces the largest disturbance.

Table 1 (right part) displays the results of the domain differentiation downstream task. DN consistently outperforms all other methods, likely indicating the involvement of the fewest tiling artifacts ($\mathcal{P}_{2}$) and the preservation of hue and color ($\mathcal{P}_{1}$). On the other hand, TIN yields the worst results, which is probably due to the large disturbances introduced by the use of global statistics.

To address the limitations of available metrics, we employed human evaluation to assess image quality (see Table 2). Three image quality challenges were conducted by forty participants, and our DN method achieved the best performance across all three challenges ($\mathcal{P}_{1}$ and $\mathcal{P}_{2}$), particularly for the Kyoto summer2autumn dataset. Additionally, we recruited eight computer vision specialists to evaluate the translation of natural images and eight pathology specialists to evaluate the results of stain transformation. Interestingly, the effectiveness of DN was more pronounced to these specialists. Furthermore, the fidelity challenge (see Fig. S8 in the supplementary material) revealed that the images generated by DN were nearly indistinguishable from real pathological images.

Table 3 presents the runtime and GPU VRAM usage for various methods. Employing operations on statistical moments generally leads to longer runtime but yields superior results. Distinctively, DN achieves faster execution in a single pass ($\mathcal{P}_{3}$) compared to KIN, with a modest increase in GPU VRAM usage. This efficiency is due to the parallel execution of prefetching and the inference branch, highlighting DN’s innovative approach to parallelism design.