Memory-Efficient Fine-Tuning Diffusion Transformers via Dynamic Patch Sampling and Block Skipping

Sunghyun Park; Jeongho Kim; Hyoungwoo Park; Debasmit Das; Sungrack Yun; Munawar Hayat; Jaegul Choo; Fatih Porikli

cs.CVMar 21, 2026

Memory-Efficient Fine-Tuning Diffusion Transformers via Dynamic Patch Sampling and Block Skipping

Authors: Sunghyun Park, Jeongho Kim, Hyoungwoo Park, Debasmit Das, Sungrack Yun, Munawar Hayat, Jaegul Choo, Fatih Porikli,

Organizations: Qualcomm AI Research · KAIST

Abstract

Diffusion Transformers (DiTs) have significantly enhanced text-to-image (T2I) generation quality, enabling high-quality personalized content creation. However, fine-tuning these models requires substantial computational complexity and memory, limiting practical deployment under resource constraints. To tackle these challenges, we propose a memory-efficient fine-tuning framework called DiT-BlockSkip, integrating timestep-aware dynamic patch sampling and block skipping by precomputing residual features. Our dynamic patch sampling strategy adjusts patch sizes based on the diffusion timestep, then resizes the cropped patches to a fixed lower resolution. This approach reduces forward & backward memory usage while allowing the model to capture global structures at higher timesteps and fine-grained details at lower timesteps. The block skipping mechanism selectively fine-tunes essential transformer blocks and precomputes residual features for the skipped blocks, significantly reducing training memory. To identify vital blocks for personalization, we introduce a block selection strategy based on cross-attention masking. Evaluations demonstrate that our approach achieves competitive personalization performance qualitatively and quantitatively, while reducing memory usage substantially, moving toward on-device feasibility (e.g., smartphones, IoT devices) for large-scale diffusion transformers.

Figures & tables

Figure 1: Memory consumption analysis during fine-tuning DiT ( i.e. , FLUX). (a) We propose DiT-BlockSkip including patch sampling and block skipping, which reducing the training memory usage significantly. (b) Comparison with baselines and DiT-BlockSkip in aspects of GPU memory usage and subject & text fidelity.

Figure 2: Overview of the proposed method. (a) The dynamic patch sampling applies different patch sizes for each diffusion timestep, enabling the model to learn both global structure and fine-grained details depending on the noise level. Cropped patches of various sizes are resized to the same fixed resolution ( $s_{min}\times s_{min}$ ) before being fed into the model. (b) Block skipping with residual features for DiT stores the $\Delta$ value between the input and output of the consecutive skipped blocks in the precomputation stage. During fine-tuning, these residual feature is added to the input of the skipped blocks, while the remaining blocks are updated end-to-end.

Figure 3: Qualitative results and mean semantic distances are reported for LoRA fine-tuned FLUX on the CustomConcept101 dataset [ 19 ] , evaluated across 30 randomly selected classes. We mask the attention scores from image (query) to text (key) within the joint attention across 14 consecutive blocks. Masking mid-level blocks causes a significant drop in similarity compared to full attention.

Method	FLUX					SANA
Method	Ratio	Resolution	DINO	CLIP-I	CLIP-T	Ratio	Resolution	DINO	CLIP-I	CLIP-T
TI	–	512 $\times$ 512	0.3384	0.6140	0.2005	–	1024 $\times$ 1024	0.4185	0.6823	0.2334
DreamBooth	–	512 $\times$ 512	0.7131	0.8122	0.3073	–	1024 $\times$ 1024	0.6270	0.7612	0.3294
LoRA	–	512 $\times$ 512	0.7324	0.8146	0.3173	–	1024 $\times$ 1024	0.7374	0.8108	0.3254
LISA	–	512 $\times$ 512	0.7387	0.8194	0.3177	–	1024 $\times$ 1024	0.5226	0.7171	0.3188
LoRA-FA	–	512 $\times$ 512	0.6017	0.7433	0.3191	–	1024 $\times$ 1024	0.4692	0.6892	0.3224

Table 1: Comparison with baselines based on FLUX and SANA. Ratio refers to skip ratio. Resolution denotes the training resolution. Inference resolution is fixed at 1024

\times

1024.

Panel of Figure 4 — Figure 4: Comparison of memory and TFLOPS.

Figure 5: Qualitative comparison with baselines on the DreamBooth dataset.

Method	Ratio	Resolution	DINO	CLIP-I	CLIP-T
Ablation Study on Dynamic Patch Sampling
LoRA	–	512 $\times$ 512	0.7324	0.8146	0.3173
+ Resize	–	256 $\times$ 256	0.7164	0.8044	0.3176
+ Patch Sampling	–	256 $\times$ 256	0.7253	0.8099	0.3196
Ablation Study on Residual Feature Precomputing
+ Block Skipping	30%	512 $\times$ 512	0.4313	0.6815	0.3047

Table 2: Ablation study on FLUX.

Metric	LoRA	Tie	Ours
Subject $\uparrow$	40.0%	23.3%	37.7%
Text $\uparrow$	29.4%	25.0%	45.6%

Table 3: User preference on subject fidelity and text fidelity.

Figure 6: Qualitative ablation results of (a) dynamic patch sampling and (b) skipped block positions.

Figure 7: Qualitative ablation results of block skipping with and without residual features.

Model-Encoder	Skip ratio	$(\mathbf{n^{},m^{}})$
FLUX-DINO	30%	(7, 10)
	40%	(13, 10)
	50%	(19, 10)
SANA-DINO	30%	(4, 2)
	40%	(4, 4)
	50%	(5, 5)

Table 4: Skip indices selected in FLUX and SANA for the CustomConcept101 datasets. Encoder denotes the pre-trained encoder to extract image embedding for semantic distance.

LIVE Prompt List	Non-LIVE Prompt List
a <*> <subject> in the jungle	a <*> <subject> in the jungle
a <*> <subject> in the snow	a <*> <subject> in the snow
a <*> <subject> on the beach	a <*> <subject> on the beach
a <*> <subject> on a cobblestone street	a <*> <subject> on a cobblestone street
a <*> <subject> on top of pink fabric	a <*> <subject> on top of pink fabric
a <*> <subject> on top of a wooden floor	a <*> <subject> on top of a wooden floor

Table 5: Evaluation prompt list we used.

Method	Skip Ratio	Training Resolution	Training Memory	Parameter Memory	Optimizer State Memory	Other Memory	Training TFLOPs	Precompute TFLOPs
Textual Inversion	–	512 $\times$ 512	42.15 GiB	31.27 GiB	0.14 GiB	10.74 GiB	44.57	–
DreamBooth	–	512 $\times$ 512	68.96 GiB	22.17 GiB	22.52 GiB	24.27 GiB	158.77	–
LoRA	–	512 $\times$ 512	35.99 GiB	22.84 GiB	1.34 GiB	11.82 GiB	41.67	–
LISA	–	512 $\times$ 512	35.98 GiB	22.84 GiB	1.34 GiB	9.33 GiB	41.67	–
LoRA-FA	–	512 $\times$ 512	32.48 GiB	22.84 GiB	0.67 GiB	8.97 GiB	41.32	–
HollowedNet	30 %	512 $\times$ 512	24.11 GiB	15.62 GiB	1.04 GiB	7.46 GiB	24.35	19.85

Table 6: Detailed training memory comparison and TFLOPs with baselines based on FLUX . Inference resolution is

1024\times 1024

Method	Skip Ratio	Training Resolution	Training Memory	Parameter Memory	Optimizer State Memory	Other Memory	Training TFLOPs	Precompute TFLOPs
Textual Inversion	–	1024 $\times$ 1024	15.31 GiB	7.86 GiB	2.20 GiB	5.25 GiB	8.91	–
DreamBooth	–	1024 $\times$ 1024	24.69 GiB	2.99 GiB	5.98 GiB	15.72 GiB	13.34	–
LoRA	–	1024 $\times$ 1024	8.35 GiB	3.50 GiB	1.03 GiB	3.83 GiB	6.41	–
LISA	–	1024 $\times$ 1024	6.67 GiB	3.50 GiB	0.10 GiB	3.07 GiB	6.00	–
LoRA-FA	–	1024 $\times$ 1024	7.40 GiB	3.50 GiB	0.51 GiB	3.39 GiB	6.16	–
HollowedNet	30 %	1024 $\times$ 1024	5.90 GiB	2.48 GiB	0.72 GiB	2.70 GiB	4.48	2.91

Table 7: Detailed training memory comparison and TFLOPs with baselines based on SANA . Inference resolution is

1024\times 1024

Method	Skip Ratio	Training Res.	Training Mem.	DINO	CLIP-I	CLIP-T
LoRA	-	512 $\times$ 512	35.99 GiB	0.6726	0.7961	0.2937
LISA	-	512 $\times$ 512	35.98 GiB	0.5985	0.7491	0.3085
HollowedNet	30 %	512 $\times$ 512	24.11 GiB	0.4378	0.6570	0.3021
HollowedNet	40 %	512 $\times$ 512	18.40 GiB	0.4568	0.6662	0.3073
HollowedNet	50 %	512 $\times$ 512	12.23 GiB	0.4242	0.6553	0.3061
Ours	30 %	256 $\times$ 256	20.78 GiB	0.6455	0.7737	0.3004

Table 8: Comparison with baselines based on FLUX in CustomConcept101 dataset. Inference resolution is

1024\times 1024

Method	Skip Ratio	Training Res.	Training Mem.	DINO	CLIP-I	CLIP-T
LoRA	-	1024 $\times$ 1024	8.35 GiB	0.6545	0.7792	0.3096
LISA	-	1024 $\times$ 1024	6.67 GiB	0.4944	0.6970	0.3213
HollowedNet	30 %	1024 $\times$ 1024	5.90 GiB	0.6459	0.7708	0.2983
HollowedNet	40 %	1024 $\times$ 1024	5.08 GiB	0.6257	0.7537	0.2932
HollowedNet	50 %	1024 $\times$ 1024	4.26 GiB	0.5828	0.7191	0.2955
Ours	30 %	512 $\times$ 512	4.29 GiB	0.6522	0.7826	0.3040

Table 9: Comparison with baselines based on SANA in CustomConcept101 dataset. Inference resolution is

1024\times 1024

Method	Training Mem.	DINO ↑	CLIP-I ↑	CLIP-T ↑
LoRA	8.35 GiB	0.7374	0.8108	0.3254
Grad. Checkpointing	5.66 GiB	0.7374	0.8122	0.3241
Partial LoRA (50%)	7.20 GiB	0.6870	0.7874	0.3261
Ours (50%)	3.10 GiB	0.7277	0.8034	0.3177

Table 10: Comparison with partial LoRA and gradient checkpointing using SANA.

Figure 8: Performance comparison across training iterations for LoRA and our method.

Method (FLUX)	Training Res.	DINO ↑	CLIP-I ↑	CLIP-T ↑
High-to-Low	$256\times 256$	0.7164	0.8093	0.3182
Random	$256\times 256$	0.7187	0.8068	0.3192
Ours	$256\times 256$	0.7253	0.8099	0.3196

Table 11: Comparison of different patch sampling using FLUX.

Method (FLUX)	Training Res.	Memory (GiB)
LoRA	512x512	35.99
Dynamic Patch Sampling	256x256	30.48
Block Skipping (30%)	512x512	24.63
Block Skipping (40%)	512x512	18.76
Block Skipping (50%)	512x512	12.85
Ours (30%)	256x256	20.78

Table 12: Memory usage with dynamic patch sampling and block skipping applied.

Figure 9: Visualization of average semantic distances for all combinations of $n$ and $m$ . Each subplot is labeled with the corresponding model, dataset, and pre-trained encoder. Since $m=k-n$ , each column shares the same total number of skipped blocks $k$ . Each value shows the average distance when cross-attention is masked on the first $n$ and the last $m$ blocks.

Figure 10: Comparison of training time between LoRA and our method including block skipping and dynamic patch sampling.

Figure 11: Additional qualitative comparison with baselines on the DreamBooth dataset.

Figure 12: Qualitative comparison with baselines on the CustomConcept101 dataset.

Memory-Efficient Fine-Tuning

Memory-efficient fine-tuning reduces the GPU, device, and communication resources needed to adapt large models while preserving task performance.

See all 61 papers

Jan 13, 2026

SnapGen++: Unleashing Diffusion Transformers for Efficient High-Fidelity Image Generation on Edge Devices

Recent advances in diffusion transformers (DiTs) have set new standards in image generation, yet remain impractical for on-device deployment due to their high computational and memory costs. In this work, we present an efficient DiT framework tailored for mobile and edge devices that achieves transformer-level generation quality under strict resource constraints. Our design combines three key components. First, we propose a compact DiT architecture with an adaptive global-local sparse attention mechanism that balances global context modeling and local detail preservation. Second, we propose an elastic training framework that jointly optimizes sub-DiTs of varying capacities within a unified supernetwork, allowing a single model to dynamically adjust for efficient inference across different hardware. Finally, we develop Knowledge-Guided Distribution Matching Distillation, a step-distillation pipeline that integrates the DMD objective with knowledge transfer from few-step teacher models, producing high-fidelity and low-latency generation (e.g., 4-step) suitable for real-time on-device use. Together, these contributions enable scalable, efficient, and high-quality diffusion models for deployment on diverse hardware.

Dongting Hu, Aarush Gupta, Magzhan Gabidolla +12

Snap Inc. · University of Melbourne · MBZUAI

Diffusion Transformer Diffusion Model Distillation Few-Step Diffusion Sampling Diffusion Model Inference Acceleration On-Device Inference

Ask PDF

Feb 19, 2026

DDiT: Dynamic Patch Scheduling for Efficient Diffusion Transformers

Diffusion Transformers (DiTs) have achieved state-of-the-art performance in image and video generation, but their success comes at the cost of heavy computation. This inefficiency is largely due to the fixed tokenization process, which uses constant-sized patches throughout the entire denoising phase, regardless of the content's complexity. We propose dynamic tokenization, an efficient test-time strategy that varies patch sizes based on content complexity and the denoising timestep. Our key insight is that early timesteps only require coarser patches to model global structure, while later iterations demand finer (smaller-sized) patches to refine local details. During inference, our method dynamically reallocates patch sizes across denoising steps for image and video generation and substantially reduces cost while preserving perceptual generation quality. Extensive experiments demonstrate the effectiveness of our approach: it achieves up to

3.52\times

and

3.2\times

speedup on FLUX-1.Dev and Wan

2.1

, respectively, without compromising the generation quality and prompt adherence.

Dahye Kim, Deepti Ghadiyaram, Raghudeep Gadde

Boston University · Amazon

Diffusion Model Inference Acceleration Diffusion-Based Image Generation Diffusion Transformer Video Diffusion Models

Ask PDF

May 13, 2026

CoReDiT: Spatial Coherence-Guided Token Pruning and Reconstruction for Efficient Diffusion Transformers

Diffusion Transformers (DiTs) deliver remarkable image and video generation quality but incur high computational cost, limiting scalability and on-device deployment. We introduce CoReDiT, a structured token pruning framework for DiTs across vision tasks. CoReDiT uses a linear-time spatial coherence score to estimate local redundancy in the latent token lattice and skips high coherence (redundant) tokens in self-attention. To maintain a dense representation and avoid visual discontinuities, we reconstruct skipped attention outputs via coherence-guided aggregation of spatially neighboring retained tokens. We further introduce a progressive, block-adaptive pruning schedule that increases pruning gradually and allocates larger budgets to blocks and denoising steps with higher redundancy. Across state-of-the-art diffusion backbones including PixArt-α and MagicDrive-V2, CoReDiT achieves up to 55% self-attention FLOPs reduction and inference speedups of 1.33x on cloud GPUs and 1.72x on mobile NPUs, while maintaining high visual quality. Notably, CoReDiT also increases on-device memory head-room, enabling higher-resolution generation.

Zhuojin Li, Hsin-Pai Cheng, Hong Cai +2

Qualcomm AI Research†

Diffusion Transformer Visual Token Pruning Diffusion Model Inference Acceleration Diffusion Model Pruning

Ask PDF

Abstract

Memory-Efficient Fine-Tuning

Explore similar work

SnapGen++: Unleashing Diffusion Transformers for Efficient High-Fidelity Image Generation on Edge Devices

DDiT: Dynamic Patch Scheduling for Efficient Diffusion Transformers

CoReDiT: Spatial Coherence-Guided Token Pruning and Reconstruction for Efficient Diffusion Transformers