Skip to content

UNet COT-FM experiments: EMD/sinkhorn ablation + architecture search #18

Draft
Justinezgh wants to merge 19 commits intou/Justinezgh/pqmassfrom
u/Justinezgh/unet_exp
Draft

UNet COT-FM experiments: EMD/sinkhorn ablation + architecture search #18
Justinezgh wants to merge 19 commits intou/Justinezgh/pqmassfrom
u/Justinezgh/unet_exp

Conversation

@Justinezgh
Copy link
Copy Markdown
Collaborator

@Justinezgh Justinezgh commented Mar 17, 2026

This PR some tasks of #14. It adds infrastructure and results for the UNet hyperparameter search:

  • cot_fm.py: add --ot_method (emd/sinkhorn_log), --exp_config YAML loader, and post-training PQMass evaluation on last and best checkpoints
  • scripts/ot_diagnostic.py: compute OT pairing quality (entropy, $E[|\Delta\theta|]$, displacement) without training to narrow eps/ot_reg candidates
  • scripts/ps_unet_analysis.py: compute power spectrum and PDF of UNet-generated maps (via ODE solve) vs N-body reference
  • configs: phase 1 (EMD, sinkhorn ot_reg=1e-3) and phase 2 (small/large arch) experiment YAMLs
  • docs: experiment tracking doc with phase rationale, results, and figures

Phase 0+1 conclusion: EMD (χ²_best=113.6) outperforms sinkhorn; phase 2 architecture ablation (small/large) in progress.

@Justinezgh Justinezgh changed the base branch from main to u/Justinezgh/pqmass March 17, 2026 03:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant