Revert: remove NeurIPS ablation checkpoints + configs

checkpoints/fm/arcu/fm_default.yaml

@@ -1,111 +0,0 @@
-run:
-  train_model: fm
-  name: rp_fm_arcu_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 20
-  dataset_size: 20000000
-  compile: model
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_arcu_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  val_conf:
-    val_frac: 0.01
-    seed: 0
-  base_conf:
-    kappa: 8.0
-    bs_frac: 0.0
-    base_dist: poles
-    scale_dist: sm_norm
-    sm_scale: 0.5
-    tanh_theta: true
-  model_conf:
-    manifold:
-    - name: euclidean
-      dim: 1
-    - name: sphere
-      dim: 3
-    - name: sphere
-      dim: 3
-    ot_coupling: true
-    ot_e_only: false
-    max_energy: 300.0
-    t_dist: sd3
-    t_dist_scale: 0.0
-    t_dist_shift: 1.3
-    loss_sc: 0.0
-    one_step_euler_fac: 0.0
-    curv_fac: 0.0
-    curv_every: 4
-    curv_eps: 0.05
-    curv_warmup: 500
-    reflow_kwargs:
-      method: midpoint
-      step_size: 0.5
-    reflow_start_epoch: 0
-    reflow_every: 4
-    t_zero_frac: 0.0
-    base_dist_loss: 0.0
-    base_pretrain_batches: 0
-    base_pretrain_bs: 16384
-    mask_conf:
-      p_forward: 1.0
-    cond_cube: false
-    canon_sym: false
-    model_args:
-      h_dim: 256
-      in_dim: 7
-      nlayers: 6
-      nhead: 8
-      dropout: 0.0
-      ff_mult: 4
-      nvtypes: 2
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      ff_swish: true
-      ff_glu: true
-      ff_no_bias: false
-      attn_flash: true
-      attn_qk_norm: true
-      attn_value_rmsnorm: true
-  opt_conf:
-    opt: muon
-    lr: 0.001
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: Amortized over argon + copper (rp_arcu_240726).

checkpoints/fm/arcu/rp_fm_arcu_240726.pt

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

checkpoints/fm/energy_1gev/fm_default.yaml

@@ -1,111 +0,0 @@
-run:
-  train_model: fm
-  name: rp_fm_energy_1gev_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 20
-  dataset_size: 20000000
-  compile: model
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_1gev_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  val_conf:
-    val_frac: 0.01
-    seed: 0
-  base_conf:
-    kappa: 8.0
-    bs_frac: 0.0
-    base_dist: poles
-    scale_dist: sm_norm
-    sm_scale: 0.5
-    tanh_theta: true
-  model_conf:
-    manifold:
-    - name: euclidean
-      dim: 1
-    - name: sphere
-      dim: 3
-    - name: sphere
-      dim: 3
-    ot_coupling: true
-    ot_e_only: false
-    max_energy: 1000.0
-    t_dist: sd3
-    t_dist_scale: 0.0
-    t_dist_shift: 1.3
-    loss_sc: 0.0
-    one_step_euler_fac: 0.0
-    curv_fac: 0.0
-    curv_every: 4
-    curv_eps: 0.05
-    curv_warmup: 500
-    reflow_kwargs:
-      method: midpoint
-      step_size: 0.5
-    reflow_start_epoch: 0
-    reflow_every: 4
-    t_zero_frac: 0.0
-    base_dist_loss: 0.0
-    base_pretrain_batches: 0
-    base_pretrain_bs: 16384
-    mask_conf:
-      p_forward: 1.0
-    cond_cube: false
-    canon_sym: false
-    model_args:
-      h_dim: 256
-      in_dim: 7
-      nlayers: 6
-      nhead: 8
-      dropout: 0.0
-      ff_mult: 4
-      nvtypes: 2
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      ff_swish: true
-      ff_glu: true
-      ff_no_bias: false
-      attn_flash: true
-      attn_qk_norm: true
-      attn_value_rmsnorm: true
-  opt_conf:
-    opt: muon
-    lr: 0.001
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/fm/energy_1gev/rp_fm_energy_1gev_240726.pt

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

checkpoints/fm/energy_1gev_no_ot/fm_default.yaml

@@ -1,111 +0,0 @@
-run:
-  train_model: fm
-  name: rp_fm_energy_1gev_no_ot_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 20
-  dataset_size: 20000000
-  compile: model
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_1gev_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  val_conf:
-    val_frac: 0.01
-    seed: 0
-  base_conf:
-    kappa: 8.0
-    bs_frac: 0.0
-    base_dist: poles
-    scale_dist: sm_norm
-    sm_scale: 0.5
-    tanh_theta: true
-  model_conf:
-    manifold:
-    - name: euclidean
-      dim: 1
-    - name: sphere
-      dim: 3
-    - name: sphere
-      dim: 3
-    ot_coupling: false
-    ot_e_only: false
-    max_energy: 1000.0
-    t_dist: sd3
-    t_dist_scale: 0.0
-    t_dist_shift: 1.3
-    loss_sc: 0.0
-    one_step_euler_fac: 0.0
-    curv_fac: 0.0
-    curv_every: 4
-    curv_eps: 0.05
-    curv_warmup: 500
-    reflow_kwargs:
-      method: midpoint
-      step_size: 0.5
-    reflow_start_epoch: 0
-    reflow_every: 4
-    t_zero_frac: 0.0
-    base_dist_loss: 0.0
-    base_pretrain_batches: 0
-    base_pretrain_bs: 16384
-    mask_conf:
-      p_forward: 1.0
-    cond_cube: false
-    canon_sym: false
-    model_args:
-      h_dim: 256
-      in_dim: 7
-      nlayers: 6
-      nhead: 8
-      dropout: 0.0
-      ff_mult: 4
-      nvtypes: 2
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      ff_swish: true
-      ff_glu: true
-      ff_no_bias: false
-      attn_flash: true
-      attn_qk_norm: true
-      attn_value_rmsnorm: true
-  opt_conf:
-    opt: muon
-    lr: 0.001
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: 'OT ablation at 1 GeV: ot_coupling=false vs energy_1gev (ot_coupling=true).'

checkpoints/fm/energy_1gev_no_ot/rp_fm_energy_1gev_no_ot_240726.pt

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

checkpoints/fm/learned_base_params/fm_default.yaml

@@ -1,111 +0,0 @@
-run:
-  train_model: fm
-  name: rp_fm_learned_base_params_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 20
-  dataset_size: 20000000
-  compile: model
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_baseline_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  val_conf:
-    val_frac: 0.01
-    seed: 0
-  base_conf:
-    kappa: 8.0
-    bs_frac: 0.0
-    base_dist: poles
-    scale_dist: sm_norm
-    sm_scale: 0.5
-    tanh_theta: true
-  model_conf:
-    manifold:
-    - name: euclidean
-      dim: 1
-    - name: sphere
-      dim: 3
-    - name: sphere
-      dim: 3
-    ot_coupling: true
-    ot_e_only: false
-    max_energy: 300.0
-    t_dist: sd3
-    t_dist_scale: 0.0
-    t_dist_shift: 1.3
-    loss_sc: 0.0
-    one_step_euler_fac: 0.0
-    curv_fac: 0.0
-    curv_every: 4
-    curv_eps: 0.05
-    curv_warmup: 500
-    reflow_kwargs:
-      method: midpoint
-      step_size: 0.5
-    reflow_start_epoch: 0
-    reflow_every: 4
-    t_zero_frac: 0.0
-    base_dist_loss: 0.0
-    base_pretrain_batches: 1000
-    base_pretrain_bs: 16384
-    mask_conf:
-      p_forward: 1.0
-    cond_cube: false
-    canon_sym: false
-    model_args:
-      h_dim: 256
-      in_dim: 7
-      nlayers: 6
-      nhead: 8
-      dropout: 0.0
-      ff_mult: 4
-      nvtypes: 2
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      ff_swish: true
-      ff_glu: true
-      ff_no_bias: false
-      attn_flash: true
-      attn_qk_norm: true
-      attn_value_rmsnorm: true
-  opt_conf:
-    opt: muon
-    lr: 0.001
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/fm/learned_base_params/rp_fm_learned_base_params_240726.pt

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

checkpoints/fm/naive_base/fm_default.yaml

@@ -1,111 +0,0 @@
-run:
-  train_model: fm
-  name: rp_fm_naive_base_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 20
-  dataset_size: 20000000
-  compile: model
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_baseline_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  val_conf:
-    val_frac: 0.01
-    seed: 0
-  base_conf:
-    kappa: 8.0
-    bs_frac: 0.0
-    base_dist: iso
-    scale_dist: logit_norm
-    sm_scale: 0.5
-    tanh_theta: true
-  model_conf:
-    manifold:
-    - name: euclidean
-      dim: 1
-    - name: sphere
-      dim: 3
-    - name: sphere
-      dim: 3
-    ot_coupling: true
-    ot_e_only: false
-    max_energy: 300.0
-    t_dist: sd3
-    t_dist_scale: 0.0
-    t_dist_shift: 1.3
-    loss_sc: 0.0
-    one_step_euler_fac: 0.0
-    curv_fac: 0.0
-    curv_every: 4
-    curv_eps: 0.05
-    curv_warmup: 500
-    reflow_kwargs:
-      method: midpoint
-      step_size: 0.5
-    reflow_start_epoch: 0
-    reflow_every: 4
-    t_zero_frac: 0.0
-    base_dist_loss: 0.0
-    base_pretrain_batches: 0
-    base_pretrain_bs: 16384
-    mask_conf:
-      p_forward: 1.0
-    cond_cube: false
-    canon_sym: false
-    model_args:
-      h_dim: 256
-      in_dim: 7
-      nlayers: 6
-      nhead: 8
-      dropout: 0.0
-      ff_mult: 4
-      nvtypes: 2
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      ff_swish: true
-      ff_glu: true
-      ff_no_bias: false
-      attn_flash: true
-      attn_qk_norm: true
-      attn_value_rmsnorm: true
-  opt_conf:
-    opt: muon
-    lr: 0.001
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/fm/naive_base/rp_fm_naive_base_240726.pt

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

checkpoints/fm/no_ot/fm_default.yaml

@@ -1,111 +0,0 @@
-run:
-  train_model: fm
-  name: rp_fm_no_ot_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 20
-  dataset_size: 20000000
-  compile: model
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_baseline_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  val_conf:
-    val_frac: 0.01
-    seed: 0
-  base_conf:
-    kappa: 8.0
-    bs_frac: 0.0
-    base_dist: poles
-    scale_dist: sm_norm
-    sm_scale: 0.5
-    tanh_theta: true
-  model_conf:
-    manifold:
-    - name: euclidean
-      dim: 1
-    - name: sphere
-      dim: 3
-    - name: sphere
-      dim: 3
-    ot_coupling: false
-    ot_e_only: false
-    max_energy: 300.0
-    t_dist: sd3
-    t_dist_scale: 0.0
-    t_dist_shift: 1.3
-    loss_sc: 0.0
-    one_step_euler_fac: 0.0
-    curv_fac: 0.0
-    curv_every: 4
-    curv_eps: 0.05
-    curv_warmup: 500
-    reflow_kwargs:
-      method: midpoint
-      step_size: 0.5
-    reflow_start_epoch: 0
-    reflow_every: 4
-    t_zero_frac: 0.0
-    base_dist_loss: 0.0
-    base_pretrain_batches: 0
-    base_pretrain_bs: 16384
-    mask_conf:
-      p_forward: 1.0
-    cond_cube: false
-    canon_sym: false
-    model_args:
-      h_dim: 256
-      in_dim: 7
-      nlayers: 6
-      nhead: 8
-      dropout: 0.0
-      ff_mult: 4
-      nvtypes: 2
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      ff_swish: true
-      ff_glu: true
-      ff_no_bias: false
-      attn_flash: true
-      attn_qk_norm: true
-      attn_value_rmsnorm: true
-  opt_conf:
-    opt: muon
-    lr: 0.001
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/fm/no_ot/rp_fm_no_ot_240726.pt

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

checkpoints/fm/no_raytrace/fm_default.yaml

@@ -1,111 +0,0 @@
-run:
-  train_model: fm
-  name: rp_fm_no_raytrace_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 20
-  dataset_size: 20000000
-  compile: model
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_non_raytraced_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  val_conf:
-    val_frac: 0.01
-    seed: 0
-  base_conf:
-    kappa: 8.0
-    bs_frac: 0.0
-    base_dist: iso_pos
-    scale_dist: sm_norm
-    sm_scale: 0.5
-    tanh_theta: true
-  model_conf:
-    manifold:
-    - name: euclidean
-      dim: 1
-    - name: sphere
-      dim: 3
-    - name: sphere
-      dim: 3
-    ot_coupling: true
-    ot_e_only: false
-    max_energy: 300.0
-    t_dist: sd3
-    t_dist_scale: 0.0
-    t_dist_shift: 1.3
-    loss_sc: 0.0
-    one_step_euler_fac: 0.0
-    curv_fac: 0.0
-    curv_every: 4
-    curv_eps: 0.05
-    curv_warmup: 500
-    reflow_kwargs:
-      method: midpoint
-      step_size: 0.5
-    reflow_start_epoch: 0
-    reflow_every: 4
-    t_zero_frac: 0.0
-    base_dist_loss: 0.0
-    base_pretrain_batches: 0
-    base_pretrain_bs: 16384
-    mask_conf:
-      p_forward: 1.0
-    cond_cube: false
-    canon_sym: false
-    model_args:
-      h_dim: 256
-      in_dim: 7
-      nlayers: 6
-      nhead: 8
-      dropout: 0.0
-      ff_mult: 4
-      nvtypes: 2
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      ff_swish: true
-      ff_glu: true
-      ff_no_bias: false
-      attn_flash: true
-      attn_qk_norm: true
-      attn_value_rmsnorm: true
-  opt_conf:
-    opt: muon
-    lr: 0.001
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/fm/no_raytrace/rp_fm_no_raytrace_240726.pt

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

checkpoints/fm/path_curvature/fm_default.yaml

@@ -1,111 +0,0 @@
-run:
-  train_model: fm
-  name: rp_fm_path_curvature_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 20
-  dataset_size: 20000000
-  compile: model
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_baseline_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  val_conf:
-    val_frac: 0.01
-    seed: 0
-  base_conf:
-    kappa: 8.0
-    bs_frac: 0.0
-    base_dist: poles
-    scale_dist: sm_norm
-    sm_scale: 0.5
-    tanh_theta: true
-  model_conf:
-    manifold:
-    - name: euclidean
-      dim: 1
-    - name: sphere
-      dim: 3
-    - name: sphere
-      dim: 3
-    ot_coupling: true
-    ot_e_only: false
-    max_energy: 300.0
-    t_dist: sd3
-    t_dist_scale: 0.0
-    t_dist_shift: 1.3
-    loss_sc: 0.0
-    one_step_euler_fac: 0.0
-    curv_fac: 0.5
-    curv_every: 4
-    curv_eps: 0.05
-    curv_warmup: 500
-    reflow_kwargs:
-      method: midpoint
-      step_size: 0.5
-    reflow_start_epoch: 0
-    reflow_every: 4
-    t_zero_frac: 0.0
-    base_dist_loss: 0.0
-    base_pretrain_batches: 0
-    base_pretrain_bs: 16384
-    mask_conf:
-      p_forward: 1.0
-    cond_cube: false
-    canon_sym: false
-    model_args:
-      h_dim: 256
-      in_dim: 7
-      nlayers: 6
-      nhead: 8
-      dropout: 0.0
-      ff_mult: 4
-      nvtypes: 2
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      ff_swish: true
-      ff_glu: true
-      ff_no_bias: false
-      attn_flash: true
-      attn_qk_norm: true
-      attn_value_rmsnorm: true
-  opt_conf:
-    opt: muon
-    lr: 0.001
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/fm/path_curvature/rp_fm_path_curvature_240726.pt

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

checkpoints/fm/path_reflow/fm_default.yaml

@@ -1,111 +0,0 @@
-run:
-  train_model: fm
-  name: rp_fm_path_reflow_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 20
-  dataset_size: 20000000
-  compile: model
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_baseline_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  val_conf:
-    val_frac: 0.01
-    seed: 0
-  base_conf:
-    kappa: 8.0
-    bs_frac: 0.0
-    base_dist: poles
-    scale_dist: sm_norm
-    sm_scale: 0.5
-    tanh_theta: true
-  model_conf:
-    manifold:
-    - name: euclidean
-      dim: 1
-    - name: sphere
-      dim: 3
-    - name: sphere
-      dim: 3
-    ot_coupling: true
-    ot_e_only: false
-    max_energy: 300.0
-    t_dist: sd3
-    t_dist_scale: 0.0
-    t_dist_shift: 1.3
-    loss_sc: 0.0
-    one_step_euler_fac: 0.0
-    curv_fac: 0.0
-    curv_every: 4
-    curv_eps: 0.05
-    curv_warmup: 500
-    reflow_kwargs:
-      method: midpoint
-      step_size: 0.5
-    reflow_start_epoch: 4
-    reflow_every: 4
-    t_zero_frac: 0.0
-    base_dist_loss: 0.0
-    base_pretrain_batches: 0
-    base_pretrain_bs: 16384
-    mask_conf:
-      p_forward: 1.0
-    cond_cube: false
-    canon_sym: false
-    model_args:
-      h_dim: 256
-      in_dim: 7
-      nlayers: 6
-      nhead: 8
-      dropout: 0.0
-      ff_mult: 4
-      nvtypes: 2
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      ff_swish: true
-      ff_glu: true
-      ff_no_bias: false
-      attn_flash: true
-      attn_qk_norm: true
-      attn_value_rmsnorm: true
-  opt_conf:
-    opt: muon
-    lr: 0.001
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/fm/path_reflow/rp_fm_path_reflow_240726.pt

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

checkpoints/fm/standard_phys_base_no_opts/fm_default.yaml

@@ -1,111 +0,0 @@
-run:
-  train_model: fm
-  name: rp_fm_baseline_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 20
-  dataset_size: 20000000
-  compile: model
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_baseline_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  val_conf:
-    val_frac: 0.01
-    seed: 0
-  base_conf:
-    kappa: 8.0
-    bs_frac: 0.0
-    base_dist: poles
-    scale_dist: sm_norm
-    sm_scale: 0.5
-    tanh_theta: true
-  model_conf:
-    manifold:
-    - name: euclidean
-      dim: 1
-    - name: sphere
-      dim: 3
-    - name: sphere
-      dim: 3
-    ot_coupling: true
-    ot_e_only: false
-    max_energy: 300.0
-    t_dist: sd3
-    t_dist_scale: 0.0
-    t_dist_shift: 1.3
-    loss_sc: 0.0
-    one_step_euler_fac: 0.0
-    curv_fac: 0.0
-    curv_every: 4
-    curv_eps: 0.05
-    curv_warmup: 500
-    reflow_kwargs:
-      method: midpoint
-      step_size: 0.5
-    reflow_start_epoch: 0
-    reflow_every: 4
-    t_zero_frac: 0.0
-    base_dist_loss: 0.0
-    base_pretrain_batches: 0
-    base_pretrain_bs: 16384
-    mask_conf:
-      p_forward: 1.0
-    cond_cube: false
-    canon_sym: false
-    model_args:
-      h_dim: 256
-      in_dim: 7
-      nlayers: 6
-      nhead: 8
-      dropout: 0.0
-      ff_mult: 4
-      nvtypes: 2
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      ff_swish: true
-      ff_glu: true
-      ff_no_bias: false
-      attn_flash: true
-      attn_qk_norm: true
-      attn_value_rmsnorm: true
-  opt_conf:
-    opt: muon
-    lr: 0.001
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/fm/standard_phys_base_no_opts/rp_fm_baseline_240726.pt

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

checkpoints/mult/arcu/mult_default.yaml

@@ -1,86 +0,0 @@
-run:
-  train_model: mult
-  name: rp_mult_arcu_240726
-  devices:
-  - 0
-  - 1
-  - 2
-  - 3
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 10
-  dataset_size: 20000000
-  compile: false
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_arcu_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  mm_conf:
-    in_dim: 11
-    cond_scalars:
-    - Density
-    - Z
-    - A
-    - Size
-    dropout: 0.05
-    h_dim: 128
-    n_layers: 6
-    n_heads: 6
-    bs: 4096
-    lr: 0.001
-    weight_decay: 0.01
-    warmup_steps: 0
-    post_emb_norm: false
-    use_abs_pos_emb: true
-    train_inverse: false
-    model_args:
-      ff_swish: true
-      ff_glu: true
-      attn_qk_norm: true
-      rotary_pos_emb: true
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      residual_attn: true
-    inv_model_args:
-      ff_swish: true
-      ff_glu: true
-      attn_qk_norm: true
-      rotary_pos_emb: true
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      residual_attn: true
-  opt_conf:
-    opt: muon
-    lr: 0.01
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: Amortized over argon + copper (rp_arcu_240726).

checkpoints/mult/arcu/rp_mult_arcu_240726.pt

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

checkpoints/mult/baseline/mult_default.yaml

@@ -1,86 +0,0 @@
-run:
-  train_model: mult
-  name: rp_mult_baseline_240726
-  devices:
-  - 0
-  - 1
-  - 2
-  - 3
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 10
-  dataset_size: 20000000
-  compile: false
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_baseline_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  mm_conf:
-    in_dim: 11
-    cond_scalars:
-    - Density
-    - Z
-    - A
-    - Size
-    dropout: 0.05
-    h_dim: 128
-    n_layers: 6
-    n_heads: 6
-    bs: 4096
-    lr: 0.001
-    weight_decay: 0.01
-    warmup_steps: 0
-    post_emb_norm: false
-    use_abs_pos_emb: true
-    train_inverse: false
-    model_args:
-      ff_swish: true
-      ff_glu: true
-      attn_qk_norm: true
-      rotary_pos_emb: true
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      residual_attn: true
-    inv_model_args:
-      ff_swish: true
-      ff_glu: true
-      attn_qk_norm: true
-      rotary_pos_emb: true
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      residual_attn: true
-  opt_conf:
-    opt: muon
-    lr: 0.01
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/mult/baseline/rp_mult_baseline_240726.pt

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

checkpoints/mult/energy_1gev/mult_default.yaml

@@ -1,86 +0,0 @@
-run:
-  train_model: mult
-  name: rp_mult_energy_1gev_240726
-  devices:
-  - 4
-  - 5
-  - 6
-  - 7
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 10
-  dataset_size: 20000000
-  compile: false
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_1gev_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  mm_conf:
-    in_dim: 11
-    cond_scalars:
-    - Density
-    - Z
-    - A
-    - Size
-    dropout: 0.05
-    h_dim: 128
-    n_layers: 6
-    n_heads: 6
-    bs: 4096
-    lr: 0.001
-    weight_decay: 0.01
-    warmup_steps: 0
-    post_emb_norm: false
-    use_abs_pos_emb: true
-    train_inverse: false
-    model_args:
-      ff_swish: true
-      ff_glu: true
-      attn_qk_norm: true
-      rotary_pos_emb: true
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      residual_attn: true
-    inv_model_args:
-      ff_swish: true
-      ff_glu: true
-      attn_qk_norm: true
-      rotary_pos_emb: true
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      residual_attn: true
-  opt_conf:
-    opt: muon
-    lr: 0.01
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/mult/energy_1gev/rp_mult_energy_1gev_240726.pt

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

checkpoints/mult/no_raytrace/mult_default.yaml

@@ -1,86 +0,0 @@
-run:
-  train_model: mult
-  name: rp_mult_no_ray_240726
-  devices:
-  - 0
-  - 1
-  - 2
-  - 3
-  precision: 32
-  matmul_precision: medium
-  dtype: float32
-  epochs: 10
-  dataset_size: 20000000
-  compile: false
-  gradient_clip_val: 1.0
-  val_check_interval: 0.2
-  limit_val_batches: 4
-  strategy: ddp
-  ckpt_dir: null
-logging:
-  comet: true
-  project: lego_pdgid
-config:
-  dl_conf:
-    lds_args:
-      data: neurips_2026/rp_non_raytraced_240726
-      frac: 1.0
-      cutoff_mev: 10
-      min_particles: 0
-    bs: 4096
-    num_workers: 16
-  mm_conf:
-    in_dim: 11
-    cond_scalars:
-    - Density
-    - Z
-    - A
-    - Size
-    dropout: 0.05
-    h_dim: 128
-    n_layers: 6
-    n_heads: 6
-    bs: 4096
-    lr: 0.001
-    weight_decay: 0.01
-    warmup_steps: 0
-    post_emb_norm: false
-    use_abs_pos_emb: true
-    train_inverse: false
-    model_args:
-      ff_swish: true
-      ff_glu: true
-      attn_qk_norm: true
-      rotary_pos_emb: true
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      residual_attn: true
-    inv_model_args:
-      ff_swish: true
-      ff_glu: true
-      attn_qk_norm: true
-      rotary_pos_emb: true
-      use_adaptive_rmsnorm: true
-      use_adaptive_layerscale: true
-      residual_attn: true
-  opt_conf:
-    opt: muon
-    lr: 0.01
-    momentum: 0.95
-    nesterov: true
-    ns_steps: 5
-    weight_decay: 0.007
-    weight_decouple: true
-    adamw_lr: 0.003
-    adamw_betas:
-    - 0.9
-    - 0.999
-    adamw_wd: 0.01
-    adamw_eps: 1.0e-08
-    scheduler:
-      cls: warmup_cosine
-      warmup_frac: 0.1
-      eta_min: 1.0e-06
-      interval: step
-  additional:
-    notes: ''

checkpoints/mult/no_raytrace/rp_mult_no_ray_240726.pt

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