CRAFT: Constrained Reward via Attention Fine-Tuning for Subject Personalization without Composed Targets

AI summary

Overview

• Research area: Subject-driven image personalization (generating new images that preserve the identity of one or more reference subjects), built on multimodal diffusion transformers (MMDiTs) and reward-based fine-tuning of diffusion models.
• Technical level: Advanced. The paper assumes familiarity with diffusion transformers, cross-attention sub-blocks, LoRA adapters, flow matching, and single-step Reward Feedback Learning (ReFL).
• One-sentence scope: CRAFT replaces expensive (reference, composed-target) paired supervision with a compact reference-only recipe that shapes where a pretrained MMDiT looks inside its own attention, then gates a pixel-level identity reward with the resulting attention masks.

Note: the paper content provided is truncated. Several sections referenced by the paper (Appendices A, B, E, F, I, J) are not included, so their contents cannot be summarized here.

What This Paper Is About

Current generalized subject-personalization systems fine-tune one pretrained MMDiT on large paired datasets of (reference, composed-target) examples, where each composed target is a synthesized image of the subject in a novel scene. Building those targets requires a multi-stage curation pipeline (LLM prompt generation, text-to-image composed-target synthesis, reference-subject extraction, VLM quality filtering, correspondence labeling) and yields 150K to over 2M paired samples per system. CRAFT's goal is to achieve strong personalization without any composed-target supervision at all, using only 10K reference images and subject masks.

Key Contributions

1. Reformulates generalized subject personalization as a reference-side reward problem. CRAFT is a single-step ReFL framework that adapts a pretrained reference-aware MMDiT through LoRA adapters trained solely on reference-side supervision, removing the need for composed-target pairs.

2. Introduces attention-level rewards realizing a "Where to look" principle. Noise-token and phrase-token attention are aligned with the correct reference subject region at a small subset of (step, block) coordinates identified by profiling the unmodified backbone.

3. Ties image-space supervision to learned attention. Per-subject masks derived from the attention gate a pixel-level identity reward, so pixel-level supervision stays consistent with the routing the attention rewards produce.

4. Demonstrates state-of-the-art results at a fraction of the data, and transferability. CRAFT reaches state-of-the-art on multiple benchmarks using only 10K reference-only instances, against the 150K to over 2M composed-target pairs prior generalized methods require, and the same recipe improves other reference-aware MMDiT backbones (UNO and UMO).

Main Findings

• Best Overall on XVerseBench with no composed-target supervision: CRAFT scores 76.47 Overall while training on 10K reference-only instances with no composed targets. The next-best prior method in the table, MOSAIC, scores 76.04 and uses 1.2M composed-target pairs.
• Single-subject gains: CRAFT reaches the best Single AVG (81.62), improving over MOSAIC by +1.57 points, and achieves the highest Single ID (84.22) and Single IP (84.23).
• Multi-subject profile: CRAFT achieves the best Multi IP (77.25) and second-best Multi AVG (71.32). Its Multi ID (61.16) is lower than MOSAIC's (69.90). The paper describes this as partly a deliberate operating point: raising the identity weight w_id recovers Multi ID to 68.83 at a small quality cost.
• Mask-free inference performs well or better: CRAFT (mask-free), which uses raw unsegmented references at inference, scores 77.80 Overall versus 76.47 for the segmented-input configuration. The paper notes masks are reward-side training annotations only, so mask-free is CRAFT's intended inference mode; this configuration is reported for reference and not included in the ranking.
• Improvement over the frozen backbone: CRAFT improves Overall by +5.43 points over the frozen FLUX.2-klein backbone (71.04 → 76.47), with the largest gains in subject identity preservation.
• All three CRAFT rewards are needed: Ablations give Overall 71.04 (backbone), 71.91 (+ R_ref), 72.37 (+ R_cons, cumulative), 73.14 (R_id alone), and 76.47 (all three rewards).
• Transfer to other backbones: Applying CRAFT to UNO raises Overall from 64.03 to 71.83; applying it to UMO raises Overall from 70.94 to 74.62. Both are marked as trained without composed-target supervision.
• User study preference: CRAFT is preferred on identity consistency (IC) at 55.2%, prompt fidelity (PF) at 62.6%, and image quality (IQ) at 34.2%. Baselines score: UMO 15.0 / 5.9 / 26.0, XVerse 15.0 / 18.3 / 16.5, MOSAIC 14.8 / 13.2 / 23.3.
• Stable results: Averaging over four samples, the run-to-run spread is small (±0.08 Overall), far below the lead over the next-best method.
• Better relational composition qualitatively: The paper reports that on prompts requiring two subjects to interact (a woman walking a corgi, an anime space ranger riding a bicycle, an Avatar wearing a cap while holding a glowing ring), CRAFT places subjects in compatible positions, orientations, and scales, while competing methods drop subjects, separate them spatially, or distort one subject.

Methodology in Plain English

CRAFT starts from the observation that a reference-aware MMDiT already produces reasonably subject-aligned attention before any fine-tuning. Rather than teaching the model where to look by showing it synthesized target images, CRAFT directly shapes that existing behavior with a lightweight reward.

Training is a single-step Reward Feedback Learning loop. The denoising prefix is rolled out without gradients up to a chosen reward step t*; then two forward passes run at that step — one on the frozen base model and one on the LoRA-adapted model, which also exposes its cross-modal attention sub-blocks. A backward pass updates only the LoRA parameters.

Three subject-specific attention handles are used per subject k: attention from noise tokens to that subject's reference tokens, noise tokens to that subject's phrase tokens, and phrase tokens to that subject's reference tokens. From these, three rewards are built:

• R_noise-ref measures the fraction of normalized noise-to-reference attention falling inside the subject's reference mask.
• R_text-ref does the same on the phrase-to-reference side, grounding each referring phrase in the reference subject region rather than generic text-encoder semantics.
• R_cons is a probabilistic soft IoU that forces the noise-grid localization derived from noise-to-reference attention to agree with the localization derived from noise-to-phrase attention, preventing a subject's attention from splitting across regions.

The noise-to-reference attention is also reused to build a per-subject noise-grid mask: it is mask-weighted, Gaussian-smoothed with σ = 2, normalized, and thresholded at 0.5. This mask gates a DINOv2 cosine-similarity identity reward computed between the generated subject region in the decoded pre-image and the reference subject region. Because the gate is derived from the very attention being optimized, the pixel-level supervision tracks the routing instead of contradicting it.

Two auxiliary terms from standard reward fine-tuning are added for prompt fidelity and aesthetics (a CLIP text-image similarity reward and an Aesthetic Predictor score), plus a velocity-space anchor L_anchor = ||v_LoRA − sg(v_base)||²₂ tying the LoRA velocity field to the base model for stability. All reward terms enter the total loss with negative signs.

The reward locus is chosen empirically rather than set arbitrarily: a subject-routing analysis on the unmodified backbone identifies the step and blocks whose noise-to-reference attention best aligns with generated subject masks. For FLUX.2-klein this yields t* = 2 and blocks {single_1, single_9, single_8}. Reading attention tensors disables Flash Attention, which is why a compact locus is used.

Implementation specifics reported: backbone FLUX.2-klein, a 9B-parameter MMDiT distilled to four denoising steps; backbone frozen; LoRA rank r = 64 attached to all attention layers; training at 1024² resolution on four NVIDIA B200 GPUs with AdamW at learning rate 2×10⁻⁶ (constant schedule, 50 warmup steps); 3,000 optimizer steps. Loss weights: w_nr = w_tr = 0.5, w_c = 1.0, w_id = 1.0, w_t = 0.3, w_a = 3×10⁻³, w_anchor = 0.5. Training data consists of isolated single-subject reference images rendered with FLUX.2 in text-to-image mode, with masks obtained via Grounded-SAM; no scene composition or composed-target synthesis is involved.

Why This Matters

Impact on research. The paper attacks a structural bottleneck in generalized personalization: the cost and coupling of composed-target curation pipelines. It shows that supervising a model's internal cross-modal routing directly can substitute for supervising synthesized output images, at roughly two orders of magnitude less training data than the largest baselines in its comparison (10K reference-only instances versus 150K to over 2M composed-target pairs). It also moves reward signals off the output image and onto the cross-modal attention sub-blocks, which prior reward fine-tuning work largely left unsupervised.

Real-world applications (as named in the paper):

• On-demand product imagery.
• Advertising creative.
• Personalized media with consistent characters across multiple scenes.
• Brand and IP asset generation that must remain identity-coherent across diverse contexts.

Industry relevance. Because CRAFT trains on reference images and masks only, it removes dependence on a particular target synthesizer and curation choice, which the paper identifies as a source of tight coupling in existing systems. Its transfer results on UNO and UMO indicate the recipe can be layered onto existing reference-aware backbones rather than requiring a new model to be trained from scratch, and its mask-free inference mode means no test-time reference segmentation is needed.

Future Directions

• Closing the multi-subject identity gap. CRAFT's Multi ID (61.16) trails MOSAIC's (69.90). The paper reports that raising w_id recovers Multi ID to 68.83 at a small quality cost, framing the current setting as an operating point rather than a ceiling — the trade-off and its full characterization are discussed in a section of the paper that is not included in the provided content.
• Broadening the reward locus. The reward locus is currently a compact set of (step, block) coordinates chosen by profiling the unmodified backbone. Whether richer or automatically selected loci further improve results — and how to make attention reading cheaper given that it disables Flash Attention — is a natural extension.
• Wider backbone transfer. The recipe is verified on FLUX.2-klein, UNO, and UMO. Whether the same "Where to look" formulation transfers to other reference-aware MMDiT families and to architectures beyond this class is left open.
• Reducing the remaining annotation requirements. Training still requires subject masks (via Grounded-SAM) on the reward side, and the main comparison uses segmented reference images at evaluation. The mask-free result (77.80 Overall) suggests further work on removing segmentation from the pipeline entirely.

Target Audience

Researchers and engineers working on diffusion-transformer-based image generation, subject-driven personalization, and reward fine-tuning of generative models. It is most useful to readers already comfortable with attention mechanics, LoRA, and ReFL-style training loops, and to practitioners who need to build or adapt personalization systems while avoiding large composed-target curation pipelines. Readers seeking a beginner-level introduction to personalization would need to consult the cited background work first.

Authors’ abstract

Subject-driven image personalization---generating new images that preserve the identity of one or several reference subjects in novel scenes---is a foundational capability for modern visual content creation. It is currently dominated by generalized methods that fine-tune a pretrained multimodal diffusion transformer (MMDiT) on hundreds of thousands to millions of paired \emph{(reference, composed-target)} examples, where each composed target is a synthesized image of the subject in a novel scene. Producing such targets demands a costly multi-stage curation pipeline---LLM-based prompt generation, T2I-based composed-target synthesis, reference-subject extraction, VLM-based quality filtering, and correspondence labeling---and tightly couples each method to a particular target synthesizer and curation choice. We introduce \emph{CRAFT} (Constrained Reward via Attention Fine-Tuning), a single-step ReFL framework that fine-tunes a pre-trained \emph{reference-aware} MMDiT via LoRA adapters using a compact reference-only data construction---$10$K reference images and subject masks, with no composed-target supervision. CRAFT realizes a \emph{Where to look} principle: attention-level rewards align noise- and phrase-token attention with the correct reference subject, and the resulting per-subject attention masks gate a pixel-level identity reward to keep image-space supervision consistent with the learned attention routing. Applied to FLUX.2-klein-9B, CRAFT achieves state-of-the-art performance on XVerseBench \rev{while using no composed-target supervision---only $10$K reference-only samples, whereas prior generalized methods require $150$K to over $2$M composed-target pairs}. The same recipe transfers to other reference-aware backbones, consistently improving performance. Project page: https://jihun999.github.io/projects/CRAFT/.