| 19 | |
| 20 | @si_module |
| 21 | class GaussianMixtureIOLayer(nn.Module): |
| 22 | class Config: |
| 23 | latent_dim: int |
| 24 | dim: int |
| 25 | num_components: int |
| 26 | |
| 27 | def __init__(self, c: Config): |
| 28 | super().__init__() |
| 29 | self.latent_dim = c.latent_dim |
| 30 | self.num_components = c.num_components |
| 31 | self.input_projection = nn.Linear(c.latent_dim, c.dim) |
| 32 | |
| 33 | self.fc_loc = nn.Linear(c.dim, c.num_components * c.latent_dim) |
| 34 | self.fc_scale = nn.Linear(c.dim, c.num_components * c.latent_dim) |
| 35 | self.fc_weight = nn.Linear(c.dim, c.num_components) |
| 36 | |
| 37 | def _square_plus(self, x): |
| 38 | return (x + T.sqrt(T.square(x) + 4)) / 2 |
| 39 | |
| 40 | def input(self, sampled_latents: T.Tensor) -> T.Tensor: |
| 41 | """Pre-sampled latents T.Tensor (B, L, Z) -> float tensor (B, L, D)""" |
| 42 | hidden = self.input_projection(sampled_latents) |
| 43 | return hidden |
| 44 | |
| 45 | def output(self, h: T.Tensor) -> Tuple[T.Tensor, T.Tensor, T.Tensor]: |
| 46 | """float tensor (B, L, D) -> Tuple of locs, scales, and weights""" |
| 47 | batch_size, seq_len, _ = h.shape |
| 48 | |
| 49 | locs = self.fc_loc(h).view(batch_size, seq_len, self.num_components, self.latent_dim) |
| 50 | scales = T.clamp(self._square_plus(self.fc_scale(h)), min=1e-6).view(batch_size, seq_len, self.num_components, self.latent_dim) |
| 51 | weights = self.fc_weight(h).view(batch_size, seq_len, self.num_components) |
| 52 | |
| 53 | return (locs, scales, weights) |
| 54 | |
| 55 | def loss(self, data, dataHat): |
| 56 | locs, scales, weights = dataHat |
| 57 | log_probs = -0.5 * T.sum( |
| 58 | (data.unsqueeze(-2) - locs).pow(2) / scales.pow(2) + |
| 59 | 2 * T.log(scales) + |
| 60 | T.log(T.tensor(2 * T.pi)), |
| 61 | dim=-1 |
| 62 | ) |
| 63 | log_weights = F.log_softmax(weights, dim=-1) |
| 64 | return -T.logsumexp(log_weights + log_probs, dim=-1) |
| 65 | |
| 66 | |
| 67 | def temp_sample(self, orig_pdist, temp): |
| 68 | locs, scales, weights = orig_pdist |
| 69 | if temp is None: |
| 70 | component_samples = locs + scales * T.randn_like(scales) |
| 71 | mixture_samples = F.gumbel_softmax(weights, hard=True) |
| 72 | sampled = (component_samples * mixture_samples.unsqueeze(-1)).sum(dim=-2) |
| 73 | elif isinstance(temp, tuple): |
| 74 | assert len(temp) == 2 |
| 75 | categorical_temp, gaussian_temp = temp |
| 76 | component_samples = locs + scales * gaussian_temp * T.randn_like(scales) |
| 77 | mixture_samples = F.gumbel_softmax(weights / categorical_temp, hard=True) |
| 78 | sampled = (component_samples * mixture_samples.unsqueeze(-1)).sum(dim=-2) |
nothing calls this directly
no outgoing calls
no test coverage detected