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>

Author: Shizoqua

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]:

monai/networks/nets/varautoencoder.py

@@ -51,6 +51,8 @@ class VarAutoEncoder(AutoEncoder):
             According to `Performance Tuning Guide <https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html>`_,
             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]:

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()

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()