| """Part D: Latent diffusion synthesizer for diffusive attention.""" |
| import jax, jax.numpy as jnp, flax.nnx as nn |
|
|
| class UNetMLP(nn.Module): |
| """U-Net-style MLP: predicts noise eps_theta(z_tau, tau, h_t).""" |
| def __init__(self, d_z: int, d_cond: int, d_hidden: int = 128, |
| n_layers: int = 3, *, rngs: nn.Rngs): |
| super().__init__() |
| self.time_embed = nn.Linear(1, d_hidden, rngs=rngs) |
| self.cond_embed = nn.Linear(d_cond, d_hidden, rngs=rngs) |
| self.input_proj = nn.Linear(d_z, d_hidden, rngs=rngs) |
| self.layers = [nn.Linear(d_hidden, d_hidden, rngs=rngs) for _ in range(n_layers)] |
| self.out_proj = nn.Linear(d_hidden, d_z, rngs=rngs) |
|
|
| def __call__(self, z, tau, h_cond): |
| t_emb = jax.nn.silu(self.time_embed(tau[..., None])) |
| c_emb = jax.nn.silu(self.cond_embed(h_cond)) |
| h = jax.nn.silu(self.input_proj(z)) |
| for layer in self.layers: |
| h = jax.nn.silu(layer(h + c_emb + t_emb)) + h |
| return self.out_proj(h) |
|
|
|
|
| class LatentDiffusionSynthesizer(nn.Module): |
| """ |
| Continuous-time latent diffusion conditioned on SSM state h_t. |
| Forward: q(z_tau|z_{tau-1}) = N(z_tau; sqrt(1-beta_tau)*z_{tau-1}, beta_tau*I) |
| Reverse: p_theta(z_{tau-1}|z_tau, h_t) = N(z_{tau-1}; mu_theta, Sigma_theta) |
| mu_theta via U-Net MLP. Time complexity: O(T * d_z^2) |
| """ |
| def __init__(self, d_z: int, d_cond: int, T: int = 10, |
| d_hidden: int = 128, n_unet_layers: int = 3, *, rngs: nn.Rngs): |
| super().__init__() |
| self.d_z = d_z |
| self.T = T |
| self.d_cond = d_cond |
| self.beta_raw = nn.Param(jnp.linspace(-4, 4, T)) |
| self.unet = UNetMLP(d_z, d_cond, d_hidden, n_unet_layers, rngs=rngs) |
| self.z_prior_proj = nn.Linear(d_cond, d_z, rngs=rngs) |
|
|
| @property |
| def beta(self): |
| return jax.nn.sigmoid(self.beta_raw.value) |
|
|
| @property |
| def alpha_cumprod(self): |
| return jnp.cumprod(1.0 - self.beta) |
|
|
| def reverse_step(self, z_tau, tau_idx, h_t, key): |
| tau = float(tau_idx + 1) / self.T |
| beta_t = self.beta[tau_idx] |
| alpha_t = 1.0 - beta_t |
| alpha_bar_t = self.alpha_cumprod[tau_idx] |
| tau_arr = jnp.full(z_tau.shape[:-1], tau) |
| eps_theta = self.unet(z_tau, tau_arr, h_t) |
| mu_theta = (1.0 / jnp.sqrt(alpha_t)) * ( |
| z_tau - beta_t / jnp.sqrt(1.0 - alpha_bar_t) * eps_theta |
| ) |
| if tau_idx > 0: |
| sigma = jnp.sqrt(beta_t) |
| z_prev = mu_theta + sigma * jax.random.normal(key, z_tau.shape) |
| else: |
| z_prev = mu_theta |
| return z_prev |
|
|
| def sample(self, h_t, key): |
| shape_prefix = h_t.shape[:-1] |
| key, sk = jax.random.split(key) |
| z = jax.random.normal(sk, (*shape_prefix, self.d_z)) |
| for tau_idx in range(self.T - 1, -1, -1): |
| key, sk = jax.random.split(key) |
| z = self.reverse_step(z, tau_idx, h_t, sk) |
| return z |
|
|
| def __call__(self, h_t, key=None): |
| if key is None: |
| key = jax.random.key(0) |
| return self.sample(h_t, key) |
|
|
