最近在调一个视频生成模型的训练流程遇到一个很典型的瓶颈分辨率从 256 提高到 512帧数从 8 提高到 16显存直接翻了几倍。最初我以为是模型参数太大后来把注意力模块单独抽出来 profiling 才发现真正的瓶颈几乎都集中在注意力矩阵的计算上。视频生成模型里的 “Token” 并不只是文本概念。画面被切成若干 patch每个 patch 是一个 token一段时间被拆成若干帧每一帧都包含一批 patch token。序列长度不是一个固定值而是由帧数、分辨率、patch size 共同决定。序列一长全局注意力就变成一种昂贵到难以承受的奢侈。这就是我看到 “Token Radius Attention for Efficient Video Generation” 这个方向时会特别留意的原因。从命名看它的核心思路非常直白与其让每个 token 去关注序列里所有其他 token不如只允许它注意自己附近某个“半径”范围内的 token把计算量真正降下来。但我先要说一句可能不太顺耳的话这类方法不是“无损优化”。它本质上是拿一部分长距离建模能力换训练效率和显存可控性。它真正解决的不是“省一点显存”这种表面问题而是让视频生成在长序列条件下变得可扩展、可训练、可部署。半径设多大、怎么设、什么时候设会出事才是实践里的核心问题。1. 先搞清楚 Token Radius Attention 到底在改哪个环节1.1 视频生成里的 token比想象中多得多Transformer 类视频生成模型里token 通常有几个来源文本条件 token来自视频描述文本数量不大通常几百到几千个。空间 patch token每一帧被切成 H×W 个 patch每个 patch 编码成一个 token。时间维度 token一些模型会加入时间位置编码或额外的帧级 token用来辅助感知运动。很多人会把视频生成的序列长度下意识类比成文本生成里的“几千个 token”。但真实视频任务里空间 token 一上来就是几千、几万。批量计算时复杂度还要再乘上 batch size。所以视频生成模型的注意力开销往往比纯文本模型高出一个数量级。1.2 全局注意力的平方复杂度是视频生成的显存黑洞标准注意力需要为每个 query token 计算它和所有 key token 的相关性。假设序列长度为 N注意力矩阵的形状是 N×N显存占用和计算量都是 O(N²)。看一个简单示意分辨率patch size每帧 token 数8 帧总 token 数注意力对数256×25616×162562048约 419 万512×51216×1610248192约 6711 万512×51216×16102416 帧 16384约 2.68 亿这组数字是示意数量不是某个模型的官方参数但趋势是确定的序列长度翻倍注意力对数量会翻四倍。更麻烦的是如果模型是扩散模型denoising 过程往往要迭代几十步每一步都会执行注意力。平方复杂度被反复放大这也是视频生成比图像生成更吃算力的核心原因。1.3 局部性假设为什么“只看附近 token”通常够用Token Radius Attention 的前提是视频内容存在明显的局部性。相邻的 patch 通常属于同一物体或同一纹理区域相邻帧之间同一空间位置对应的内容往往是连续的。因此大量相距很远的 token它们之间的注意力权重本来就很小计算它们大概率是一种浪费。不过这个假设不是没有例外。真正实践时要特别注意两类情况快速运动一个物体在几帧内从画面左侧跑到右侧固定的小半径可能看不到对应关系。全局语义关联远处出现一个关键物体名称或者整体场景信息需要跨越大范围传递。这两个例外恰好说明“局部性”是一个强先验并不是所有视频内容都完美满足。所以 Token Radius Attention 才需要配合“半径”这个可配置变量而不是简单地把全局注意力一刀切掉。2. “半径”不是单一数值三种典型设计2.1 空间半径在帧内约束注意力范围空间半径是最容易理解的一种设计。每一帧内部token 只关注距离自己一定空间范围内的其他 token。比如每个 patch 只关注周围 3×3、5×5 或 7×7 邻域。这种方式的优点是实现简单mask 可以离线计算。因为图像坐标固定两个 token 之间的空间距离是确定的不需要动态推断。在这种配置下注意力几乎变成一个“动态权重版本的卷积”感受野固定但权重由注意力机制计算而不是由卷积核生成。如果一帧内有 N 个 token空间半径为 k×k 邻域注意力复杂度可以降到 N×K其中 K 是邻域 token 数通常远小于 N。2.2 时间半径让注意力拥有跨帧能力时间半径则让 token 可以关注前后若干帧的对应位置或邻域。实现时需要为每个 token 记录它在视频中的坐标(t, h, w)再根据时间维度的差值限制候选 key。一个最简单的规则是允许当前帧 token 关注时间坐标在[t - r_t, t r_t]范围内的 token。更精细的做法还会限制跨帧时只能关注对应空间位置附近而不是把整帧都放开。很多初学者容易把时间半径理解成序列窗口。这并不准确。序列窗口只看 token 的排列顺序而时间半径要理解视频的帧结构。在视频里token 顺序可能因为数据预处理被打乱真正可靠的判断依据始终是 token 的实际坐标。2.3 语义半径从固定范围到动态候选语义半径是更进阶的一种设计。它不再用固定阈值决定注意力范围而是让模型根据当前 token 的特征动态选出若干候选 key。这可以理解成每个 query 维护一个“语义邻居列表”列表之外的内容直接不看。动态语义半径的表达能力更强但工程复杂度也明显更高选择候选 key 的模块会带来额外计算和显存。每个 query 的候选数量如果不一致批处理时要做 padding。不规则索引访问对底层 kernel 提出了更高要求。如果你还在学习这个方向我不建议从动态半径起步。固定空间半径加固定时间半径的组合已经足够搭出一个有效 baseline。2.4 不同半径组合时的取舍在实际项目里半径设计不是互斥的常见组合有几种半径设计实现成本表达能力常见风险空间半径低中等无法跨帧运动变化只能靠帧内信息时间半径中较高固定时间窗口可能漏掉快速位移空间 时间固定半径中较高超参变多需要逐个验证动态语义半径高高候选选择分支可能抵消收益从工程经验看我更建议先用“空间半径 时间半径”的组合并且把半径设为可配置项。不要一上来就把三个维度全部打开组合超参的调参空间会迅速膨胀。注意不要一上来就把半径设计成完全动态先用固定半径跑通流程再考虑加语义候选分支。否则出了问题很难定位是半径逻辑出错还是候选选择模块出错。3. 工程化之前先想清楚坐标、mask 和注意力 kernel3.1 token 坐标映射是第一步无论选择哪种半径动手写 attention 之前都要先做一件事建立 token 坐标表。视频输入通常是(B, C, T, H, W)patch 化之后会变成类似(B, L, D)的结构。要算半径不能只看 L而要恢复出每个 token 对应的(t, h, w)。下面是一个常见的坐标构建示例写法不是特定模型的官方实现import torch T, H, W 4, 8, 8 # 4 帧每帧 8x8 网格 grid torch.meshgrid( torch.arange(T), torch.arange(H), torch.arange(W), indexingij ) # 返回 3 个 shape(T, H, W) 的坐标张量 coords torch.stack(grid, dim-1).reshape(-1, 3) # coords 的形状是 (L, 3)L T * H * W # 每一行是 (t, h, w)有了coords才能继续算两个 token 之间的距离也才能判断它们是否在指定半径内。3.2 用 mask 还是用稀疏索引限制注意力范围工程上主要有两种主流实现方式。mask 方式生成一个 L×L 的布尔矩阵在 softmax 之前把不允许的位置置为负无穷。这种实现最直观也最容易调试。但底层如果仍然是 dense attention它并不会真正减少 L² 的内存开销只是把部分位置“捂起来”不参与计算。稀疏索引方式为每个 query 只保留允许的 key 索引在 kernel 内部也只计算这些 key。这样可以真正省掉被裁剪的注意力对数但对框架和 kernel 有更高要求。实际操作中我建议先用 mask 方式验证正确性因为逻辑直观、方便打印。等 baseline 确认没问题再换用稀疏 kernel 提升效率。第一次实现就上稀疏 kernel排错难度会明显增加。3.3 mask 构建示例假设radius表示允许的欧氏距离阈值mask 可以这样构建# coords: (L, 3)每一行是 (t, h, w) d torch.cdist(coords.float(), coords.float()) mask d radius # (L, L)这个写法非常简洁但要注意cdist在 L 很大时内存占用会很高只适合小规模验证。大规模场景下更合理的做法是按坐标差值直接计算局部邻域而不是先构造稠密距离矩阵。是否包含自身 token取决于设计。通常建议包含否则每个 token 连自己的信息都看不到。3.4 半径不是越小越好消融实验是关键从直觉看半径越小计算量越小。但半径过小会带来一个典型问题每个 token 只能看到自己和极少数邻居信息量不足训练会不收敛或生成质量明显劣化。所以消融实验的设计非常重要。建议从这几步开始固定全局注意力作为 baseline。保持其他模块不变只替换注意力模块。依次加入空间半径、时间半径。记录训练 loss、显存峰值、单 step 耗时、验证集指标。不要在实验里同时改 patch size、batch size 和学习率。否则损失曲线一波动很难分辨是半径导致的还是其他配置导致的。4. 自己动手实现时建议按这个排查链路走4.1 先分现象显存、速度、收敛、画质遇到问题后的第一步不是改代码而是先确认现象属于哪一类显存没降甚至比全局注意力还高。训练速度没有变快反而更慢。loss 不收敛或收敛后明显变差。生成的视频出现块状感、闪烁、运动不连贯。不同现象对应完全不同的排查方向。如果你还没有区分现象就急着调参数很容易在原地打转。4.2 五层排查输入、环境、参数、实现、工具边界按顺序排查比较高效不要跳层。第一层看现象。是 loss 变成 NaN还是训练稳定但不下降是显存真没降还是只降了前向、反向却爆了第二层看输入。token 坐标 mapping 是否正确。最容易出问题的点是把不同帧的 token 坐标当成同一帧导致时间半径失效。第三层看环境。依赖版本和 attention kernel 是否支持自定义 mask。很多框架在开启加速路径后会直接忽略用户传入的 mask。第四层看参数。半径是否设得太小比如半径等于 0 或 1。每个 token 只能看自己和极少数邻居时训练很难收敛。第五层看工具边界。确认稀疏 kernel 对 mask 的解释方式和你预期是否一致。有的 kernel 把 mask 理解为“保留位”有的理解为“屏蔽位”反过来就会出问题。4.3 常见误判和对应处理现象可能原因处理建议显存没降mask 没有真正裁剪矩阵换稀疏 kernel或改用局部索引loss 不收敛mask 可能没有进入注意力计算打印 mask 使用日志检查加速路径速度反而变慢候选索引的 gather/scatter 开销过大减少候选集或改为固定邻域生成画面有块状感空间半径过小增大空间半径做对照实验运动不连贯时间半径过小或没有跨帧增加时间半径观察变化注意不要在同一轮实验里同时改多个变量。先保持其他模块不动只改注意力范围才能判断收益和损失到底来自哪里。5. 从单次实验到可复用流程一个三步法5.1 第一步固定变量先跑最小样例用最小配置把流程跑通。建议使用小分辨率、小帧数、小 batch。比如 2 帧、每帧 32×32、patch size 16这样 token 数量很少方便打印 mask 和注意力权重。这一步的目的不是训练好模型而是确认几个基本事实token 坐标 mapping 是否正确。mask 形状是否和注意力输入匹配。前向和反向能否正常执行。渲染不同的半径值注意力模式是否符合直觉。单次跑通只说明流程没有断。真正麻烦的是批量任务、异常重试和长期维护。所以这个最小样例不要太快丢掉后面排查问题还会反复用到。5.2 第二步只替换注意力模块做对照实验在已有模型里只替换注意力模块。这是最干净的实验设计。因为其他模块没有变化注意力改动带来的差异更容易被观察到。建议至少跑三个配置配置说明全局注意力baseline不设任何半径空间半径注意力只看帧内邻域空间半径 时间半径注意力同时限制空间和时间范围记录指标包括训练 loss、验证 loss、单 step 耗时、显存峰值、生成样本的视觉质量。一个常见结果会是局部注意力比全局注意力 loss 高一点但显存和耗时下降明显。这是正常的。关键问题不是“高了多少”而是“这个交换是否匹配你的项目目标”。5.3 第三步把超参和日志沉淀成可复用记录做完实验后把配置保存下来。一个简单做法是用 JSON 记录超参和结果{ model: video_diffusion_baseline, radius_spatial: 7, radius_temporal: 2, patch_size: 16, batch_size: 4, frames: 8, peak_memory_mb: 18200, step_time_ms: 340, val_loss: 0.132 }这一步看起来不起眼但长期价值很大。当项目从实验走向产品化时重新实验、调参、回归都依赖这些记录。没有记录后面大概率会重复踩同一个坑。6. 这类设计的边界与长期价值6.1 适合谁、不适合谁Token Radius Attention 并不是所有视频生成任务的银弹。它有自己的适用边界。适合的场景视频分辨率高、帧数长全局注意力完全跑不动。需要在单卡或低显存设备上完成训练或推理。视频内容以静态或慢速运动为主局部性较强。不适合或需要更谨慎的场景视频中有大量全局性信息需求比如在一段画面里必须从很远的地方引入关键信息。目标检测、多物体交互等对远距离关系敏感的任务需要先验证。对生成质量要求极高且算力充足的场景不一定非要牺牲全局注意力。总之这个方案的价值高度依赖任务本身的信号结构。脱离数据特征谈“省显存”很容易在下一轮实验里被质量损失打脸。6.2 和 Flash Attention、Deformable Attention 等方法的关系很多人会把所有“省显存”的 attention 方法混在一起但实际上差异很大。Flash Attention 通过底层 kernel 重新设计减少显存带宽占用但注意力语义没有变化依然是所有 token 互相看。它属于“用更省的方式做同一件事”。Token Radius Attention 改变的是注意力语义从“所有 token 互相看”变成“只在一定范围内互相看”。它属于“用不同的方式描述信息交互范围”。Deformable Attention 更接近动态 token 选择可以看作语义半径这条路线的一种实现。理解这个区别能帮你避免一个误区不要把所有 attention 优化都当成同类。先问自己到底是想在不改变语义的情况下省资源还是想改变语义来控制信息交互半径。方向不同工程做法完全不同。6.3 核心价值注意力预算的重新分配回到文章开头的主判断。Token Radius Attention 对视频生成最大的意义不是某个具体公式而是把注意力从“无差别的全量关系建模”变成“有预算的关系建模”。视频是时间和空间的组合体真正的语义变化往往发生在局部。让每个 token 在有限半径内做更细致的动态加权比让它去扫描所有 token 更符合视频信号的特征。这种先验换来的效率提升在长视频、高分辨率、多次 denoising 迭代的场景下会累积成非常可观的收益。但如果要长期使用光改注意力模块还不够。还需要补上资源消耗统计、显存峰值 profiling、推理阶段的半径切换逻辑以及覆盖质量、速度和资源三个维度的评估指标。这是一个逐步工程化的过程不是一次实验就能完成的。我在最初的显存排查里花了很长时间纠结“为什么 N×N 注意力矩阵这么大”。后来想明白了一个更基础的问题不是矩阵大而是我根本没有裁剪过信息范围。Token Radius Attention 让我重新认识到视频生成的注意力优化不是盲目寻找更快的 kernel而是要设计“每个 token 应该从多远的地方获取信息”。这个半径本质上是模型对视频结构的一种先验约束。如果你也在做类似尝试我建议下一步别急着把半径设大或设小。先写一个最小样例把 token 坐标和 mask 打印出来看看在 2 帧、4 帧、8 帧下注意力模式是否符合预期。然后保持其他模块不变加一个全局注意力作为对照。有了这个基础后续再谈工程优化和部署才有真正可靠的依据。