lrm/flux/docs/architecture.md

Flux LRM Architecture

Data Path

1. Dataset returns pairwise examples:

• caption
• image_0, image_1
• preference labels label_0, label_1
• timestep tensor for each image in pair

2. Tokenization:

• tokenizer -> input_ids for CLIP text encoder
• tokenizer_2 -> input_ids_2 for T5 text encoder

3. Image preprocessing:

• Resize/crop/normalize to model input resolution.

Model Path

1. Text branch

• CLIP text encoder produces pooled prompt embeddings.
• T5 encoder produces sequence prompt embeddings.
• Pooled CLIP output is projected to reward text embedding space.

2. Image branch

• VAE encodes images into latent tensors.
• Noise is sampled and mixed with latents using flow-style sigma schedule.
• Latents are packed into Flux token format.
• Flux transformer predicts token outputs conditioned on text embeddings.
• Token outputs are pooled and projected to reward image embedding space.

3. Reward scoring

• Normalize text/image embeddings.
• Pairwise logits are computed with learnable temperature (logit_scale).
• Criterion computes pairwise preference loss.

Training Control

• Task orchestrates train/eval loops and metric reporting.
• Accelerator manages distributed setup, mixed precision, checkpointing.
• Hydra config composes model/dataset/criterion/task/optimizer/scheduler groups.

Notes

• Flux schnell guidance is typically 0.0.
• Dual text encoders are trainable by default (unless explicitly frozen).