Upload qandc/variance_compensation.py

qandc/variance_compensation.py

@@ -0,0 +1,90 @@
+"""
+Variance Compensation (VC) for Exposure Bias Correction.
+Second key contribution of Q&C paper (Section 3.2, Eq 9-12).
+Corrects variance shifts when quantization + cache are combined.
+"""
+import torch
+import torch.nn as nn
+import numpy as np
+
+class VarianceCompensation:
+    """Full VC with analytical K_t from Eq 12."""
+    def __init__(self, n_timesteps, n_channels, device="cuda"):
+        self.n_timesteps = n_timesteps
+        self.n_channels = n_channels
+        self.device = device
+        self.K = torch.ones(n_timesteps, n_channels, device=device)
+        self._numerator = torch.zeros(n_timesteps, n_channels, device=device)
+        self._denominator = torch.zeros(n_timesteps, n_channels, device=device)
+        self._count = torch.zeros(n_timesteps, device=device)
+        self._calibrated = False
+
+    def accumulate(self, t_idx, x_fp, x_qc):
+        reduce_dims = (0, 2, 3) if x_fp.dim() == 4 else (0,)
+        mu = x_qc.mean(dim=reduce_dims)
+        x_fp_c = x_fp - mu.view(1, -1, 1, 1) if x_fp.dim() == 4 else x_fp - mu
+        x_qc_c = x_qc - mu.view(1, -1, 1, 1) if x_qc.dim() == 4 else x_qc - mu
+        term1 = (x_fp_c * x_qc_c).mean(dim=reduce_dims)
+        x_fp_safe = x_fp.clone(); x_fp_safe[x_fp_safe.abs() < 1e-6] = 1e-6
+        term2 = (x_qc_c / x_fp_safe).mean(dim=reduce_dims)
+        dterm1 = (x_qc_c ** 2).mean(dim=reduce_dims)
+        dterm2 = ((x_qc_c ** 2) / (x_fp_safe ** 2)).mean(dim=reduce_dims)
+        self._numerator[t_idx] += (term1 + term2).to(self.device)
+        self._denominator[t_idx] += (dterm1 + dterm2).to(self.device)
+        self._count[t_idx] += 1
+
+    def calibrate(self):
+        for t in range(self.n_timesteps):
+            if self._count[t] > 0:
+                self.K[t] = torch.clamp(self._numerator[t] / torch.clamp(self._denominator[t], min=1e-8), 0.5, 2.0)
+        self._calibrated = True
+        print(f"VC calibrated: K range [{self.K.min():.4f}, {self.K.max():.4f}]")
+
+    def apply(self, x, t_idx):
+        if not self._calibrated: return x
+        K_t = self.K[t_idx]
+        if x.dim() == 4:
+            mu = x.mean(dim=(0, 2, 3), keepdim=True)
+            return mu + K_t.view(1, -1, 1, 1).to(x.device) * (x - mu)
+        else:
+            mu = x.mean(dim=0, keepdim=True)
+            return mu + K_t.view(1, -1).to(x.device) * (x - mu)
+
+class SimpleVarianceCompensation:
+    """Simplified VC using variance ratio estimation. Applies only in later denoising stages."""
+    def __init__(self, n_timesteps, n_channels, apply_after_fraction=0.5, device="cuda"):
+        self.n_timesteps = n_timesteps
+        self.n_channels = n_channels
+        self.apply_after_fraction = apply_after_fraction
+        self.device = device
+        self.var_ratio = torch.ones(n_timesteps, n_channels, device=device)
+        self._fp_var_sum = torch.zeros(n_timesteps, n_channels, device=device)
+        self._qc_var_sum = torch.zeros(n_timesteps, n_channels, device=device)
+        self._count = torch.zeros(n_timesteps, device=device)
+        self._calibrated = False
+
+    def accumulate(self, t_idx, x_fp, x_qc):
+        reduce_dims = (0, 2, 3) if x_fp.dim() == 4 else (0,)
+        self._fp_var_sum[t_idx] += x_fp.var(dim=reduce_dims).to(self.device)
+        self._qc_var_sum[t_idx] += x_qc.var(dim=reduce_dims).to(self.device)
+        self._count[t_idx] += 1
+
+    def calibrate(self):
+        for t in range(self.n_timesteps):
+            if self._count[t] > 0:
+                fp_var = self._fp_var_sum[t] / self._count[t]
+                qc_var = torch.clamp(self._qc_var_sum[t] / self._count[t], min=1e-8)
+                self.var_ratio[t] = torch.clamp(torch.sqrt(fp_var / qc_var), 0.5, 2.0)
+        self._calibrated = True
+        print(f"Simple VC calibrated: ratio [{self.var_ratio.min():.4f}, {self.var_ratio.max():.4f}]")
+
+    def apply(self, x, t_idx):
+        if not self._calibrated: return x
+        if t_idx < int(self.n_timesteps * self.apply_after_fraction): return x
+        ratio = self.var_ratio[t_idx]
+        if x.dim() == 4:
+            mu = x.mean(dim=(0, 2, 3), keepdim=True)
+            return mu + ratio.view(1, -1, 1, 1).to(x.device) * (x - mu)
+        else:
+            mu = x.mean(dim=0, keepdim=True)
+            return mu + ratio.view(1, -1).to(x.device) * (x - mu)