Patch VP logSNR stability in exported full_capacitor code

full_capacitor/encoder.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 
 import torch
+import torch.nn.functional as F
 from torch import Tensor, nn
 
 from .fcdm_block import FCDMBlock
@@ -30,15 +31,15 @@ class EncoderPosterior:
 
     @property
     def alpha(self) -> Tensor:
-        """VP signal coefficient: sqrt(sigmoid(logsnr)), computed in float32."""
+        """VP signal coefficient computed stably in float32."""
         logsnr_fp32 = self.logsnr.to(torch.float32)
-        return torch.sigmoid(logsnr_fp32).sqrt()
+        return torch.exp(0.5 * F.logsigmoid(logsnr_fp32))
 
     @property
     def sigma(self) -> Tensor:
-        """VP noise coefficient: sqrt(sigmoid(-logsnr)), computed in float32."""
+        """VP noise coefficient computed stably in float32."""
         logsnr_fp32 = self.logsnr.to(torch.float32)
-        return torch.sigmoid(-logsnr_fp32).sqrt()
+        return torch.exp(0.5 * F.logsigmoid(-logsnr_fp32))
 
     def mode(self) -> Tensor:
         """Posterior mode in token space: alpha * mean, computed in float32."""
@@ -130,7 +131,7 @@ class Encoder(nn.Module):
             mean, logsnr = projection.chunk(2, dim=1)
             mean = self.norm_out(mean)
             logsnr_fp32 = logsnr.to(torch.float32)
-            alpha = torch.sigmoid(logsnr_fp32).sqrt()
+            alpha = torch.exp(0.5 * F.logsigmoid(logsnr_fp32))
             return (alpha * mean.to(torch.float32)).to(dtype=mean.dtype)
         z = self.norm_out(projection)
         return z