Add ptpl-xattn-pma-p05 (PTPL-XAttn-PMA-B (p=0.5))

ptpl-xattn-pma-p05/README.md

@@ -0,0 +1,20 @@
+---
+license: apache-2.0
+tags:
+  - histopathology
+  - diffusion
+  - spatial-transcriptomics
+  - icml-2026-sd4h-workshop
+---
+
+# PTPL-XAttn-PMA-B (p=0.5)
+
+EMA-only inference weights for the **PTPL-XAttn-PMA-B (p=0.5)** row reported in the
+ICML 2026 SD4H workshop submission *Transcriptomics-Conditioned Virtual Tissue
+Synthesis via Diffusion Transformers*.
+
+- **Source checkpoint**: `step_2617000_ema.pt`
+- **Architecture**: see `training_config.yaml` in this folder.
+- **License**: Apache-2.0.
+
+See the umbrella repo README at `stmdit-anon/stmdit-checkpoints` for usage.

ptpl-xattn-pma-p05/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0505ed6ab7ffaa4c5dded254187e17b54bec2a98f5ef9f8434450d7d35356df
+size 630725612

ptpl-xattn-pma-p05/training_config.yaml

@@ -0,0 +1,119 @@
+# Training Configuration - PTPL-XAttn-PMA-B (Pooling by Multihead Attention)
+# ===========================================================================
+# Pooling by Multihead Attention: 32 learned seeds attend over CF tokens.
+# Trained on PTPL features (DeepSpot predicted GE, corrected normalization).
+# Uses best dropout from ablation study: p=0.5 (weights 25/25/25/25).
+#
+# Usage:
+#   run-training configs/training.yaml --lightning
+
+output_dir: "/cluster/work/grlab/projects/projects2025-virtual-tissue-gen/scratch/10x_TuPro/PixCell-PTPL-XAttn/training/ptpl-xattn-pma-B"
+device: "cuda"
+
+# ============================================================================
+# MODEL
+# ============================================================================
+
+model:
+  type: "pixart_ge_xattn"
+  variant: "B"  # 130M params: depth=12, hidden=768, heads=12
+  ge_encoder_type: "cancerfoundation"
+  ge_hidden_dim: 512
+  cf_model_dir: "/cluster/home/pvlachas/leomed-home/pretrained_model_weights/cancer-foundation"
+  cf_freeze_backbone: true
+  ge_token_source: "pma"
+  ge_num_tokens: 32
+  ge_xattn_fusion: "xattn_only"
+
+# ============================================================================
+# DATA — PTPL features (DeepSpot predicted GE, same VAE/UNI as Visium)
+# ============================================================================
+
+data:
+  features_dir: "/cluster/work/grlab/projects/projects2025-virtual-tissue-gen/scratch/10x_TuPro-PTPL/feat-extraction/features_train"
+  load_gene_expression: true
+  load_gene_expression_binned: true
+  num_workers: 8
+  pin_memory: true
+  val_split: 0.1
+
+# ============================================================================
+# DIFFUSION
+# ============================================================================
+
+diffusion:
+  timesteps: 1000
+  beta_schedule: "linear"
+  image_size: 256
+  latent_size: 32
+
+# ============================================================================
+# TRAINING
+# ============================================================================
+
+training:
+  batch_size: 32
+  batch_size_val: 32
+  gradient_accumulation_steps: 4  # effective batch = 128
+  num_epochs: 1000
+  seed: 42
+  gradient_clip: 0.01
+  ema_rate: 0.9999
+
+  optimizer:
+    lr: 2e-5
+    weight_decay: 0.01
+    betas: [0.9, 0.999]
+
+  scheduler:
+    warmup_steps: 1000
+    min_lr_ratio: 0.1
+
+  # Best dropout from ablation: p=0.5 (25/25/25/25)
+  classifier_free_guidance:
+    conditioning_schedule:
+      - mask: [uni, ge]   # full conditioning (UNI + GE active)
+        weight: 25
+      - mask: [ge]        # GE only (UNI dropped)
+        weight: 25
+      - mask: [uni]       # UNI only (GE dropped)
+        weight: 25
+      - mask: []           # unconditional (both dropped)
+        weight: 25
+
+  convergence:
+    monitor_timestep_range: [900, 1000]
+    patience: 10
+    min_epochs: 50
+
+# ============================================================================
+# DISTRIBUTED
+# ============================================================================
+
+distributed:
+  precision: "bf16-mixed"
+  compile_model: true
+
+# ============================================================================
+# CHECKPOINT
+# ============================================================================
+
+checkpoint:
+  save_every: 1000
+  resume: null
+
+# ============================================================================
+# LOGGING
+# ============================================================================
+
+logging:
+  log_every: 100
+  validate_every: 0
+  gpu_monitor: true
+  gpu_monitor_interval: 60.0
+  sample_every_epochs: 10
+  sample_every_steps: 0
+  num_samples: 16
+  sample_guidance_scale: 3.0
+  sample_num_steps: 20
+  sample_vae_path: "/cluster/home/pvlachas/leomed-home/pretrained_model_weights/stability-ai-stable-diffusion-3-5-large/models--stabilityai--stable-diffusion-3.5-large/snapshots/ceddf0a7fdf2064ea28e2213e3b84e4afa170a0f/vae"