torchcell.models.diffusion_decoder.DiffusionDecoder

class DiffusionDecoder(input_dim: int, hidden_dim: int = 128, output_dim: int = 1, num_layers: int = 4, num_heads: int = 8, dropout: float = 0.1, norm: str = 'layer', num_timesteps: int = 1000, mlp_ratio: float = 4.0, beta_schedule: str = 'cosine', beta_start: float = 0.0001, beta_end: float = 0.02, cosine_s: float = 0.008, sampling_steps: int = 50, parameterization: str = 'x0')[source]

Bases: Module

Diffusion decoder with cross-attention conditioning on graph embeddings.

forward_diffusion(x_0: Tensor, t: Tensor, noise: Tensor | None = None) → tuple[Tensor, Tensor][source]

Add noise to data for training.

Parameters:

  • x_0 – Original data [batch, 1]

  • t – Timesteps [batch]

  • noise – Optional noise [batch, 1]

Returns:

Noisy data noise: Noise that was added

Return type:

x_t

denoise(x_t: Tensor, context: Tensor, t: Tensor, predict_x0: bool = True) → Tensor[source]

Predict x0 or noise given noisy data and conditioning.

Parameters:

  • x_t – Noisy data [batch, 1]

  • context – Graph embeddings [batch, hidden_dim]

  • t – Timesteps [batch]

  • predict_x0 – If True, predict x0 directly; if False, predict noise

Returns:

Predicted x0 or noise [batch, 1]

sample(context: Tensor, num_samples: int | None = None, device: device | None = None, sampling_steps: int | None = None) → Tensor[source]

Sample from the diffusion model using DDIM-style sampling.

Parameters:

  • context – Graph embeddings [batch, hidden_dim]

  • num_samples – Number of samples (defaults to batch size)

  • device – Device to sample on

  • sampling_steps – Number of denoising steps (defaults to 50 for faster sampling)

Returns:

Sampled phenotypes [batch, 1]

loss(x_0: Tensor, context: Tensor, predict_x0: bool = True, t_mode: str = 'random') → Tensor[source]

Compute diffusion loss using x0-prediction or noise prediction.

Parameters:

  • x_0 – Original phenotype values [batch, 1]

  • context – Graph embeddings [batch, hidden_dim]

  • predict_x0 – If True, predict x0 directly (better for sparse data)

  • t_mode – Timestep sampling mode: - “zero” or “t0”: Always use t=0 (pure reconstruction, no noise) - “partial” or “small”: Sample from [0, num_timesteps/10) - “full” or “random”: Sample from full range [0, num_timesteps) (default)

Returns:

Loss value