EasyControl: Plug‑and‑Play DiT Control with Arbitrary Aspect Ratios and Accelerated Inference

Overview

EasyControl is a new conditional generation paradigm for Diffusion Transformers (DiT). It treats each condition as an independent condition branch that is injected into a pretrained DiT model via a low‑rank LoRA module, enabling seamless integration with custom models and efficient multi‑condition fusion.

Key Innovations

Lightweight Condition‑Injection LoRA Module : isolates condition signals, applies low‑rank projection only to the condition branch, and freezes text and noise branches.

Position‑Aware Training Paradigm (PATP) : normalizes input conditions to a fixed resolution and uses position‑aware interpolation to keep spatial consistency, allowing arbitrary aspect ratios.

Causal Attention + KV Cache : combines causal attention with a KV‑cache that pre‑computes condition features at timestep 0, reducing repeated computation and cutting inference delay.

Condition‑Injection LoRA

In the Transformer, each token passes through query, key, and value projections before self‑attention. Standard QKV matrices are shared across text, noise, and condition branches, which limits condition representation. EasyControl inserts a LoRA layer into the condition branch only, yielding the modified QKV as shown in the following equation: Q_c = Q + A_c·B_c where A_c and B_c are low‑rank matrices learned for the condition branch, while the text and noise branches remain unchanged.

Causal Attention Mechanisms

Two causal attention variants are introduced:

Conditional Causal Attention : enforces that condition tokens cannot attend to text or noise tokens during training, preventing unwanted information flow.

Mutual Causal Attention : allows condition tokens to attend to each other but isolates them from text/noise queries, enabling multi‑condition inference without cross‑condition interference.

Both mechanisms use custom attention masks to control the flow of information, as illustrated in the mask equations (images).

Position‑Aware Training Paradigm

To support arbitrary resolutions, control signals are down‑sampled to a target size using Position‑Aware Interpolation (PAI) . PAI interpolates positional encodings alongside pixel values, preserving pixel‑level alignment. The scaling factors for height ( s_h) and width ( s_w) are computed as:

s_h = H_{target} / H_{orig},  s_w = W_{target} / W_{orig}

Each block in the resized condition image is mapped back to its original location using a linear mapping, ensuring spatial consistency.

Loss Function

The training objective is a flow‑matching loss: L = \|\epsilon_\theta(x_t, c) - \epsilon\|^2 where x_t denotes the noisy image at time t, c the condition, \theta the model parameters, and \epsilon the ground‑truth noise.

Efficient Inference via KV‑Cache

Because the condition branch is independent of diffusion timesteps, its Key‑Value pairs are pre‑computed at t=0 and cached. During subsequent steps, the cached KV pairs are reused, eliminating the need to recompute condition features for each of the N diffusion steps.

Experiments

Implementation details : EasyControl builds on the FLUX.1‑dev pretrained DiT, training on 4×A100 80 GB GPUs with batch size 1, learning rate 1e‑4 for 100 k steps. Inference uses flow‑matching sampling with 25 steps.

Evaluation setup : Four settings are compared – (1) single‑condition generation, (2) single‑condition adaptation with a custom model, (3) multi‑condition integration, and (4) resolution adaptability. Metrics include inference time, parameter count, controllability, generation quality, and text‑image alignment.

Results :

Inference latency is reduced by ~30% thanks to KV‑cache; the full model runs in 16.3 s (single condition) vs. 38.5 s without PATP and KV‑cache.

Parameter count stays at 15 M, far smaller than ControlNet’s 3 B.

Qualitative comparisons show EasyControl maintains text consistency and produces higher‑quality images across Canny, depth, OpenPose, and thematic controls.

Quantitative tables (Table 1) confirm the speed‑quality trade‑off advantage.

Ablation study demonstrates that removing any component (CIL, PATP, or causal mutual attention) degrades multi‑condition generalization, high‑resolution fidelity, or control precision.

Conclusion

EasyControl delivers an efficient and flexible framework for conditional diffusion generation. By combining a lightweight condition‑injection LoRA, a position‑aware training paradigm, and causal attention with KV‑cache, it resolves the efficiency and flexibility bottlenecks of existing DiT‑based control methods while supporting arbitrary resolutions and plug‑and‑play compatibility with community‑customized models.

References

[1] EasyControl: Adding Efficient and Flexible Control for Diffusion Transformer