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

SnapGen++: Unleashing Diffusion Transformers for Efficient High-Fidelity Image Generation on Edge Devices
Jan 13, 2026Dongting Hu, Aarush Gupta, Magzhan Gabidolla +12Diffusion Transformer Diffusion Model Distillation
Snap Inc. · University of Melbourne · MBZUAI
DDiT: Dynamic Patch Scheduling for Efficient Diffusion Transformers
Feb 19, 2026Dahye Kim, Deepti Ghadiyaram, Raghudeep GaddeDiffusion Model Inference Acceleration Diffusion-Based Image Generation
Boston University · Amazon
CoReDiT: Spatial Coherence-Guided Token Pruning and Reconstruction for Efficient Diffusion Transformers
May 13, 2026Zhuojin Li, Hsin-Pai Cheng, Hong Cai +2Diffusion Transformer Visual Token Pruning
Qualcomm AI Research†

Jan 13, 2026J/K move · Enter open · S save

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

Dongting Hu, Aarush Gupta, Magzhan Gabidolla, Arpit Sahni+11

Snap Inc. · University of Melbourne · MBZUAI

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

Open paper AskPDF arXiv

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.

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