Fix posterior VP interpolation to use float32 precision

README.md

@@ -17,6 +17,7 @@ library_name: fcdm_diffae
 
 | Date | Change |
 |------|--------|
+| 2026-04-08 | Fix posterior VP interpolation to use float32 precision (was using model dtype) |
 | 2026-04-07 | Rename package `capacitor_diffae` → `fcdm_diffae`, class `FCDMDiffAE`; encode() now returns whitened latents, decode() dewhitens internally |
 | 2026-04-06 | Initial release |
 

fcdm_diffae/encoder.py

@@ -30,24 +30,28 @@ class EncoderPosterior:
 
     @property
     def alpha(self) -> Tensor:
-        """VP signal coefficient: sqrt(sigmoid(logsnr))."""
-        return torch.sigmoid(self.logsnr).sqrt()
+        """VP signal coefficient: sqrt(sigmoid(logsnr)), computed in float32."""
+        logsnr_fp32 = self.logsnr.to(torch.float32)
+        return torch.sigmoid(logsnr_fp32).sqrt()
 
     @property
     def sigma(self) -> Tensor:
-        """VP noise coefficient: sqrt(sigmoid(-logsnr))."""
-        return torch.sigmoid(-self.logsnr).sqrt()
+        """VP noise coefficient: sqrt(sigmoid(-logsnr)), computed in float32."""
+        logsnr_fp32 = self.logsnr.to(torch.float32)
+        return torch.sigmoid(-logsnr_fp32).sqrt()
 
     def mode(self) -> Tensor:
-        """Posterior mode in token space: alpha * mean."""
-        return self.alpha.to(dtype=self.mean.dtype) * self.mean
+        """Posterior mode in token space: alpha * mean, computed in float32."""
+        return (self.alpha * self.mean.to(torch.float32)).to(dtype=self.mean.dtype)
 
     def sample(self, *, generator: torch.Generator | None = None) -> Tensor:
-        """Sample from posterior: alpha * mean + sigma * eps."""
-        eps = torch.randn_like(self.mean, generator=generator)  # type: ignore[call-overload]
-        alpha = self.alpha.to(dtype=self.mean.dtype)
-        sigma = self.sigma.to(dtype=self.mean.dtype)
-        return alpha * self.mean + sigma * eps
+        """Sample from posterior: alpha * mean + sigma * eps, computed in float32."""
+        mean_fp32 = self.mean.to(torch.float32)
+        eps = torch.randn(
+            mean_fp32.shape, device=mean_fp32.device, dtype=torch.float32,
+            generator=generator,
+        )
+        return (self.alpha * mean_fp32 + self.sigma * eps).to(dtype=self.mean.dtype)
 
 
 class Encoder(nn.Module):
@@ -123,7 +127,8 @@ class Encoder(nn.Module):
         if self.bottleneck_posterior_kind == "diagonal_gaussian":
             mean, logsnr = projection.chunk(2, dim=1)
             mean = self.norm_out(mean)
-            alpha = torch.sigmoid(logsnr).sqrt().to(dtype=mean.dtype)
-            return alpha * mean
+            logsnr_fp32 = logsnr.to(torch.float32)
+            alpha = torch.sigmoid(logsnr_fp32).sqrt()
+            return (alpha * mean.to(torch.float32)).to(dtype=mean.dtype)
         z = self.norm_out(projection)
         return z