We propose an image restoration algorithm that can control the perceptual quality and/or the mean square error (MSE) of any pre-trained model, trading one over the other at test time. Our algorithm is few-shot: Given about a dozen images restored by the model, it can significantly improve the perceptual quality and/or the MSE of the model for newly restored images without further training. Our approach is motivated by a recent theoretical result that links between the minimum MSE (MMSE) predictor and the predictor that minimizes the MSE under a perfect perceptual quality constraint. Specifically, it has been shown that the latter can be obtained by optimally transporting the output of the former, such that its distribution matches the source data. Thus, to improve the perceptual quality of a predictor that was originally trained to minimize MSE, we approximate the optimal transport by a linear transformation in the latent space of a variational auto-encoder, which we compute in closed-form using empirical means and covariances. Going beyond the theory, we find that applying the same procedure on models that were initially trained to achieve high perceptual quality, typically improves their perceptual quality even further. And by interpolating the results with the original output of the model, we can improve their MSE on the expense of perceptual quality. We illustrate our method on a variety of degradations applied to general content images of arbitrary dimensions.

翻译：我们提出一种图像复原算法，能够控制任何预训练模型的感知质量和/或均方误差（MSE），并在测试时实现两者之间的权衡。我们的算法具有少样本特性：给定约十几张由模型复原的图像，无需进一步训练即可显著提升模型对新复原图像的感知质量和/或MSE。本方法的动机源于近期一项连接最小均方误差（MMSE）预测器与完美感知质量约束下最小化MSE预测器的理论成果。具体而言，研究表明后者可通过前者的输出进行最优传输获得，使其分布与源数据匹配。因此，为改进原本以最小化MSE为目标训练的预测器的感知质量，我们通过变分自编码器隐空间中的线性变换来近似最优传输，该变换可利用经验均值与协方差以闭式解计算。超越理论范畴，我们发现对最初以高感知质量为训练目标的模型应用相同流程，通常能进一步提升其感知质量。通过将处理结果与模型原始输出进行插值，我们可以在牺牲感知质量的前提下改善其MSE。我们在任意尺寸的通用内容图像上应用多种退化类型，验证了本方法的有效性。