Source-linked AI summary
USO: Unified Style and Subject-Driven Generation via Disentangled and Reward Learning
Shaojin Wu, Mengqi Huang, Yufeng Cheng, Wenxu Wu, Jiahe Tian, Yiming Luo, Fei Ding, Qian He
TL;DR
Style-driven and subject-driven generation are commonly treated as separate objectives, despite both requiring content–style disentanglement. USO unifies them with cross-task triplets, disentangled training, style reward learning, and USO-Bench, achieving state-of-the-art results across subject, style, and joint tasks.
Problem
Existing methods treat style-driven and subject-driven generation as isolated tasks, limiting cross-task disentanglement of relevant and irrelevant visual features.
Source: §1 · pp. 3–4Method
USO combines cross-task triplet curation, progressive style-alignment and content–style disentanglement training, style reward learning, and a unified benchmark.
Source: §1 · p. 4, §3.3.1 · p. 7Results
USO achieves state-of-the-art results on subject-driven, style-driven, and joint style-subject-driven tasks, attaining the highest CLIP-T, DINO, and CSD scores.
Source: §1 · p. 4Takeaways & Limitations
Cross-task co-disentanglement supports one customization model that handles individual tasks and free-form combinations with subject consistency, style fidelity, and text controllability.
Source: §1 · p. 4
Abstract
from arXiv · showhide
Existing literature typically treats style-driven and subject-driven generation as two disjoint tasks: the former prioritizes stylistic similarity, whereas the latter insists on subject consistency, resulting in an apparent antagonism. We argue that both objectives can be unified under a single framework because they ultimately concern the disentanglement and re-composition of content and style, a long-standing theme in style-driven research. To this end, we present USO, a Unified Style-Subject Optimized customization model. First, we construct a large-scale triplet dataset consisting of content images, style images, and their corresponding stylized content images. Second, we introduce a disentangled learning scheme that simultaneously aligns style features and disentangles content from style through two complementary objectives, style-alignment training and content-style disentanglement training. Third, we incorporate a style reward-learning paradigm denoted as SRL to further enhance the model's performance. Finally, we release USO-Bench, the first benchmark that jointly evaluates style similarity and subject fidelity across multiple metrics. Extensive experiments demonstrate that USO achieves state-of-the-art performance among open-source models along both dimensions of subject consistency and style similarity. Code and model: https://github.com/bytedance/USO
1 Introduction
USO argues that style-driven and subject-driven generation can be unified through cross-task co-disentanglement, because each task learns complementary inclusion and exclusion of visual features. It implements this idea with curated triplets, disentangled training, style rewards, and joint evaluation, reporting state-of-the-art results across both task dimensions.
Problem
- The underlying challenge is selecting task-relevant features from noisy visual references while excluding irrelevant style or subject information.Style-driven generation needs style features, whereas subject-driven generation needs subject appearance features.Source: §1 · pp. 1, 3
Method
- The cross-task framework constructs triplets and trains USO with progressive style alignment, content–style disentanglement, and style reward learning.The triplets contain style references, de-stylized subject references, and stylized subject results.Source: §1 · p. 4
Motivation
- USO identifies isolated task-specific disentanglement as a limitation and proposes cross-task co-disentanglement to mutually improve style-driven and subject-driven generation.Style-driven generation must retain style while excluding subject appearance, whereas subject-driven generation requires the opposite separation.Source: §1 · pp. 3–4
Evaluation
- USO-Bench provides a unified benchmark for evaluating subject-driven, style-driven, and joint style-subject-driven customization.The benchmark is presented as the first evaluation framework tailored to cross-task customization.Source: §1 · p. 4
Results
- USO achieves state-of-the-art results on subject-driven, style-driven, and joint style-subject-driven tasks, with the highest reported CLIP-T, DINO, and CSD scores.The model also handles individual tasks and free-form combinations while exhibiting subject consistency, style fidelity, and text controllability.Source: §1 · p. 4
2 Related Work
Prior work studies style transfer and subject-driven generation largely through task-specific architectures, datasets, or feature-injection strategies. These lines of work increasingly suggest that reference-guided generation tasks may share a common in-context foundation, motivating their joint treatment.
2.1 Style Transfer
- Style-transfer methods use pretrained encoders, attention manipulation, or layer-aware injection to preserve style while limiting content leakage.Several approaches improve disentanglement through coupled data or time- and layer-aware feature injection, but analyses may depend on specific architectures.Source: §2.1 · p. 4
- USO's triplet curation framework generates both layout-preserved and layout-shifted triplets to support deeper cross-task content–style disentanglement.This contrasts with prior triplet work focused on style-driven generation and layout-preserved data.Source: Figure 3 · p. 5
2.2 Subject-Driven generation
- Subject-driven generation conditions image synthesis on text instructions and reference subjects, with recent transformer-based methods using shared attention and in-context learning.Related systems extend references to layout control, multiple images, or simple style transfer.Source: §2.2 · p. 5
- The literature increasingly indicates that diverse reference-guided generation tasks can be unified within transformer in-context frameworks.This observation motivates investigating whether jointly addressing tasks can provide mutual benefits.Source: §2.2 · p. 5
3 Methodology
USO unifies subject-driven and style-driven generation through cross-task triplets, disentangled conditional encoders, style alignment, and style reward learning. Its framework combines semantic and visual conditioning so subjects and styles can be recomposed across tasks.
3.1 Preliminary
- USO extends multimodal DiT conditioning by concatenating text tokens, noisy latents, and condition tokens so representations interact across their respective spaces.Source: §3.1 · p. 5
3.2 Cross-Task Triplet Curation Framework
- Cross-task triplets include layout-preserved and layout-shifted examples, enabling deeper content–style disentanglement than prior layout-preserved style-driven data.The dataset is built from stylized pairs and complementary stylization and de-stylization experts, with filtering for style similarity and subject consistency.Source: §3.2 · pp. 5–6
3.3 Unified Customization Framework (USO)
- USO unifies subject-driven and style-driven generation so arbitrary subjects can be recombined with arbitrary styles beyond rigid layout preservation.Source: §3.3 · pp. 6–7
3.3.1 Style Alignment Training.
- Style alignment replaces VAE-only reference encoding with SigLIP semantic features and a hierarchical projector that captures both high-level structure and fine-grained style details.Only the projector is trained in this stage, aligning extracted style features with the pretrained text-to-image model’s distribution.Source: §3.3.1 · p. 7
3.3.2 Content-Style Disentanglement Training.
- Content-style disentanglement uses separate encoders and triplet training to reduce content leakage while supporting both subject-driven and style-driven tasks.The design allows style-image details to be excluded when they are not relevant to the requested task.Source: §3.3.2 · p. 7
3.3.3 Style Reward Learning
- Style Reward Learning supplements flow matching with a style-similarity reward computed by a VLM-based filter or CSD model, then combines the reward and pretraining losses.SRL alternates gradient-free inference with reward back-propagation, using L = LPre + λLSRL after the specified training step.Source: §3.3.3 · pp. 7–8
4 Experiments
USO-Bench evaluates subject consistency, style similarity, and text–image alignment across subject-driven, style-driven, and joint generation. USO achieves leading results across these settings and receives top user-study performance on both subject-driven and style-driven tasks.
4.1 Experiments Setting
- USO-Bench contains 6000 subject-driven, 7040 style-driven, and 29500 joint-task samples from 50 content and 50 style references.Source: §4.1 · p. 11
- Evaluation measures subject consistency with CLIP-I and DINO, style similarity with CSD, and text–image alignment with CLIP-T.Source: §4.1 · p. 11
4.2 Experimental Results
- USO achieves the highest DINO and CLIP-I scores and a leading CLIP-T score for subject-driven generation on USO-Bench.Source: Table 1 · p. 11, §4.2 · p. 12
- USO achieves the highest CSD and CLIP-T scores among style-driven approaches, including abstract material and Pixar-style references.Source: Table 1 · p. 11, §4.2 · p. 12
- For joint style-subject-driven generation, USO achieves the highest CSD and CLIP-T scores and substantially outperforms all baselines.Source: Table 2 · p. 11, §4.2 · p. 12
- A user study finds USO performs best on both tasks across text fidelity, visual appeal, consistency or similarity, and overall quality.Source: §4.2 · p. 12, Figure 9 · p. 12
5 Ablation Study
Ablations show that style reward learning, style alignment training, disentangled encoders, and the hierarchical projector each contribute to USO’s performance. Removing or replacing these components degrades style, subject, identity, or text-related outcomes.
5.1 Effect of Style Reward Learning (SRL).
- SRL improves style similarity in style-driven and joint tasks and improves subject and identity consistency in subject-driven examples.Source: §5.1 · p. 13, Figure 10 · p. 13
- Removing SRL sharply lowers CSD, CLIP-I, and CLIP-T, while SRL also improves texture quality and reduces plastic artifacts.Source: §5.1 · p. 13, Table 3 · p. 13
5.2 Effect of Style Alignment Training (SAT).
- Removing SAT degrades subject-driven CLIP-T and joint-task CSD, while weakening the oil-painting style in the cheetah example.Source: §5.2 · pp. 13–14, Table 4 · p. 13, Figure 11 · p. 14
5.3 Effect of Disentangled Encoder (DE).
- Replacing disentangled encoders with one VAE harms nearly every metric, reduces the cheetah’s stylization, and weakens the man’s identity.Source: §5.3 · p. 14
5.4 Effect of Hierarchical Projector.
- The hierarchical projector obtains the highest CSD and a leading CLIP-T score, substantially benefiting style-alignment training.Source: §5.4 · p. 14
6 Conclusion
USO unifies subject-driven, style-driven, and joint style-subject-driven generation through cross-task co-disentanglement. It combines triplet curation, style alignment, content–style disentanglement, and style-reward learning, and achieves state-of-the-art results across these tasks.
- USO combines triplet curation, progressive style-alignment training, content–style disentanglement training, and style-reward learning in one framework.Source: §6 · p. 14
- USO-Bench provides task-specific and joint evaluation, while USO achieves state-of-the-art subject consistency, style fidelity, and text controllability.Source: §6 · p. 14
F.1.1 Implementation Details.
The implementation uses distinct training stages and reports their main optimization settings. The section also includes a USO-Bench example figure.
F.1.1 Implementation Details.
- Style-alignment training uses 23,000 steps at batch size 16, learning rate 8e−5, resolution 768, and reward steps S = 16,000.Source: §F.1.1 · p. 18
- Content–style disentanglement training uses 21,000 steps at batch size 64, learning rate 8e−5, resolution 1024, and reward steps S = 18,000.Source: §F.1.1 · p. 18
- LoRA rank 128 is used throughout the reported training configuration.Source: §F.1.1 · p. 18
- Figure 12 provides examples from USO-Bench.Source: Figure 12 · p. 18
F.1.2 Details of USO-Bench.
USO-Bench jointly evaluates subject-driven, style-driven, and style-subject-driven generation by combining prompt types with style references and varying layout conditions.
F.1.2 Details of USO-Bench.
- USO-Bench covers subject-driven, style-driven, and joint style-subject-driven generation through paired prompts and style-reference images.Joint samples are formed through the Cartesian product of subject-driven prompts and style references.Source: §F.1.2 · p. 18
- The benchmark separates joint-generation prompts into layout-shifted and layout-preserved variants.Source: §F.1.2 · p. 18
F.2 More Results
On DreamBench, USO achieves the strongest reported subject-consistency results, while additional qualitative results show broad control over subjects, styles, and their joint composition.
F.2.1 Quantitative Evaluation on DreamBench [25].
- USO achieves the highest CLIP-I and DINO scores on DreamBench, while its CLIP-T score of 0.317 trails the top result of 0.318.The evaluation uses 4,500 image groups across all subjects, with six generated images per prompt.Source: §F.2.1 · p. 18, Table 5 · p. 19
F.2.2 Additional Results.
- USO extracts task-relevant content features while maintaining subject consistency across diverse prompts and tasks.The qualitative results span subject-driven generation, instruction-based stylized editing, and identity preservation.Source: §F.2.2 · p. 19
- USO captures fine-grained brushwork, material textures, and abstract artistic styles beyond simple color transfer.Source: §F.2.2 · p. 19
- USO combines arbitrary subjects with arbitrary styles in both layout-preserving and layout-shifting generations.Source: §F.2.2 · p. 19