lrm/lrm_sana/docs/architecture.md

SANA LRM Architecture

System Overview

This implementation reuses the modular trainer architecture from flux:

• scripts/train.py orchestrates load/prepare/train/eval/save.
• accelerators/* manage distributed runtime, logging, and checkpoints.
• datasets/* build pairwise preference batches with timestep controls.
• models/* provide latent-space text-image feature extraction.
• criterions/* compute pairwise preference loss.
• tasks/* compute inference probabilities and accuracy metrics.

Data Flow

1. Dataset loads pairwise examples from pickapic-anonymous/pickapic_v1.
2. Each sample yields prompt tokens, image pair tensors, labels, and timestep pair.
3. Criterion concatenates image pairs and passes text/image/timestep to the model.
4. Model returns normalized text/image embeddings in shared projection space.
5. Criterion computes pairwise logits and weighted preference loss.
6. Task evaluation converts scores to probabilities and computes accuracy.

SANA Model Path

1. Text branch:

• Load SANA tokenizer/text encoder(s) by checkpoint profile.
• Build prompt conditioning embeddings for transformer cross-attention.
• Build pooled text representation for reward embedding projection.

2. Image branch:

• Preprocess RGB image to configured resolution.
• Encode image with SANA VAE to latents.
• Apply scheduler-based noise at selected timesteps.
• Run SANA transformer on noisy latents with text conditioning.
• Pool transformer latent outputs to image representation.

3. Reward projection:

• Project text and image representations to common dimension.
• L2-normalize features.
• Compute scaled similarity using learnable logit_scale.

Training Controls

• Mixed precision and gradient accumulation from accelerator config.
• Distributed feature gather for pairwise loss consistency.
• Validation/test evaluation at configured intervals.
• Periodic checkpointing and best-metric tracking.

Logging and Outputs

• Output root: logs/lrm/{project_name}/{run_name}.
• Saved artifacts:
  • config snapshot
  • train log
  • periodic checkpoints
  • final model state

Variant-Aware Profile Design

A model profile selects:

• pretrained checkpoint id
• default image size (512 or 1024)
• tokenizer/text path behavior
• optional special overrides for memory or precision

This keeps one codebase while enabling all four requested SANA checkpoints.