机器学习笔记(17.6): FloQ

可能更好的阅读体验 Jeefy's Blog: FloQ

所以 floq 其实就是将 FQL 中的 Q network 变成了用 flow-matching “多步推理” 的密集监督的网络。

核心建模

其核心的建模就是：

\[Q(s, a) = z(0) + \frac 1 K \sum_{i = 1}^K v_{\theta}\left(\left.\frac i K, \psi_{\theta}\left(\left. \frac {i - 1} K, z \right\vert s, a\right) \right\vert s, a\right) \]

换成递推式：

\[z(t + 1) \leftarrow z(t) + \frac 1 K v_{\theta}\left(\left. \frac t K, z(t) \right\vert s, a\right) \]

那么根据 flow-matching 最开始需要用一点随机的噪声，这里就用的是：

\[\begin{aligned} z &\sim {\rm Unif}[l, u] \\ u &= Q_{max} \\ u - l &= {\kappa} \cdot (Q_{max} - Q_{min}) & \kappa \text{根据训练稳定性越大越好} \end{aligned} \]

这个 \(u-l\) 太小了，那么就会退化成直线轨迹，如果太大了，就会导致训练困难（积分不稳定）

损失也很简单：

\[{\mathcal L}_{floq}(\theta) = {\mathbb E}_{z, t} \left[ \left\| v_\theta(t, z(t) | s, a) - \frac {y(s, a) - z}{1 - 0} \right\|^2 \right] \]

---

网络的输入

• \(s\)：状态向量

• \(a\)：动作向量

• \(z_j\)：插值中间值，使用 HL-Gauss 编码。

• \(t\)：时间值，采用 Fourier 嵌入。

• \(\Delta\) HL-Gauss

就是将 \(z_j\) 这个标量，通过在数轴上放置一个 \(N(z_j, \sigma^2)\) 的高斯分布，然后区域求概率后，归一化而来。

\({\rm erf}(x) = \frac 2 {\sqrt \pi} \int_0^x e^{-t^2} {\rm d}t\) 也就是误差函数，可以对应高斯分布的累计概率函数
若 \(X \sim N(\mu, \sigma^2)\) 则 \(P(X \le x) = \Phi(\frac {x - \mu}{\sigma}) = \frac 1 2 [1 + {\rm erf}(\frac {x - \mu}{\sqrt 2 \sigma})]\)
论文中采用 \(N = 50, \sigma=16\)

取 \({\rm bin\_width} = \frac {Q_{max} - Q_{min}} {N - 1}\)，那么：

\[{\rm bin\_prob}_k = {\rm erf} \left(\frac {s_{k + 1} - z_j}{\sqrt 2 \sigma} \right) - {\rm erf} \left(\frac {s_{k} - z_j}{\sqrt 2 \sigma} \right) \]

最后对 \(\rm bin\_prob\) 归一化一下。

需要注意，\(Q_{max}, Q_{min}\) 是全局的最大最小，不是基于某一个局面的最大最小。
另外，会适当的将 \(Q_{max}, Q_{min}\) 向外扩张，留出余量，减少 \(z_j\) 跑出这个区间的可能。

• \(\Delta\) Fourier Embedding

\[\gamma(t) = [\sin(2\pi b_1 t), \cos (2\pi b_1 t), \cdots, \sin(2\pi b_{32} t), \cos(2\pi b_{32}t)] \in {\mathbb R}^{64} \]

所以为什么要这样编码？

1. 数值稳定

做过 DQN 的人都知道，Q 函数的值域可以非常大，所以直接将 \(z_j\) 按照 Q 的实际值输入，可能导致 MLP 激活过大梯度爆炸之类的问题。而使用了 HL-Gauss，可以保证分量有界，而且变化的是分布，而非数值尺度。
并且输入了更多维度的 \(z_j\)，平滑，且信息更丰富，网络对于这个信息的 “注意力” 会更强。

2. 谱偏差理论：Rahaman et al., ICML 2019《On the Spectral Bias of Neural Networks》；Tancik et al., NeurIPS 2020《Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains》

3. 分类 vs 回归的漂移鲁棒性：Farebrother et al. 2024《Stop Regressing: Training Value Functions via Classification for Scalable Deep RL》

---

解决难题？

这个框架提供了高密度的监督信号，类似于 LLM 的 CoT。这里的每一步的速度场，其实就和 CoT 非常类似，让模型进行多步迭代和思考。但是更本质的叙述应该是 “高密度信号监督” 带来的表示能力提升。

这就是论文中所述的 test-time scaling in RL。通过更多计算获取更强能力。

但是为什么这样会更好，这仍然是个黑盒。论文中的一个论述逻辑是：

直线轨迹 ⟹ \(v\) 无需依赖 \((z, t)\) ⟹ 与 monolithic 网络同构；弯曲轨迹 ⟹ \(v\) 必须消费 \((z,t)\) 输入 ⟹ 迭代计算有真实容量。

这个真实容量，赋予了网络逐步修正前面步的偏差的能力，让逐步的积分变成串行的误差修正过程。

但是为什么存在随着 K 出现了性能下降的问题？这个可能可以归结于积分步骤过多导致的过拟合，以及 TD 学习（FQL 框架下）固有的动力学不稳定的问题（Q 网络的分布漂移）

另一个实验是，如果训练时积分多，但是推理时积分少，性能也会非常的接近，说明这个 test-time scaling 的收益主要在于训练时的监督信号让 Q 值的估计更加精确。这个和 LLM 的 scaling 有所区别。

还有一个实验是，如果只监督 \(t = 0\) 呢？单步积分的表现就已经 \(14\% \to 49\%\) 了，多步积分到 \(53\%\)（HM-Large）。作者解释为 对不同噪声回归同一目标构成了更鲁棒的监督信号，从而提高了整体收益。

---

更进一步：《What Does Flow Matching Bring To TD Learning?》用更多的实验说明了这个机制。