提出SAFE-DiT框架，通过智能分离冗余注意力掩码，实现高分辨率扩散模型加速与内存优化。
SAFE-DiT: Semantics-Aware Fast-path Execution for High-Resolution Diffusion Transformers
- 将注意力掩码分为可移除与需保留两类，仅清除无语义的冗余掩码
- 在1024²下提速2.69倍，2560²下提速5.09倍，内存从94.1GB降至27.9GB
- 无需训练即可部署，适合高分辨率图像生成与资源受限场景
高分辨率扩散变压器推理存在显著的空间冗余，但多数自适应实现将区域计算编码为注意力掩码，意外导致缩放点积注意力（SDPA）偏离FlashAttention快速路径。我们识别出这一可避免的系统瓶颈为掩码诱导调度税（MIDT），其随潜在序列长度增长而加剧。提出SAFE-DiT，一种无需训练的语义感知快速路径执行框架，将精确掩码剔除与基于近似的空间调度分离。SAFE-DiT仅移除确保证据来源的图像自注意力掩码（此类掩码在注意力日志上引起行级常数偏移），保留语义承载掩码如文本填充掩码，并通过提示条件的标记分块、带全局上下文的选择性状态更新及周期性上下文刷新实现空间自适应。称此纯加速配置为SAFE-Core，敏感度加权的无分类器引导则单独报告为SAFE-DiT+SW。在评估的PyTorch SDPA栈中，冗余掩码使长序列注意力比无掩码路径慢4.1至5.8倍。在Lumina-Next上，SAFE-DiT在1024²下实现2.69倍端到端加速，在2560²下达5.09倍，2560²峰值内存从94.1GB降至27.9GB，使3072²生成成为可能。配对指标、组件消融及盲测人类评估支持SAFE-Core与密集快速路径基线在视觉上无差异，而SAFE-DiT+SW提供独立的提示对齐操作点且不重新引入空间自注意力掩码。代码已开源。
原文摘要 · Abstract (English)
High-resolution Diffusion Transformer (DiT) inference contains substantial spatial redundancy, but many spatially adaptive implementations encode regional computation as attention masks, which can inadvertently move scaled dot-product attention (SDPA) away from FlashAttention fast paths. We identify this avoidable systems bottleneck as Mask-Induced Dispatch Tax (MIDT) and show that it grows with latent sequence length. We introduce SAFE-DiT, a training-free Semantics-Aware Fast-path Execution framework that separates exact mask elision from approximation-based spatial scheduling. SAFE-DiT removes only provenance-certified image self-attention masks that induce a row-wise constant shift in attention logits, preserves semantics-bearing masks such as text-padding masks, and realizes spatial adaptation through prompt-conditioned token partitioning, selective state updates with global context, and periodic context refresh. We call this acceleration-only configuration SAFE-Core and report sensitivity-weighted classifier-free guidance separately as SAFE-DiT+SW. On the evaluated PyTorch SDPA stack, redundant masks make long-sequence attention $4.1\times$ to $5.8\times$ slower than the mask-free path. On Lumina-Next, SAFE-DiT achieves $2.69\times$ end-to-end acceleration at $1024^2$ resolution and $5.09\times$ at $2560^2$, reduces peak memory at $2560^2$ from 94.1 to 27.9 GB, and enables $3072^2$ generation when dense inference runs out of memory. Paired metrics, component ablations, and a blinded human study support visual non-inferiority of SAFE-Core to the dense fast-path baseline, while SAFE-DiT+SW provides a separate prompt-alignment operating point without reintroducing spatial self-attention masks. Code is available at https://github.com/xuanhuayin/SAFE-DiT.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。