Skip to content

fix: store z_mu instead of sampled posterior in latent precompute - #52

Merged
Can-Zhao merged 3 commits into
NVIDIA-Medtech:mainfrom
Can-Zhao:fix/latent-z-mu
Oct 7, 2026
Merged

Can-Zhao merged 3 commits into
NVIDIA-Medtech:mainfrom
Can-Zhao:fix/latent-z-mu

Conversation

@Can-Zhao

@Can-Zhao Can-Zhao commented Oct 6, 2026

Copy link
Copy Markdown
Contributor

encode_stage_2_inputs() returns z_mu + sigma*eps. With weak KL regularisation (kl_weight~1e-7), the noise term accounts for ~27% of latent variance and acts as an irreducible floor on the diffusion L1 objective. The released checkpoints were trained on z_mu (encode()[0]); this aligns the script with that practice.

Fixes #45.

encode_stage_2_inputs() returns z_mu + sigma*eps. With weak KL
regularisation (kl_weight~1e-7), the noise term accounts for ~27% of
latent variance and acts as an irreducible floor on the diffusion L1
objective. The released checkpoints were trained on z_mu (encode()[0]);
this aligns the script with that practice.

Fixes NVIDIA-Medtech#45.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@greptile-apps

greptile-apps Bot commented Oct 6, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerConfidence Score: 5/5

[High risk] Changes how training data embeddings are computed and stored.

The PR appears safe to merge; no outstanding findings remain.

Summary

The PR changes latent precomputation to store the autoencoder’s posterior mean instead of a sampled latent, and updates the embedding instructions to match.

  • The README now prominently advises regenerating embeddings created before the fix.

Reviews (3) · Last reviewed commit: "fix: split long line in data/README.md t..."

Comment thread scripts/diff_model_create_training_data.py
Comment thread scripts/diff_model_create_training_data.py
Can-Zhao and others added 2 commits October 6, 2026 11:02
data/README.md: replace encode_stage_2_inputs with encode(x)[0]
(posterior mean z_mu) and note that old embeddings should be
regenerated.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@Can-Zhao
Can-Zhao merged commit dba166e into NVIDIA-Medtech:main Oct 7, 2026
2 checks passed
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.

Fine-tuning rflow-mr on rectal T2

1 participant