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

TL;DR

SEGA is a training-free, inference-only modification to pretrained diffusion transformers (DiTs) that lets them generate well past their training resolution — up to 4096² on FLUX and Qwen-Image. Existing training-free recipes combine RoPE extrapolation (YaRN, NTK) with a uniform attention-temperature scale; SEGA replaces the uniform scale with a per-RoPE-component, per-axis, per-sample, per-denoising-step dynamic scale derived from an FFT of the current latent. Low-energy spectral bands receive larger scaling (to recover under-resolved structure) and high-energy bands receive smaller scaling (to avoid amplifying detail that’s already there), eliminating the structure-vs-detail trade-off that fixed-scale methods face.

Key claims

Existing training-free high-resolution recipes (YaRN, DyPE, UltraImage, ScaleDiff, I-Max, HiFlow) apply a content-agnostic uniform scale across RoPE components, which forces a trade-off between preserving global structure and recovering fine detail [Abstract, project page].
SEGA dynamically rescales each rotary dimension by a scale that depends on the latent’s spatial-frequency content at the current denoising step — varying across dimensions, height/width axes, samples in the batch, and timesteps [project page].
The mapping is computed by reshaping the latent into its 2D spatial layout, applying an FFT per axis, then routing spectral energy onto the RoPE dimensions that encode each band, so the per-dimension scale is directly indexed by the energy in that band [project page Method].
SEGA is fully training-free: no new networks, no learned parameters, and the modification is layered on top of existing extrapolation methods (YaRN, DyPE, UltraImage, ScaleDiff, I-Max, HiFlow) rather than replacing them [project page].
SEGA is demonstrated on both FLUX and Qwen (Qwen-Image) DiT backbones generating up to 4096×4096, paired with each of YaRN, DyPE, UltraImage, ScaleDiff, I-Max, and HiFlow [project page gallery].
“SEGA consistently improves high-resolution synthesis across multiple target resolutions, outperforming state-of-the-art training-free baselines” [Abstract] — concrete metrics are in the paper body.

Method

The setup is a pretrained DiT with axial RoPE, run at a target resolution higher than its training resolution. SEGA leaves training and architecture untouched and inserts a single dynamic scaling step into each transformer block’s attention at each denoising step.

At step $t$ , the current noisy latent is reshaped into its 2D spatial layout (H × W tokens). For each axis (height, width) separately, SEGA takes the FFT of the latent along that axis, computes per-frequency-band energy, and maps this energy onto the RoPE dimensions for that axis — exploiting the fact that each rotary dimension $i$ encodes a specific spatial frequency $\omega_i$ . The per-dimension scale is then a function of the energy in band $i$ : low-energy bands get a larger scale (stretch attention reach in that band to fill in structure the model would have under-resolved), high-energy bands get a smaller scale (keep detail that’s already there from being over-amplified). The scale also varies per sample in the batch, since two different prompts may have very different spectral signatures at the same target resolution.

The modification is layered on top of an existing RoPE-extrapolation method (e.g. YaRN) and an existing attention-temperature recipe (e.g. UltraImage). Fixed-scale methods can be viewed as the SEGA degenerate case where the spectral-energy mapping is replaced by a constant. The galleries show 4096² outputs for FLUX and Qwen-Image pairings with each of YaRN, DyPE, UltraImage, ScaleDiff, I-Max, and HiFlow.

Results

The paper claims consistent improvements over the SOTA training-free baselines (YaRN, DyPE, UltraImage, ScaleDiff, I-Max, HiFlow) across multiple target resolutions on both FLUX and Qwen-Image, with qualitative comparisons at 4096×4096 [project page gallery, Abstract]. Specific FID / IS / metric numbers and resolution sweeps are reported in the paper body; the project page does not surface them.

Why it’s interesting

SEGA sits at the intersection of three threads on the wiki. First, it complements Native-Resolution Image Synthesis (NiT) at the training-free end: NiT is the cleanest existence proof that retraining with native-resolution packing + axial 2D-RoPE delivers SOTA across resolutions, but it requires building a new model from scratch; SEGA is what you do when you have a frozen pretrained DiT (FLUX, Qwen) and want 4K outputs at inference time only. The two papers are not in tension — they’re at opposite ends of the cost/quality lever.

Second, it lands cleanly into the wiki’s growing RoPE-pathology cluster. RoPE Distinguishes Neither Positions Nor Tokens in Long Contexts, Provably proves a content-independent failure floor for RoPE in long contexts, and Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings (PoPE) responds with a re-parameterisation that decouples content and position. SEGA addresses a different, content-dependent failure mode of RoPE — that uniform extrapolation can’t be right because different prompts have different spectral signatures at any given resolution — and patches it without retraining. The polar-positional-embedding line and the spectral-energy-attention line are independent levers and could plausibly stack.

Third, it slots into Classifier-Free Guidance variants as another dynamic, per-sample, per-step inference-time recipe — the same general design pattern as VLM-Guided Adaptive Negative Prompting (VLM-Guided Adaptive Negative Prompting for Creative Generation) and Dynamic CFG (Dynamic Classifier-Free Diffusion Guidance via Online Feedback) — except the lever is RoPE-component attention scaling rather than guidance strength. The recurring pattern is: read a signal off the current latent (or VLM judgment of it), use it to set a knob per-sample per-step. SEGA’s specific signal — spectral energy mapped onto RoPE dimensions — is novel.