Boosting Resolution Generalization of Diffusion Transformers with

Randomized Positional Encodings

Liang Hou1 * , Cong Liu1,2 * † , Mingwu Zheng1 , Xin Tao1‡ , Pengfei Wan1 , Di Zhang1 , Kun Gai1
1
Kling Team, Kuaishou Technology
2
Southeast University
{lianghou96, jiangsutx}@[Link] {liucong08, zhengmingwu, wanpengfei, zhangdi08}@[Link]

Abstract diffusion frameworks, making Transformer-based diffusion

arXiv:2503.18719v2 [[Link]] 7 Jan 2026

architectures one of the central focuses of modern diffusion

Resolution generalization in image generation tasks enables
the production of higher-resolution images with lower training model research (Lu et al. 2024; Ma et al. 2024; Chen et al.
resolution overhead. However, a key obstacle for diffusion 2023a, 2024b). However, existing image generation models
transformers in addressing this problem is the mismatch be- are typically trained at a specific resolution to produce high-
tween positional encodings seen at inference and those used quality images only at that resolution. Scaling these models
during training. Existing strategies such as positional encod- directly to higher resolutions usually incurs a multiplicative
ings interpolation, extrapolation, or hybrids, do not fully re- increase in training cost, which becomes prohibitive when
solve this mismatch. In this paper, we propose a novel two- computation and data resources are limited. This situation
dimensional randomized positional encodings, namely RPE- calls for models with genuine cross-resolution generalization
2D, that prioritizes the order of image patches rather than ability, such that, even when trained solely on low-resolution
their absolute distances, enabling seamless high- and low-
images, they can still generate high-quality images at higher
resolution generation without training on multiple resolutions.
Concretely, RPE-2D independently samples positions along resolutions, thereby avoiding the substantial costs associated
the horizontal and vertical axes over an expanded range dur- with conventional high-resolution training.
ing training, ensuring that the encodings used at inference lie A number of approaches have been proposed to address
within the training distribution and thereby improving reso- resolution generalization of image generation. The first line of
lution generalization. We further introduce a simple random
resize-and-crop augmentation to strengthen order modeling
work (He et al. 2023; Du et al. 2024; Lu et al. 2024; Teterwak
and add micro-conditioning to indicate the applied cropping et al. 2019; Yang et al. 2019) focuses on enhancing network
pattern. On the ImageNet dataset, RPE-2D achieves state-of- architectures, but often leads to complex designs that are
the-art resolution generalization performance, outperforming tightly coupled to specific frameworks or training pipelines.
competitive methods when trained at 2562 and evaluated at A second line of work (Jin et al. 2023) improves extrapolation
3842 and 5122 , and when trained at 5122 and evaluated at by modifying the attention mechanism to account for changes
7682 and 10242 . RPE-2D also exhibits outstanding capabili- in attention entropy. However, it largely overlooks a key
ties in low-resolution image generation, multi-stage training bottleneck: the one-to-one correspondence introduced by
acceleration, and multi-resolution inheritance. positional encodings (PEs), which enables Transformers to
perceive positional information but simultaneously constrains
Introduction the resolution generalization capacity of DiTs. A third line of
Diffusion models (Ho, Jain, and Abbeel 2020; Nichol and methods (Zhuo et al. 2024; Lu et al. 2024; Peng et al. 2023;
Dhariwal 2021; Song, Meng, and Ermon 2020; Song et al. NTK 2024) explicitly targets the limitations that PEs impose
2020) have effectively replaced traditional generative models on generalization, proposing interpolation-, extrapolation-, or
such as variational autoencoders (Kingma 2013) and gen- hybrid-based schemes. Yet these methods remain bounded by
erative adversarial networks (Goodfellow et al. 2014) as the intrinsic extrapolation limits of the underlying PEs and
the predominant paradigm in the field of image generation do not fully close the PE gap between training and inference.
due to their strong generative performance (Dhariwal and In this work, we revisit resolution generalization in im-
Nichol 2021; Rombach et al. 2022). Diffusion Transform- age generation from the perspective of PEs. We argue that
ers (DiTs) (Peebles and Xie 2023) further demonstrate that the fundamental reason existing methods perform poorly at
Transformers (Vaswani 2017) can be effectively scaled within resolution extrapolation is that many PEs required at test
* These authors contributed equally. time have never “truly” appeared during training, leading to
†
This work was conducted during the author’s internship at Kling a systematic distributional mismatch of PEs between train-
Team, Kuaishou Technology. ing and inference. To fundamentally alleviate this issue, we
‡
Corresponding author. posit that all PEs used at test time should, in a statistical
Copyright © 2026, Association for the Advancement of Artificial sense, be “covered” by the sampling process during training.
Intelligence ([Link]). All rights reserved. Guided by this principle, and inspired by the success of one-
dimensional randomized positional encodings (RPE-1D) in cating that RPE-2D is broadly compatible with existing DiT
handling length extrapolation in large language models (Ru- architectures. Beyond upward resolution extrapolation, RPE-
oss et al. 2023), we propose RPE-2D, a two-dimensional, 2D also supports downward resolution generation, accelerates
training-based randomized positional encoding framework multi-stage training when fine-tuning to higher resolutions,
tailored for resolution generalization in image generation. and enables flexible multi-resolution inheritance, highlight-
In contrast to conventional approaches that attempt to ex- ing its practical value for scalable diffusion transformers.
tend positions along fixed coordinate axes, RPE-2D performs
random sampling over a larger two-dimensional grid while
only enforcing consistency of order along the horizontal
Related Work
and vertical axes. As a result, all PEs required during high- Length Generalization in Languages Models
resolution inference can be regarded as statistically covered
by the random sampling process at training time. This re- A significant stride in extrapolation has been achieved with
frames an out-of-domain extrapolation task as an in-domain ALIBI (Press, Smith, and Lewis 2021), a method that em-
interpolation problem and models them in a unified man- ploys local attention to reinforce the model’s ability to cap-
ner via random selection, thereby avoiding any additional ture local dependencies within the data. This is crucial as
training overhead. Conceptually, each image can be regarded it allows the model to maintain a more refined understand-
as a cropped, resized, or geometrically transformed view ing of the data’s structure, thereby improving the quality of
of a larger latent canvas. This is fundamentally different [Link] notable approach is the NTK (NTK
from the one-dimensional textual sequences processed by 2024), which adjusts the frequency components of the po-
language models, where it is natural to assume a uniform sition encodings. This method is designed to preserve the
step size between adjacent tokens along the sequence and high-frequency information during the extrapolation process,
to use equally spaced positional encodings to represent their ensuring a more accurate representation of the data’s charac-
order. In contrast, in two-dimensional visual settings, differ- teristics. YaRN (Peng et al. 2023) is an innovative approach
ent views correspond to different regions and scales of the that extends the context window of large language models
same underlying canvas. From this perspective, using exactly efficiently. It does so by modifying the attention mechanism
the same, equally spaced positional encodings to model all to handle longer sequences without the need for fine-tuning,
such views introduces unnecessary constraints on positional thus maintaining a consistent level of performance across
modeling. The design of RPE-2D is precisely motivated by various lengths of input data. The concept of random posi-
this observation: by assigning randomized two-dimensional tion encoding (Ruoss et al. 2023) has also gained traction,
positional encodings, it aims to weaken the model’s reliance offering a more natural and elegant solution to the challenge
on specific positional intervals and instead encourage it to of handling longer sequences during prediction. This method
exploit positional order, which is an essential factor that has has been shown to be effective not only in language models
often been overlooked in prior work. but also in non-language models, where the generation of
Concretely, RPE-2D performs without-replacement ran- images or other data types requires a broader context un-
dom sampling along the horizontal and vertical axes of a derstanding. Attention Masking is another strategy that has
predefined maximal grid, followed by sorting the sampled proven effective in language models, which are inherently
indices in ascending order to construct a two-dimensional local in nature. By ”forcing” the model to focus on a limited
set of random positions. At test time, we instead adopt a number of tokens, it can effectively manage the increased
deterministic, equidistant sampling strategy to achieve better complexity during prediction. However, its applicability to
generalization in expectation. To further enhance the model’s non-language models is still under exploration .
ability to capture positional order, we introduce a data aug-
mentation strategy that combines random resizing and crop- Resolution Generalization in Diffusion Models
ping, and employ micro-conditioning to explicitly inject the
corresponding cropping and resizing information. This al- In the realm of computer vision, extrapolation techniques
lows the model to preserve the topological structure of im- have been pivotal in advancing the capabilities of models
ages while relying more on positional order than on precise to generate images and predict video sequences beyond the
distances. In addition, we incorporate attention scaling and limits of their training data. The development of FiT (Lu et al.
timestep shifting strategies during inference to alleviate per- 2024) and LuminaNext (Zhuo et al. 2024) has showcased
formance degradation caused by changes in attention entropy the potential of local attention mechanisms in enhancing the
and signal-to-noise ratio when sampling at high resolutions. performance of image generation models. Local attention fo-
We empirically validate RPE-2D on ImageNet at both cuses on specific regions within an image, allowing for more
2562 and 5122 training resolutions. When trained at 2562 detailed and accurate generation of high-resolution images. In
and evaluated at 3842 and 5122 , as well as trained at 5122 addition to these, there are modifications to the network struc-
and evaluated at 7682 and 10242 , RPE-2D consistently out- ture, such as attention scale (Jin et al. 2023), neighborhood
performs strong positional-encoding extrapolation baselines, attention (Hassani et al. 2023), and KV-compression (Chen
demonstrating state-of-the-art resolution generalization per- et al. 2024a). In summary, while current methods have made
formance under all evaluation settings. Moreover, integrating limited improvements in extrapolation capabilities, they still
RPE-2D with different PE families maintains or improves fail to address the fundamental issue of the position encoding
in-distribution image quality at the training resolution, indi- gap between training and prediction.
 1 2 … w … W

1,1 1,3 … 1,w … 1,W

1 1 2 … w
2 3 w W
3,1 3,3 … 3,w … 3,W
2
1 1,2 1,3 1,w 1,W … … … … … …
… 1 1,2 1,3 1,w 1,W

2 4,1 4,3 4,w 4,W h,1 h,3 … h,w … h,W

h 4 4,1 4,3 4,w 4,W

… h,2 h,3 h,w h,W … … … … … …

… h h,2 h,3 h,w h,W

h H,2 h,3 h,w H,W H,1 H,3 … H,w … H,W

H H H,2 h,3 h,w H,W

Training Inference

Figure 1: Illustration of RPE-2D for training and inference. During training (left), row and column indices are randomly sampled
without replacement from the maximal grid H × W and sorted to form a set of 2D positions matching the training resolution.
During inference (right), a deterministic, approximately equidistant grid matching the inference resolution is used.

Preliminary The full d-dimensional rotation matrix is block-diagonal,

Positional Encodings composed of such 2 × 2 rotation blocks.
Sinusoidal PE Positional encodings (PEs) (Vaswani 2017) 2D Positional Encodings
play a significant role in Transformer-based sequence mod-
eling, as they inject positional information into token rep- For image-like data with a two-dimensional structure, PEs
resentations to compensate for the order-agnostic nature of are typically extended to 2D by composing two independent
self-attention. A widely used choice is the sinusoidal PE, 1D PEs along the horizontal and vertical axes. Taking 2D
which adds to each token embedding xm ∈ Rd at position RoPE as an example, consider the query qx1 ,y1 at spatial
m ∈ {1, 2, . . . , L} a positional vector PE(m) := pm ∈ Rd , position (x1 , y1 ) and the key kx2 ,y2 at position (x2 , y2 ). The
where d ∈ N+ is the embedding dimension. Its components corresponding 2D rotary matrix can be written as
are defined as  
Rx2 −x1 0
PE(m, 2i) := pm,2i = sin(mθi ), (1) Rx2 −x1 , y2 −y1 = , (5)
0 Ry2 −y1
PE(m, 2i + 1) := pm,2i+1 = cos(mθi ), (2)
where Rx2 −x1 and Ry2 −y1 are 1D RoPE rotation matrices
where i ∈ {0, 1, . . . , d/2 − 1} and θi = b−2i/d is the fre- along the horizontal and vertical directions, respectively, and
quency associated with the i-th pair of dimensions, with base the full 2D rotation is realized as a block-diagonal composi-
b ∈ R+ . tion of the two. This construction naturally adapts RoPE to
RoPE Rotary positional encoding (RoPE) (Su et al. 2024) 2D grids while preserving its relative-position property along
is a form of relative PE that has shown strong length gener- each axis.
alization and has become a preferred choice in both modern
LLMs and DiTs. Instead of adding a positional vector, RoPE Method
applies a position-dependent rotation to the query and key 2D Randomized Positional Encodings
vectors in self-attention. Let qm ∈ Rd and kn ∈ Rd denote
the query and key at positions m and n, respectively, and let We consider resolution generalization in image generation,
f denote the attention function. RoPE modifies f as where a model is trained only at a low resolution due to
computational constraints but is expected to generate images
f (qm , kn , m, n) = (Rm qm )⊤ (Rn kn ) at higher resolutions at test time. Let htrain , wtrain ∈ N+
= q⊤ ⊤ ⊤
m Rm Rn kn = qm Rn−m kn , (3) denote the spatial size of the training images (or VAE la-
where Rm and Rn are rotation matrices that depend on the tents), and htest , wtest ∈ N+ that of the test images, with
absolute positions, and Rn−m := R⊤ htest > htrain and wtest > wtrain . Under such resolution ex-
m Rn depends only on
the relative offset (n − m). For a single 2D subspace (a pair trapolation, many positional encodings required at test time
of channels), the relative rotation matrix takes the form inevitably lie outside the range seen during training.
  NTK (NTK 2024) and YaRN (Peng et al. 2023) extend
cos((n − m)θi ) − sin((n − m)θi ) the usable context range by combining interpolation and ex-
Rn−m = , (4)
sin((n − m)θi ) cos((n − m)θi ) trapolation, but they do not resolve a fundamental issue: the
positional encoding associated with each token differs be- Data Augmentation and Micro-Conditioning
tween training and inference. Inspired by one-dimensional To further enhance the model’s ability to perceive the order
randomized positional encodings (RPE-1D) (Ruoss et al. of image patches, we jointly apply resize and crop opera-
2023) in LLMs, we reinterpret resolution extrapolation in tions to transform the “collected” high-resolution images into
image generation as an interpolation problem and propose low-resolution inputs suitable for training. The resize opera-
2D Randomized Positional Encodings (RPE-2D). The core tion encourages the model to capture global structure, while
idea is to ensure that all positional encodings used at test the crop operation prompts it to attend to local details. Im-
time are statistically covered by the training-time sampling portantly, the low-resolution images produced by these two
process. By randomly assigning positions to image patches operations are kept at the same spatial resolution. To address
in a structured manner, every test-time position lies within the issue of image incompleteness introduced by cropping,
the training distribution, thereby improving robustness to we design a micro-conditioning mechanism. We first upsam-
positional shifts. ple each low-resolution image in the training set to a high
As illustrated in Fig. 1, RPE-2D extends RPE-1D, orig- resolution (if necessary) and record its base resolution as
inally designed for text, to a two-dimensional setting suit- coriginal = (horiginal , woriginal ). During each training iteration,
able for images. A naive extension would be to flatten the we then randomly select start and end coordinates from a set
htrain × wtrain patches into a 1D sequence and sample of cropping options (including the no-crop case correspond-
positions from a longer 1D range of length HW , where ing to global resizing) and crop the base image accordingly,
H > htest > htrain and W > wtest > wtrain are hyper- yielding crop coordinates ccrop = (ctop , cleft , cdown , cright ). The
parameters. However, such flattening ignores the inherent cropped region is subsequently resized to a target resolution
2D structure of images and entangles horizontal and verti- cresize = (htarget , wtarget ), where htarget × wtarget matches the
cal neighbors in an unnatural way, leading to distorted dis- desired training resolution. These three types of conditioning
tances along the two axes. For 2D image data, the horizontal information are injected into the model via adaLN (Xu et al.
and vertical axes are naturally decoupled. RPE-2D there- 2019). Concretely, each component is independently embed-
fore performs independent randomized position sampling ded using Fourier feature encoding (Tancik et al. 2020), and
along each axis. Formally, at each training step we sam- the resulting embeddings are concatenated into a single vec-
ple, without replacement, index sets X ⊂ {1, 2, . . . , H} and tor. We then add this vector to the DiT (Peebles and Xie
Y ⊂ {1, 2, . . . , W } such that |X | = htrain and |Y| = wtrain . 2023) timestep embedding, thereby providing the model with
We then sort them in ascending order, X = {x1 , . . . , xhtrain } explicit information about the original resolution, cropping
with x1 < x2 < · · · < xhtrain and Y = {y1 , . . . , ywtrain } pattern, and final resize configuration.
with y1 < y2 < · · · < ywtrain . The 2D random position set is
constructed via the Cartesian product Training-Free Sampling Strategy
X × Y = {(x, y) | x ∈ X , y ∈ Y}. (6) Attention Scale In addition to the changes in PEs, reso-
lution extrapolation inevitably leads to an increase in the
For the patch at training index (i, j), where 1 ≤ i ≤ htrain number of image patches, creating another inconsistency be-
and 1 ≤ j ≤ wtrain , its randomized positional encoding is tween testing and training. Since attention is scale-dependent,
defined as this dependency arises from the fact that the entropy of at-
RPE(i, j) := PE(xi , yj ) ∈ Rd , tention changes as the number of patches increases (Jin et al.
2023). We also attempt to use the proposed scaling factor to
with (xi , yj ) ∈ X × Y and PE(·, ·) denoting any 2D posi- mitigate the variations in attention entropy,
tional encoding function (e.g., SinPE or RoPE). This con-  
struction preserves the monotonic order along each axis and logn m ⊤
induces a consistent 2D grid structure: along any fixed row, Attention(Q, K, V) = softmax √ QK V, (7)
d
vertical coordinates are aligned, and along any fixed column,
horizontal coordinates are aligned, while the actual intervals where m = htest × wtest and n = htrain × wtrain represent the
between sampled positions vary across training steps and number of patches during testing and training, respectively.
thus prevent the model from memorizing specific lengths. Timestep Shift When generating large images with diffu-
At test time, RPE-2D uses deterministic and approxi- sion models, the increase in resolution leads to an increase
mately equidistant positions. Given a maximal grid of size in the signal-to-noise ratio (SNR) of the noise schedule used
H × W , we choose x1 = 1, xhtest = H and y1 = 1, in training (Hoogeboom, Heek, and Salimans 2023). There-
ywtest = W ,with spacings
 xi+1 − xi = H/htest and fore, it is necessary to adjust the inference timestep spacing
yj+1 − yj = W/wtest . In this way, all test-time positions during sampling to maintain the SNR as much as possible.
lie within the support of the randomized training positions Specifically, We follow SD3 (Esser et al. 2024) to map the
while covering the full spatial extent of the maximal grid, time tep tn ∈ {1, 2, . . . , T } for n patches in training to the
effectively turning resolution extrapolation into interpolation timestep tm ∈ {1, 2, . . . , T } for m patches in inference to
over a shared 2D positional range. Our RPE-2D training approximate the same level of SNR,
paradigm is orthogonal to the specific choice of positional $ p m tn %
encoding and can be applied on top of both sinusoidal PEs n × T
tm = pm
t × T. (8)
n −1 × T
and RoPE; we empirically validate this compatibility in our 1+ n

experiments (see Table 3).

 ImageNet 256 × 256
Method 384 × 384 512 × 512
FID↓ sFID↓ IS↑ Precision↑ Recall↑ FID↓ sFID↓ IS↑ Precision↑ Recall↑
PI 18.87 41.59 260.97 0.8312 0.0602 30.64 57.76 159.72 0.676 0.055
Ext 15.95 37.79 374.74 0.8970 0.0636 28.35 54.77 232.41 0.607 0.181
NTK 16.56 35.92 375.28 0.9203 0.0686 27.88 49.8 227.45 0.619 0.177
YaRN 16.97 26.08 264.34 0.7864 0.1050 19.13 34.31 253.16 0.749 0.151
RPE-2D 15.63 14.40 385.67 0.9631 0.1174 17.95 18.23 348.99 0.849 0.181

Table 1: Comparison of RPE-2D with different methods on resolution extrapolation trained on ImageNet 256 × 256.

ImageNet 512 × 512

Method 768 × 768 1024 × 1024
FID↓ sFID↓ IS↑ Precision↑ Recall↑ FID↓ sFID↓ IS↑ Precision↑ Recall↑
PI 27.24 68.63 150.53 0.8441 0.373 38.64 92.49 139.38 0.7066 0.361
Ext 20.57 53.65 223.39 0.8184 0.462 45.77 116.54 180.16 0.5698 0.499
NTK 21.58 46.11 225.00 0.7883 0.452 32.90 73.69 216.12 0.6304 0.579
YaRN 55.21 75.84 62.93 0.5547 0.513 50.65 79.04 103.38 0.6010 0.461
RPE-2D 20.45 40.46 271.05 0.8271 0.512 25.40 47.18 192.05 0.8109 0.535

Table 2: Comparison of RPE-2D with different methods on resolution extrapolation trained on ImageNet 512 × 512.

Experiments RoPE (Su et al. 2024). It is worth noting that all competitors
Experimental Setup are implemented with two-dimensional positional encodings:
in particular, NTK and YaRN are extended to their 2D RoPE-
Training Settings We follow DiT (Peebles and Xie style versions by applying Eq. (5) following FiT (Lu et al.
2023) using ImageNet-2562 and ImageNet-5122 as training 2024). Starting from weights trained on ImageNet at 2562
datasets, employing the DiT-XL/2 network architecture while for 400k iterations, we extrapolate to 3842 and 5122 . As
keeping the other training hyper-parameters unchanged1 . On shown in Table 1, RPE-2D achieves state-of-the-art metrics
the ImageNet-2562 , we trained the model from scratch using at both 3842 and 5122 , reducing the previous best sFID of
the proposed random position encoding for 400k iterations, 34.31 obtained by YaRN to 18.23, and thus substantially
and compared it with the baseline. Subsequently, we apply improving resolution extrapolation. The qualitative results
the weights obtained from training ImageNet2562 for 400k in Fig. 2 further indicate that RPE-2D maintains superior
iterations to ImageNet5122 for an additional 800k iterations, visual quality at 5122 , suggesting that our method effectively
and compared the resolution extrapolation results with the pushes the practical extrapolation range beyond that of YaRN
baseline method. and NTK.
Evaluation Metrics Following DiT, we use FID (Heusel We further fine-tune the models on ImageNet at 5122 for
et al. 2017), sFID (Nash et al. 2021), IS (Salimans et al. an additional 800k iterations and extrapolate to 7682 and
2016), and Precision/Recall (Kynkäänniemi et al. 2019) as 10242 . As reported in Table 2, RPE-2D continues to obtain
the quantitative evaluation metrics in the experiments. We the best overall performance at higher resolutions, while PI
also follow the default value of 4.0 for “cfg-scale” in the achieves competitive results on the precision metric. Fig. 2
[Link] file in the DiT official code. presents qualitative comparisons between our method and
baseline approaches: only RPE-2D consistently preserves
Comparisons both global structure and fine details across resolutions. In
contrast, methods such as NTK and YaRN, which combine
RPE-2D is a training approach for positional encodings rather interpolation and extrapolation, tend to exhibit structural ar-
than a specific encoding form, making it theoretically compat- tifacts, whereas PI, as a purely interpolation-based method,
ible with any type of positional encoding. We apply RPE-2D often suffers from noticeable detail loss.
to both SinPE and RoPE, and the results in Table 3 show that
combining either PE with RPE-2D consistently improves Ablation Studies
performance.
We then compare RPE-2D with PI (Chen et al. 2023b), We conduct ablation studies on the main components of
extrapolation (Ext), NTK (NTK 2024), and YaRN (Peng RPE-2D: (i) the random resize-and-crop augmentation with
et al. 2023) for resolution extrapolation, all built on top of micro-conditioning (Cond-Aug), (ii) attention scaling, and
(iii) timestep shifting.
1
[Link] Our Cond-Aug treats the collected low-resolution training
 Figure 2: Qualitative results of RPE-2D against different positional encoding extrapolation methods at different resolutions.

images as resized or cropped views of larger images, where exhibits resolution generalization in both upward and down-
different sampling intervals and starting points correspond ward directions.
to resizing scale factors and cropping coordinates. As shown
in Table 4, Cond-Aug reduces the FID from 20.78 to 19.12 Multi-Stage Training Acceleration Because RPE-2D en-
and improves the IS from 293.96 to 325.76, yielding substan- ables high-quality high-resolution generation from models
tial gains at extrapolated resolutions and indicating that it trained at lower resolutions, a natural application is to facili-
strengthens the modeling of positional order. tate multi-stage, multi-resolution training. We take a model
In addition, the carefully designed attention scaling and pre-trained on ImageNet at 2562 , fine-tune it at a resolu-
timestep shifting further improve IS while significantly re- tion of 5122 , and compare the loss convergence of standard
ducing sFID, as reported in Table 4, demonstrating their ef- RoPE against RoPE equipped with our randomized positional
fectiveness when combined with RPE-2D for high-resolution encoding scheme. As illustrated in Fig. 4, the model with
sampling. RPE-2D starts from a lower loss and converges more rapidly,
which is beneficial for staged training of large-scale models.
Applications
Low-Resolution Image Generation As shown in Fig. 3, Conclusion
RPE-2D can not only generate images with higher resolu- This work investigates the resolution generalization prob-
tions than those used for training, but also synthesize images lem in diffusion transformers from the perspective of posi-
at lower resolutions, e.g., generating 1282 images when the tional encodings (PEs). Previous approaches have not fully
training resolution is 2562 . This demonstrates that RPE-2D addressed the inconsistency of PEs between training and test-
 ImageNet 256 × 256 ImageNet 512 × 512
Method
FID↓ sFID↓ IS↑ Precision↑ Recall↑ FID↓ sFID↓ IS↑ Precision↑ Recall↑
SinPE 17.82 11.33 359.32 0.927 0.149 14.75 10.02 277.92 0.8341 0.143
RoPE 17.30 11.07 366.19 0.956 0.176 14.17 7.16 313.73 0.8677 0.176
SinPE w/ RPE-2D 17.33 11.21 358.43 0.931 0.155 14.53 8.72 302.55 0.8324 0.145
RoPE w/ RPE-2D 16.92 11.02 362.14 0.959 0.170 14.09 6.61 398.35 0.8399 0.225

Table 3: Comparison of RPE-2D applied on absolute position encoding (SinPE) and relative position encoding (RoPE).

ImageNet 256 × 256

Method 256 × 256 512 × 512
FID↓ sFID↓ IS↑ Precision↑ Recall↑ FID↓ sFID↓ IS↑ Precision↑ Recall↑
RPE-2D 17.17 11.19 358.17 0.9579 0.1683 20.78 27.29 293.96 0.8155 0.155
+ Cond-Aug 17.09 11.13 359.55 0.9587 0.1687 19.12 23.91 325.76 0.8399 0.163
+ Attention Scale 17.02 11.05 361.73 0.9578 0.1693 18.33 19.17 347.82 0.8469 0.177
+ Timestep Shift 16.92 11.02 362.14 0.9591 0.1703 17.95 18.23 348.99 0.8493 0.181

Table 4: Ablation study on components of RPE-2D. Ablation components are progressively integrated in sequence.

Figure 4: Training loss and FID curves of RoPE and RPE-2D.

ing. We propose RPE-2D, ensuring that the PEs during testing

are all trained. By modeling the position orders among im-
age patches rather than their absolute distances, our method
bridge the gap between training and testing. Additionally,
we propose random data augmentation further enhance the
model’s ordering modeling while reducing its dependency on
the exact number of tokens. To address the potential issue of
image incompleteness caused by random data augmentation,
we also introduce micro-conditioning, enabling the model to
perceive the specific augmentation methods applied. During
Figure 3: Generated images at different resolutions, including high-resolution inference, we also employ attention scaling
128×128, 256×256, 512×512, 768×768, and 1024×1024, and timestep shifting to address issues related to attention
where the model is trained only at resolutions of 256 × 256 entropy increase and signal-to-noise ratio mismatch. Exper-
and 512 × 512. imental results on ImageNet-256/512 demonstrate that our
proposed method significantly outperforms existing compet-
ing approaches in the resolution generalization problem.
 References Jin, Z.; Shen, X.; Li, B.; and Xue, X. 2023. Training-free
2024. Ntk-aware Scaled Rope Allows Llama Models to Have diffusion model adaptation for variable-sized text-to-image
Extended (8k+) Context Size Without Any Fine-tuning and synthesis. Advances in Neural Information Processing Sys-
Minimal Perplexity Degradation. Accessed: 2024-4-10. tems, 36: 70847–70860.
Chen, J.; Ge, C.; Xie, E.; Wu, Y.; Yao, L.; Ren, X.; Wang, Kingma, D. P. 2013. Auto-encoding variational bayes. arXiv
Z.; Luo, P.; Lu, H.; and Li, Z. 2024a. Pixart-sigma: Weak-to- preprint arXiv:1312.6114.
strong training of diffusion transformer for 4k text-to-image Kynkäänniemi, T.; Karras, T.; Laine, S.; Lehtinen, J.; and
generation. arXiv preprint arXiv:2403.04692. Aila, T. 2019. Improved precision and recall metric for as-
Chen, J.; Yu, J.; Ge, C.; Yao, L.; Xie, E.; Wu, Y.; Wang, Z.; sessing generative models. Advances in neural information
Kwok, J.; Luo, P.; Lu, H.; et al. 2023a. Pixart-alpha: Fast processing systems, 32.
training of diffusion transformer for photorealistic text-to- Lu, Z.; Wang, Z.; Huang, D.; Wu, C.; Liu, X.; Ouyang, W.;
image synthesis. arXiv preprint arXiv:2310.00426. and Bai, L. 2024. Fit: Flexible vision transformer for diffu-
Chen, S.; Wong, S.; Chen, L.; and Tian, Y. 2023b. Extend- sion model. arXiv preprint arXiv:2402.12376.
ing context window of large language models via positional Ma, N.; Goldstein, M.; Albergo, M. S.; Boffi, N. M.; Vanden-
interpolation. arXiv preprint arXiv:2306.15595. Eijnden, E.; and Xie, S. 2024. Sit: Exploring flow and
Chen, S.; Xu, M.; Ren, J.; Cong, Y.; He, S.; Xie, Y.; Sinha, diffusion-based generative models with scalable interpolant
A.; Luo, P.; Xiang, T.; and Perez-Rua, J.-M. 2024b. GenTron: transformers. arXiv preprint arXiv:2401.08740.
Diffusion Transformers for Image and Video Generation.
Nash, C.; Menick, J.; Dieleman, S.; and Battaglia, P. W.
In Proceedings of the IEEE/CVF Conference on Computer
2021. Generating images with sparse representations. arXiv
Vision and Pattern Recognition, 6441–6451.
preprint arXiv:2103.03841.
Dhariwal, P.; and Nichol, A. 2021. Diffusion models beat
gans on image synthesis. Advances in neural information Nichol, A. Q.; and Dhariwal, P. 2021. Improved denoising
processing systems, 34: 8780–8794. diffusion probabilistic models. In International conference
on machine learning, 8162–8171. PMLR.
Du, R.; Chang, D.; Hospedales, T.; Song, Y.-Z.; and Ma, Z.
2024. Demofusion: Democratising high-resolution image Peebles, W.; and Xie, S. 2023. Scalable diffusion models with
generation with no $$$. In Proceedings of the IEEE/CVF transformers. In Proceedings of the IEEE/CVF International
Conference on Computer Vision and Pattern Recognition, Conference on Computer Vision, 4195–4205.
6159–6168. Peng, B.; Quesnelle, J.; Fan, H.; and Shippole, E. 2023. Yarn:
Esser, P.; Kulal, S.; Blattmann, A.; Entezari, R.; Müller, J.; Efficient context window extension of large language models.
Saini, H.; Levi, Y.; Lorenz, D.; Sauer, A.; Boesel, F.; et al. arXiv preprint arXiv:2309.00071.
2024. Scaling rectified flow transformers for high-resolution Press, O.; Smith, N. A.; and Lewis, M. 2021. Train short,
image synthesis. In Forty-first International Conference on test long: Attention with linear biases enables input length
Machine Learning. extrapolation. arXiv preprint arXiv:2108.12409.
Goodfellow, I.; Pouget-Abadie, J.; Mirza, M.; Xu, B.; Warde- Rombach, R.; Blattmann, A.; Lorenz, D.; Esser, P.; and Om-
Farley, D.; Ozair, S.; Courville, A.; and Bengio, Y. 2014. mer, B. 2022. High-resolution image synthesis with latent dif-
Generative adversarial nets. Advances in neural information fusion models. In Proceedings of the IEEE/CVF conference
processing systems, 27. on computer vision and pattern recognition, 10684–10695.
Hassani, A.; Walton, S.; Li, J.; Li, S.; and Shi, H. 2023. Ruoss, A.; Delétang, G.; Genewein, T.; Grau-Moya, J.;
Neighborhood attention transformer. In Proceedings of Csordás, R.; Bennani, M.; Legg, S.; and Veness, J. 2023.
the IEEE/CVF Conference on Computer Vision and Pattern Randomized positional encodings boost length generaliza-
Recognition, 6185–6194. tion of transformers. arXiv preprint arXiv:2305.16843.
He, Y.; Yang, S.; Chen, H.; Cun, X.; Xia, M.; Zhang, Y.; Salimans, T.; Goodfellow, I.; Zaremba, W.; Cheung, V.; Rad-
Wang, X.; He, R.; Chen, Q.; and Shan, Y. 2023. Scalecrafter: ford, A.; and Chen, X. 2016. Improved techniques for training
Tuning-free higher-resolution visual generation with diffu- gans. Advances in neural information processing systems,
sion models. In The Twelfth International Conference on 29.
Learning Representations.
Song, J.; Meng, C.; and Ermon, S. 2020. Denoising diffusion
Heusel, M.; Ramsauer, H.; Unterthiner, T.; Nessler, B.; and
implicit models. arXiv preprint arXiv:2010.02502.
Hochreiter, S. 2017. Gans trained by a two time-scale update
rule converge to a local nash equilibrium. Advances in neural Song, Y.; Sohl-Dickstein, J.; Kingma, D. P.; Kumar, A.; Er-
information processing systems, 30. mon, S.; and Poole, B. 2020. Score-based generative model-
Ho, J.; Jain, A.; and Abbeel, P. 2020. Denoising diffusion ing through stochastic differential equations. arXiv preprint
probabilistic models. Advances in neural information pro- arXiv:2011.13456.
cessing systems, 33: 6840–6851. Su, J.; Ahmed, M.; Lu, Y.; Pan, S.; Bo, W.; and Liu, Y. 2024.
Hoogeboom, E.; Heek, J.; and Salimans, T. 2023. simple Roformer: Enhanced transformer with rotary position embed-
diffusion: End-to-end diffusion for high resolution images. ding. Neurocomputing, 568: 127063.
In International Conference on Machine Learning, 13213– Tancik, M.; Srinivasan, P.; Mildenhall, B.; Fridovich-Keil,
13232. PMLR. S.; Raghavan, N.; Singhal, U.; Ramamoorthi, R.; Barron, J.;
and Ng, R. 2020. Fourier features let networks learn high
frequency functions in low dimensional domains. Advances
in neural information processing systems, 33: 7537–7547.
Teterwak, P.; Sarna, A.; Krishnan, D.; Maschinot, A.; Be-
langer, D.; Liu, C.; and Freeman, W. T. 2019. Boundless:
Generative adversarial networks for image extension. In
Proceedings of the IEEE/CVF International Conference on
Computer Vision, 10521–10530.
Vaswani, A. 2017. Attention is all you need. Advances in
Neural Information Processing Systems.
Xu, J.; Sun, X.; Zhang, Z.; Zhao, G.; and Lin, J. 2019. Un-
derstanding and improving layer normalization. Advances in
neural information processing systems, 32.
Yang, Z.; Dong, J.; Liu, P.; Yang, Y.; and Yan, S. 2019. Very
long natural scenery image prediction by outpainting. In
Proceedings of the IEEE/CVF international conference on
computer vision, 10561–10570.
Zhuo, L.; Du, R.; Xiao, H.; Li, Y.; Liu, D.; Huang, R.; Liu,
W.; Zhao, L.; Wang, F.-Y.; Ma, Z.; et al. 2024. Lumina-Next:
Making Lumina-T2X Stronger and Faster with Next-DiT.
arXiv preprint arXiv:2406.18583.