xenium-senescence-checkpoints

Xenium Senescence Checkpoints

UNet (42M params) trained with OT-CFM for cell trajectory generation.

Model Architecture

• UNet_CCDM_LatentCond: dim=64, dim_mults=[1,2,2,4,8], channels=1
• Input: 64×64 grayscale cell image
• Conditioning: senescence score (sinusoidal) + VAE latent (128-dim) + cell type embedding
• Output: velocity field v_pred (1, 64, 64)

Checkpoints

v8_{tissue}/model-50000.pt    # V8 baseline (50K steps, constant LR 5e-5)
c25_20k/model-{step}.pt       # C25 (cosine LR 2e-4, contrastive, 20K steps)

Checkpoint Format

ckpt = torch.load("model-50000.pt")
# Keys: 'model' (state_dict), 'ema' (EMA state_dict), 'step' (int)

Loading

from unet_ccdm_latent_cond import UNet_CCDM_LatentCond

unet = UNet_CCDM_LatentCond(
    dim=64, dim_mults=[1,2,2,4,8], cond_emb_dim_init=128,
    channels=1, num_cell_types=10, dropout=0.0
)
ckpt = torch.load("model-50000.pt", map_location="cpu")
unet.load_state_dict(ckpt['model'])

Training Config (C25 — best FID)

accelerate launch --multi_gpu --num_processes 8 train_flow_matching.py \
  --image_start --batch_size 512 --train_steps 20000 \
  --lr 2e-4 --lr_schedule cosine --warmup_steps 2000 --lr_min 1e-6 \
  --weight_decay 0.01 --contrastive_weight 0.1 \
  --latent_drop_prob 0.5 --noise_prob 0.2 --noise_level 0.2 \
  --num_cell_types 10 --auto_resume

Results (Lung)