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

Problem Statement

Diffusion Transformers (DiTs) have established themselves as the dominant architecture for text-to-image (T2I) generation. However, their performance significantly degrades when generating images at resolutions beyond their training range (typically 1024² to 2048²). Extrapolating to higher resolutions often results in blurred textures, repetitive patterns, and structural breakdowns.

Existing training-free solutions attempt to mitigate this by modifying inference-time attention behavior, primarily through Rotary Position Embedding (RoPE) extrapolation combined with attention scaling. Current methods, such as YaRN, apply a uniform and content-agnostic scaling across all RoPE components. This approach treats short-wavelength components (governing fine-grained texture) identically to long-wavelength components (shaping global structure). The paper argues that this static uniformity induces a fundamental trade-off: strategies that preserve global structure often fail to recover fine details, and vice versa. Furthermore, the spectral distribution of the latent representation evolves dynamically during the denoising process and varies across different image contents, rendering static scaling suboptimal.

Methodology: SEGA

The authors introduce SEGA (Spectral-Energy Guided Attention), a training-free method that dynamically scales attention across RoPE components based on the latent's spatial-frequency structure at each denoising step.

Core Insight

SEGA is motivated by the observation that RoPE components are coupled to specific spatial frequencies. Instead of applying a fixed scaling factor, SEGA derives per-dimension scaling magnitudes by analyzing the energy distribution of the current intermediate latent.

Algorithmic Workflow

1. Spectral Analysis: At each denoising step, the intermediate latent $Z$ is transformed from the spatial domain to the frequency domain using a 2D Fast Fourier Transform (FFT).

   • Axis-wise Profiles: The 2D power spectrum is marginalized over orthogonal axes to obtain 1D energy profiles ($E_H, E_W$) for horizontal and vertical frequencies.
   • Radial Profile: The power spectrum is averaged within concentric rings to obtain a rotation-invariant profile ($E_{iso}$) representing energy distribution across spatial scales.

2. Per-Dimension Correction ($s_d$):

   • For each RoPE dimension $d$, the corresponding wavelength is identified to map it to a specific frequency band in the axis-wise profile.
   • The log-energy of this band is standardized.
   • A non-linearity (tanh) is applied, and the result is zero-centered to ensure a strict redistribution of scaling. Dimensions associated with low-energy bands receive stronger scaling to preserve positional discrimination, while high-energy bands receive weaker scaling to avoid over-amplification.

3. Global Amplitude Factor ($\sigma$):

   • The radial profile is reduced to a single scalar using spectral flatness (Wiener entropy).
   • This factor controls the strength of the scaling. If the latent's spectrum is flat (noise-dominated), $\sigma \to 0$, suppressing the correction. As structural content emerges and the spectrum becomes less flat, $\sigma \to 1$, applying the correction at full strength.

4. Final Scaling Formula:
   The final scaling factor $m_d$ for a dimension $d$ is computed as:
   $$m_d = m_{ref} \cdot (1 - \sigma \cdot s_d)$$
   Where $m_{ref}$ is a reference scale determined by the resolution ratio (using a power-law formulation). This allows SEGA to adapt continuously to the latent's spectral content, sharpening attention at under-resolved frequencies and softening it at over-emphasized ones.

Key Contributions

• Dynamic, Content-Aware Scaling: Unlike prior uniform scaling methods, SEGA adapts attention scaling to the specific spectral structure of the latent at every denoising step.
• Resolution of the Structure-Detail Trade-off: By differentiating between RoPE components based on their associated frequency bands and energy levels, SEGA simultaneously improves global structural coherence and fine-detail fidelity.
• Training-Free and Architecture-Agnostic: The method requires no fine-tuning, no learnable parameters, and no architectural changes. It integrates directly into standard RoPE-based pipelines (demonstrated on Flux and Qwen).
• Superior Performance at Ultra-High Resolutions: SEGA enables stable generation at resolutions exceeding 36 million pixels (e.g., 6144²), a regime where other methods often fail completely.

Experimental Results

The authors evaluated SEGA on Flux and Qwen architectures across multiple high-resolution settings (up to 6144²) using the Aesthetic-4K dataset and a custom zero-shot benchmark.

• Quantitative Performance: SEGA consistently outperformed state-of-the-art baselines (including NTK, YaRN, DyPE, UltraImage, and multi-stage guidance methods like HiFlow and I-Max) across all metrics: ImageReward (IR), PickScore (PS), CLIP Score (CS), MUSIQ (MSQ), and FID. Notably, at 4096², SEGA achieved the best scores across every evaluated metric.
• Qualitative Improvements: Visual comparisons demonstrated that SEGA effectively prevents common failure modes such as repetitive tiling artifacts, structural collapse, and loss of semantic clarity. It preserves complex geometries and fine textures (e.g., fur, fabric details) that other methods blur or distort.
• Robustness: The method maintained stability across non-square aspect ratios and extreme extrapolation factors (up to 36× the training area), whereas competing methods showed significant performance degradation as resolution increased.

Significance and Claims

The paper claims that SEGA represents a minimal yet effective solution for unlocking the high-resolution generation capabilities of pre-trained DiTs. By shifting from static, heuristic-based scaling to a dynamic, spectral-energy-guided approach, the authors address a fundamental limitation in current extrapolation strategies.

The authors emphasize that their work provides a "spectral perspective" on attention scaling, revealing a coupling between RoPE dimensions and spatial frequencies that can be exploited to guide inference-time adaptation. They position SEGA not as a replacement for training-based approaches, but as a powerful, plug-and-play enhancement that allows existing models to generalize far beyond their native training resolutions without the computational cost of retraining or the complexity of multi-stage pipelines. The method is presented as a step toward better understanding and controlling the behavior of attention mechanisms in resolution extrapolation tasks.