我自己的原文哦~ https://blog.51cto.com/whaosoft/12750420
#傅里叶特征 (Fourier Feature）与核回归
位置编码背后的理论解释
本文探讨了位置编码背后的理论基础，特别是傅里叶特征（Fourier Feature）与核回归（Kernel Regression）的联系，并解释了如何通过这些理论来增强神经网络对高频信息的学习能力。
最近我在看位置编码最新技术时，看到了一个叫做 "NTK-aware" 的词。我想：「"NTK"是什么？Next ToKen （下一个词元）吗？为什么要用这么时髦的缩写？」看着看着，我才发现不对劲。原来，NTK 是神经网络理论里的一个概念，它从 kernel regression 的角度解释了神经网络的学习方法。基于 NTK 理论，有人解释了位置编码的理论原理并将其归纳为一种特殊的 Fourier Feature （傅里叶特征）。这么多专有名词一下就把我绕晕了，我花了几天才把它们之间的关系搞懂。
在这篇文章里，我主要基于论文_Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains_ （后文简称为「傅里叶特征论文」），介绍傅里叶特征这一概念。为了讲清这些理论的发展脉络，我会稍微讲一下 NTK 等理论概念。介绍完傅里叶特征后，我还会讲解它在其他方法中的应用。希望读完本文后，读者能够以这篇论文为基点，建立一个有关位置编码原理的知识网络，以从更深的层次来思考新的科研方向。
用 MLP 表示连续数据
我们先从一个具体的任务入手，直观体会傅里叶特征能够做些什么事。
我们知道，神经网络，哪怕是最简单的多层感知机（MLP），都有着很强的泛化能力：训练完毕后，对于训练集里完全没见过的输入，网络也能给出很正确的输出。特别地，如果新输入恰好和训练集的某个输入很近，那么它的输出也会和对应的训练集输出很近；随着新输出与训练集输入的距离不断增加，新输出也会逐渐变得不同。这反映了神经网络的连续性：如果输入的变化是连续的，那么输出的变化也是连续的。
基于神经网络的这一特性，有人想到：我们能不能用神经网络来表示连续数据呢？比如我想表达一张处处连续的图像，于是我令神经网络的输入是(x, y) 表示的二维坐标，输出是 RGB 颜色。之后，我在单张图像上过拟合这个 MLP。这样，学会表示这张图像后，哪怕输入坐标是分数而不是整数，神经网络也能给出一个颜色输出。
这种连续数据有什么好处呢？我们知道，计算机都是以离散的形式来存储数据的。比如，我们会把图像拆成一个个像素，每个像素存在一块内存里。对于图像这种二维数据，计算机的存储空间还勉强够用。而如果想用密集的离散数据表达更复杂的数据，比如 3D 物体，计算机的容量就捉襟见肘了。但如果用一个 MLP 来表达 3D 物体的话，我们只需要存储 MLP 的参数，就能获取 3D 物体在任何位置的信息了。
这就是经典工作神经辐射场 (Neural Radiance Field, NeRF) 的设计初衷。NeRF 用一个 MLP 拟合 3D 物体的属性，其输入输出如下图所示。我们可以用 MLP 学习每个 3D 坐标的每个 2D 视角处的属性（这篇文章用的属性是颜色和密度）。根据这些信息，利用某些渲染算法，我们就能重建完整的 3D 物体。
上述过程看起来好像很简单直接。但在 NeRF 中，有一个重要的实现细节：必须给输入加上位置编码，MLP 才能很好地过拟合连续数据。这是为什么呢？让我们先用实验复现一下这个现象。
MLP 拟合连续图像实验
为了快速复现和位置编码相关的问题，我们简单地用一个 MLP 来表示图像：MLP 的输入是 2D 坐标，输出是此处的三通道 RGB 颜色。我为这篇博文创建一个 GitHub 文件夹 https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/FourierFeature ，该实验的 Notebook 代码在文件夹的image_mlp.ipynb 中，欢迎大家 clone 项目并动手尝试。
一开始，我们先导入库并可视化要拟合的图片。
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.io import read_image, ImageReadMode
from torchvision.transforms.functional import to_pil_imagefrom tqdm import tqdm
from einops import rearrangedef viz_image(pt_img: torch.Tensor):pil_img = to_pil_image(pt_img)display(pil_img)input_image = read_image('misuzu.png', ImageReadMode.RGB)
input_image = input_image.to(torch.float32) / 255
input_image = input_image.unsqueeze(0)
input_image = F.interpolate(input_image, (256, 256), mode='bilinear')
viz_image(input_image[0])
我们再定义一个 MLP 类。稍后我们会并行地传入二维坐标。具体来说，我们会将输入定义为一个[1, 2, H, W] 形状的数据，其中通道数 2 表示(i, j) 格式的坐标。由于输入是以图像的形式并行输入的，我们可以用 的 2D 卷积来表示二维数据上的并行 MLP。所以在下面这个 MLP 里，我们只用到 卷积、激活函数、归一化三种层。按照傅里叶特征论文的官方示例，网络最后要用一个 Sigmoid 激活函数调整输出的范围。
class MLP(nn.Module):def __init__(self, in_c, out_c=3, hiden_states=256):super().__init__()self.mlp = nn.Sequential(nn.Conv2d(in_c, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),nn.Conv2d(hiden_states, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),nn.Conv2d(hiden_states, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),nn.Conv2d(hiden_states, out_c, 1), nn.Sigmoid())def forward(self, x):return self.mlp(x)之后我们来定义训练数据。在一般的任务中，输入输出都是从训练集获取的。而在这个任务中，输入是二维坐标，输出是图像的颜色值。输出图像input_image 我们刚刚已经读取完毕了，现在只需要构建输入坐标即可。我们可以用下面的代码构建一个[1, 2, H, W] 形状的二维网格，grid[0, :, i, j] 处的数据是其坐标(i, j) 本身。当然，由于神经网络的输入一般要做归一化，所以我们会把原本0~H 和0~W 里的高宽坐标缩放都到0~1。最终grid[0, :, i, j]==(i/H, j/W)。
H, W = input_image.shape[2:]h_coord = torch.linspace(0, 1, H)
w_coord = torch.linspace(0, 1, W)
grid = torch.stack(torch.meshgrid([h_coord, w_coord]), -1).permute(2, 0, 1).unsqueeze(0)准备好一切后，我们就可以开始训练了。我们初始化模型model 和优化器optimizer，和往常一样训练这个 MLP。如前所述，这个任务的输入输出非常直接，输入就是坐标网格grid，目标输出就是图片input_image。每训练一段时间，我们就把当前 MLP 拟合出的图片和误差打印出来。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = MLP(2).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
n_loops = 400
input_image = input_image.to(device)
grid = grid.to(device)
for epoch in tqdm(range(n_loops)):output = model(grid)loss = F.l1_loss(output, input_image)optimizer.zero_grad()loss.backward()optimizer.step()if epoch % 100 == 0 or epoch == n_loops - 1:viz_image(output[0])print(loss.item())运行代码，大致能得到如下输出。可以看到，从一开始，图像就非常模糊。
不过，如果我们在把坐标输入进网络前先将其转换成位置编码——一种特殊的傅里叶特征，那么 MLP 就能清晰地拟合出原图片。这里我们暂时不去关注这段代码的实现细节。
class FourierFeature(nn.Module):def __init__(self, in_c, out_c, scale):super().__init__()fourier_basis = torch.randn(in_c, out_c // 2) * scaleself.register_buffer('_fourier_basis', fourier_basis)def forward(self, x):N, C, H, W = x.shapex = rearrange(x, 'n c h w -> (n h w) c')x = x @ self._fourier_basisx = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)x = 2 * torch.pi * xx = torch.cat([torch.sin(x), torch.cos(x)], dim=1) return xfeature_length = 256
model = MLP(feature_length).to(device)
fourier_feature = FourierFeature(2, feature_length, 10).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
n_loops = 400
for epoch in tqdm(range(n_loops)):x = fourier_feature(grid)output = model(x)loss = F.l1_loss(output, input_image)optimizer.zero_grad()loss.backward()optimizer.step()if epoch % 100 == 0 or epoch == n_loops - 1:viz_image(output[0])print(loss.item())
prev_output = output
简单地对比一下，此前方法的主要问题是 MLP 无法拟合高频的信息（如图块边缘），只能生成模糊的图像。而使用位置编码后，MLP 从一开始就能较好地表示高频信息。可见，问题的关键在于如何让 MLP 更好地拟合数据的高频信息。
接下来，我们来从一个比较偏理论的角度看一看论文是怎么分析位置编码在拟合高频信息中的作用的。
核回归
傅里叶特征论文使用了神经正切核（Nerual Tangent Kernel, NTK）来分析 MLP 的学习规律，而 NTK 又是一种特殊的核回归 (Kernel Regression) 方法。在这一节里，我会通过代码来较为仔细地介绍核回归。下一节我会简单介绍 NTK。
和神经网络类似，核回归也是一种数学模型。给定训练集里的输入和输出，我们建立这样一个模型，用来拟合训练集表示的未知函数。相比之下，核回归的形式更加简单，我们有更多的数学工具来分析其性质。
核回归的设计思想来源于我们对于待拟合函数性质的观察：正如我们在前文的分析一样, 要用模型拟合一个函数时, 该模型在训练数据附近最好是连续变化的。离训练集输入越近, 输出就要和其对应输出越近。基于这种想法，核回归直接利用和所有数据的相似度来建立模型：假设训练数据为 , 我们定义了一个计算两个输入相似度指标 , 那么任意输入 的输出为:
也就是说，对于一个新输入 ，我们算它和所有输入 的相似度 ，并把相似度归一化。最后的输出 是现有 的相似度加权和。
这样看来，只要有了相似度指标，最终模型的形式也就决定下来了。我们把这个相似度指标称为「核」。至于为什么要把它叫做核，是因为这个相似度指标必须满足一些性质，比如非负、对称。但我们这里不用管那么多，只需要知道核是一种衡量距离的指标，决定了核就决定了核回归的形式。
我们来通过一个简单的一维函数拟合实验来进一步熟悉核回归。该实验代码在项目文件夹下的kernel_regression.ipynb 中。
先导入库。
%matplotlib inlineimport numpy as np
import matplotlib.pyplot as plt再创建一个简单的非线性函数，做为我们的拟合目标。这个函数就是一个简单的周期为 2 的正弦函数乘上线性函数 。我们可以简单可视化一下函数在 之间的图像。
def func(x):return np.sin(np.pi * x) * (1 - x)xs = np.linspace(-1, 1, 100)
ys = func(xs)
plt.plot(xs, ys)
plt.show()
基于这个函数，我们等间距地选一些点做为训练数据。
sample_x = np.linspace(-1, 1, 10)
sample_y = func(sample_x)
plt.scatter(sample_x, sample_y)
plt.show()
有了数据后，我们来用核回归根据数据拟合这个函数。在决定核回归时，最重要的是决定核的形式。这里我们用正态分布的概率密度函数来表示核，该核唯一的超参数是标准差，需要我们根据拟合结果手动调整。标准差为1 的标准正态分布核的图像如下所示。由于最后要做归一化，正态分布密度函数的系数被省略掉了。
def kernel_func(x_ref, x_input, sigma=1):return np.exp(-(x_input-x_ref)**2 / (2 * sigma**2))xs = np.linspace(-1, 1, 100)
ys = kernel_func(0, xs)
plt.plot(xs, ys)
plt.show()
可以从图像中看出，离某输入越近（假设该输入是0），那么相似度就越高。这符合我们对于相似度函数的要求。
有了核函数后，我们就直接得到了模型。根据核回归模型计算结果的函数为kernel_regression。函数参数xs, ys 表示训练数据，x_input 表示测试时用的输入坐标，sigma 是核回归的超参数。
假设有n 个训练样本，有m 个测试输入，那么我们要计算每个测试输入对每个训练输入的n * m 个相似度，这些相似度会存到矩阵weight 里。为此，我们需要对xs 和x_input 做一些形状变换，再用上面定义的核函数kernel_func 求出每对相似度。有了相似度后，我们根据公式计算点乘结果weight_dot 及归一化系数weight_sum，并最终计算出核回归的结果res。
基于这个函数，我们可以将测试输入定义成[-1, 1] 上一些更密集的坐标，并用上面定义好的 10 个样本做为训练集，得到核回归的结果。
def kernel_regression(xs, ys, x_input, sigma=1):# xs: [n, ]# ys: [n, ]# x_input: [m, ]N = xs.shape[0]xs = np.expand_dims(xs, 1)ys = np.expand_dims(ys, 1)x_input = np.expand_dims(x_input, 0)x_input = np.repeat(x_input, N, 0)weight = kernel_func(xs, x_input, sigma) # [n, m]weight_sum = np.sum(weight, 0)weight_dot = weight.T @ ysweight_dot = np.squeeze(weight_dot, 1)res = weight_dot / weight_sumreturn ressigma = 1
xs = np.linspace(-1, 1, 100)
ys = kernel_regression(sample_x, sample_y, xs, sigma)
plt.title(f'sigma = {sigma}')
plt.plot(xs, ys)
plt.show()我们可以通过修改sigma 来得到不同的拟合效果。以下是我的一些结果：
可以看出，标准差越小，模型倾向于过拟合；随着标准差变大，曲线会逐渐平缓。我们需要不断调整超参数，在过拟合和欠拟合之间找到一个平衡。这种现象很容易解释：正态分布核函数的标准差越小，意味着每个训练数据的影响范围较小，那么测试样本更容易受到少数样本的影响；标准差增大之后，各个训练样本的影响开始共同起作用，我们拟合出的函数也越来越靠近正确的函数；但如果标准差过大，每个训练样本的影响都差不多，那么模型就什么都拟合不了了。
从实验结果中，我们能大致感受到核回归和低通滤波很像，都是将已知数据的平均效果施加在未知数据上。因此，在分析核回归的时候，往往会从频域分析核函数。如果核函数所代表低通滤波器的带宽 （bandwidth）越大，那么剩下的高频信息就更多，核回归也更容易拟合高频信息较多的数据。
神经正切核
那么，核回归是怎么和神经网络关联起来的呢？有研究表明，在一些特殊条件下，MLP 的最终优化结果可以用一个简单的核回归来表示。这不仅意味着我们可以神奇地提前预测梯度下降的结果，还可以根据核回归的性质来分析神经网络的部分原理。这种能表示神经网络学习结果的核被称为神经正切核（NTK）。
这些特殊条件包括 MLP 无限宽、SGD 学习率的学习率趋近 0 等。由于这些条件和实际神经网络的配置相差较远，我们难以直接用核回归预测复杂神经网络的结果。不过，我们依然可以基于这些理论来分析和神经网络相关的问题。傅里叶特征的分析就是建立在 NTK 上的。
NTK 的形式为
其中, 是参数为 的神经网络, 为内积运算。简单来看, 这个式子是说神经网络的核回归中，任意两个向量间的相似度等于网络对参数的偏导的内积的期望。基于 NTK，我们可以分析出很多神经网络的性质, 比如出乎意料地, 神经网络的结果和随机初始化的参数无关, 仅和网络结构和训练数据有关。
在学习傅里叶特征时, 我们不需要仔细研究这些这些理论, 而只需要知道一个结论: 一般上述 NTK 可以写成标量函数 , 也就是可以先算内积再求偏导。这意味用核回归表示神经网络时, 真正要关心的是输入间的内积。别看 NTK 看起来那么复杂, 傅里叶特征论文其实主要就用到了这一个性质。
为了从理论上讲清为什么 MLP 难以拟合高频，作者还提及了很多有关 NTK 的分析，包括一种叫做谱偏差（spectral bias）的现象：神经网络更容易学习到数据中的低频特征。可能作者默认读者已经熟悉了相关的理论背景，这部分论述经常会出现逻辑跳跃，很难读懂。当然，不懂这些理论不影响理解傅里叶特征。我建议不要去仔细阅读这篇文章有关谱偏差的那一部分。
正如我们在前文的核回归实验里观察到的，核回归模型能否学到高频取决于核函数的频域特征。因此，这部分分析和 NTK 的频域有关。对这部分内容感兴趣的话可以去阅读之前有关谱偏差的论文。
傅里叶特征的平移不变性
在上两节中，我们花了不少功夫去认识谱回归和 NTK。总结下来，其实我们只需要搞懂两件事：
- 神经网络最终的收敛效果可以由简单的核回归决定。而核回归重点是定义两个输入之间的相似度指标（核函数）。
- 表示神经网络的核回归相似度指标是 NTK，它其实又只取决于两个输入的内积。
根据这一性质，我们可以部分解释为什么在文章开头那个 MLP 拟合连续图像的实验中，位置编码可以提升 MLP 拟合高频信息的能力了。这和位置输入的特性有关。
当 MLP 的输入表示位置时, 我们希望模型对输入位置具有平移不变性。比如我们现在有一条三个样本组成的句子 。当我们同时改变句子的位置信息时, 比如将句子的位置改成 时, 网络能学出完全一样的东西。但显然不对输入位置做任何处理的话， 和 对神经网络来说是完全不同的意思。
而使用位置编码的话，情况就完全不同了。假如输入数据是二维坐标 ，我们可以用下面的式子建立一个维度为 的位置编码:
其中 是系数, 是一个投影矩阵, 用于把原来 2 D 的位置变成一个更长的位置编码。当然, 由于位置编码中既要有 也要有 , 所以最终的位置编码长度为 。
根据我们之前的分析，NTK 只取决于输入间的内积。算上位置编码后，一对输入位置 的内积为:
而根据三角函数和角公式可知：
这样，上面那个内积恰好可以写成：
上式完全由位置间的相对距离决定。上式决定了 NTK，NTK 又决定了神经网络的学习结果。所以，神经网络的收敛结果其实完全取决于输入间的相对距离，而不取决于它们的绝对距离。也因此，位置编码使得 MLP 对于输入位置有了平移不变性。
加入位置编码后，虽然 MLP 满足了平移不变性，但这并不代表 MLP 学习高频信息的能力就变强了。平移不变性能给我们带来什么好处呢？作者指出，当满足了平移不变性后，我们就能手动调整 NTK 的带宽了。回想一下我们上面做的核回归实验，如果我们能够调整核的带宽，就能决定函数是更加高频（尖锐）还是更加低频（平滑）。这里也是同理，如果我们能够调大 NTK 的带宽，让它保留更多高频信息，那么 MLP 也就能学到更多的高频信息。
作者在此处用信号处理的知识来分析平移不变性的好处，比如讲了新的 NTK 就像一个重建卷积核 （reconstruction filter），整个 MLP 就像是在做卷积。还是由于作者省略了很多推导细节，这部分逻辑很难读懂。我建议大家直接记住推理的结论：平移不变性使得我们能够调整 NTK 的带宽，从而调整 MLP 学习高频的能力。
那我们该怎么调整 NTK 的带宽呢？现在的新 NTK 由下面的式子决定：
为了方便分析, 我们假设 和 都是一维实数。那么, 如果我们令 的话:
这个式子能令你想到什么? 没错, 就是傅里叶变换。 较大的项就表示 NTK 的高频分量。我们可以通过修改前面的系数 来手动调整 NTK 的频域特征。我们能看到，位置编码其实就是在模拟傅里叶变换，所以作者把位置编码总结为傅里叶特征。
作者通过实验证明我们可以手动修改 NTK 的频谱。实验中, 作者令 。 表示位置编码只有第一项： 。不同 时 NTK 的空域和频域示意图如下所示。可以看出, 令 时, 即傅里叶特征所有项的系数都为 1 时, NTK 的高频分量不会衰减。这也意味着 MLP 学高频信息和低频信息的能力差不多。
随机傅里叶特征
现在我们已经知道傅里叶特征的公式是什么, 并知道如何设置其中的参数 了。现在, 还有一件事我们没有决定：该如何设置傅里叶特征的长度 呢？
既然我们说傅里叶特征就是把输入的位置做了一次傅里叶变换, 那么一般来讲, 傅里叶特征的长度应该和原图像的像素数一样。比如我们要表示一个 的图像, 那么我们就需要令 表示不同方向上的频率： 。但这样的话, 神经网络的参数就太多了。可不可以令 更小一点呢?
根据之前的研究Random features for large-scale kernel machines 表明, 我们不需要密集地采样傅里叶特征, 只需要稀疏地采样就行了。具体来说, 我们可以从某个分布随机采样 个频率 来, 这样的学习结果和密集采样差不多。当然, 根据前面的分析, 我们还是令所有系数 。在实验中, 作者发现, 从哪种分布里采样都无所谓, 关键是 的采样分布的标准差, 因为这个标准差决定了傅里叶特征的带宽, 也决定了网络拟合高频信息的能力。实验的结果如下:
我们可以不管图片里 是啥意思, 只需要知道 是三组不同的实验就行。虚线是密集采样傅里叶特征的误差，它的结果反映了一个「较好」的误差值。令人惊讶的是，不管从哪种分布里采样 , 最后学出来的网络误差都差不多。问题的关键在于采样分布的标准差。把标准差调得够好的话, 模型的误差甚至低于密集采样的误差。
也就是说，虽然我们花半天分析了位置编码和傅里叶变换的关系，但我们没必要照着傅里叶变换那样密集地采样频率，只需要随机选一些频率即可。当然，这个结论只对 MLP 拟合连续数据的任务有效，和 Transformer 里的位置编码无关。
代码实现随机傅里叶特征
现在，我们可以回到博文开头的代码，看一下随机傅里叶特征是怎么实现的。
class FourierFeature(nn.Module):def __init__(self, in_c, out_c, scale):super().__init__()fourier_basis = torch.randn(in_c, out_c // 2) * scaleself.register_buffer('_fourier_basis', fourier_basis)def forward(self, x):N, C, H, W = x.shapex = rearrange(x, 'n c h w -> (n h w) c')x = x @ self._fourier_basisx = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)x = 2 * torch.pi * xx = torch.cat([torch.sin(x), torch.cos(x)], dim=1) return xfeature_length = 256
model = MLP(feature_length).to(device)
fourier_feature = FourierFeature(2, feature_length, 10).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
n_loops = 400
for epoch in tqdm(range(n_loops)):x = fourier_feature(grid)output = model(x)loss = F.l1_loss(output, input_image)optimizer.zero_grad()loss.backward()optimizer.step()if epoch % 100 == 0 or epoch == n_loops - 1:viz_image(output[0])print(loss.item())
prev_output = output傅里叶特征通过类FourierFeature 实现。其代码如下：
class FourierFeature(nn.Module):def __init__(self, in_c, out_c, scale):super().__init__()fourier_basis = torch.randn(in_c, out_c // 2) * scaleself.register_buffer('_fourier_basis', fourier_basis)def forward(self, x):N, C, H, W = x.shapex = rearrange(x, 'n c h w -> (n h w) c')x = x @ self._fourier_basisx = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)x = 2 * torch.pi * xx = torch.cat([torch.sin(x), torch.cos(x)], dim=1) return x构造函数里的fourier_basis 表示随机傅里叶特征的频率，对应论文公式里的，scale 表示采样的标准差。初始化好了随机频率后，对于输入位置x，只要按照公式将其投影到长度为out_c / 2 的向量上，再对向量的每一个分量求sin, cos 即可。按照之前的分析，我们令所有系数 为，所以不需要对输出向量乘系数。
傅里叶特征在 StyleGAN3 里的应用
傅里叶特征最经典的应用就是 NeRF 这类过拟合连续数据任务。除此之外，傅里叶特征另一次大展身手是在 StyleGAN3 中。
StyleGAN3 希望通过平滑地移动生成网络的输入来使输出图片也发生对应的移动。为此，StyleGAN3 将生成网络的输入定义为频域上的一个有限带宽图像信号：根据信号处理知识，我们能够将有限带宽信号转换成空域上无限连续的信号。也就是说，不管输入的分辨率（采样率）多低，我们都能够平滑地移动输入图片。StyleGAN3 借助随机傅里叶特征来实现这样一个频域图像。
以下代码选自 StyleGAN3 中傅里叶特征的构造函数。这个函数的关键是随机生成一些频率固定，但方向可以不同的傅里叶频率。函数先随机采样了一些频率，再将它们归一化，最后乘上指定的带宽bandwidth，保证所有频率大小相等。
class SynthesisInput(torch.nn.Module):def __init__(self,w_dim, # Intermediate latent (W) dimensionality.channels, # Number of output channels.size, # Output spatial size: int or [width, height].sampling_rate, # Output sampling rate.bandwidth, # Output bandwidth.):super().__init__()self.w_dim = w_dimself.channels = channelsself.size = np.broadcast_to(np.asarray(size), [2])self.sampling_rate = sampling_rateself.bandwidth = bandwidth# Draw random frequencies from uniform 2D disc.freqs = torch.randn([self.channels, 2])radii = freqs.square().sum(dim=1, keepdim=True).sqrt()freqs /= radii * radii.square().exp().pow(0.25)freqs *= bandwidthphases = torch.rand([self.channels]) - 0.5而在使用这个类获取网络输入时，和刚刚的 MLP 实现一样，我们会先生成一个二维坐标表格grid 用于查询连续图片每一处的颜色值，再将其投影到各个频率上，并计算新向量的正弦函数。
这段代码中，有两块和我们自己的实现不太一样。第一，StyleGAN3 允许对输入坐标做仿射变换（平移和旋转）。仿射变换对坐标的影响最终会转化成对三角函数相位phases 和频率freqs 的影响。第二，在计算三角函数时，StyleGAN3 只用了正弦函数，没有用余弦函数。
def forward(self, ...):...# Transform frequencies.phases = ...freqs = ...# Construct sampling grid.theta = torch.eye(2, 3, device=w.device)theta[0, 0] = 0.5 * self.size[0] / self.sampling_ratetheta[1, 1] = 0.5 * self.size[1] / self.sampling_rategrids = torch.nn.functional.affine_grid(theta.unsqueeze(0), [1, 1, self.size[1], self.size[0]], align_corners=False)# Compute Fourier features.x = (grids.unsqueeze(3) @ freqs.permute(0, 2, 1).unsqueeze(1).unsqueeze(2)).squeeze(3) # [batch, height, width, channel]x = x + phases.unsqueeze(1).unsqueeze(2)x = torch.sin(x * (np.pi * 2))x = x * amplitudes.unsqueeze(1).unsqueeze(2)...# Ensure correct shape.x = x.permute(0, 3, 1, 2) # [batch, channel, height, width]return x我们在 MLP 拟合连续图像的实验里复现一下这两个改动。首先是二维仿射变换。给定旋转角theta 和两个方向的平移tx, ty，我们能够构造出一个 的仿射变换矩阵。把它乘上坐标[x, y, 1] 后，就能得到仿射变换的输出。我们对输入坐标grid 做仿射变换后得到grid_ext，再用grid_ext 跑一遍傅里叶特征和 MLP。
N, C, H, W = grid.shape
tx = 50 / H
ty = 0
theta = torch.tensor(torch.pi * 1 / 8)
affine_matrix = torch.tensor([[torch.cos(theta), -torch.sin(theta), tx],[torch.sin(theta), torch.cos(theta), ty],[0, 0, 1]
]
).to(device)
grid_ext = torch.ones(N, 3, H, W).to(device)
grid_ext[:, :2] = grid.clone()
grid_ext = grid_ext.permute(0, 2, 3, 1)
grid_ext = (grid_ext @ affine_matrix.T)
grid_ext = grid_ext.permute(0, 3, 1, 2)[:, :2]x = fourier_feature(grid_ext)
output = model(x)
viz_image(output[0])在示例代码中，我们可以得到旋转 45 度并向下平移 50 个像素的图片。可以看到，变换成功了。这体现了连续数据的好处：我们可以在任意位置对数据采样。当然，由于这种连续数据是通过过拟合实现的，在训练集没有覆盖的坐标处无法得到有意义的颜色值。
之后，我们来尝试在傅里叶特征中只用正弦函数。我们将投影矩阵的输出通道数从out_c / 2 变成out_c，再在forward 里只用sin 而不是同时用sin, cos。经实验，这样改了后完全不影响重建质量，甚至由于通道数更多了，重建效果更好了。
class FourierFeature(nn.Module):def __init__(self, in_c, out_c, scale):super().__init__()fourier_basis = torch.randn(in_c, out_c) * scaleself.register_buffer('_fourier_basis', fourier_basis)def forward(self, x):N, C, H, W = x.shapex = rearrange(x, 'n c h w -> (n h w) c')x = x @ self._fourier_basisx = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)x = 2 * torch.pi * xx = torch.sin(x)return xStyleGAN3 论文并没有讲为什么只用sin，网上也很少有人讨论傅里叶特征的实现细节。我猜傅里叶特征并不是非得和傅里叶变换完全对应，毕竟它只是用来给神经网络提供更多信息，而没有什么严格的意义。只要把输入坐标分解成不同频率后，神经网络就能很好地学习了。
只用sin 而不是同时用sin, cos 后，似乎我们之前对 NTK 平移不变的推导完全失效了。但是，根据三角函数的周期性可知，只要是把输入映射到三角函数上后，网络主要是从位置间的相对关系学东西。绝对位置对网络来说没有那么重要，不同的绝对位置只是让所有三角函数差了一个相位而已。只用sin 的神经网络似乎也对绝对位置不敏感。为了证明这一点，我把原来位于[0, 1] 间的坐标做了一个幅度为10 的平移。结果网络的误差几乎没变。
for epoch in tqdm(range(n_loops)):x = fourier_feature(grid + 10)output = model2(x)loss = F.l1_loss(output, input_image)optimizer.zero_grad()loss.backward()optimizer.step()根据这些实验结果，我感觉是不是从 NTK 的角度来分析傅里叶特征完全没有必要？是不是只要从直觉上理解傅里叶特征的作用就行了？按我的理解，傅里叶特征在真正意义在于显式把网络对于不同频率的关注度建模出来，从而辅助网络学习高频细节。
总结
在这篇博文中，我们学习了傅里叶特征及其应用，并顺带了解其背后有关核回归、NTK 的有关理论知识。这些知识很杂乱，我来按逻辑顺序把它们整理一下。
为了解释为什么 NeRF 中的位置编码有效，傅里叶特征论文研究了用 MLP 拟合连续数据这一类任务中如何让 MLP 更好地学到高频信息。论文有两大主要结论：
- 通过从 NTK 理论的分析，位置编码其实是一种特殊的傅里叶特征。这种特征具有平移不变性。因此，神经网络就像是在对某个输入信号做卷积。而我们可以通过调整傅里叶特征的参数来调整卷积的带宽，也就是调整网络对于不同频率的关注程度，从而使得网络不会忽略高频信息。
- 傅里叶特征的频率不需要密集采样，只需要从任意一个分布随机稀疏采样。影响效果的关键是采样分布的标准差，它决定了傅里叶特征的带宽，也就决定了网络是否能关注到高频信息。
除了过拟合连续数据外，傅里叶特征的另一个作用是直接表示带宽有限信号，以实现在空域上的连续采样。StyleGAN3 在用傅里叶特征时，允许对输入坐标进行仿射变换，并且计算特征时只用了正弦函数而不是同时用正弦、余弦函数。这表明有关 NTK 的理论分析可能是没有必要的，主要说明问题的还是实验结果。
傅里叶特征论文仅研究了拟合连续数据这一类问题，没有讨论 Transformer 中位置编码的作用。论文中的一些结论可能无法适用。比如在大模型的位置编码中，我们还是得用密集的sin, cos 变换来表示位置编码。不过，我们可以依然借助该论文中提到的理论分析工具，来尝试分析所有位置编码的行为。
只通过文字理解可能还不太够，欢迎大家尝试我为这篇博客写的 Notebook，通过动手做实验来加深理解。https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/FourierFeature
#让模型预见分布漂移
动态系统颠覆性设计引领时域泛化新革命本研究提出了一种方法，能够在领域数据分布持续变化的动态环境中，基于随机时刻观测的数据分布，在任意时刻生成适用的神经网络。
下图展示了模型在领域数据随时间发生旋转和膨胀时的泛化表现。通过在一些随机时间点（蓝色标记点）的观测，模型可以在任意时刻生成适用的神经网络，其决策边界始终与数据分布保持协调一致。
01 摘要
在实际应用中，数据集的数据分布往往随着时间而不断变化，预测模型需要持续更新以保持准确性。时域泛化旨在预测未来数据分布，从而提前更新模型，使模型与数据同步变化。
然而，传统方法假设领域数据在固定时间间隔内收集，忽视了现实任务中数据集采集的随机性和不定时性，无法应对数据分布在连续时间上的变化。此外，传统方法也难以保证泛化过程在整个时间流中保持稳定和可控。
为此，本文提出了连续时域泛化任务，并设计了一个基于模型动态系统的时域泛化框架 Koodos，使得模型在连续时间中与数据分布的变化始终保持协调一致。Koodos 通过库普曼算子将模型的复杂非线性动态转化为可学习的连续动态系统，同时利用先验知识以确保泛化过程的稳定性和可控性。
实验表明，Koodos 显著超越现有方法，为时域泛化开辟了全新的研究方向。
02 论文信息
论文链接：
https://arxiv.org/pdf/2405.16075
开源代码：
https://github.com/Zekun-Cai/Koodos/
OpenReview：
https://openreview.net/forum?id=G24fOpC3JE
我们在代码库中提供了详细的逐步教程，涵盖了 Koodos 的实现、核心概念的解读以及可视化演示：
https://github.com/Zekun-Cai/Koodos/blob/main/Tutorial_for_Koodos.ipynb
整个教程流程紧凑，十分钟即可快使掌握 Koodos 的使用方法，力荐尝试！
03 情景导入
在实际应用中，训练数据的分布通常与测试数据不同，导致模型在训练环境之外的泛化能力受限。领域泛化（Domain Generalization, DG）作为一种重要的机器学习策略，旨在学习一个能够在未见目标领域中也保持良好表现的模型。
近年来研究人员发现，在动态环境中，领域数据（Domain Data）分布往往具有显著的时间依赖性，这促使了时域泛化（Temporal Domain Generalization, TDG）技术的快速发展。
时域泛化将多个领域视为一个时间序列而非一组独立的静态个体，利用历史领域预测未来领域，从而实现对模型参数的提前调整，显著提升了传统 DG 方法的效果。
然而，现有的时域泛化研究集中在“离散时间域”假设下，即假设领域数据在固定时间间隔（如逐周或逐年）收集。基于这一假设，概率模型被用于预测时域演变，例如通过隐变量模型生成未来数据，或利用序列模型（如 LSTM）预测未来的模型参数。
然而在现实中，领域数据的观测并不总是在离散、规律的时间点上，而是随机且稀疏地分布在连续时间轴上。例如，图 1 展示了一个典型的例子——基于推文数据进行社交媒体舆情预测。
与传统 TDG 假设的领域在时间轴上规律分布不同，实际中我们只能在特定事件（如总统辩论）发生时获得一个域，而这些事件的发生时间并不固定。同时，概念漂移（Concept Drift）在时间轴上发生，即领域数据分布随着时间不断演变：如活跃用户增加、新交互行为形成、年龄与性别分布变化等。
理想情况下，每个时态域对应的预测模型也应随时间逐渐调整，以应对这种概念漂移。最后，由于未来的域采集时间未知，我们希望可以泛化预测模型到未来任意时刻。
▲ 图1：连续时域泛化示意图。图中展示了通过推文训练分类模型进行舆情预测。其中训练域仅能在特定政治事件（如总统辩论）前后采集。我们希望通过这些不规律时间分布的训练域来捕捉分布漂移，并最终使模型能够推广到任意未来时刻。
事实上，领域分布在连续时间上的场景十分常见，例如：
事件驱动的数据采集：仅在特定事件发生时采集领域数据，事件之间没有数据。
流数据的随机观测：领域数据在数据流的任意时间点开始或结束采集，而非持续进行。
离散时态域但缺失：尽管领域数据基于离散时间点采集，但部分时间节点的领域数据缺失。
为了应对这些场景中的模型泛化，我们提出了“连续时域泛化”（Continuous Temporal Domain Generalization, CTDG）任务，其中观测和未观测的领域均分布于连续时间轴上随机的时间点。CTDG 关注于如何表征时态领域的连续动态，使得模型能够在任意时间点实现稳定、适应性的调整，从而完成泛化预测。
04 核心挑战
CTDG 任务的挑战远超传统的 TDG 方法。CTDG 不仅需要处理不规律时间分布的训练域，更重要的是，它旨在让模型泛化到任意时刻，即要求在连续时间的每个点上都能精确描述模型状态。
而 TDG 方法则仅关注未来的单步泛化：在观测点优化出当前模型状态后，只需将其外推一步即可。这使得 CTDG 区别于 TDG 任务：CTDG 的关键在于如何在连续时间轴上同步数据分布和模型参数的动态演变，而不是仅局限于未来某一特定时刻的模型表现。
具体而言，与 TDG 任务相比，CTDG 的复杂性主要来自以下几个尚未被充分探索的核心挑战：
如何建模数据动态并同步模型动态：CTDG 要求在连续时间轴上捕捉领域数据的动态，并据此同步调整模型状态。然而，数据动态本身难以直接观测，需要通过观测时间点来学习。此外，模型动态的演变过程也同样复杂。理解数据演变如何驱动模型演变构成了 CTDG 的首要挑战。
如何在高度非线性模型动态中捕捉主动态：领域数据的预测模型通常依赖过参数化（over-parametrized）的深度神经网络，模型动态因此呈现出高维、非线性的复杂特征。这导致模型的主动态嵌藏在大量潜在维度中。如何有效提取并将这些主动态映射到可学习的空间，是 CTDG 任务中的另一重大挑战。
如何确保长期泛化的稳定性和可控性：为实现未来任意时刻的泛化，CTDG 必须确保模型的长期稳定性。此外，在许多情况下，我们可能拥有数据动态的高层次先验知识。如何将这些先验知识嵌入 CTDG 的优化过程中，进而提升泛化的稳定性和可控性，是一个重要的开放性问题。
05 技术方法
5.1 问题定义
在 CTDG 中，一个域 表示在时间 采集的数据集，由实例集 组成，其中 和 分别为特征值，目标值和实例数。我们重点关注连续时间上的渐进性概念漂移，表示为领域数据的条件概率分布 随时间平滑变化。
在训练阶段，模型接收一系列在不规律时间点 上收集的观测域 ，其中每个时间点 是定义在连续时间轴 上的实数，且满足 $t_1<t_2<\ldots<t_t$ 。<="" p="">
在每个 上，模型学习到领域数据 的预测函数 ，其中 表示 时刻的模型参数。CTDG 的目标是建模参数的动态变化，以便在任意给定时刻 上预测模型参数 ，从而得到泛化模型 。
在后续部分中，我们使用简写符号 、、 和 ，分别表示在时间 上的 、 、 和 。
5.2 设计思路
我们的方法通过模型与数据的同步、动态简化表示，以及高效的联合优化展开。具体思路如下：
1. 同步数据和模型的动态：我们证明了连续时域中模型参数的连续性，而后借助神经微分方程（Neural ODE）建立模型动态系统，从而实现模型动态与数据动态的同步。
2. 表征高维动态到低维空间：我们将高维模型参数映射到一个结构化的库普曼空间（Koopman Space）中。该空间通过可学习的低维线性动态来捕捉模型的主要动态。
3. 联合优化模型与其动态：我们将单个领域的模型学习与各时间点上的连续动态进行联合优化，并设计了归纳偏置的约束接口，通过端到端优化保证泛化的稳定性和可控性。
▲ 模型设计
5.3 解决方案
Step 1. 数据动态建模与模型动态同步
分布变化的连续性假设：我们首先假设数据分布在时间上具有连续演化的特性，即条件概率分布 随时间平滑变化, 其演化规律可由一个函数 所描述的动态系统刻画。尽管真实世界中的渐进概念漂移可能较为复杂，但因概念漂移通常源于底层的连续过程（如自然、生物、物理、社会或经济因素），这一假设不失普适性。
分布变化引发的模型参数连续演化：基于上述假设，模型的函数功能空间应随数据分布变化同步调整。我们借助常微分方程来描述这一过程：
由此可推导出模型参数的演化满足：
其中, 是 对 的雅可比矩阵。
这一结果表明，如果数据分布的演化在时间上具有连续性，那么的演化过程也具有连续性，即模型参数会随数据分布的变化而平滑调整。上式为建立了一个由微分方程描述的模型动态系统。
模型动态系统学习：由于数据动态 的具体形式未知, 直接求解上述微分方程并不可行。为此, 我们引入一个由神经网络定义的连续动态系统, 用可学习的函数 描述模型参数 的变化。
通过鼓励模型动态和数据动态之间的拓扑共轭（Topological Conjugation）关系使 逼近真实动态。具体而言, 拓扑共轭要求通过泛化获得的模型参数与直接训练得到的参数保持一致。为此, 我们设定以下优化目标, 以学习 的参数 :
其中, 通过在时刻 的领域上直接训练获得, 则表示从时间 通过动态 演变至 的泛化参数:
通过这一优化过程，我们建立了模型动态与数据动态之间的同步机制。借助动态函数，我们可以在任意时刻精确求解模型的状态。
Step 2. 通过库普曼算子简化模型动态
非线性动态线性化
在实际任务中, 预测模型通常依赖于过参数化的深度神经网络, 使得模型动态 呈现为在高维空间中纠缠的非线性动态。直接对 建模不仅计算量大，且极易导致泛化不稳定。
然而, 受数据动态 的支配, 而数据动态通常是简单、可预测的。这意味着在过参数化空间中，模型的主动态（Principal Dynamics）可以在适当转换的空间内进行更易于管理的表示。
受此驱动，我们引入库普曼理论（Koopman Theory）来简化复杂的模型动态。库普曼理论在保持动态系统特征的同时将复杂的非线性动态线性化。
具体而言, 我们定义一个库普曼嵌入函数 , 将原始的高维参数空间映射到一个低维的库普曼空间中:
其中, 表示库普曼空间中的低维表示。通过库普曼算子 , 我们可以在线性空间中刻画 的动态：
一旦获得了简化的动态表示，我们可以在库普曼空间中更新模型参数，而后将其反映射回原始参数空间：
最终，通过库普曼算子的引入，我们实现了对模型动态的简化，保证了泛化过程的稳健性。
Step 3. 联合优化与先验知识结合
模型及其动力学的联合优化：我们对多个组件同时施加约束确保模型能稳定泛化，其包含以下关键项：
- 预测准确性：通过最小化预测误差，使预测模型在每个观测时间点都能准确预测实际数据。
- 泛化准确性：通过最小化预测误差，使泛化模型在每个观测时间点都能准确预测实际数据。
- 重构一致性：确保模型参数在原始空间与库普曼空间之间的转换具有一致性。
- 动态保真性：约束库普曼空间的动态行为，使得映射后的空间符合预期的动态系统特征。
- 参数一致性：确保泛化模型参数映射回原始空间后与预测模型参数保持一致。
利用库普曼算子评估和控制泛化过程：引入库普曼理论的另一优势在于，我们可以通过库普曼算子的谱特性来评估模型的长期稳定性。此外，还可以在库普曼算子中施加约束来控制模型的动态行为。
1. 系统稳定性评估
通过观察库普曼算子的特征值，可以判断系统是否稳定：
- 若所有特征值实部为负，系统会稳定地趋向于一个平衡状态。
- 若存在特征值实部为正，系统将变得不稳定，模型在未来可能会崩塌。
- 若特征值实部为零，系统可能表现出周期性行为。通过分析这些特征值的分布，我们可以预测系统的长期行为，识别模型在未来是否可能出现崩溃的风险。
2. 泛化过程约束
我们可以通过对库普曼算子施加显式约束来调控模型的动态行为。例如：
- 周期性约束：当数据动态为周期性时，可将库普曼算子设为反对称矩阵，使其特征值为纯虚数，从而使模型表现出周期性行为。
- 低秩近似：将表示为低秩矩阵，有助于控制模型的自由度，避免过拟合到次要信息。
通过这些手段，我们不仅提高了泛化的长期稳定性，还增强了模型在特定任务中的可控性。
06 实验
6.1 实验设置
为验证算法效果，我们使用了合成数据集和多种真实世界场景的数据集：
合成数据集：包括 Rotated 2-Moons 和 Rotated MNIST 数据集，通过在连续时间区间内随机生成时间戳，并对 Moons 和 MNIST 数据按时间戳逐步旋转生成连续时域。
真实世界数据集：
- 事件驱动数据集 Cyclone：基于热带气旋的卫星图像预测风力强度，气旋发生日期对应连续时域。
- 流数据集 Twitter 和 House：分别从任意时间段抽取推文和房价数据流构成一个领域，多次随机抽取形成连续时域。
- 不规则离散数据集 Yearbook：人像图片预测性别，从 84 年中随机抽取 40 年数据作为连续时域。
6.2 实验结果与分析
定量分析
我们首先对比了 Koodos 方法与各基线方法的定量性能。表 1 显示，Koodos 方法在所有数据集上展现了显著的性能提升。
在合成数据集上，Koodos 能够轻松应对持续的概念漂移，而所有基线方法在这种场景下全部失效。
在真实世界数据集上，尽管某些基线方法（如 CIDA、DRAIN 和 DeepODE）在少数场景中略有表现，但其相较于简单方法（如 Offline）的改进非常有限。相比之下，Koodos 显著优于所有现有方法，彰显出在时域泛化任务中考虑分布连续变化的关键作用。
▲ 实验结果
定性分析
决策边界：为直观展示泛化效果，我们在 Rotated 2-Moons 数据集上进行了决策边界的可视化。该任务具有极高难度：模型需在 0 到 35 秒左右的 35 个连续时域上训练，随后泛化到不规律分布在 35 到 50 秒的 15 个测试域。而现有方法通常只能泛化至未来的一个时域（T+1），且难以处理不规律的时间分布。图 3 从 15 个测试域中选取了 7 个进行可视化。结果清晰地表明，基线方法在应对连续时域的动态变化时表现不足。随着时间推进，决策边界逐渐偏离理想状态。尤其是最新的 DRAIN 方法（ICLR23）在多步泛化任务中明显失效。
相比之下，Koodos 在所有测试域上展现出卓越的泛化能力，始终保持清晰、准确的决策边界，与实际数据分布变化高度同步。这一效果突显了 Koodos 在时域泛化任务中的革命性优势。
▲ 图3：2-Moons 数据集决策边界的可视化（紫色和黄色表示数据区域，红线表示决策边界）。从上到下比较了两种基线方法和 Koodos；从左到右显示了部分测试域（15 选 7，所有测试域的分布在时间轴上用红点标记）。
模型演变轨迹：为更深入地分析模型的泛化能力，我们通过 t-SNE 降维，将不同方法的模型参数的演变过程（Model Evolution Trajectory）在隐空间中可视化（图 4）。
可以看出，Koodos 的轨迹呈现出平滑而有规律的螺旋式上升路径，从训练域平滑延伸至测试域。这一轨迹表明，Koodos 能够在隐空间中有效捕捉数据分布的连续变化，并随时间自然地扩展泛化。
相比之下，基线模型的轨迹在隐空间中缺乏清晰结构，随着时间推移，逐渐出现明显的偏离，未能形成一致的动态模式。
▲ 图4：模型状态在隐空间中的时空轨迹。Koodos 展现出与数据动态和谐同步的模型动态。
时域泛化的分析与控制：在 Koodos 模型中，库普曼算子为分析模型动态提供了有效手段。我们对 Koodos 在 2-Moons 数据集上分析表明，库普曼算子的特征值在复平面上分布在稳定区和不稳定区，这意味着 Koodos 在中短期内能稳定泛化，但在极长时间的预测上将会逐渐失去稳定性，偏离预期路径（图 5b）。
为提升模型的稳定性，我们通过将库普曼算子配置为反对称矩阵（即Koodos版本），确保所有特征值为纯虚数，使模型具有周期性稳定特性。在这一配置下，Koodos展现出高度一致的轨迹，即使在长时间外推过程中依然保持稳定和准确，证明了引入先验知识对增强模型稳健性的效果（图 5c）。
，时长00:23
▲ 图5：非受控和受控条件下的极长期泛化预测模型轨迹。a：部分训练域数据；b：不受控，模型最终偏离预期；c：受控，模型始终稳定且准确。
▲ 图5：非受控和受控条件下的极长期泛化预测模型轨迹。a：部分训练域数据；b：不受控，模型最终偏离预期；c：受控，模型始终稳定且准确。
07 结论
我们设计了一种基于模型连续动态系统的时域泛化方法，能够在数据域随时间逐渐演变的环境中，实现泛化模型的稳定性与可控性。未来，我们计划从多个方向进一步拓展这一技术的应用：
生成式模型扩展：时域泛化与生成式模型任务有天然的关联，Koodos 所具备的泛化能力能够为神经网络生成技术带来新的可能。
非时态泛化任务：Koodos 的应用并不局限于时域泛化，它也可以适用于其他分布变化的任务中。我们计划探索其在非时态领域的应用。
大模型集成：我们将探索时域泛化在大模型中的集成，帮助 LLM 在复杂多变的分布中保持鲁棒性和稳定性。
我们对时域泛化任务在未来的广阔应用前景充满期待。如有任何问题或合作意向，欢迎联系我们！
邮箱： caizekun@csis.u-tokyo.ac.jp
GitHub： https://github.com/Zekun-Cai/Koodos/
Paper： https://arxiv.org/pdf/2405.16075
#Scaling Laws for Precision 解读
本文探讨了模型量化对性能的影响，并提供了关于训练时量化和后训练量化的实用建议。文章强调了在不同训练精度下，如何平衡模型性能和量化损失，以及在实际应用中选择合适的量化策略的重要性。
来自链接 https://zhuanlan.zhihu.com/p/6848989432
原文
- https://arxiv.org/abs/2411.04330
前置知识：
scaling law:
- Training Compute-Optimal Large Language Models(Chinchilla scaling law)
个人讨厌晦涩难懂+无法应用于实际场景的"装逼结论"，因此先按照自己的理解帮大家rephrase一下论文的主要发现(in plain language)：
首先，这是一篇研究精度(precision)、参数量(parameters)和训练数据量(tokens)之间关系的重要论文。
1. 关于后训练量化(Post-Training Quantization, PTQ):1.1 基本概念
- 指的是pretrain以较高精度(bf16)进行，结束后再量化到更低精度（如int4）
1.2 结论1
模型预训练的trained_token/parameter比率越高，预训练结束后，使用PTQ带来的性能下降就越大。这里作者没写明白有误导性！！！实际上这个结论指的是：
- 我们都知道PTQ一定会带来性能下降（PTQ后，valid loss相比pretrain之后会上升)，这个下降可以用
- 论文提出了预测这个下降值的公式：
- 其中:
- 训练数据量D越大,PTQ带来的损失越大(正相关)
- 参数量N越大,PTQ带来的损失越小(负相关)
- 量化后的精度Ppost越低,损失增加越多(负指数关系)
- N: 参数量
- D: 训练token数
- : PTQ后的精度
- γγγ: 拟合常数
- 这个公式告诉我们:
- 注意，δPTQ还有一种完整形式（section 5) 同时考虑了训练精度和推理精度（继续往后看）：
- 那么如果你必须进行PTQ，那么对于同样参数量大小的模型，被训更多token的模型的 δPTQ 会比喂更少数据的模型要大。但最终loss的绝对数量是多少并不一定，因为即便 δPTQ 这个正数会让loss上升（性能下降），但模型终归被训了更多数据，这么一抵消可能loss还是会下降。相当于两只无形的手（数据量的上升带来的loss下降、PTQ带来的loss上升）在掰手腕；给定模型参数量和固定的精度，具体谁能掰过谁会有一个打平手的cutoff数据量。
- 举例子，如果你要固定70B模型参数量并pretrain时候采用bf16，并且pretrain后要PTQ到int4。那么采用两种数据量：
- a) 用10B token训出来模型
- b) 5B token训出来的模型
那么一定是a)情况的 δPTQ 更大，但最终PTQ结束之后的loss的数值是多少就不一定了。
因此作者也在原文中提到了**there exists an amount of pretraining data beyond which additional data is actively harmful to performance at inference-time (see top-left, Figure 2)，也就是给定你要进行PTQ，那么对于你的实验设置，总有一个cutoff的数据量，称之为临界的数据量 Dcrit ，超过这个量后继续训练会导致PTQ后性能下降。这个临界点并不是说超过后训练数据就“有害”，而是说在进行PTQ后，性能的提升可能会被性能的下降所抵消。因此，在实际应用中，需要权衡训练数据量与模型量化后的性能。
论文给出了计算这个临界点的公式:
其他结论
- 在某些情况下，过度训练（more tokens）反而会让PTQ后的模型性能变差
- 更大的模型在相同的token/parameter比率下，对PTQ更鲁棒
- 对于固定大小的数据集，增加模型参数量可以提高PTQ的鲁棒性
- 这种规律在不同的PTQ方法中都存在（论文验证了GPTQ、AWQ和RTN三种方法）
训练精度的影响
- 训练时使用较低精度的模型在PTQ时性能下降较小
- 如果你知道模型最终需要被量化到很低的精度(比如4bit)，那么在训练时就使用相对较低的精度(比如8bit)可能比使用高精度(比如16bit)更好，因为这样可以让模型在训练过程中就适应量化噪声。
- 实话说这个结论初看有点脱裤子放屁，因为太符合直觉了（bushi）。用脚想想就知道【训练用int8然后量化到int4】肯定比【训练用bf16然后量化到int4】要好，原文section 5：models trained in lower precision are more robust to post-train quantization in the sense of incurring lower degradation.
- 这也解释了为什么一些较新的大语言模型倾向于使用BF16而不是FP32来训练，因为这不仅可以节省计算资源，还可能让模型在后续量化时表现更好
1.3 PTQ造成loss degradation的深入分析1.3.1 两个竞争效应（section 5）
在分析PTQ对模型性能的影响时，论文发现了两个相互竞争的效应：
- Robustification效应
- 低精度训练会让模型更适应量化噪声
- 这使得模型在后续PTQ时更加鲁棒
- 可以理解为模型学会了如何在噪声环境中运作
- Overtraining效应
- 低精度训练会降低模型的有效参数量()，这意味着模型在相同的数据量下“看起来”参数量更少，从而在PTQ时对参数量化的敏感性增加
- 因为 和 成正比, 较低的Neff理论上会导致更大的性能下降: (section 5这边第一次读还以为写错了）。作者说的 实际上应该参考公式 9 变为 ，随着 的增加, 确实增加, 也就是成正比。说明白点就是低精度训练会下降Neff, 也就是一个 模型的可能有效的参数只有 10 B , 然后 变大, 然后根据section 3 的公式就会造成更大的degradation)
- 这个效应与Robustification效应相反
在实践中，Robustification效应通常占主导，这就是为什么低精度训练的模型在PTQ时表现更好。
1.3.2 精度阈值效应
一个重要发现是，当精度低于5-bit时，PTQ带来的性能下降会急剧增加：
- 在高精度区间（如8-bit以上），D/N比率的增加对性能的影响相对温和
- 在5-bit以下，即使很小的D/N比率增加也可能导致显著的性能下降
- 这个发现对实践中选择量化精度有重要指导意义-- 在实际应用中，应避免将模型量化到低于5-bit的精度，除非有特定的需求和相应的优化技术支持
1.3.3 理论解释
论文在附录中提供了两个可能的理论解释：
Sharpness假说
- 模型在训练过程中会逐渐变得更"sharp"-- 随着训练的进行，模型的损失函数变得更加“尖锐”（sharp），即梯度和Hessian矩阵的特征值增加，这导致模型对参数扰动更加敏感。因此，PTQ带来的参数量化噪声会对尖锐的损失函数产生更大的影响。
- Sharp的模型对参数扰动更敏感
- 这种敏感性会随着训练的进行而增加
- 这解释了为什么过度训练可能导致更大的PTQ降质
分层学习假说
- 模型通过分层方式学习特征-- 模型通过逐步学习更复杂的特征，这些特征依赖于之前学习的基础特征。量化噪声影响基础特征，会级联地影响到更高层次的复杂特征，从而导致整体性能的下降。
- 早期学习基础特征，后期学习复杂特征
- 复杂特征依赖于基础特征的准确性
- 当基础特征受到量化噪声影响时，会对依赖它们的复杂特征造成级联效应
- 这解释了为什么训练时间越长，模型对量化越敏感
2. 关于训练时量化（Training-time Quantization）
2.1 基本概念
论文中将训练时量化分为两种情况：
- 仅量化权重（Quantization-Aware Training, QAT）：只将模型的权重量化到低精度，其他部分保持高精度，以适应推理阶段的低精度环境。
- 全面量化（Low-precision Training）：同时量化模型的权重、激活值和注意力计算（即键-值缓存），以减少计算资源需求。
注意：这里的权重指模型中所有线性层（Linear layers）的权重矩阵，包括：
- Transformer 中的所有投影矩阵（例如 query、key、value 的投影权重）；
- 嵌入层（Embedding layers）权重矩阵；
- 最终输出层的权重矩阵。
但在论文的实验中未对嵌入层（Embedding layer）进行量化。
量化实现细节：
- 论文遵循了 FP8 训练的标准规范（Micikevicius et al., 2022）；
- 权重采用 按通道(per-channel) 量化；
- 激活值采用 按张量(per-tensor) 量化；
- 对于后训练量化（PTQ），主要针对模型权重进行量化。
2.2 核心发现
权重、激活值和注意力的量化效果是独立且可乘的，这一点非常关键。
论文提出了“有效参数量 Neff effective parameter count）”的概念。简而言之， Neff 代表了模型在低精度下的“真实有效”参数量。在低精度训练时，模型的实际参数量 N会被折减为较低的 Neff ，这有助于评估模型在低精度量化下的性能损失。
基本形式：
完整形式（全面量化）：
其中：
- N：模型的实际参数量；
- Pw ：权重精度；
- Pa：激活值精度；
- Pkv ：注意力计算精度；
- γw、γa、γkv ：各部分的敏感度系数，反映了模型对不同量化精度的适应性。
举个例子，在相同的计算预算下，有两种方案：
- a) 使用 16-bit 精度训练较小的模型；
- b) 使用 8-bit 精度训练较大的模型（参数量约为前者的 2 倍）。
根据论文的 Neff 分析，第二种方案通常更优，因为：
- 增加的参数量带来的性能提升超过了精度降低造成的损失；
- 8-bit 精度已接近论文中发现的计算最优精度（7-8 bits）；
- 低精度训练可以在相同的计算预算下处理更多的数据。
最优训练精度的计算：论文发现，在一般情况下，最优的训练精度为 7-8 bits。这意味着当前常用的 16-bit（BF16）训练精度其实存在冗余。但如果追求极低精度（例如 4-bit 以下），则需要不成比例地增加模型大小才能维持性能。
但是，如果模型大小被固定（例如受限于硬件资源），情况会有所不同：
- 此时，最优训练精度会随着训练数据量的增加而提高。具体来说，最优精度与训练数据量和参数量的比值成对数关系，即：
最优精度训练数据量参数量最优精度∝log(训练数据量参数量)（见论文 Section 4.3.3）
2.3 训练成本分析
训练成本的计算公式如下：
其中：
- C：计算成本；
- N ：模型参数量；
- D ：训练 token 数；
- P ：训练精度；
- 6/16：标准化系数（基于 Chinchilla 成本模型）。
这意味着什么？ 举个例子：假设你的计算预算是固定的，希望训练一个模型，有两种选择：
- 使用 16-bit 精度训练一个 35B 参数量的模型；
- 使用 8-bit 精度训练一个 70B 参数量的模型。
根据论文的发现，第二种方案可能更优，因为增加的参数量带来的性能提升超过了精度降低带来的损失。
2.4 实践建议
如果计算预算有限：
- 优先选择 7-8 bit 的训练精度，并利用节省下来的资源增加模型参数量；
- 避免使用低于 4-bit 的训练精度，因为这需要大幅增加模型大小才能维持性能（见论文 Section 4.3.2）。
如果模型大小受限：
- 在需要处理更大量数据时，提高训练精度；
- 例如，当 token/parameter 比率超过 1000 时，建议使用 8-bit 以上的精度；
- 在高 token/parameter 比率下，避免使用低于 6-bit 的训练精度（见论文 Section 4.3.3）。
各部分的精度选择：
- 权重（Weights）在极低精度（3-bit）下仍能保持稳定；
- 激活值（Activations）和注意力计算（KV-cache）在低于 3-bit 时可能会出现不稳定；
- 这种差异可能与量化方式有关（权重采用按通道量化，激活值采用按张量量化），而不一定是固有特性。
3.限制与未来研究方向
3.1 固定的模型架构
这篇论文采用了固定的Transformer++架构，以便在一个可控的环境中分析精度、参数量和数据量之间的关系。然而，在实际应用中，低精度训练通常会伴随着模型架构的调整。例如，一些先进的低精度训练方法可能会引入特殊的正则化技术或优化策略，以减轻低精度带来的负面影响。因此，论文的结论主要适用于固定架构的情况，尚未在经过优化的低精度架构中进行验证。
3.2 计算成本与系统开销
虽然理论上，降低训练精度（比如从16-bit降到8-bit）可以按比例减少计算需求，但在实际操作中，由于系统开销和硬件实现的限制，精度降低所带来的性能提升通常低于理论预期。例如，某些硬件可能无法高效支持极低精度（如4-bit以下）的计算，导致实际的加速效果有限。此外，不同精度下的数据移动和存储优化表现也可能有所不同，这进一步影响了低精度训练的实际效率。
3.3 仅关注验证损失，缺乏下游任务评估
论文主要关注于训练过程中的验证损失（validation loss）作为性能评估指标，而没有对下游任务的具体表现进行评估。尽管验证损失是衡量模型性能的重要指标，但不同任务对模型精度和量化的敏感性可能存在差异。
3.4 实验规模的限制
虽然论文中训练了多达17亿(17B)参数的模型，并使用了高达26B tokens的数据集，但这些规模相对较小，与当前最先进的大规模语言模型（如数百亿甚至千亿参数级别）相比仍有差距。因此，论文的scaling law在更大规模模型上的适用性尚未得到验证。
4. 量化方法的多样性
这篇论文主要关注于整数类型的量化方法，并通过GPTQ、AWQ和RTN等方法进行了验证。然而，浮点类型的量化方法（如FP8、FP4）在实际应用中也具有重要意义，尤其是在某些硬件平台上具有更好的支持和性能表现。不同量化方法在引入量化噪声和影响模型性能方面可能存在显著差异，因此，未来的研究应涵盖更多种类的量化方法，以全面理解量化对模型性能的影响。
5. 数据集和训练策略的单一性
这篇论文使用了Dolma V1.7数据集，并采用了特定的训练策略和超参数设置。不同的数据集和训练策略可能会影响模型对量化的敏感性。例如，某些数据集可能具有更高的复杂性或多样性，导致模型在低精度下表现出不同的鲁棒性。因此，未来的研究应在多样化的数据集和训练配置下进行，以验证缩放规律的普适性。