From 528243de2ce1efde2933c097dc9500697a07f6d6 Mon Sep 17 00:00:00 2001 From: Lanre Shittu <136805224+Shizoqua@users.noreply.github.com> Date: Fri, 19 Jun 2026 20:23:56 +0100 Subject: [PATCH 1/3] fix(networks): return mu from VAE reparameterize at inference VarAutoEncoder.reparameterize added the standard deviation to mu at eval time (`std.add_(mu)` with no noise term), so inference returned mu + std instead of the posterior mean. At inference the latent code should be mu; the random term belongs only to training (the reparameterization trick). Return mu directly when not training, and compute mu + eps * std out-of-place otherwise. VarFullyConnectedNet.reparameterize had the identical bug and is fixed the same way. Adds regression tests asserting eval is deterministic and equals mu while training stays stochastic. Fixes #8413. Signed-off-by: Lanre Shittu <136805224+Shizoqua@users.noreply.github.com> --- monai/networks/nets/fullyconnectednet.py | 11 +++++----- monai/networks/nets/varautoencoder.py | 11 +++++----- tests/networks/nets/test_fullyconnectednet.py | 21 ++++++++++++++++++ tests/networks/nets/test_varautoencoder.py | 22 +++++++++++++++++++ 4 files changed, 55 insertions(+), 10 deletions(-) diff --git a/monai/networks/nets/fullyconnectednet.py b/monai/networks/nets/fullyconnectednet.py index be179e5b59c..8c331fb38d2 100644 --- a/monai/networks/nets/fullyconnectednet.py +++ b/monai/networks/nets/fullyconnectednet.py @@ -172,12 +172,13 @@ def decode_forward(self, z: torch.Tensor, use_sigmoid: bool = True) -> torch.Ten return x def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: + if not self.training: + # At inference the latent code is the posterior mean; the random + # term is only added during training (the reparameterization trick). + return mu std = torch.exp(0.5 * logvar) - - if self.training: # multiply random noise with std only during training - std = torch.randn_like(std).mul(std) - - return std.add_(mu) + eps = torch.randn_like(std) + return mu + eps * std def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: mu, logvar = self.encode_forward(x) diff --git a/monai/networks/nets/varautoencoder.py b/monai/networks/nets/varautoencoder.py index 0674094aa78..9a8110aae9e 100644 --- a/monai/networks/nets/varautoencoder.py +++ b/monai/networks/nets/varautoencoder.py @@ -142,12 +142,13 @@ def decode_forward(self, z: torch.Tensor, use_sigmoid: bool = True) -> torch.Ten return x def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: + if not self.training: + # At inference the latent code is the posterior mean; the random + # term is only added during training (the reparameterization trick). + return mu std = torch.exp(0.5 * logvar) - - if self.training: # multiply random noise with std only during training - std = torch.randn_like(std).mul(std) - - return std.add_(mu) + eps = torch.randn_like(std) + return mu + eps * std def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: mu, logvar = self.encode_forward(x) diff --git a/tests/networks/nets/test_fullyconnectednet.py b/tests/networks/nets/test_fullyconnectednet.py index 863d1399a9a..0eca66a0f2e 100644 --- a/tests/networks/nets/test_fullyconnectednet.py +++ b/tests/networks/nets/test_fullyconnectednet.py @@ -64,6 +64,27 @@ def test_vfc_shape(self, input_param, input_shape, expected_shape): result = net.forward(torch.randn(input_shape).to(device))[0] self.assertEqual(result.shape, expected_shape) + def test_vfc_reparameterize_eval_returns_mu(self): + # At eval the latent code must equal mu (deterministic); at train it must + # be stochastic. Same #8413 reparameterize bug as VarAutoEncoder. + net = VarFullyConnectedNet( + in_channels=10, out_channels=10, latent_size=30, encode_channels=(15, 20, 25), decode_channels=(15, 20, 25) + ).to(device) + data = torch.randn(3, 10).to(device) + + with eval_mode(net): + _, mu1, _, z1 = net(data) + _, _, _, z2 = net(data) + self.assertTrue(torch.allclose(z1, mu1)) + self.assertTrue(torch.allclose(z1, z2)) + + net.train() + with torch.no_grad(): + _, mu_t, _, zt1 = net(data) + _, _, _, zt2 = net(data) + self.assertFalse(torch.allclose(zt1, mu_t)) + self.assertFalse(torch.allclose(zt1, zt2)) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_varautoencoder.py b/tests/networks/nets/test_varautoencoder.py index 459c537c555..d24ece16594 100644 --- a/tests/networks/nets/test_varautoencoder.py +++ b/tests/networks/nets/test_varautoencoder.py @@ -122,6 +122,28 @@ def test_script(self): test_data = torch.randn(2, 1, 32, 32) test_script_save(net, test_data, rtol=1e-3, atol=1e-3) + def test_reparameterize_eval_returns_mu(self): + # At eval the latent code must equal mu (deterministic, no noise added); + # at train it must be stochastic. Regression test for #8413, where eval + # returned mu + std. + net = VarAutoEncoder( + spatial_dims=2, in_shape=(1, 32, 32), out_channels=1, latent_size=4, channels=(4, 8), strides=(2, 2) + ).to(device) + data = torch.randn(2, 1, 32, 32).to(device) + + with eval_mode(net): + _, mu1, _, z1 = net(data) + _, _, _, z2 = net(data) + self.assertTrue(torch.allclose(z1, mu1)) + self.assertTrue(torch.allclose(z1, z2)) + + net.train() + with torch.no_grad(): + _, mu_t, _, zt1 = net(data) + _, _, _, zt2 = net(data) + self.assertFalse(torch.allclose(zt1, mu_t)) + self.assertFalse(torch.allclose(zt1, zt2)) + if __name__ == "__main__": unittest.main() From b7415af9d32d88ece98598371f796e545d2c6f50 Mon Sep 17 00:00:00 2001 From: Lanre Shittu <136805224+Shizoqua@users.noreply.github.com> Date: Wed, 5 Aug 2026 03:09:17 +0100 Subject: [PATCH 2/3] docs(networks): add docstrings to VAE reparameterize methods and tests Signed-off-by: Lanre Shittu <136805224+Shizoqua@users.noreply.github.com> --- monai/networks/nets/fullyconnectednet.py | 12 ++++++++++++ monai/networks/nets/varautoencoder.py | 12 ++++++++++++ tests/networks/nets/test_fullyconnectednet.py | 7 +++++-- tests/networks/nets/test_varautoencoder.py | 7 ++++--- 4 files changed, 33 insertions(+), 5 deletions(-) diff --git a/monai/networks/nets/fullyconnectednet.py b/monai/networks/nets/fullyconnectednet.py index 8c331fb38d2..d7354d4bd24 100644 --- a/monai/networks/nets/fullyconnectednet.py +++ b/monai/networks/nets/fullyconnectednet.py @@ -172,6 +172,18 @@ def decode_forward(self, z: torch.Tensor, use_sigmoid: bool = True) -> torch.Ten return x def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: + """Sample a latent code using the reparameterization trick. + + At inference (eval mode) the posterior mean is returned directly. During + training, returns ``mu + eps * std`` with ``eps ~ N(0, I)``. + + Args: + mu: Posterior mean, shape ``(batch, latent_size)``. + logvar: Log-variance of the posterior, same shape as ``mu``. + + Returns: + Sampled latent code, same shape as ``mu``. + """ if not self.training: # At inference the latent code is the posterior mean; the random # term is only added during training (the reparameterization trick). diff --git a/monai/networks/nets/varautoencoder.py b/monai/networks/nets/varautoencoder.py index 9a8110aae9e..23c8aa37dca 100644 --- a/monai/networks/nets/varautoencoder.py +++ b/monai/networks/nets/varautoencoder.py @@ -142,6 +142,18 @@ def decode_forward(self, z: torch.Tensor, use_sigmoid: bool = True) -> torch.Ten return x def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: + """Sample a latent code using the reparameterization trick. + + At inference (eval mode) the posterior mean is returned directly. During + training, returns ``mu + eps * std`` with ``eps ~ N(0, I)``. + + Args: + mu: Posterior mean, shape ``(batch, latent_size)``. + logvar: Log-variance of the posterior, same shape as ``mu``. + + Returns: + Sampled latent code, same shape as ``mu``. + """ if not self.training: # At inference the latent code is the posterior mean; the random # term is only added during training (the reparameterization trick). diff --git a/tests/networks/nets/test_fullyconnectednet.py b/tests/networks/nets/test_fullyconnectednet.py index 0eca66a0f2e..b1ee3155852 100644 --- a/tests/networks/nets/test_fullyconnectednet.py +++ b/tests/networks/nets/test_fullyconnectednet.py @@ -65,8 +65,11 @@ def test_vfc_shape(self, input_param, input_shape, expected_shape): self.assertEqual(result.shape, expected_shape) def test_vfc_reparameterize_eval_returns_mu(self): - # At eval the latent code must equal mu (deterministic); at train it must - # be stochastic. Same #8413 reparameterize bug as VarAutoEncoder. + """A VFC latent code is deterministic at eval (equals mu) and stochastic at train. + + Regression test for the #8413 reparameterize bug, which returned ``mu + std`` + at inference instead of ``mu``. + """ net = VarFullyConnectedNet( in_channels=10, out_channels=10, latent_size=30, encode_channels=(15, 20, 25), decode_channels=(15, 20, 25) ).to(device) diff --git a/tests/networks/nets/test_varautoencoder.py b/tests/networks/nets/test_varautoencoder.py index d24ece16594..ce1e2e50e7f 100644 --- a/tests/networks/nets/test_varautoencoder.py +++ b/tests/networks/nets/test_varautoencoder.py @@ -123,9 +123,10 @@ def test_script(self): test_script_save(net, test_data, rtol=1e-3, atol=1e-3) def test_reparameterize_eval_returns_mu(self): - # At eval the latent code must equal mu (deterministic, no noise added); - # at train it must be stochastic. Regression test for #8413, where eval - # returned mu + std. + """A VarAutoEncoder latent code is deterministic at eval (equals mu) and stochastic at train. + + Regression test for #8413, where eval returned ``mu + std`` instead of ``mu``. + """ net = VarAutoEncoder( spatial_dims=2, in_shape=(1, 32, 32), out_channels=1, latent_size=4, channels=(4, 8), strides=(2, 2) ).to(device) From a0b471d06e53205dffb1724f91e5c57c8b3ab724 Mon Sep 17 00:00:00 2001 From: Lanre Shittu <136805224+Shizoqua@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:53:47 +0100 Subject: [PATCH 3/3] fix(nets): preserve backward compat for VAE reparameterize at inference Address review: instead of changing the default inference behaviour (which returned mu + std), gate the new behaviour behind a use_mean_at_inference constructor argument. The default stays mu + std for backward compatibility; use_mean_at_inference=True returns the posterior mean at eval, matching the issue #8413 request. Both VarAutoEncoder and VarFullyConnectedNet get the new argument, with regression tests covering both the new opt-in and the preserved default. Signed-off-by: Lanre Shittu <136805224+Shizoqua@users.noreply.github.com> --- monai/networks/nets/fullyconnectednet.py | 14 ++++++--- monai/networks/nets/varautoencoder.py | 14 ++++++--- tests/networks/nets/test_fullyconnectednet.py | 28 ++++++++++++++++- tests/networks/nets/test_varautoencoder.py | 30 ++++++++++++++++++- 4 files changed, 76 insertions(+), 10 deletions(-) diff --git a/monai/networks/nets/fullyconnectednet.py b/monai/networks/nets/fullyconnectednet.py index d7354d4bd24..9c22e0a8e49 100644 --- a/monai/networks/nets/fullyconnectednet.py +++ b/monai/networks/nets/fullyconnectednet.py @@ -103,6 +103,8 @@ class VarFullyConnectedNet(nn.Module): act: activation type and arguments. Defaults to PReLU. bias: whether to have a bias term in linear units. Defaults to True. adn_ordering: order of operations in :py:class:`monai.networks.blocks.ADN`. + use_mean_at_inference: whether to return the posterior mean (rather than ``mu + std``) + as the latent code during inference. Defaults to False for backward compatibility. Examples:: @@ -122,11 +124,13 @@ def __init__( act: tuple | str | None = Act.PRELU, bias: bool = True, adn_ordering: str | None = None, + use_mean_at_inference: bool = False, ) -> None: super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.latent_size = latent_size + self.use_mean_at_inference = use_mean_at_inference self.encode = nn.Sequential() self.decode = nn.Sequential() @@ -174,8 +178,10 @@ def decode_forward(self, z: torch.Tensor, use_sigmoid: bool = True) -> torch.Ten def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: """Sample a latent code using the reparameterization trick. - At inference (eval mode) the posterior mean is returned directly. During - training, returns ``mu + eps * std`` with ``eps ~ N(0, I)``. + At inference (eval mode), if ``use_mean_at_inference`` is enabled, the posterior + mean is returned directly. Otherwise, ``mu + std`` is returned, matching the + original behaviour. During training, returns ``mu + eps * std`` with + ``eps ~ N(0, I)``. Args: mu: Posterior mean, shape ``(batch, latent_size)``. @@ -184,12 +190,12 @@ def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor Returns: Sampled latent code, same shape as ``mu``. """ - if not self.training: + if not self.training and self.use_mean_at_inference: # At inference the latent code is the posterior mean; the random # term is only added during training (the reparameterization trick). return mu std = torch.exp(0.5 * logvar) - eps = torch.randn_like(std) + eps = torch.randn_like(std) if self.training else torch.ones_like(std) return mu + eps * std def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: diff --git a/monai/networks/nets/varautoencoder.py b/monai/networks/nets/varautoencoder.py index 23c8aa37dca..d6d6cba92f1 100644 --- a/monai/networks/nets/varautoencoder.py +++ b/monai/networks/nets/varautoencoder.py @@ -51,6 +51,8 @@ class VarAutoEncoder(AutoEncoder): According to `Performance Tuning Guide `_, if a conv layer is directly followed by a batch norm layer, bias should be False. use_sigmoid: whether to use the sigmoid function on final output. Defaults to True. + use_mean_at_inference: whether to return the posterior mean (rather than ``mu + std``) + as the latent code during inference. Defaults to False for backward compatibility. Examples:: @@ -90,9 +92,11 @@ def __init__( dropout: tuple | str | float | None = None, bias: bool = True, use_sigmoid: bool = True, + use_mean_at_inference: bool = False, ) -> None: self.in_channels, *self.in_shape = in_shape self.use_sigmoid = use_sigmoid + self.use_mean_at_inference = use_mean_at_inference self.latent_size = latent_size self.final_size = np.asarray(self.in_shape, dtype=int) @@ -144,8 +148,10 @@ def decode_forward(self, z: torch.Tensor, use_sigmoid: bool = True) -> torch.Ten def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: """Sample a latent code using the reparameterization trick. - At inference (eval mode) the posterior mean is returned directly. During - training, returns ``mu + eps * std`` with ``eps ~ N(0, I)``. + At inference (eval mode), if ``use_mean_at_inference`` is enabled, the posterior + mean is returned directly. Otherwise, ``mu + std`` is returned, matching the + original behaviour. During training, returns ``mu + eps * std`` with + ``eps ~ N(0, I)``. Args: mu: Posterior mean, shape ``(batch, latent_size)``. @@ -154,12 +160,12 @@ def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor Returns: Sampled latent code, same shape as ``mu``. """ - if not self.training: + if not self.training and self.use_mean_at_inference: # At inference the latent code is the posterior mean; the random # term is only added during training (the reparameterization trick). return mu std = torch.exp(0.5 * logvar) - eps = torch.randn_like(std) + eps = torch.randn_like(std) if self.training else torch.ones_like(std) return mu + eps * std def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: diff --git a/tests/networks/nets/test_fullyconnectednet.py b/tests/networks/nets/test_fullyconnectednet.py index b1ee3155852..9d5efa7870f 100644 --- a/tests/networks/nets/test_fullyconnectednet.py +++ b/tests/networks/nets/test_fullyconnectednet.py @@ -71,7 +71,12 @@ def test_vfc_reparameterize_eval_returns_mu(self): at inference instead of ``mu``. """ net = VarFullyConnectedNet( - in_channels=10, out_channels=10, latent_size=30, encode_channels=(15, 20, 25), decode_channels=(15, 20, 25) + in_channels=10, + out_channels=10, + latent_size=30, + encode_channels=(15, 20, 25), + decode_channels=(15, 20, 25), + use_mean_at_inference=True, ).to(device) data = torch.randn(3, 10).to(device) @@ -88,6 +93,27 @@ def test_vfc_reparameterize_eval_returns_mu(self): self.assertFalse(torch.allclose(zt1, mu_t)) self.assertFalse(torch.allclose(zt1, zt2)) + def test_vfc_reparameterize_default_keeps_original_behaviour(self): + """By default, eval returns ``mu + std`` (deterministic) for backward compatibility. + + The default must preserve the pre-#8413 behaviour: at inference the standard + deviation is added to the mean without random noise. + """ + net = VarFullyConnectedNet( + in_channels=10, + out_channels=10, + latent_size=30, + encode_channels=(15, 20, 25), + decode_channels=(15, 20, 25), + ).to(device) + data = torch.randn(3, 10).to(device) + + with eval_mode(net): + _, mu1, _, z1 = net(data) + _, _, _, z2 = net(data) + self.assertFalse(torch.allclose(z1, mu1)) + self.assertTrue(torch.allclose(z1, z2)) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_varautoencoder.py b/tests/networks/nets/test_varautoencoder.py index ce1e2e50e7f..193ed631135 100644 --- a/tests/networks/nets/test_varautoencoder.py +++ b/tests/networks/nets/test_varautoencoder.py @@ -128,7 +128,13 @@ def test_reparameterize_eval_returns_mu(self): Regression test for #8413, where eval returned ``mu + std`` instead of ``mu``. """ net = VarAutoEncoder( - spatial_dims=2, in_shape=(1, 32, 32), out_channels=1, latent_size=4, channels=(4, 8), strides=(2, 2) + spatial_dims=2, + in_shape=(1, 32, 32), + out_channels=1, + latent_size=4, + channels=(4, 8), + strides=(2, 2), + use_mean_at_inference=True, ).to(device) data = torch.randn(2, 1, 32, 32).to(device) @@ -145,6 +151,28 @@ def test_reparameterize_eval_returns_mu(self): self.assertFalse(torch.allclose(zt1, mu_t)) self.assertFalse(torch.allclose(zt1, zt2)) + def test_reparameterize_default_keeps_original_behaviour(self): + """By default, eval returns ``mu + std`` (deterministic) for backward compatibility. + + The default must preserve the pre-#8413 behaviour: at inference the standard + deviation is added to the mean without random noise. + """ + net = VarAutoEncoder( + spatial_dims=2, + in_shape=(1, 32, 32), + out_channels=1, + latent_size=4, + channels=(4, 8), + strides=(2, 2), + ).to(device) + data = torch.randn(2, 1, 32, 32).to(device) + + with eval_mode(net): + _, mu1, _, z1 = net(data) + _, _, _, z2 = net(data) + self.assertFalse(torch.allclose(z1, mu1)) + self.assertTrue(torch.allclose(z1, z2)) + if __name__ == "__main__": unittest.main()