Stein's Unbiased Risk Estimator (SURE) loss and Conjugate Gradient

[cxlcl]

Description

Based on the discussion topic here, we implemented the Conjugate-Gradient algorithm for linear operator inversion, and Stein's Unbiased Risk Estimator (SURE) [1] loss for ground-truth-date free diffusion process guidance that is proposed in [2] and illustrated in the algorithm below:

The Conjugate-Gradient (CG) algorithm is used to solve for the inversion of the linear operator in Line-4 in the algorithm above, where the linear operator is too large to store explicitly as a matrix (such as FFT/IFFT of an image) and invert directly. Instead, we can solve for the linear inversion iteratively as in CG.

The SURE loss is applied for Line-6 above. This is a differentiable loss function that can be used to train/giude an operator (e.g. neural network), where the pseudo ground truth is available but the reference ground truth is not. For example, in the MRI reconstruction, the pseudo ground truth is the zero-filled reconstruction and the reference ground truth is the fully sampled reconstruction. The reference ground truth is not available due to the lack of fully sampled.

Reference
[1] Stein, C.M.: Estimation of the mean of a multivariate normal distribution. Annals of Statistics 1981 [paper link]
[2] B. Ozturkler et al. SMRD: SURE-based Robust MRI Reconstruction with Diffusion Models. MICCAI 2023
[paper link]

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[ericspod]

@marksgraham you may want to coordinate here with other diffusion work being integrated into core.

[KumoLiu]

Sorry @ericspod, I accidentally deleted your comment, I just wanted to delete the pending one from me which duplicate with yours.

[KumoLiu]

Thanks for the PR, would be great add some user guide in the tutorial repo perhaps after all components added. Could you please help fix the CI error?

[marksgraham]

@marksgraham you may want to coordinate here with other diffusion work being integrated into core.

Thanks @ericspod. This PR looks complementary - @cxlcl let me know if you want to discuss when you get onto adding the generation step work mentioned in your original issue. MONAI Generative is currently being ported into core so we need to make sure these additions work together :)

[cxlcl]

@marksgraham you may want to coordinate here with other diffusion work being integrated into core.

Thanks @ericspod. This PR looks complementary - @cxlcl let me know if you want to discuss when you get onto adding the generation step work mentioned in your original issue. MONAI Generative is currently being ported into core so we need to make sure these additions work together :)

@marksgraham For the generation step, I was planning to use the two components in this PR (CG and SURE loss) to replace the ones in the original implementation in Reference [2]. This is done locally and I will upload it to the research repo of MONA as an example of how to use CG and SURE during the sampling stage. Will let you know once it's done. So that we may start from there.

[cxlcl]

Thanks for the PR, would be great add some user guide in the tutorial repo perhaps after all components added. Could you please help fix the CI error?

Will work on that. Thanks for the suggestion on making a suer guide. Is there a guidance on making a tutorial or exemplar tutorial I can follow?

[KumoLiu]

Is there a guidance on making a tutorial or exemplar tutorial I can follow?

Here is a contribution guide for tutorial repo.
https://github.com/Project-MONAI/tutorials/blob/main/CONTRIBUTING.md

[cxlcl]

Is there a guidance on making a tutorial or exemplar tutorial I can follow?

Here is a contribution guide for tutorial repo. https://github.com/Project-MONAI/tutorials/blob/main/CONTRIBUTING.md

One quick question related to the data used for the tutorial: can we include it in the tutorial itself, or shall we include it anywhere else (such as a link) and then use wget to download it?

[ericspod]

If you can include the data directly in the notebook (ie. in a cell) you can do that, otherwise please host it elsewhere such that it can be downloaded through the browser or with a tool like gdown that's readily available.

[cxlcl]

Hi, maybe I didn't quite get how the script ./runtest.sh decide if a the code format is ready to go, and reformatted it if needed. But after I tried to ran ./runtest.sh --autofix after getting this CI error, I got this CI error.

It turns out that the ./runtest.sh --autofix did the changes similar to the following to multiple files:

before (in: monai/data/decathlon_datalist.py)

@overload
def _compute_path(base_dir: PathLike, element: PathLike, check_path: bool = False) -> str:
    ...

after

@overload
def _compute_path(base_dir: PathLike, element: PathLike, check_path: bool = False) -> str: ...

These changes are applied to multiple files and lead to the following error in the CI test [link to the error message]

/home/runner/work/MONAI/MONAI/monai/data/decathlon_datalist.py:27:1: E704 multiple statements on one line (def)
/home/runner/work/MONAI/MONAI/monai/data/decathlon_datalist.py:31:1: E704 multiple statements on one line (def)
/home/runner/work/MONAI/MONAI/monai/utils/dist.py:53:1: E704 multiple statements on one line (def)
/home/runner/work/MONAI/MONAI/monai/utils/dist.py:57:1: E704 multiple statements on one line (def)
/home/runner/work/MONAI/MONAI/monai/utils/dist.py:61:1: E704 multiple statements on one line (def)
/home/runner/work/MONAI/MONAI/monai/utils/misc.py:106:1: E704 multiple statements on one line (def)
/home/runner/work/MONAI/MONAI/monai/utils/misc.py:110:1: E704 multiple statements on one line (def)
7     E704 multiple statements on one line (def)
7
Check failed!
Please run auto style fixes: ./runtests.sh --autofix
Error: Process completed with exit code 1.

As suggested in the guideline, I've tried to run ./runtests.sh --autofix locally before pushing. But after that, when I run ./runtests.sh --codeformat, I got the same error above.

Maybe I'm missing something, like the configuration for autofix?

[KumoLiu]

Hi @cxlcl, it looks like it was introduced by the new release from black.
psf/black#4173
We should ignore E704. cc @ericspod

[cxlcl]

Hi @cxlcl, it looks like it was introduced by the new release from black. psf/black#4173 We should ignore E704. cc @ericspod

Thanks for the reply. Is there anything I can do from my end to pass the CI test to ignore E704? I guess it is a server-side modification for testing.

[KumoLiu]

Hi @cxlcl, no.
I have created a PR to fix it. #7422
You can directly pull the latest change, thanks.

[cxlcl]

Hi @cxlcl I went through this and the tutorial with a number of comments. It looks good in general though there's a few things to address. Please do still coordinate with @marksgraham if there's anything to align with our other diffusion work. Thanks!

Hi @ericspod , thanks a lot for the comments and suggestions. I did some modifications based on the suggestions in the new commit. Please let me know if there's any modification needed. Thanks!

[ericspod]

Looks good now. There's a minor PEP8 complaint but otherwise it's fine.

[KumoLiu]

Overall looks good to me, thanks!

[KumoLiu]

/build

[KumoLiu]

/build