dit_prior

dit_prior

Denoising Score-Matching based Transformer Prior for variational models.

Implements a DiT-style transformer prior that learns p(z|s) through Denoising Score-Matching with either diffusion or flow matching, where z is the latent variable and s is the conditioning (observations).

DiTPrior

DiTPrior(latent_dimension, embedding_dimension, number_of_heads, number_of_layers, feedforward_dimension, device, observation_horizon=1, algorithm_type=value, sigma=0.0, ode_solver=value, timestep_sampler=BETA.value, logit_mean=0.0, logit_std=1.0, beta_alpha=1.5, beta_beta=1.0, max_timestep=0.999, num_train_timesteps=100, num_inference_steps=10, beta_start=0.0001, beta_end=0.02, beta_schedule=value, scheduler_type=value, prediction_type=value, clip_sample=False, variance_type=None, dropout=0.1, normalization_type=value, attention_type=value, number_of_key_value_heads=None, activation=value, use_gating=True, exclude_keys=None, prior_target_key=value, latent_standardization_enabled=True, latent_standardization_eps=1e-06, latent_standardization_max_batches=None, require_fitted_latent_standardization=False)

Bases: PriorLatentEncoder

DiT prior supporting diffusion and flow matching.

Uses a DiT transformer network where the noisy latent z is treated as a CLS token appended to observation tokens. The transformer attends bidirectionally across all tokens, is gated by the timestep embedding (Adaptive LayerNorm Zero), and the final CLS representation is projected to predict noise (diffusion) or velocity (flow matching).

Parameters:

Name	Type	Description	Default
latent_dimension	int	Dimension of latent variable z.	required
embedding_dimension	int	Hidden dimension of the transformer.	required
number_of_heads	int	Number of attention heads.	required
number_of_layers	int	Number of DiT decoder layers.	required
feedforward_dimension	int	Dimension of the feedforward network.	required
device	str	Device to place prior on.	required
observation_horizon	int	Observation history size.	1
algorithm_type	str	Algorithm type ("diffusion" or "flow_matching").	value
sigma	float	Noise level for flow matching (0 = deterministic OT).	0.0
ode_solver	str	ODE solver for flow matching ("euler", "heun", or "rk4").	value
num_train_timesteps	int	Number of diffusion timesteps during training.	100
num_inference_steps	int	Number of denoising/integration steps.	10
beta_start	float	Starting beta for noise schedule (diffusion).	0.0001
beta_end	float	Ending beta for noise schedule (diffusion).	0.02
beta_schedule	str	Type of noise schedule (diffusion).	value
scheduler_type	str	Diffusion scheduler type.	value
prediction_type	str	What diffusion model predicts (epsilon, sample, velocity).	value
clip_sample	bool	Whether to clip samples during diffusion.	False
variance_type	str | None	Variance type for DDPM scheduler.	None
dropout	float	Dropout rate.	0.1
normalization_type	str	Type of adaptive normalization layer	value
activation	str	Activation function name.	value
use_gating	bool	Whether to use AdaLN-Zero gating in DiT layers.	True
exclude_keys	list[str] | None	Keys to exclude from observations.	None
prior_target_key	str	Posterior output key used as denoising target.	value
latent_standardization_enabled	bool	Whether to standardize DiT target latents.	True
latent_standardization_eps	float	Numerical epsilon used in latent standardization.	1e-06
latent_standardization_max_batches	int | None	Maximum train batches to scan when fitting latent standardization stats. None scans the full train loader.	None
require_fitted_latent_standardization	bool	Whether missing latent stats should raise.	False

Source code in src/versatil/models/decoding/latent/prior/dit_prior.py

def __init__(
    self,
    latent_dimension: int,
    embedding_dimension: int,
    number_of_heads: int,
    number_of_layers: int,
    feedforward_dimension: int,
    device: str,
    observation_horizon: int = 1,
    algorithm_type: str = DenoisingAlgorithm.FLOW_MATCHING.value,
    sigma: float = 0.0,
    ode_solver: str = ODESolver.EULER.value,
    timestep_sampler: str = TimestepSampler.BETA.value,
    logit_mean: float = 0.0,
    logit_std: float = 1.0,
    beta_alpha: float = 1.5,
    beta_beta: float = 1.0,
    max_timestep: float = 0.999,
    num_train_timesteps: int = 100,
    num_inference_steps: int = 10,
    beta_start: float = 0.0001,
    beta_end: float = 0.02,
    beta_schedule: str = BetaSchedule.SQUAREDCOS_CAP_V2.value,
    scheduler_type: str = SchedulerType.DDIM.value,
    prediction_type: str = PredictionType.EPSILON.value,
    clip_sample: bool = False,
    variance_type: str | None = None,
    dropout: float = 0.1,
    normalization_type: str = NormalizationType.LAYER_NORM.value,
    attention_type: str = AttentionType.MULTI_HEAD.value,
    number_of_key_value_heads: int | None = None,
    activation: str = ActivationFunction.SILU.value,
    use_gating: bool = True,
    exclude_keys: list[str] | None = None,
    prior_target_key: str = LatentKey.POSTERIOR_MU.value,
    latent_standardization_enabled: bool = True,
    latent_standardization_eps: float = 1e-6,
    latent_standardization_max_batches: int | None = None,
    require_fitted_latent_standardization: bool = False,
):
    super().__init__(latent_dimension=latent_dimension, device=device)
    self.embedding_dimension = embedding_dimension
    self.observation_horizon = observation_horizon
    self.algorithm_type = algorithm_type
    self.num_train_timesteps = num_train_timesteps
    self.num_inference_steps = num_inference_steps
    self.exclude_keys = exclude_keys or []
    if prior_target_key not in {
        LatentKey.POSTERIOR_MU.value,
        LatentKey.POSTERIOR_LATENT.value,
    }:
        raise ValueError(
            f"Unsupported DiTPrior prior_target_key: {prior_target_key}. "
            f"Expected one of "
            f"{[LatentKey.POSTERIOR_MU.value, LatentKey.POSTERIOR_LATENT.value]}"
        )
    self.prior_target_key = prior_target_key
    self.timestep_sampling_config = TimestepSamplingConfig(
        sampler=timestep_sampler,
        logit_mean=logit_mean,
        logit_std=logit_std,
        beta_alpha=beta_alpha,
        beta_beta=beta_beta,
        max_timestep=max_timestep,
    )
    self.latent_standardizer = LatentStandardizer(
        latent_dimension=latent_dimension,
        enabled=latent_standardization_enabled,
        epsilon=latent_standardization_eps,
        require_fitted=require_fitted_latent_standardization,
    )
    if (
        latent_standardization_max_batches is not None
        and latent_standardization_max_batches <= 0
    ):
        raise ValueError(
            "latent_standardization_max_batches must be positive when set, "
            f"got {latent_standardization_max_batches}."
        )
    self.latent_standardization_max_batches = latent_standardization_max_batches

    if algorithm_type == DenoisingAlgorithm.FLOW_MATCHING.value:
        self.flow_matcher = ConditionalFlowMatcher(sigma=sigma)
        self.ode_solver = ode_solver
        self.noise_scheduler = None
    elif algorithm_type == DenoisingAlgorithm.DIFFUSION.value:
        self.flow_matcher = None
        self.ode_solver = None
        self.prediction_type = prediction_type
        scheduler_config = DiffusionSchedulerConfig(
            scheduler_type=scheduler_type,
            num_train_timesteps=num_train_timesteps,
            num_inference_steps=num_inference_steps,
            beta_start=beta_start,
            beta_end=beta_end,
            beta_schedule=beta_schedule,
            prediction_type=prediction_type,
            clip_sample=clip_sample,
            variance_type=variance_type,
        )
        self.noise_scheduler = create_noise_scheduler(scheduler_config)
    else:
        raise ValueError(
            f"Unknown algorithm_type: {algorithm_type}. "
            f"Expected one of {[e.value for e in DenoisingAlgorithm]}"
        )

    self.timestep_embed = SinusoidalPositionalEncoding1D(
        embedding_dimension=embedding_dimension,
        position_source=PositionSource.SCALAR.value,
        precompute_encodings=False,
    )
    self.timestep_mlp = nn.Sequential(
        nn.Linear(embedding_dimension, embedding_dimension),
        nn.SiLU(),
        nn.Linear(embedding_dimension, embedding_dimension),
    )
    self.latent_input_proj = nn.Linear(latent_dimension, embedding_dimension)
    self.final_layer = FinalPredictionLayer(
        hidden_dimension=embedding_dimension,
        output_dimension=latent_dimension,
        activation=activation,
    )
    temporal_positional_encoding = None
    if self.observation_horizon > 1:
        temporal_positional_encoding = LearnedPositionalEncoding1D(
            embedding_dimension=embedding_dimension
        )
    self.input_builder = TransformerInputBuilder(
        embedding_dimension=embedding_dimension,
        spatial_positional_encoding_layer=SinusoidalPositionalEncoding2D(
            embedding_dimension=embedding_dimension, normalize=True
        ),
        flat_positional_encoding_layer=SinusoidalPositionalEncoding1D(
            embedding_dimension=embedding_dimension,
            maximum_sequence_length=1000,
        ),
        temporal_positional_encoding_layer=temporal_positional_encoding,
    )
    self.decoder = ConditionalBidirectionalDecoder(
        number_of_layers=number_of_layers,
        embedding_dimension=embedding_dimension,
        conditioning_dimension=embedding_dimension,
        number_of_heads=number_of_heads,
        feedforward_dimension=feedforward_dimension,
        dropout=dropout,
        activation=activation,
        normalization_type=normalization_type,
        attention_type=attention_type,
        number_of_key_value_heads=number_of_key_value_heads,
        positional_encoding_type=None,  # Handled externally by input_builder
        use_cross_attention=False,
        use_gating=use_gating,
        use_final_normalization=False,  # FinalPredictionLayer has its own AdaNorm
    )

    def _init_module(module: nn.Module) -> None:
        if getattr(module, "_is_modulation_layer", False):
            return
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)

    self.latent_input_proj.apply(_init_module)
    self.timestep_mlp.apply(_init_module)
    self.to(torch.device(device))

timestep_sampler property

timestep_sampler

Configured flow-matching timestep sampler name.

logit_mean property

logit_mean

Configured logit-normal sampler mean.

logit_std property

logit_std

Configured logit-normal sampler standard deviation.

beta_alpha property

beta_alpha

Configured beta sampler alpha parameter.

beta_beta property

beta_beta

Configured beta sampler beta parameter.

max_timestep property

max_timestep

Configured maximum sampled timestep.

get_auxiliary_output_keys

get_auxiliary_output_keys()

DiT prior outputs denoising predictions and targets.

Source code in src/versatil/models/decoding/latent/prior/dit_prior.py

def get_auxiliary_output_keys(self) -> set[str]:
    """DiT prior outputs denoising predictions and targets."""
    return {
        LatentKey.PRIOR_LATENT.value,
        LatentKey.PRIOR_PREDICTION.value,
        LatentKey.PRIOR_TARGET.value,
    }

get_callbacks

get_callbacks(experiment_config)

Provide DiT prior training callbacks.

Parameters:

Name	Type	Description	Default
experiment_config	ExperimentConfig	Experiment-level callback configuration.	required

Returns:

Type	Description
list	Prior target standardization callback when enabled, otherwise no callbacks.

Source code in src/versatil/models/decoding/latent/prior/dit_prior.py

def get_callbacks(self, experiment_config: ExperimentConfig) -> list:
    """Provide DiT prior training callbacks.

    Args:
        experiment_config: Experiment-level callback configuration.

    Returns:
        Prior target standardization callback when enabled, otherwise no callbacks.
    """
    if not self.latent_standardizer.enabled:
        return []
    return [
        PriorTargetStandardizationCallback(
            max_batches=self.latent_standardization_max_batches
        )
    ]

build_training_target

build_training_target(posterior_output)

Select and detach the configured posterior target for DiT training.

Parameters:

Name	Type	Description	Default
posterior_output	dict[str, Tensor]	Posterior encoder output dictionary.	required

Returns:

Type	Description
Tensor	Detached latent target selected by prior_target_key.

Source code in src/versatil/models/decoding/latent/prior/dit_prior.py

def build_training_target(
    self, posterior_output: dict[str, torch.Tensor]
) -> torch.Tensor:
    """Select and detach the configured posterior target for DiT training.

    Args:
        posterior_output: Posterior encoder output dictionary.

    Returns:
        Detached latent target selected by ``prior_target_key``.
    """
    return posterior_output[self.prior_target_key].detach()

forward

forward(target_latents, observations)

Compute denoising predictions for training.

Parameters:

Name	Type	Description	Default
target_latents	Tensor | None	Clean latent samples from posterior (B, latent_dim).	required
observations	dict[str, Tensor]	Dictionary of conditioning features.	required

Returns:

Type	Description
dict[str, Tensor]	Dictionary with LatentKey.PRIOR_PREDICTION.value and LatentKey.PRIOR_TARGET.value.

Source code in src/versatil/models/decoding/latent/prior/dit_prior.py

def forward(
    self,
    target_latents: torch.Tensor | None,
    observations: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Compute denoising predictions for training.

    Args:
        target_latents: Clean latent samples from posterior (B, latent_dim).
        observations: Dictionary of conditioning features.

    Returns:
        Dictionary with LatentKey.PRIOR_PREDICTION.value and LatentKey.PRIOR_TARGET.value.
    """
    if target_latents is None:
        raise ValueError(
            "DiTPrior.forward() requires target_latents for denoising "
            "training. Use sample_prior() for inference."
        )
    standardized_target_latents = self.latent_standardizer.standardize(
        target_latents
    )
    batch_size = standardized_target_latents.shape[0]
    device = standardized_target_latents.device
    filtered_obs = self._filter_observations(observations)

    if self.algorithm_type == DenoisingAlgorithm.FLOW_MATCHING.value:
        noise = torch.randn_like(standardized_target_latents)  # (B, latent_dim)
        sampled_time = sample_timesteps_from_config(
            config=self.timestep_sampling_config,
            batch_size=batch_size,
            device=device,
        )  # (B,)
        (
            time,
            interpolated_latent,
            target_velocity,
        ) = self.flow_matcher.sample_location_and_conditional_flow(
            x0=noise, x1=standardized_target_latents, t=sampled_time
        )  # time: (B,), interpolated_latent: (B, latent_dim), target_velocity: (B, latent_dim)
        timestep_embedding = self._get_timestep_embedding_continuous(time)  # (B, D)
        predicted_velocity = self._predict_from_tokens(
            noisy_latent=interpolated_latent,
            observations=filtered_obs,
            timestep_embedding=timestep_embedding,
        )  # (B, latent_dim)
        return {
            LatentKey.PRIOR_PREDICTION.value: predicted_velocity,
            LatentKey.PRIOR_TARGET.value: target_velocity,
        }

    timesteps = sample_random_timesteps(
        batch_size=batch_size,
        num_train_timesteps=self.num_train_timesteps,
        device=device,
    )  # (B,)
    noisy_latents, noise = add_noise_to_tensor(
        clean=standardized_target_latents,
        noise_scheduler=self.noise_scheduler,
        timesteps=timesteps,
    )  # noisy_latents: (B, latent_dim), noise: (B, latent_dim)
    timestep_embedding = self._get_timestep_embedding(timesteps)  # (B, D)
    model_output = self._predict_from_tokens(
        noisy_latent=noisy_latents,
        observations=filtered_obs,
        timestep_embedding=timestep_embedding,
    )  # (B, latent_dim)
    if self.prediction_type == PredictionType.EPSILON.value:
        target = noise
    elif self.prediction_type == PredictionType.SAMPLE.value:
        target = standardized_target_latents
    elif self.prediction_type == PredictionType.VELOCITY.value:
        target = self.noise_scheduler.get_velocity(
            sample=standardized_target_latents, noise=noise, timesteps=timesteps
        )
    else:
        raise ValueError(
            f"Unknown prediction_type: {self.prediction_type}. "
            f"Expected one of {[e.value for e in PredictionType]}"
        )
    return {
        LatentKey.PRIOR_PREDICTION.value: model_output,
        LatentKey.PRIOR_TARGET.value: target,
    }

sample_prior

sample_prior(batch_size, observations=None)

Sample latent variable from learned prior.

Parameters:

Name	Type	Description	Default
batch_size	int	Number of samples to generate.	required
observations	dict[str, Tensor] | None	Dictionary of conditioning features.	None

Returns:

Type	Description
Tensor	Sampled latent embeddings (batch_size, latent_dim).

Source code in src/versatil/models/decoding/latent/prior/dit_prior.py

def sample_prior(
    self,
    batch_size: int,
    observations: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Sample latent variable from learned prior.

    Args:
        batch_size: Number of samples to generate.
        observations: Dictionary of conditioning features.

    Returns:
        Sampled latent embeddings (batch_size, latent_dim).
    """
    device = next(self.parameters()).device
    if observations is None:
        observations = {}
    filtered_obs = self._filter_observations(observations)
    z = torch.randn(
        batch_size, self.latent_dimension, device=device
    )  # (B, latent_dim)

    if self.algorithm_type == DenoisingAlgorithm.FLOW_MATCHING.value:

        def velocity_fn(
            current_latent: torch.Tensor, current_time: torch.Tensor
        ) -> torch.Tensor:
            """Predict the flow velocity at a continuous timestep."""
            timestep_embedding = self._get_timestep_embedding_continuous(
                current_time
            )  # (B, D)
            return self._predict_from_tokens(
                noisy_latent=current_latent,
                observations=filtered_obs,
                timestep_embedding=timestep_embedding,
            )  # (B, latent_dim)

        standardized_sample = integrate_ode(
            z_init=z,
            velocity_fn=velocity_fn,
            num_steps=self.num_inference_steps,
            solver=self.ode_solver,
        )  # (B, latent_dim)
        return self.latent_standardizer.unstandardize(standardized_sample)

    setup_inference_timesteps(self.noise_scheduler, self.num_inference_steps)
    for t in self.noise_scheduler.timesteps:
        timesteps = torch.full(
            (batch_size,), t, device=device, dtype=torch.long
        )  # (B,)
        timestep_embedding = self._get_timestep_embedding(timesteps)  # (B, D)
        model_output = self._predict_from_tokens(
            noisy_latent=z,
            observations=filtered_obs,
            timestep_embedding=timestep_embedding,
        )  # (B, latent_dim)
        z = self.noise_scheduler.step(
            model_output=model_output, timestep=t, sample=z
        ).prev_sample  # (B, latent_dim)
    return self.latent_standardizer.unstandardize(z)