原论文：arxiv 2503.19062v1

引言

Color Transfer with Modulated Flows，这是一篇笔者个人非常喜欢的工作。这篇文章的亮点在于逻辑非常清晰，闭环严密，同时思想非常优雅自然，摆脱了很多工程上的trivial。工作所关注的Color Transfer，实际上就是这样一种问题：给定一张content图像和一张style图像，如何在保持content图像纹理，细节，结构等语义信息全都不变的情况下，将style的颜色风格迁移到content图像上？形式化的讲，这是一个分布迁移问题：每张图像的所有像素点的RGB值，都可以看作是分布在该图像RGB色彩空间上的随机变量，Color Transfer的目标就是将content颜色分布上的随机变量$X_0$加以变换$T$，得到$X_1 = T(X_0)$，使得$X_1$的分布与style图像的颜色分布尽可能接近。

初步想法：Flow Matching及其问题

该问题可以被表述为分布之间的迁移，但这种迁移并不是唯一的：存在无数种转化后的content图像，它们都满足style的颜色分布要求。举个例子，假如style的颜色分布是50%红色，50%蓝色，那么content图像既可以是一半纯红，一半纯蓝，也可以是四分之一红色，四分之一蓝色，剩下的一半是绿色（当然，这两种在视觉上肯定是不佳的），也可以是多种颜色的平滑过渡（视觉上效果更好）。除了分布匹配之外，我们还希望转化后的视觉效果尽可能好，因此该问题目标被重新表述如下：

  1. content转化后的颜色分布与style尽量匹配
  2. 转化后的content颜色过渡尽可能平滑，自然，合理

这样，该问题就可以自然套入最优传输（Optimal Transport）的框架中。假设 content 中有很多红色、蓝色，style 中有很多橙色、青色。我们希望把颜色“搬过去”，但不要乱搬。比如浅红最好变成接近的暖色，而不是突然变成深蓝。于是定义一个 cost：$c(x,y)=|x-y|^2$，然后寻找一个映射$T$，使得$\mathbb{E}_{x\sim \mu_c}[|x-T(x)|^2]$ 尽量小，同时满足：$T\mu_c=\mu_s$。直观的讲，OT最小化传输/转移代价，在期望意义下，原content分布中的每个点都被转移至style分布中最接近的点（不是贪心思想），这就保证了转移后图像色彩的平滑性。

OT未必存在解析解，并且求解困难。因此我们自然想到利用Flow Matching/Rectified Flow方法，用以近似求解OT分布迁移问题：采样得到大量(content, style)像素点对，然后利用FM方法学习从content颜色分布到style颜色分布的映射$T$，最后把content的每个像素都输入这个映射，整个content就被推到了style的颜色分布。

几点说明：

  1. 这里的采样样本全都是单个像素点而非整个图像。这是因为我们关注的是RGB色彩转移，而不是结构/风格的语义转移，单个像素点只包含色彩信息，而如果加上CNN，考虑局部结构，如纹理，边缘等，就会引入结构语义信息，这不是我们关注的。
  2. 每个分布间的映射可以看作FLow，也可以看作Transport plan/coupling $\pi(x,y)$，即联合概率分布。刚开始我们可以假设这个coupling是trivial的，即彼此独立$\pi_{trivial}(x,y)=\mu_c(x)\mu_s(y)$，当然这个coupling会导致传输代价很大。然后我们通过rectified flow学习到一个新的flow$T_{1\text{-rectified}}:X_0\mapsto X_1$，得到coupling$\pi_{1\text{-rectified}}(X_0,X_1) = \pi_0(X_0)\times \delta(X_1-T_{1\text{-rectified}}(X_0)).$，这可以大幅减小传输代价。

然而，这种flow matching有个问题：每拿到一对新的content-style图，就需要重新训练一个FM模型，并且不同的FM模型参数不同，不能复用，假如有N张图，就需要训练$O(N^2)$个模型。也就是说，FM模型是对输入特化的，不能泛化到其他图像对，导致开销很大。于是我们需要一个通用的FM模型，它可以在训练阶段学习到所有图像对的颜色分布迁移规律，在测试阶段直接输入新的content-style图像对，就能得到转化后的content图像。

从数学洞察到普适思想：求同存异

对于一维的$X_0$和$X_1$及其累积分布函数CDF$F_0$和$F_1$，最优传输映射$T^*$可以通过CDF的逆函数来求解：

推广到3维，这一公式不精确成立，但可以提供一个很好的近似。此外，在高维下CDF的逆函数不一定存在，但FM构造出的ODE是可逆的，也就是说我们依旧可以通过FM模型从target分布倒退回source分布，也就实现了逆CDF的效果。再次审视这个公式，可以写成$T^*(x) = F_1^{-1}(Y)$，其中$Y = F_0(x)$。由于$Y$是$X_0$的CDF，服从$[0,1]$上的均匀分布，因此我们得到一个很强的启发：可以把$Y$看作是一个通用的中间变量，它不依赖于具体的source和target分布。也就是说，所有的source分布都可以映射到这个均匀分布的latent变量上，而所有的target分布都可以从这个中间变量映射回去。于是我们就有了一个普适的思想：求同存异。

这样，我们就可以只构造N个FM模型：N个source分布到中间变量的映射（不用构造逆映射，因为FM本身就是可逆的），就可以实现对任意source-target图像对的颜色迁移，而不需要为每一对图像都训练一个FM模型。但这样依然是依赖于输入的，我们仍然需要训练N个模型。那么我们能否进一步把N个模型合并为一个模型呢？

工程化的实践：参数化FLow

上面提到，我们通过FM模型构造两个分布之间的映射$T$，实际上，这个模型可以相当简单，因为输入仅仅是4维向量：(R, G, B, t)，其中t是时间步。作者仅使用了一个2层MLP作为该FM模型，该模型实际上就对应了一个唯一的flow。既然这样的FM模型需要训练很多个，而它本身又非常简单，那么我们是否可以把这个flow参数化，让另一个大模型去学习该FM模型的参数呢？这样就可以把N个FM模型合并为一个大模型了。于是作者提出了Modulated Flow的概念：通过一个大模型（如CNN）去学习FM模型的参数，这个大模型的输入是图像，输出该图像到均匀分布的flow参数。论文中把这个大模型叫做编码器，这样，只需要分别输入新的content和style图像，就能得到对应的FM模型，从而实现颜色迁移。这里的 flow 网络参数不是正常训练出来的，而是被 encoder 预测出来的。所以它叫 modulated flow：flow 的权重由 encoder 调制/生成。

训练分两个阶段：
alt text
alt text

推理：
alt text

说明：作者发现，使用这种encoder会导致FM模型输出的target分布与真实目标分布略有偏移，但也提供了隐式正则化，降低了ModFlow相较于从头训练的Rectified Flow的平均 Lipschitz 常数，从⽽产⽣更少视觉伪影。

alt text

一点思考

  • 为什么这种做法对于Color Transfer问题是可行的？为什么能保留空间内容？
  • 因为该模型的推理过程本质上是在pixel而非整个image的层面进行的，这就保证了没有引入局部空间信息。此外，pixel space的rectified flow也保证了这一变换符合OT性质，进而保证了平滑性。

结语

这篇工作有两个亮点：分布迁移中引入中间态的优雅高角度观点与学习参数的有益工程启发。是一篇有价值的好文章。