Logic
74 lines
13 KiB
Markdown
74 lines
13 KiB
Markdown
本文的主要idea是使用imf和attnres来加速流匹配的推理速度，同时保持相近的推理质量，imf来自于计算机视觉中的he kaiming的研究Improved mean flows: On the challenges of fastforward generative models,之前是用于流匹配图像生成模型的改进，其前述相关研究如下
|
||
## MeanFlow [@geng2025meanflow]
|
||
Diffusion的生成模型是基于一步一步的方式生成的，这显然不够快，对flow matching中的核心ODE公式
|
||
$$\frac{dx_t}{dt} = v(x_t, t)$$
|
||
这里提一下在看这篇文章时突然产生的对微分算子$d$的理解，$x_t$表示的是某一时刻的具体位置，我们想知道这一时刻的瞬时速度，要根据速度的计算公式进行计算，速度的计算公式是位移时间比，但$x_t$表示的是某一时刻的具体位置，不是一段时间的位移，因此我们使用微分算子表示在$x_t$附近取很小的一段位移除以很小的一段时间，这里一定注意加入微分算子前后的含义变化，对$t$来说加入微分算子前表示的是具体的时刻$t$，而加入后表示的是一段时间$t$。如果取的足够小，通过这种方式计算出来的一小段内的速度就约等于瞬时速度，这是微分算子的作用。
|
||
回到ODE公式本身，可以看出这个公式其实计算的是每一时刻的瞬时速度，而原本的flow matching需要多步采样就来自这个时刻的瞬时性，我们得到瞬时速度后是通过一种近似采样的方式向前移动，是把这个瞬时速度视为一段时间内的平均速度，这也是为什么一开始的模型采样步数多效果就好，因为采样步数太少会导致这种估计非常不准确，使得路径偏移严重。为了解决这种问题，Rectified Flow可译为重整流直接将路径假设为一条直线，注意此处的直线是一种假设，我们只是希望模型能够学会按照这条直线去走，为了能够让模型学会尽可能的按照直线去走，我们在计算损失时让模型预测的速度尽可能接近直线路径的速率，另一方面还通过重整修正，让训练好的但预测路径还不那么直的模型继续学习更直的路线（通过让第一次训练好的模型从任意起点出发得到一个目标点后学习二者间的直线来重整）。
|
||
尽管如此，实际实验中发现，模型学习的路线仍然不可能是我们预期中的直线，因此，这种假设可能本身存在一定的问题，真实的分布转移路径可能并不直，那是否有办法让模型可以学习到真实的平均速度而不是基于路径直线这一假设的平均速度呢？
|
||
回到平均速度的定义，平均速度是位移除以时间，那么任意两个时间点之间的平均速度是
|
||
$$\bar{v}(z_t, r, t) = \frac{1}{t-r} \int_r^t v(z_\tau, \tau)d\tau$$
|
||
我们训练模型去拟合这个平均速度，但写出损失函数后可以发现，我们并不知道这个平均速度的真实值（因为没有假设条件了，因此路径可能是任意的）。但MeanFlow提出，可以通过构建平均速度和瞬时速度的关系来获得真实的平均速度。推导如下
|
||
$$\begin{aligned} & \bar{v}(z_t, r, t) = \frac{1}{t-r} \int_r^t v(z_\tau, \tau) d\tau \\ \Leftrightarrow \quad & (t - r)\bar{v}(z_t, r, t) = \int_r^t v(z_\tau, \tau) d\tau \\ \Leftrightarrow \quad & \frac{d}{dt} \left[ (t - r)\bar{v}(z_t, r, t) \right] = \frac{d}{dt} \int_r^t v(z_\tau, \tau) d\tau \\ \Leftrightarrow \quad & \bar{v}(z_t, r, t) + (t - r)\frac{d}{dt}\bar{v}(z_t, r, t) = v(z_t, t) \\ \Leftrightarrow \quad & \bar{v}(z_t, r, t) = v(z_t, t) - (t - r)\frac{d}{dt}\bar{v}(z_t, r, t) \end{aligned} \tag{1}$$
|
||
其中
|
||
$$\frac{d}{dt}\bar{v}(z_t, r, t) = \frac{d\bar{v}}{dz} \cdot \frac{dz}{dt} + \frac{d\bar{v}}{dr} \cdot \frac{dr}{dt} + \frac{d\bar{v}}{dt} \cdot \frac{dt}{dt} = \frac{d\bar{v}}{dz} \cdot v(z_t, t) + \frac{d\bar{v}}{dt}$$
|
||
$$= jvp(\bar{v}, (z, r, t), (v, 0, 1))$$
|
||
最后可得
|
||
$$\bar{v}(z_t, r, t) = v(z_t, t) - (t - r)jvp(\bar{v}, (z, r, t), (v, 0, 1)) \tag{2}$$
|
||
这个表达式是自举的，即右式中包含左式，左右相关。式中的瞬时速度$v$我们用条件速度$\epsilon-x_0$代替，结合如下的损失函数
|
||
$$L = \mathbb{E}\|\bar{v}_\theta(z_t, r, t) - sg(v(z_t, t) - (t - r)jvp(\bar{v}_\theta, (z, r, t), (v, 0, 1)))\|_2^2$$
|
||
我们可以发现，这个损失函数表示的是一种自我纠正，即我们已知起始和终止位置的情况下的条件瞬时速度始终是$v$，模型预测的平均速度是空间中的宏观流，表示在$z$位置从时间$r$到$t$，平均的流向和平均流向的变化率是什么样的，我们用这个特定样本的瞬时速度校正这个平均速度应该是什么样的，和模型预测的平均速度的差值，最终希望的是模型在这一点的平均速度能够逆推出该特定样本的特定瞬时速度。
|
||
我们此时会考虑这样一个问题，既然我们想要让模型预测平均速度，为何不直接让模型去拟合直线平均速度$\epsilon-x$，而是费劲进行表达式的推导呢。
|
||
原因在于上式中的瞬时速度，从模型最终推理的角度来说，应该是边缘速度场，我们现在考虑的是真实的生成过程中的瞬时速度和平均速度的关系，而边缘速度场不可能是平坦的，因此不能将平均速度场简单的认为和条件速度一致。训练算法如下
|
||
此处模型学会的平均速度不再是一个固定值，而是虽然变化但其变化能符合我们特定样本的条件瞬时速度的平均速度。
|
||
尽管如此，这部分理解起来仍然怪怪的，并不是十分有说服力。这大概也是后续improved meanflow中提出改进的原因。
|
||
可得到训练算法：
|
||
|
||
其中和flow matching的主要区别在于加入了雅可比向量积的计算，引入了约16%的额外计算量，但得到的好处是生成时只需一步生成。
|
||
|
||
## iMF [@gengImprovedMeanFlows2025]
|
||
本篇文章是对MeanFlow的改进工作，MeanFlow的推导过程中对平均速度表达式中的瞬时速度都直接使用条件速度替代，实际上原表达式中的所有瞬时速度都应该是边际速度（边缘概率）而不应该是条件速度，在Flow Matching训练时使用条件速度的原因是条件速度的期望和边际速度二者的梯度是相同的，仅相差常数。
|
||
但在式2中，jvp内的瞬时速度如果使用条件速度，jvp表达式是一个复杂的非线性表达式，此处的瞬时速度本应该使用边际速度，但这里使用条件速度进行替代，而非线性表达式的特性导致该表达式的期望的梯度并不和使用边际速度时的梯度相同，这就导致训练不稳定。
|
||
因此可以对式1进行简单变形，得到瞬时速度和平均速度之间的关系，如下
|
||
$$\mathbf{v}(z_t) = \mathbf{u}(z_t) + (t - r)\frac{d}{dt}\mathbf{u}(z_t) $$
|
||
$$\mathbf{V}_\theta(z_t) \triangleq \mathbf{u}_\theta(z_t) + (t - r)JVP_{sg}(\mathbf{u}_\theta; \mathbf{v}_\theta) \quad (12)$$
|
||
这里$v$和$t$使用同一个网络预测，时刻记住meanflow中预测平均速度的网络$u$的输入是$t,r,z_t$三个参数。当$r=t$时，网络预测的就相当于瞬时速度，因此使用同一个网络预测两遍分别预测$v$和$t$，根据式12得到修正后的瞬时速度，让这个修正后的瞬时速度尽可能接近在直线路径假设下的条件速度$\epsilon-x$。
|
||
这里我们已经如果能让网络直接预测$v_\theta$为什么还要通过式12来计算这个瞬时速度呢，这里要结合我们的训练过程考虑，我们的最终目的仍然和MeanFlow一致，即让网络可以预测平均速度，这里我们看似是在计算瞬时速度的损失，实际上该瞬时速度是由式12的右式得来的，而且jvp不参与梯度计算即其中模型预测的瞬时速度$v_\theta$不会产生梯度更新，梯度更新仅发生在前面的$u_\theta(z_t)$中，因此我们实际上是在要求网络学会预测平均速度$u_\theta$。
|
||
通过这种改进，我们可以充分利用原本的Flow Matching的训练稳定性。
|
||
作者还发现，通过去掉adaLN-zero，但使用将同一个条件token重复多次并作为序列的一部分的形式来进行建模，可以达到和使用adaLN-zero一样的效果。
|
||
|
||
但之前都是用于计算机视觉，我们将其用于机器人的模仿学习和动作生成，之前的类似研究是 Mean-Flow based One-Step Vision-Language-Action 为CVPR2026。但这篇文章只使用了meanflow，我们的improved meanflow改进了meanflow中存在的一些理论问题，同时获得更稳定的训练效果，除此以外我们增加了attnres用于解决模型训练过程中的梯度消失和深层网络特征非常相似的问题。attnres的理论来源于大模型训练过程Residual connections [12] with PreNorm [60] are standard in modern LLMs, yet they accumulate
|
||
all layer outputs with fixed unit weights. This uniform aggregation causes uncontrolled hidden-state
|
||
growth with depth, progressively diluting each layer’s contribution [27]. We propose Attention
|
||
Residuals (AttnRes), which replaces this fixed accumulation with softmax attention over preceding
|
||
layer outputs, allowing each layer to selectively aggregate earlier representations with learned, input-
|
||
dependent weights. To address the memory and communication overhead of attending over all
|
||
preceding layer outputs for large-scale model training, we introduce Block AttnRes, which partitions
|
||
layers into blocks and attends over block-level representations, reducing the memory footprint while
|
||
preserving most of the gains of full AttnRes. Combined with cache-based pipeline communication
|
||
and a two-phase computation strategy, Block AttnRes becomes a practical drop-in replacement for
|
||
standard residual connections with minimal overhead.
|
||
Scaling law experiments confirm that the improvement is consistent across model sizes, and ablations
|
||
validate the benefit of content-dependent depth-wise selection. We further integrate AttnRes into
|
||
the Kimi Linear architecture [69] (48B total / 3B activated parameters) and pre-train on 1.4T tokens,
|
||
where AttnRes mitigates PreNorm dilution, yielding more uniform output magnitudes and gradient
|
||
distribution across depth, and improves downstream performance across all evaluated tasks 。我的概述如下 这里我们换另外一种写法，它能让我们看出更深刻的东西。先记 $\boldsymbol{y}_t = \boldsymbol{f}_t(\boldsymbol{x}_{t-1})$，那么有 $\boldsymbol{x}_t = \boldsymbol{x}_{t-1} + \boldsymbol{y}_t$，约定 $\boldsymbol{y}_0 = \boldsymbol{x}_0$，那么易得 $\boldsymbol{x}_t = \boldsymbol{y}_0 + \boldsymbol{y}_1 + \cdots + \boldsymbol{y}_t$，于是它可以等价地写成
|
||
$$
|
||
\boldsymbol{y}_{t+1} = \boldsymbol{f}_{t+1}(\boldsymbol{y}_0 + \boldsymbol{y}_1 + \cdots + \boldsymbol{y}_t) \tag{2}
|
||
$$
|
||
即从 $\boldsymbol{y}$ 的视角看，Residuals是将 $\boldsymbol{y}_0, \boldsymbol{y}_1, \cdots, \boldsymbol{y}_t$ 等权求和作为 $\boldsymbol{f}_{t+1}$ 的输入来得到 $\boldsymbol{y}_{t+1}$，那么一个自然的推广就是换成加权求和：
|
||
$$
|
||
\boldsymbol{y}_{t+1} = \boldsymbol{f}_{t+1}\left(\sum_{s=0}^t a_{t+1,s} \boldsymbol{y}_s\right) \qquad \text{where} \quad a_{t,s} \ge 0, \quad \sum_{s=0}^t a_{t+1,s} = 1 \tag{3}
|
||
$$
|
||
此时我们可以发现Hyper-connection就是这种加权求和的一种表现形式。但Hyper-connection中的$H$都由$tanh$激活后连乘得来，这导致该值有爆炸或者坍缩的风险。
|
||
Deepseek的mHC通过引入对H的交替归一化使其满足双随机性，双随机性即矩阵的所有行列和都为1,而且双随机性具有乘法的封闭性，即双随机矩阵相乘还是双随机矩阵，这样使得H的连乘不会出现爆炸和坍缩的问题。
|
||
回过头去考虑式3,既然是加权求和，那么一个权重的设置方法就是使用注意力机制，让权重$a$等于注意力分数，于是就有了AttnRes（苏神的作品）。其形式数学上看起来并不困难（RoPE也是这样）。
|
||
$$
|
||
a_{t+1,s} \propto \exp(\boldsymbol{w}_{t+1} \cdot \boldsymbol{y}_s) \tag{6}
|
||
$$
|
||
其中 $\boldsymbol{w}_t$ 是一个可训练的向量参数，即直接以一个数据无关的静态向量为Q、而K、V都是 $\boldsymbol{y}_s$ 去做Softmax Attention，这便是第一版AttnRes。随后通过一些实验得到给K多加个RMSNorm的操作，能取得比较稳定的收益，这构成了终版的AttnRes形式
|
||
$$
|
||
a_{t+1,s} \propto \exp(\boldsymbol{w}_{t+1} \cdot \text{RMSNorm}(\boldsymbol{y}_s)) \tag{7}
|
||
$$
|
||
|
||
我们研究用于对比的模型是Diffusion Policy，和ACT，在我自己的roboimi仿真环境下 有sim_transfer和pocket_insert两个仿真环境，在这个两个仿真环境下我使用imf-attnres和diffusion policy进行了很多实验。之前的实验都是使用swanlab记录的，实验以socket-insert开头 在runs目录下 可以看到之前的实验记录
|