← 返回 2026-05-22

SEGA:面向扩散 Transformer 分辨率外推的频谱能量引导注意力 SEGA: Spectral-Energy Guided Attention for Resolution Extrapolation in Diffusion Transformers

Javad Rajabi, Kimia Shaban, Koorosh Roohi, David B. Lindell, Babak Taati 📅 2026-05-21 👍 41 2026-07-13 08:36
RoPE外推 免训练方法 扩散Transformer 频谱分析 高分辨率图像生成

按潜变量频谱能量动态调节RoPE注意力缩放,实现免训练高分辨率文生图。

前置知识

RoPE(旋转位置编码)

通过对 query 和 key 向量施加成对的二维旋转来编码相对位置,每个维度对应一个独立旋转频率 θ_d,使注意力点积只依赖相对距离;在扩散 Transformer 中常按高、宽两条轴分别施加。

SEGA 的全部工作都建立在对 RoPE 各频率维施加差异化缩放之上,不懂 RoPE 就无法理解它改了什么。

Diffusion Transformer(DiT)

用自注意力替换 U-Net 的卷积主干,把潜变量展平为 token 序列逐层去噪,是当前文生图的主流架构,本文的 Flux 与 Qwen 都属此类,其分辨率上限由训练分辨率决定。

理解为何 token 网格随分辨率扩张会让 RoPE 偏移落入分布外,从而引发注意力稀释。

YaRN 与注意力缩放

大模型长上下文外推法,结合 NTK 频率调整与对 logit 的统一缩放因子 τ(s)=0.1 ln s+1,用来补偿外推引起的注意力熵塌缩,后被图像生成沿用,是 SEGA 要取代的基线。

SEGA 的核心论点正是『统一缩放不合理』,必须先明白 YaRN 怎么做才能看懂改进在哪。

长度外推 / NTK

当序列超出训练长度时 RoPE 的相对偏移落入分布外。NTK 通过放大基频 b′=b·s^{D/(D−2)} 重分配各维频率,使外推位置仍保持位置判别力,是 SEGA 默认挂载的外推主干。

SEGA 不扩展位置范围,必须叠加在 NTK 等外推方法之上,二者分工不同。

二维 FFT 与频谱平坦度

二维快速傅里叶变换把潜变量从空间域转到频域,输出功率谱表征各空间频率的能量分布;频谱平坦度(Wiener 熵,几何均值/算术均值)衡量能量是集中还是分散。

SEGA 用 FFT 提取的轴向能量与频谱平坦度分别驱动逐维修正和全局幅度,是方法的信号源。

研究动机

当前主流的文生图扩散 Transformer(如 Flux、Qwen)大多只在 1024² 到 2048² 分辨率范围内训练,一旦在推理时外推到更高分辨率(如 4096² 甚至 6144²),质量会急剧下降,出现模糊纹理、重复平铺、结构崩塌等问题。根因在于这些模型依赖旋转位置编码 RoPE 提供空间先验,而外推时相对位置偏移远超训练分布,注意力权重在扩张的 token 网格上被过度稀释,丧失了对位置偏移的判别力。从长上下文语言建模迁移来的 YaRN 等方法虽用 RoPE 外推加统一注意力缩放 τ(s)=0.1 ln s+1 部分缓解了稀释,但由于对所有频率维施加同一个常数缩放,无法兼顾全局结构与细节纹理:偏强会破坏粗结构、偏弱又救不回细节。

本文的目标是本文目标是提出一个完全免训练、内容自适应的注意力缩放方法,在不引入任何可学参数、不修改模型结构的前提下,让预训练 DiT 稳定外推到远超训练范围的分辨率(最高到约 3800 万像素的 6144²),同时保持全局结构连贯与细粒度纹理保真,并在 ImageReward、PickScore、CLIP Score、MUSIQ、FID 等多项指标上全面超越现有的直接推理和多阶段引导方法,尤其要在 4096² 分辨率下每个指标都拿到最优。

与已有工作不同的是,作者抓住了一个被忽视的关键事实:RoPE 的不同维度天然对应不同空间波长(频率),短波长维控制精细纹理、长波长维塑造全局结构。因此缩放理应按频率分量区别对待,而非一刀切。更进一步,潜变量本身的频谱能量分布并非静态——它随去噪步骤从噪声态演化到结构态,也随图像内容(如雾中湖泊与喧闹集市)而显著不同。SEGA 的切入角度正是『让缩放跟着潜变量实时的频谱能量走』:低能量频段加强缩放以保位置判别,高能量频段减弱缩放以避免过增强,并用频谱熵控制整体强度。这一信号此前没有任何方法用作引导。

核心方法

直觉上,SEGA 像一个会听频谱的『注意力调音师』:在每个去噪步,它先对当前潜变量做一次二维傅里叶变换,看清此刻能量分布在不同空间频率上,再据此为 RoPE 的每一个频率维单独拧一个缩放旋钮——能量贫乏的频段拧大、能量充沛的频段拧小。技术路线分两步:先把潜变量从空间域转频域,提取轴向能量剖面 E_a 和旋转无关的径向剖面 E_iso;再把这两个剖面换算成『逐维修正 s_d』和『全局幅度 σ』,与一个由分辨率比决定的参考缩放 m_ref 组合成最终的逐维缩放 $m_d = m_{ref}\cdot(1-\sigma\cdot s_d)$,作用于旋转后的 query、key。整个过程每步重算、内容相关、零参数。

核心创新是把『RoPE 维 ↔ 空间频率』这一耦合关系变成可计算的引导信号。每个 RoPE 维 d 的波长 $T_d=2\pi/\theta_d$ 唯一对应一条空间频带;SEGA 用该频带在潜变量功率谱里的对数能量来决定这一维的缩放方向——低于均值的频段得到 $s_d<0$ 的负修正(即增大缩放 $m_d$),高于均值的得到 $s_d>0$(减小缩放),且所有维的修正强制零和 $\sum_d s_d=0$,保证只重新分配、不整体偏移。与 YaRN、DyPE、UltraImage 的本质区别在于:它们要么用与分辨率绑定的固定函数、要么用预定义的时间步调度、要么只在少数『主导频率』上做稀疏离散修正,而 SEGA 在全部 RoPE 维和两条空间轴上给出一条连续的、每步每样本都不同的缩放曲线,真正做到了内容自适应。

方法步骤详情

输入是当前潜变量 $Z\in\mathbb{R}^{N\times C}$ 与分辨率比 $s$。①把 Z 还原成 $H\times W\times C$,沿通道求平均并去均值得零中心图 $\tilde{M}$,做一次二维 FFT 得功率谱。②对每条轴 $a\in\{H,W\}$ 沿正交轴边际化得 1D 能量剖面 $E_a$;同时在同心环内平均得旋转不变的径向剖面 $E_{iso}$。③对 RoPE 维 d 用波长 $T_d$ 定位 $E_a$ 中频带,取对数能量 $\hat{E}_d$ 标准化为 $z_d=(\hat{E}_d-\mu)/\nu$,过 tanh 减均值得 $s_d$,天然零和、负值对应低能量维。④由 $E_{iso}$ 算频谱平坦度 SF 得 $\sigma=1-SF^{\gamma}$,噪声态趋 0、结构态趋 1(取 $\gamma=1.5$)。⑤合成 $m_{ref}=(R_{target}/R_{train})^{\kappa}$,$m_d=m_{ref}(1-\sigma s_d)$,按 $f_{SEGA}=m_d\cdot f_{RoPE}$ 作用于旋转后的 query/key,每步执行(取 $\kappa=0.08$)。

技术新颖性

新颖性体现在三点。第一,诊断角度新:首次把『注意力缩放应该频率敏感』这一观点用潜变量实测频谱来落地,而非依赖 YaRN 的常数或 UltraImage 的主导频率启发式。第二,机制连续且零和:tanh 加减均值保证缩放在各维间重新分配而不偏移整体均值,比 UltraImage 的离散逐维递归修正更平滑、覆盖全部维度。第三,双重自适应:轴向剖面让高、宽两轴独立缩放,全局幅度 σ 让强度随去噪自然演化(早期抑制、晚期全开),这是 DyPE 那种预设时间步调度做不到的。此外它还是正交可组合的——可与 NTK、YaRN、DyPE 任一外推主干叠加,消融实验已证实。

SEGA scaling maps at 4096²
Figure 3: SEGA scaling maps at 4096²
Impact on Attention Evolution
Figure 4: Impact on Attention Evolution

实验结果

在 Flux 与 Qwen 两个主干、四种分辨率上 SEGA 全面领先。4096² 上 Flux 最具说服力:ImageReward 1.26、PickScore 23.18、CLIP Score 29.22、MUSIQ 45.72、FID 150.05、FIDp 51.28,每项全优,而 ScaleDiff 的 FIDp 高达 76.14、YaRN 高达 98.52。Qwen 上 4096² 同样全指标最优,IR 达 1.51。零样本基准验证泛化:Qwen+SEGA 的 IR 1.58、PS 23.86,逼近甚至部分超越闭源 Nano Banana 2(IR 1.37、PS 23.43)。极端分辨率下优势最显著:6144² 上 Flux 的 FID 232.18,远好于 DyPE 的 274.82、UltraImage 的 290.75,基线 Base 崩到 453.97,此时相对 1024² 已是 36 倍面积外推。消融证明三组件缺一不可:NTK 加固定缩放 FID 162.92,SEGA 降到 150.05;换 YaRN、DyPE 主干仍有效。注意力熵分析显示 SEGA 最接近训练分布内 1024² 基线,真正修复了注意力结构。

Comparison of SEGA against SOTA baselines on Flux
Table 1: Comparison of SEGA against SOTA baselines on Flux
Quantitative comparison on Qwen
Table 2: Quantitative comparison on Qwen
Ablation study on Flux at 4096×4096
Table 3: Ablation study on Flux at 4096×4096
查看结构化数据
任务指标本文基线提升
Flux 4096² 图像保真度 FID↓ 150.05 ScaleDiff 153.05 / YaRN 160.48 相对最强基线 ScaleDiff 降低约 3 个点,相对 YaRN 降低约 10 个点
Flux 4096² 细节保真度 FIDp↓ 51.28 YaRN 98.52 / ScaleDiff 76.14 相对 YaRN 降低约 47 个点,细节纹理优势极大
Qwen 4096² 综合质量 ImageReward↑ 1.51 DyPE 0.97 提升约 0.54,全指标领先
极端分辨率 6144² 保真度 FID↓ 232.18 DyPE 274.82 / UltraImage 290.75 降低约 42 点,36 倍面积外推下优势最显著
零样本基准 4096²(Qwen 主干) ImageReward↑ 1.58 闭源 Nano Banana 2 为 1.37 超越闭源商用模型约 0.21

局限与改进

作者在附录 D 坦承三点局限:SEGA 只调制旋转嵌入的幅值,并不扩展 RoPE 的位置范围,因此必须挂载在某个长度外推方法(实验中用 NTK)之上,会部分继承其结构先验;其次,虽能上到 8192²,但在最极端外推比下感知质量仍会下降,此时瓶颈是模型本身容量而非位置编码;第三,SEGA 自身计算开销可忽略,但生成多兆像素图像的二次方注意力成本仍然昂贵。我的额外观察:①方法的两个超参 γ=1.5、κ=0.08 是经验调出的,跨分辨率是否都最优未充分验证;②评价高度依赖 ImageReward、PickScore 等 learned metric,可能偏向某类审美,缺乏大规模人工评测;③仅在 Flux、Qwen 两类主干验证,对其他 RoPE 变体(如 3D 旋转)的普适性未触及。

独立分析的弱点

独立看,第一个弱点是统计鲁棒性:每步都对潜变量做 FFT 并标准化能量,当某轴 token 极少(极端长宽比)时统计量 μ、ν 可能不稳定,建议加小常数正则或改用稳健分位数。第二个是理论解释偏经验:把『低能量→强缩放』的因果链条用直觉论证,缺少注意力熵与频谱能量之间的解析关系,可建立简化模型证明最优缩放正比于 1/能量。第三个是超参敏感性:κ、γ 与 m_ref 形式(幂律 vs 对数)在不同分辨率区间表现分化,Table 9 显示 6144² 上幂律 FID 232 对比对数 269,存在手调痕迹,可改为按分辨率自适应或对这两个标量做可微搜索(仍保持免训练推理)。第四,缺乏人类偏好研究,建议补充众包 pairwise 评测以排除 learned metric 的偏差。

未来方向

作者明确提到频谱视角可推广到视频及其他模态的分辨率或时长外推。基于成果我延伸三个方向:①把『RoPE 维 ↔ 频率』的引导机制用于视频扩散的长时序外推,时间轴频谱可作为天然信号;②将频谱能量引导与多阶段引导(HiFlow、I-Max)结合,先稳结构再细化纹理,可能进一步突破上限;③探索线性注意力或稀疏注意力的二次方成本替代,让 8K、16K 真正实用化,因为算力才是极端分辨率的真正瓶颈;④把频谱自适应思想做成可学的轻量 adapter,在保留免训练推理选项的同时获得更强的内容建模。作者也呼吁学界把频谱视角作为修改注意力行为的新范式。

复现评估

复现门槛中等偏低。方法本身零参数、纯推理,公式清晰且超参公开(γ=1.5、κ=0.08、幂律 m_ref)。实验在 Flux(dev.Krea checkpoint)和 Qwen 上跑,bfloat16 精度、NVIDIA H100 GPU,遵循官方推理设置,每一步都施加 SEGA、不做 warmup。评价用了公开数据集 Aesthetic-4K 和自建零样本集(200 条 LLM 生成 prompt,文中描述但未明确开源),指标多为标准库(FID、CLIP Score、ImageReward、PickScore、MUSIQ)。项目页给出(rajabi2001.github.io/sega/2026),但论文为 2026-05 Preprint,代码与权重是否完整开源未在正文确认。算力方面,6144² 单图在 H100 上仍耗时可观,是复现极端分辨率实验的主要门槛,但 4096² 及以下普通研究者可负担。