AlignedGen: Style-Consistent DiT Generation

• AlignedGen is a training-free framework that enforces consistent artistic style in DiT-based text-to-image generation without modifying model architectures.
• It resolves the DiT-specific positional collision by employing Shifted Position Embedding to remap spatial coordinates and enable effective attention sharing.
• Advanced attention sharing—including selective token sharing, controllable key scaling, and layer-selective application—ensures style is transferred while preserving prompt-dependent content.

AlignedGen is a training-free framework for enforcing style consistency across images generated by Diffusion Transformer (DiT) text-to-image models, especially Flux, SD3, and SD3.5, without model fine-tuning or architectural changes. It is motivated by the observation that diffusion models can generate high-quality images from rich textual prompts yet struggle to maintain a consistent artistic style across images conditioned on the same style prompt. The framework attributes a central DiT failure mode to conflicting positional signals under naïve attention sharing, and addresses it through Shifted Position Embedding (ShiftPE), an Advanced Attention Sharing (AAS) suite, and an efficient query-key-value feature extraction procedure for external image references, while maintaining precise text-to-image alignment (Zhang et al., 21 Sep 2025).

1. Scope, objective, and problem setting

AlignedGen targets a specific deployment problem in generative imaging: producing multiple images whose semantic content varies according to different prompts while their style remains stable. The stated goals are threefold: to ensure high visual style consistency across images generated from the same style prompt, to maintain faithful alignment with potentially differing text prompts, and to enable plug-and-play conditioning on external user-provided images as style references (Zhang et al., 21 Sep 2025).

The method is situated within the broader transition from U-Net-based diffusion systems to DiT-based systems. Prior training-free style-consistency approaches are described as constrained to the U-Net architecture, and naïvely porting their attention-sharing mechanism to DiT is reported to fail. AlignedGen therefore addresses not merely style transfer in the abstract, but style-consistent generation under the architectural conditions created by unified multimodal attention and rotary positional embeddings in DiT. This architectural focus is central to the method’s design rather than an implementation detail (Zhang et al., 21 Sep 2025).

A recurrent practical distinction in the underlying paper is between style preservation and content preservation. AlignedGen is designed so that style is shared while semantic content remains prompt-dependent. This differs from methods whose primary purpose is to preserve an object identity or a personalized concept from a reference image. In AlignedGen, prompt controllability is treated as a co-equal requirement rather than as a secondary property.

2. DiT-specific failure of naïve attention sharing

A central claim of AlignedGen is that naïve attention sharing fails in DiT because of conflicting positional signals from improper position embeddings. In DiT, rotary position embeddings encode spatial position for image tokens while image and text tokens participate in unified multimodal attention. When attention is shared between a reference and a target image without modifying position assignment, reference and target tokens at the same spatial location receive the same positional embedding. The paper terms the resulting pathology a positional collision (Zhang et al., 21 Sep 2025).

This positional collision has three reported consequences. First, the attention mechanism becomes spatially locked, so correspondence is forced primarily between identical coordinates in the reference and target. Second, the target image exhibits content leakage, copying actual reference content rather than only its style. Third, prompt controllability degrades, to the point that naïve sharing on DiT can produce nearly identical images irrespective of prompt. The reported failure modes are mode collapse, content leakage, and loss of prompt controllability (Zhang et al., 21 Sep 2025).

These observations directly challenge a common assumption that attention sharing is architecture-agnostic. Within the AlignedGen formulation, the issue is not that attention sharing is intrinsically ineffective, but that in DiT its interaction with RoPE changes the geometry of cross-image correspondence. This reframes the problem from “how to share attention” to “how to share attention under non-colliding positional structure.”

3. Shifted Position Embedding

Shifted Position Embedding, or ShiftPE, is the mechanism introduced to remove positional collision. Its key idea is to assign reference and target images non-overlapping coordinate spaces in positional encoding. If the target image uses position $(i, j)$ on an image grid of width $w$, the reference image is remapped to $(i, j+w)$. The paper describes this as placing the reference image beside the target in a larger virtual grid (Zhang et al., 21 Sep 2025).

This remapping changes the positional indices fed to RoPE without changing model parameters. ShiftPE is applied when processing the query, key, and value matrices in the DiT multimodal attention blocks. Because the reference and target no longer share identical spatial coordinates, attention is not forced into one-to-one same-location correspondence. The reported effect is that attention maps spread wider and are no longer concentrated only on matching positions, which prevents content copying, sustains style transfer, and preserves text controllability (Zhang et al., 21 Sep 2025).

ShiftPE is therefore a positional reindexing strategy rather than a learned module. Its role in the overall framework is foundational: the AAS suite is introduced only after this positional conflict is resolved. The method’s logic is sequential—first remove the geometric artifact created by improper positional overlap, then exploit cross-image attention for style consistency.

4. Advanced Attention Sharing

On top of ShiftPE, AlignedGen introduces Advanced Attention Sharing, a suite of three techniques: Selective Attention Sharing (SAS), controllable style consistency via key scaling, and layer-selective application (Zhang et al., 21 Sep 2025).

Selective Attention Sharing (SAS) restricts sharing to key and value features derived from image tokens. Because DiT uses multimodal attention over both image and text tokens, sharing all keys and values allows the reference prompt to contaminate the target’s semantic content. SAS therefore excludes text-derived tokens from cross-image injection. The stated effect is that style is shared while semantic content remains faithful to each prompt, and computational overhead is also reduced. The paper further states that only the image parts of $K$ and $V$ are shared or combined across reference and target images, using AdaIN and concatenation as described in its Equations (4)–(6) (Zhang et al., 21 Sep 2025).

Controllable style consistency via key scaling introduces a scalar parameter $\lambda$ that multiplies the keys derived from the reference image. The reported intuition is that a larger $\lambda$ increases the cosine similarity between target queries and reference keys, thereby increasing the pull exerted by the reference style in the attention softmax. This makes style strength tunable at inference time. The paper gives the corresponding construction as

$$\mathbf{K}^{F} = \mathrm{Concat}(\mathbf{K}^{tar}_{txt}, \hat{\mathbf{K}}^{tar}_{img}, \lambda \cdot \mathbf{K}^{ref}_{img}).$$

It also reports that too high a value harms prompt faithfulness, while too low a value under-transfers style (Zhang et al., 21 Sep 2025).

Layer-selective application reflects the empirical claim that not all DiT layers should participate in style alignment. Applying attention sharing to all layers can degrade content because layers specialize in different generation aspects. The paper reports that best performance is achieved when attention sharing is applied to mid/high DiT layers, with Flux.1-dev using layers 19 through 57 in the ablation study. This is presented as the regime that transfers style without harming structural or semantic accuracy (Zhang et al., 21 Sep 2025).

Taken together, these three components specify a constrained form of attention sharing. The framework does not share indiscriminately; it shares only selected modalities, at selected depths, with controllable intensity. This is the principal mechanism by which AlignedGen reconciles style consistency with text-to-image alignment.

5. Conditioning on external style references

AlignedGen extends beyond generation-from-style-prompt by supporting arbitrary user-supplied images as style references. For this purpose, it introduces an efficient feature extraction pipeline that reuses the frozen DiT inference process rather than training a dedicated encoder (Zhang et al., 21 Sep 2025).

The reported procedure has four stages. First, the reference image is encoded into the VAE latent space. Second, Gaussian noise is added at a given diffusion timestep $t$. Third, the noisy latent is passed through the frozen DiT model at the same timestep used during normal generation. Fourth, intermediate $Q$, $w$0, and $w$1 features are extracted from the model’s attention layers. These features are then cached and used for style alignment in subsequent generation passes (Zhang et al., 21 Sep 2025).

This mechanism broadens the applicability of the framework in two ways. It decouples style conditioning from the need to generate a reference image within the same system, and it preserves the training-free character of the approach. The paper explicitly presents the resulting workflow as seamless incorporation of external images as style references. It also reports qualitative results for conditioning on human images.

6. Experimental evaluation

The reported experimental setup uses Flux.1-dev as the backbone and evaluates on 100 prompt sets against StyleAligned, B-Lora, IP-Adapter, StyleShot, and CSGO. Three metrics are used: text alignment $w$2, defined as CLIP cosine similarity between image and prompt; style consistency $w$3, defined as average pairwise cosine similarity over CLIP image features within an image set; and DINO consistency $w$4, defined analogously using DINO self-supervised features (Zhang et al., 21 Sep 2025).

Method	$w$5
AlignedGen	0.282 / 0.740 / 0.554
StyleAligned	0.277 / 0.728 / 0.512
IP-Adapter	0.278 / 0.737 / 0.522
B-Lora	0.259 / 0.625 / 0.271
StyleShot	0.276 / 0.693 / 0.511
CSGO	0.280 / 0.686 / 0.450

On these metrics, AlignedGen is reported to achieve the highest $w$6 and $w$7 while also achieving the highest $w$8. The paper interprets this as evidence that stronger style consistency need not require sacrificing text alignment within the DiT setting it studies (Zhang et al., 21 Sep 2025).

The qualitative comparisons follow the same pattern. Vanilla Flux is described as producing high image quality with widely varying style. StyleAligned is reported to improve style consistency on SDXL but to fail on Flux, producing mode collapse, duplicate images, and content leakage. AlignedGen is reported to produce diverse content, consistently styled images, and prompt-faithful outputs, including when conditioned on human images (Zhang et al., 21 Sep 2025).

The user study involved 40 participants and reports average ranking, where lower is better, for style/image alignment. The values given are 1.20 for AlignedGen, 3.66 for StyleAligned, and 2.04–5.25 for the other baselines. The ablation studies further report that removing SAS lowers text-prompt alignment and style consistency, that mid/later layers are crucial for style transfer, and that $w$9 exposes a style-strength continuum with an explicit prompt-faithfulness trade-off (Zhang et al., 21 Sep 2025).

The framework is also reported to generalize architecturally to SD3, SD3.5, and other DiT-based models, and to be plug-and-play with methods such as FluxControl and Dreambooth. This indicates that the method is intended as an inference-time augmentation layer for DiT pipelines rather than as a backbone-specific retraining procedure.

7. Relation to adjacent work and terminological distinctions

AlignedGen should be distinguished from the separate system AlignGen, which addresses personalized image generation through cross-modality prior alignment rather than style consistency across multiple generated images. AlignGen introduces a learnable token $(i, j+w)$0, a Deviation Extraction Module, a selective cross-modal attention mask, and a training strategy with random reference image dropout and random concept name substitution, using FLUX.1 DEV as the base model (Lin et al., 28 May 2025). Its target problem is prompt-reference misalignment in zero-shot personalized generation, where generation tends to bias toward the textual prior and lose reference content (Lin et al., 28 May 2025).

The distinction is substantive. AlignedGen is concerned with consistent artistic style across a set of outputs and with making attention sharing work on DiT through ShiftPE and AAS. AlignGen is concerned with aligning textual and visual priors so that a personalized concept from a reference image is preserved under prompt variation. Because the names are close, conflating them obscures the fact that the two methods solve different alignment problems under different operational definitions of “alignment” (Zhang et al., 21 Sep 2025, Lin et al., 28 May 2025).

Within the narrower lineage of training-free style-consistency methods, AlignedGen’s principal contribution is its claim that DiT requires a different treatment from U-Net. The framework’s core thesis is that the unification of image and text tokens and the use of rotary positional embeddings make naïve attention sharing insufficient. In that sense, AlignedGen is best understood not as a generic style-transfer heuristic, but as a DiT-specific inference framework for style-aligned generation.