SEGA: Spectral-Energy Guided Attention for Resolution Extrapolation in Diffusion Transformers¶

会议: NeurIPS 2026
arXiv: 2605.22668
OpenReview: https://openreview.net/forum?id=IUujREirW0
代码: https://rajabi2001.github.io/sega/
领域: 图像生成
关键词: 扩散Transformer, 分辨率外推, 旋转位置编码, 频域能量引导, 免训练高分辨率生成

一句话总结¶

针对 DiT 在推理期外推至超训练分辨率时注意力稀释与频带缩放失衡的问题，提出免训练的 SEGA 方法，通过潜变量频谱分析自适应调节各 RoPE 维度的注意力缩放尺度，在无需微调的前提下显著改善 4K/超高分辨率生成的全局结构与细节质感。

研究背景与动机¶

基于 Diffusion Transformers（DiT）的文生图模型（如 Flux、Qwen-Image）取得了卓越的生成质量，但其有效生成能力往往受限于训练阶段固定的图像分辨率（通常为 \(1024^2\) 至 \(2048^2\)）。当在推理阶段直接将模型外推至更高分辨率（例如 \(4096^2\) 甚至千万像素级别）时，图像质量会急剧恶化，频繁出现结构瓦解、纹理模糊以及重复多头/多物体等严重瑕疵。直接微调高分辨率模型不仅面临高分辨率图像配对数据匮乏的问题，更受限于二次方级暴涨的序列注意力计算开销，因而无需微调的免训练（training-free）外推方法成为迫切需求。

在 DiT 架构中，外推崩溃的核心根源在于旋转位置编码（RoPE）。当 token 序列沿二维网格大幅扩展时，查询与键之间的相对位置偏移显著超出训练分布范围，导致注意力权重过度稀释，丧失了在极长距离下的空间分辨能力。现有免训练方案多借鉴大语言模型中的长文本外推策略（如 NTK-aware、YaRN），通过重标定基频并辅以全局静态注意力缩放来抵消稀释效应。然而，二维图像的 RoPE 各维度天然对应着不同的空间波长：长波长分量控制全局轮廓与宏观语义布局，而短波长分量则捕捉高频局部纹理。统一、静态的缩放因子无法调和这一矛盾——缩放过大会造成高频过冲与结构断裂，缩放不足则会导致注意力弥散与纹理过度平滑。更为关键的是，潜变量的频谱能量分布在去噪全流程中动态演进（从早期白噪声主导转变为晚期结构主导），且在不同内容复杂度的图像间差异悬殊，静态超参数根本无法适配多变的内容分布。

本文的切入视角是建立潜变量空间频谱能量与 RoPE 各频率维度之间的双向映射：RoPE 各维度既然与特定空间波长严格绑定，就应由当前去噪时刻潜变量在对应频带的真实能量高低来动态裁决其注意力锐化程度。核心 idea：提出免训练的频谱能量引导注意力（SEGA），在去噪每一步对当前潜变量执行轻量级快速傅里叶变换，利用轴向能量剖面和径向频谱平坦度自适应计算各 RoPE 维度的缩放调制因子，实现内容感知且动态平衡的高分辨率平滑外推。

方法详解¶

整体框架¶

SEGA 是一种即插即用、完全免训练的推理期增强方法，无缝集成于采用 RoPE 的 DiT 注意力模块内部。其核心思想是在去噪时间步 \(t\) 直接从当前潜变量特征中提取频域先验，从而动态构造逐维度、方向解耦的位置编码缩放因子。

整个处理管线由三个串行阶段构成：首先对输入潜变量进行通道平均与空间均值中心化，通过单次二维快速傅里叶变换（2D FFT）分别提炼表征空间方向性能量的轴向剖面（Axis-wise profile）与反映整体结构聚集度的径向剖面（Radial profile）；其次根据轴向能量分布计算各 RoPE 维度的零和重分布校正量，并以基于维纳熵的频谱平坦度作为全局门控因子调控调制幅度；最后将自适应调制器与基准分辨率外推尺度融合，直接对 RoPE 变换后的查询与键向量进行维度级赋权，从而重塑注意力权重矩阵的局部聚焦性。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["中间潜变量 Z<br/>(H×W×C 空间网格)"] --> B["阶段1：轻量潜变量频谱分析<br/>2D FFT 提取轴向与径向能量剖面"]
    B --> C["阶段2：逐维度 RoPE 能量调制<br/>波长对齐映射 + 维纳熵动态门控"]
    C --> D["阶段3：频域自适应注意力缩放<br/>修正 Q/K 旋转嵌入并执行注意力计算"]
    D --> E["高分辨率去噪输出<br/>兼顾宏观结构一致性与微观纹理细节"]

关键设计¶

1. 轻量潜变量频谱分析：方向解耦与尺度聚合

针对二维图像高分辨率生成的各向异性（例如宽幅全景图与高耸构图的频域分布截然不同），统一标量缩放无法针对高度与宽度轴分别优化。SEGA 在每个去噪步从隐藏状态 \(\mathbf{Z} \in \mathbb{R}^{H \times W \times C}\) 出发，首先在通道维度上求均值并去除空间均值，得到零中心化二维图 \(\tilde{\mathbf{M}} \in \mathbb{R}^{H \times W}\)。随后执行单次 2D FFT 得到功率谱 \(|\mathcal{F}_{\text{2D}}[\tilde{\mathbf{M}}]|^2\)，并在此基础上并行提取两类互补的频域特征：一方面沿正交轴进行边缘积分边缘化，提炼出分别对应高度轴与宽度轴的 1D 轴向能量剖面 \(\mathcal{E}_H\) 和 \(\mathcal{E}_W\)；另一方面在同心圆环区域内进行径向能量平均，提取出丢弃方向性但表征全向空间尺度能量衰减特性的径向剖面 \(\mathcal{E}_{\text{iso}}\)。整个变换无需反向传播或可学习参数，单次 FFT 带来的推理时间增量不足 1%。

2. 逐维度 RoPE 能量调制：对齐波长与零和重分布

为了克服传统方法在全部旋转维度上施加固定恒等权重的缺陷，SEGA 建立了 RoPE 维度与物理空间波长的精确对齐机制。对于轴向 RoPE 的第 \(d\) 个频率维度，其固有波长为 \(T_d = 2\pi / \theta_d\)。SEGA 将该波长映射至对应轴向能量剖面 \(\mathcal{E}_a\) 的具体频带，读取其对数能量 \(\hat{E}_d^{(a)}\) 并进行跨维度 \(z\)-score 归一化：

\[z_d^{(a)} = \frac{\hat{E}_d^{(a)} - \mu^{(a)}}{\nu^{(a)}}\]

为了避免对注意力的整体能量造成非受控偏移，SEGA 引入非线性双曲正切函数，构建了严格满足零和约束（\(\sum_d s_d^{(a)} = 0\)）的重分布校正量：

\[s_d^{(a)} = \tanh(z_d^{(a)}) - \mathbb{E}[\tanh(z^{(a)})]\]

当某个维度所代表的波长在潜变量中处于低能量状态时，\(s_d^{(a)} < 0\)，系统会对其施加更强的缩放以恢复该频段的空间分辨力；反之，对能量已然过剩的频段减弱缩放强度，彻底消除了局部伪影与过度锐化。

3. 维纳熵动态门控：感知去噪演进阶段

去噪轨迹的不同阶段对位置引导的敏感度各不相同。去噪初期的潜变量充斥着高斯白噪声，频谱几乎平坦，此时强行施加频带特化调整容易把随机噪声模式固化为错误宏观布局；而在去噪后期，图像骨架清晰确立，谱能量高度聚集，正需要全强度的频域引导来雕琢微观纹理。SEGA 利用径向剖面 \(\mathcal{E}_{\text{iso}}\) 计算谱平坦度（即维纳熵，几何平均值与算术平均值之比）：

\[\mathrm{SF}(\mathcal{E}_{\text{iso}}) = \frac{\exp\left(\frac{1}{n_{\text{bins}}} \sum_{b=0}^{n_{\text{bins}}-1} \ln \mathcal{E}_{\text{iso}}[b]\right)}{\frac{1}{n_{\text{bins}}} \sum_{b=0}^{n_{\text{bins}}-1} \mathcal{E}_{\text{iso}}[b]} \in (0, 1]\]

进而定义自适应幅度门控标量 \(\sigma\)：

\[\sigma = 1 - \left[\mathrm{SF}(\mathcal{E}_{\text{iso}})\right]^\gamma \in [0, 1]\]

超参数 \(\gamma \ge 1\)（实验中设为 1.5）控制门控开启的响应敏感度。当处于早期高平坦度状态时，\(\mathrm{SF} \to 1\) 使得 \(\sigma \to 0\)，调制自动退化为安全保守的基准模式；随着结构逐步显现，\(\sigma \to 1\)，各频带自适应缩放无缝切换至全效调节状态。

4. 频域自适应注意力缩放：端到端集成

最终的 RoPE 变换结合基准分辨率外推比 \(m_{\text{ref}} = (R_{\text{target}} / R_{\text{train}})^\kappa\)（其中 \(\kappa = 0.08\)）与动态调制项，形成统一的维度级缩放因子：

\[m_d^{(a)} = m_{\text{ref}} \cdot \mathcal{M}_d^{(a)}(\mathbf{Z}) = m_{\text{ref}} \cdot \left(1 - \sigma \cdot s_d^{(a)}\right)\]

在实际执行注意力计算时，该缩放因子直接作用在各位置的旋转特征向量上：

\[\boldsymbol{f}^{\text{SEGA}}(\mathbf{x}, n, d) = m_d^{(a)} \cdot \boldsymbol{f}^{\text{RoPE}}(\mathbf{x}, n, d)\]

这一等价实现避免了对高维注意力权重矩阵进行显式的大规模逐元素乘法，既完美继承了 FlashAttention 等底层硬件加速算子的兼容性，又将自适应频域引导直接注入到 Query-Key 的点积交互过程中。

损失函数 / 训练策略¶

SEGA 属于完全免训练（training-free）的直接推理增强算法，不需要任何反向传播、权重参数微调或辅助提示词优化。在实现中，默认基于强化的 2D 修正版 NTK-aware（\(b' = b \cdot s^{2D/(D-2)}\)）作为基础频率重分配框架，超参数全局固定为 \(\gamma = 1.5\) 与 \(\kappa = 0.08\)，对于任意文本提示、宽高比以及外推分辨率均无需单独微调。

实验关键数据¶

主实验¶

论文在 Aesthetic-4K 基准上针对两大代表性开源 DiT 模型（Flux 与 Qwen-Image）开展了跨分辨率全方位评测，涵盖多种极限尺寸（\(2048 \times 4096\)、\(4096 \times 2048\)、\(3072 \times 3072\)、\(4096 \times 4096\)）。评测指标覆盖人类审美偏好（ImageReward、PickScore、HPSv2）、语义图文对齐度（CLIP Score）、无参考图像质量（MUSIQ、CLIP-IQA）以及保真度分布距离（FID 与基于 patch 的 \(\text{FID}_p\)）。下表展示了在 Flux 架构下两种极端分辨率（\(2048 \times 4096\) 非对称矩形与 \(4096 \times 4096\) 极限正方形）的代表性结果：

分辨率与方法	IR↑	PS↑	CS↑	MSQ↑	FID↓	\(\text{FID}_p\)↓
\(2048 \times 4096\)
Base (Flux 直接外推)	0.39	21.64	27.32	52.96	158.86	67.31
HiFlow (多阶段流引导)	1.14	22.61	27.72	44.30	152.24	69.23
I-Max (投影流外推)	1.11	22.71	27.60	32.58	157.82	68.21
ScaleDiff (多尺度级联)	1.17	22.84	29.03	55.84	157.62	72.21
YaRN (LLM启发均匀缩放)	0.97	22.63	28.48	52.30	156.97	83.35
DyPE (动态时间步外推)	1.10	22.90	28.87	53.35	159.81	85.12
UltraImage (主频平移对齐)	0.82	22.41	28.57	55.40	157.33	59.19
SEGA (本文方法)	1.21	22.91	29.18	53.65	151.93	64.54
\(4096 \times 4096\)
Base (Flux 直接外推)	-0.72	20.31	25.34	26.52	183.33	121.93
HiFlow	1.26	23.17	28.40	33.37	155.56	60.70
I-Max	1.26	23.10	28.45	26.36	151.74	64.06
YaRN	0.88	22.21	28.30	42.87	160.48	98.52
DyPE	1.01	22.56	28.79	43.23	156.21	97.88
UltraImage	0.61	21.74	28.16	43.61	167.04	63.91
SEGA (本文方法)	1.26	23.18	29.22	45.73	150.05	51.28

在 Qwen-Image 上的评测同样印证了 SEGA 的优越性：在 \(4096^2\) 极端分辨率下，SEGA 的 ImageReward 达到 1.51（相较 DyPE 的 0.97 与 UltraImage 的 0.81 呈断层式领先），MUSIQ 达到 45.03，全方位刷新直接推理范式下的性能上限。

消融实验¶

为了严格分离频域指导各设计组件的真实贡献，在 Flux 模型 \(4096 \times 4096\) 分辨率下展开了系统消融实验，定量对比了引导剖面维度、门控机制及基准外推策略的影响：

配置 / 变体	IR↑	PS↑	HPS↑	CS↑	MSQ↑	CQA↑	FID↓
NTK + 固定统一缩放 (Fixed Scaling)	0.66	22.10	0.270	28.38	42.21	0.690	162.92
SEGA 仅轴向能量 (Axis-only, 无维纳熵门控)	1.15	23.05	0.290	28.86	44.08	0.701	152.98
SEGA 仅全局门控 (Global-only, 无轴向波长对齐)	1.13	22.89	0.286	28.81	43.55	0.671	153.58
替换基准外推：YaRN + SEGA	0.94	22.21	0.280	28.23	44.36	0.707	159.80
替换基准外推：DyPE + SEGA	1.18	22.78	0.287	28.84	45.18	0.720	154.51
SEGA 完整方案 (NTK + Axis + Global)	1.26	23.25	0.291	29.22	45.72	0.725	150.05

关键发现¶

动态频域引导对极高分辨率至关重要：当从固定缩放切换为动态频域引导时，ImageReward 从 0.66 骤升至 1.26，FID 从 162.92 显著降至 150.05，证实了传统语言模型中单标量注意力缩放用于高分辨率视觉生成时存在的内在缺陷。
轴向与全局双路剖面相辅相成：单独启用轴向波长能量调制（Axis-only）或单纯依赖谱平坦度门控（Global-only）虽然均能带来明显增益，但两者融合时才能在去噪早期避免白噪声扰动，并在晚期提供精细的方向性频带补偿。
基频重分配选择的决定性影响：将 SEGA 架构搭建于 2D 修正版 NTK 之上效果最强。YaRN 的分段频率混合机制在二维网格下会打破几何空间连续性，即使叠加 SEGA 也会削弱全局布局质量（FID 升至 159.80）。

亮点与洞察¶

建立 RoPE 旋转维度与空间频域波长的严密对偶关系：创造性地将二维空间傅里叶分析引入旋转位置编码的注意力缩放过程，打破了以往“盲目缩放注意力矩阵”的经验主义做法，理论连贯且机制自然。
维纳熵作为天然去噪进程感知器：巧妙借助频谱平坦度（Wiener Entropy）直接捕捉扩散潜变量从无序白噪声到结构化图像的相变过程，无需人工硬编码复杂的非线性时间步调度衰减函数。
对复杂宽高比与极高分辨率的强大自适应性：方向解耦的轴向积分解构使得长宽比极度失衡的非正方形图像（如 \(2048 \times 4096\)）依然能获得精准匹配的高低频注意力校正，具备向视频扩散模型扩展的优异通用潜力。

局限与展望¶

计算瓶颈仍受限于长序列注意力开销：作为直接推理算法，虽然避免了复杂级联推理，但在超高分辨率（如 \(4096^2\)）下潜变量 token 数量巨大，全局完全注意力的内存占用和推理延迟依然较高。
超深去噪步的频域累积漂移：在去噪步数非常大（如 >100 步）时，频谱计算可能对极高频局部模式出现轻微过敏感现象，未来可探索在频域引入自适应低通平滑或小波多分辨率分析。
向视频与三维场景外推的迁移拓展：二维轴向能量分解自然契合空间网格，未来可进一步拓展至三维时空张量，解决视频 DiT 中长序列时间轴与超高清空间轴的联合外推挑战。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从频域能量分布对偶视角破解 DiT 在 RoPE 机制下的分辨率外推难题，设计优雅精妙。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 Flux 与 Qwen 两大前沿开源基准，跨越四种极端分辨率与十余项评测指标，消融论证严密扎实。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述逻辑严谨，数学公式与物理概念映射清晰，图表呈现详尽直观。
价值: ⭐⭐⭐⭐⭐ 零微调开销、即插即用，为大规模 DiT 落地工业级高分辨率图像生成提供了极具实用价值的推理范式。