Poor performance when testing the pretrained knee MRI checkpoint on knee test dataset

[Graeme-Lee]

Hi, thanks for releasing this great benchmark for inverse problems.

I tried to use the provided pretrained knee checkpoint to test on the knee MRI test dataset, but the results I obtained are much poorer than those reported in the benchmark table. I am not sure whether I made a mistake in the evaluation setup, so I would really appreciate your help.

Below is the sampling config I used:

name: multi-coil-mri
prior: exps/pretrain/MRI-knee.pt
model:
  _target_: inverse_problems.multi_coil_mri.MultiCoilMRI
  sigma_noise: 0.0
  total_lines: 320
  acceleration_ratio: 4
  pattern: random
  mask_seed: 0

data:
  _target_: training.dataset.MultiCoilMRILMDBData
  root: data/knee_test_lmdb
  image_size: [320, 320]
  id_list: null
  simulated_kspace: false

evaluator:
  _target_: eval.MRI

exp_dir: exps/inference/multi-coil-mri-knee

The overall metrics I obtained are：

Final metric results: {'psnr': 24.378959970290268, 'psnr_std': 1.7808092531148152, 'ssim': 0.41473705892660667, 'ssim_std': 0.07551422425698978, 'data misfit': 32.24088266880017, 'data misfit_std': 15.172199847063636}

These numbers seem much worse than the reported results, so I am wondering whether I may have missed something in the testing procedure.

Could you please help clarify:

whether this config is correct for reproducing the knee MRI results,
whether the pretrained checkpoint matches this exact setting,
whether simulated_kspace: false is the correct option here, and
whether there are any additional preprocessing or evaluation details required?

Thank you very much.

[devzhk]

Hi,

Could you also share which sampling algorithm you used, and what hyperparameters were used for that algorithm?

One thing to note is that the effective hyperparameter ranges for the same algorithm can vary significantly across different inverse problems. We mention this in the README as well:

we recommend starting with the optimized hyperparameters provided in Table 12 of our paper as a reference point.

[Graeme-Lee]

Hello, I use DAPS, and the default daps.yaml file:

name: DAPS
method:
  _target_: algo.daps.DAPS
  annealing_scheduler_config:
    num_steps: 200
    sigma_max: 100
    sigma_min: 0.1
    sigma_final: 0
    schedule: 'linear'
    timestep: 'poly-7'

  diffusion_scheduler_config:
    num_steps: 5
    sigma_min: 0.01
    sigma_final: 0
    schedule: 'linear'
    timestep: 'poly-7'

  lgvd_config:
    num_steps: 100
    lr: 1e-4
    tau: 0.01
    lr_min_ratio: 0.01

[devzhk]

Hi,

Can you try the following hyperparameters?

name: DAPS
method:
  _target_: algo.daps.DAPS
  annealing_scheduler_config:
    num_steps: 200
    sigma_max: 100
    sigma_min: 0.1
    sigma_final: 0
    schedule: 'linear'
    timestep: 'poly-7'

  diffusion_scheduler_config:
    num_steps: 5
    sigma_min: 0.01
    sigma_final: 0
    schedule: 'linear'
    timestep: 'poly-7'

  lgvd_config:
    num_steps: 100
    lr: 1.52e-5
    tau: 4.77e-3
    lr_min_ratio: 1.0

See table 12 in the paper for the best hyperparameters for each problem. The effective hyperparameter ranges vary a lot across different problems.

[Graeme-Lee]

Hello, thanks for your previous suggestion.

However, the result is still worse even when I use your recommended hyperparameters:

Metric results: {'psnr': 21.624613572121053, 'ssim': 0.266681276002673, 'data misfit': 42.707229021069686}

For my test, I used:

• multi-coil-mri-knee-lmdb.yaml as the problem config
• mri-knee-mvue.yaml as the pretrain config

I kept them as default without additional changes.

Could you please clarify the exact settings used in your experiments, especially:

• the problem config,
• the pretrain config,
• whether any fields need to be modified from the default values, and
• whether there are any additional preprocessing or evaluation details needed to reproduce the reported results?

Thanks a lot for your help.

[devzhk]

Thanks for reporting this and for sharing the configuration. We’ll take a closer look and get back to you early next week.

[zihuiwu]

Hi @Graeme-Lee, thank you for your interest in our work!

Here is the command I used for running DAPS on the MRI problem:

python main.py problem=multi-coil-mri-knee pretrain=mri-knee-mvue algorithm=daps algorithm.method.lgvd_config.lr=1.52e-5 algorithm.method.lgvd_config.tau=4.77e-3

To align our results, could you try to run the validation data (knee_val_lmdb on this page) with the following command?

python main.py problem=multi-coil-mri-knee-val-lmdb pretrain=mri-knee-mvue algorithm=daps algorithm.method.lgvd_config.lr=1.52e-5 algorithm.method.lgvd_config.tau=4.77e-3

where multi-coil-mri-knee-val-lmdb.yaml should exist in configs/problem and look like

name: multi-coil-mri
prior: /path/to/MRI-knee.pt                                                     <--- update this path
model:
  _target_: inverse_problems.multi_coil_mri.MultiCoilMRI
  sigma_noise: 0.0
  total_lines: 320
  acceleration_ratio: 4
  pattern: random
  mask_seed: 0

data:
  _target_: training.dataset.MultiCoilMRILMDBData
  root: /path/to/knee_val_lmdb/                                                 <--- update this path
  image_size: [320, 320]
  id_list: null
  simulated_kspace: false

evaluator:
  _target_: eval.MRI

exp_dir: exps/inference/multi-coil-mri-knee

I just tried to run this and got the following results:

[2026-05-11 19:48:01,201][utils.helper][INFO] - Final metric results: {'psnr': 29.849226174674616, 'psnr_std': 1.1903719405732942, 'ssim': 0.7442389600442286, 'ssim_std': 0.08567056824187594, 'data_fit': 31.148046281627803, 'data_fit_std': 24.5761403334592}...

The full resolved config generated in the experiment directory should look like

tf32: true
inference: true
num_samples: 1
compile: false
seed: 0
wandb: false
exp_name: default
pretrain:
  data:
    _target_: training.dataset.MultiCoilMRILMDBData
    root: /tmp/zwu2/datasets/fastmri_multicoil_knee_train_lmdb
    image_size:
    - 320
    - 320
    mvue_only: true
  model:
    _target_: models.precond.EDMPrecond
    model_type: DhariwalUNet
    img_resolution: 320
    img_channels: 2
    label_dim: 0
    model_channels: 128
    channel_mult:
    - 1
    - 1
    - 1
    - 2
    - 2
    attn_resolutions:
    - 16
    num_blocks: 1
    dropout: 0.0
  loss:
    _target_: training.loss.EDMLoss
  train:
    lr: 0.0001
    batch_size: 64
    num_steps: 100000
    warmup_steps: 500
    ema_decay: 0.9999
    grad_clip: 0.0
    ema_halflife_nimg: 204800
    ema_rampup_ratio: 0.05
    tf32: true
    resume: None
    num_workers: 1
  log:
    exp_dir: exps/pretrain
    log_every: 10
    save_every: 5000
    exp_name: EDM-MRI-KNEE-MVUE
    wandb: false
    project: EDM-MRI-KNEE-MVUE
    group: training
problem:
  name: multi-coil-mri
  prior: /path/to/MRI-knee.pt                                                     <--- your model path
  model:
    _target_: inverse_problems.multi_coil_mri.MultiCoilMRI
    sigma_noise: 0.0
    total_lines: 320
    acceleration_ratio: 4
    pattern: random
    mask_seed: 0
  data:
    _target_: training.dataset.MultiCoilMRILMDBData
    root: /path/to/knee_val_lmdb/                                                 <--- your data path
    image_size:
    - 320
    - 320
    id_list: null
    simulated_kspace: false
  evaluator:
    _target_: eval.MRI
  exp_dir: exps/inference/multi-coil-mri-knee
algorithm:
  name: DAPS
  method:
    _target_: algo.daps.DAPS
    annealing_scheduler_config:
      num_steps: 200
      sigma_max: 100
      sigma_min: 0.1
      sigma_final: 0
      schedule: linear
      timestep: poly-7
    diffusion_scheduler_config:
      num_steps: 5
      sigma_min: 0.01
      sigma_final: 0
      schedule: linear
      timestep: poly-7
    lgvd_config:
      num_steps: 100
      lr: 1.52e-05
      tau: 0.00477
      lr_min_ratio: 0.01

Hope this helps and please let me know if you run into any problem!

[Graeme-Lee]

Hello, while reproducing your results, I found a mismatch in the parameters. In the experimental setup, lgvd_config.lr_min_ratio is 0.01, but Table 12 reports a value of 1.0. I would appreciate it if you could confirm the intended setting.