ConceptPrism: Concept Disentanglement in Personalized Diffusion Models via Residual Token Optimization

1. Problem Statement

Personalized Text-to-Image (T2I) generation aims to create user-specific concepts (e.g., "my doll," "my style") from a few reference images. The core challenge is concept disentanglement: separating the target concept (shared features across images) from irrelevant residual information (background, lighting, specific poses).

• The Trade-off: Without effective disentanglement, models face a conflict between concept fidelity (how well the generated image matches the reference) and text alignment (how well it follows the user's prompt).
  • If a model captures too much information to ensure high fidelity, it often memorizes unwanted attributes (e.g., a specific background), causing the generated image to conflict with new text prompts.
  • If it focuses only on the text, it fails to capture the unique visual details of the user's concept.
• Limitations of Existing Methods: Current approaches rely heavily on manual guidance (e.g., linguistic cues, segmentation masks, or auxiliary image sets) to define what to learn. These methods are often incomplete (failing to capture intricate details), cumbersome, and lack scalability across different model architectures.

2. Methodology: ConceptPrism

ConceptPrism introduces a novel framework that extracts shared features exclusively through cross-image comparison, eliminating the need for external manual guidance. The method operates in two main stages:

A. Token Definition

The visual information in a set of $N$ reference images $\{x^{(i)}\}$ is decomposed into:

1. Target Concept ($C_{target}$): Shared visual features across all images.
2. Image-wise Residuals ($C^{(i)}_{residual}$): Unique, image-specific information (background, pose, lighting).

The model assigns distinct text tokens to represent these:

• Target Token ($t_{target}$): A single token to capture $C_{target}$.
• Residual Tokens ($\{t^{(i)}_{residual}\}$): A set of tokens, one for each image, to capture $C^{(i)}_{residual}$.

B. Dual-Loss Optimization (Token Optimization Stage)

The tokens are jointly optimized using two complementary objectives:

1. Reconstruction Loss ($L_{rec}$):

   • Ensures that the combination of the target token and the specific residual token can faithfully reconstruct the corresponding reference image.
   • Formula: Minimizes the noise prediction error when conditioning on $c^{(i)} = \text{"}t_{target} \text{ with } t^{(i)}_{residual}\text{"}$.

2. Exclusion Loss ($L_{excl}$):

   • The Core Innovation: This loss forces the residual tokens to be uninformative regarding the shared target concept.
   • It minimizes the KL divergence between the generation distribution conditioned on a residual token (using a different image $x^{(j)}$ as a proxy for the target concept) and the unconditional distribution (null condition).
   • Mechanism: By forcing the residual token to discard shared features when reconstructing other images, an "information vacuum" is created. This compels the target token to absorb the shared concept to satisfy the reconstruction objective.
   • Formula: Minimizes the distance between noise predictions conditioned on $t^{(i)}_{residual}$ vs. the null condition $\emptyset$ for images $j \neq i$.

Initialization Strategy:

• Residual Tokens: Initialized using embeddings from descriptive sentences (8–32 words) generated by a Vision-Language Model (VLM) describing the whole scene. This ensures they start with general image context.
• Target Token: Initialized randomly. The optimization process naturally drives it to capture the commonalities.

C. Concept Disentangled Fine-Tuning

Once the tokens are optimized:

• The model is fine-tuned using LoRA (Low-Rank Adaptation) on the cross-attention layers.
• The learned tokens ($t_{target}$ and $t^{(i)}_{residual}$) are used as text conditions.
• This allows the model to focus strictly on the target concept during generation, preventing the entanglement of residual details.

3. Key Contributions

1. Novel Disentanglement Paradigm: ConceptPrism is the first method to automatically identify shared concepts via cross-image comparison without relying on manual guidance, masks, or auxiliary encoders.
2. Exclusion Loss: A technical solution that isolates target features by constraining the content of residual tokens, creating an information vacuum that forces the target token to learn the shared concept.
3. Scalability: The method is lightweight (token optimization + LoRA fine-tuning) and architecture-agnostic, working effectively on both Stable Diffusion and Diffusion Transformers (DiT).

4. Experimental Results

The method was evaluated on DreamBench (30 subjects) and other datasets, comparing against baselines like DreamBooth, Custom Diffusion, DisenBooth, and ELITE.

• Quantitative Performance:
  • Text Alignment (CLIP-T): Achieved 0.357, outperforming DreamBooth (0.313) and matching Custom Diffusion.
  • Concept Fidelity (DINO): Achieved 0.210, significantly higher than all baselines (e.g., DreamBooth: 0.182, DisenBooth: 0.170).
  • Trade-off: ConceptPrism achieves the best balance, significantly improving fidelity without sacrificing text alignment, whereas other methods typically degrade one metric to improve the other.
• Qualitative Analysis:
  • Discrete Objects: Successfully generates the target object in diverse scenes without copying background elements (a common failure in DreamBooth).
  • Abstract Styles: Captures complex visual styles (e.g., "sunset," "swirl") that linguistic cues alone fail to represent, outperforming encoder-based methods.
• Efficiency:
  • Training time is comparable to DreamBooth (477s vs. ~434s) and significantly faster than DisenBooth (928s).
  • Requires fewer fine-tuning steps (120 steps) to reach optimal quality compared to DreamBooth (400 steps).
• Robustness:
  • Token Length: Performance is stable even with a single residual token ($n=1$).
  • Initialization: The method is robust to the length of the initial description for residual tokens.
  • Reference Diversity: Unintended commonalities (e.g., a shared background color) can be mitigated by adding a single counterexample image.

5. Significance

ConceptPrism addresses a fundamental bottleneck in personalized generation: the reliance on manual supervision to separate concepts. By automating this process through residual token optimization and exclusion loss, it enables:

• Higher Fidelity: Capturing intricate visual details that manual labels often miss.
• Better Generalization: Generating the concept in novel contexts without inheriting unwanted attributes.
• Broader Applicability: A scalable solution that works across diverse model architectures (SD, DiT) and concept types (objects, styles, poses) without requiring complex external tools or encoders.

This work establishes a new standard for concept disentanglement, proving that shared visual concepts can be precisely isolated through internal model dynamics rather than external guidance.