dit_block_transformer

dit_block_transformer

'DiT Block' architecture inspired by the original implementation by Peebles and Xie, adapted for robotics in "The Ingredients for Robotic Diffusion Transformers" by Dasari et al.

The implementation by Dasari had several bugs/differences from the original DiT paper, which have been corrected here.

The architecture consists of an encoder-decoder transformer where the encoder processes pooled observation features, and the decoder generates action tokens conditioned on both the timestep embedding and the encoder output mean.

References

https://arxiv.org/html/2410.10088v1 https://github.com/SudeepDasari/dit-policy/blob/main/data4robotics/models/diffusion.py#L282 https://arxiv.org/abs/2212.09748 https://github.com/facebookresearch/DiT/blob/main/models.py

DiTBlock

DiTBlock(number_of_encoder_layers, number_of_decoder_layers, embedding_dimension, number_of_heads, number_of_key_value_heads=None, feedforward_dimension=None, dropout=0.1, attention_dropout=0.0, activation=value, normalization_type=value, attention_type=value, positional_encoding_type=None, maximum_sequence_length=2048, maximum_decoder_length=256, timestep_embedding_dimension=256, bias=True, normalization_epsilon=1e-06, use_gating=True, initializer_range=0.02)

Bases: Module

DiT-Block paper network architecture.

The encoder processes input tokens bidirectionally and pools output to a single vector. The decoder generates output tokens conditioned on (timestep + pooled encoder output) via AdaLN.

Shape notation

B: batch size S: encoder sequence length T: decoder sequence length D: embedding dimension

Initialize the DiffusionTransformer.

Parameters:

Name	Type	Description	Default
number_of_encoder_layers	int	Number of encoder layers.	required
number_of_decoder_layers	int	Number of decoder layers.	required
embedding_dimension	int	Hidden dimension of the transformer.	required
number_of_heads	int	Number of attention heads.	required
number_of_key_value_heads	int | None	Number of Key/Values heads (for Group Query Attention).	None
feedforward_dimension	int | None	Feedforward network hidden dimension.	None
dropout	float	Dropout rate.	0.1
attention_dropout	float	Dropout rate for attention.	0.0
activation	str	Activation function.	value
normalization_type	str	Type of normalization.	value
attention_type	str	Type of attention.	value
positional_encoding_type	str | None	Type of positional encoding for encoder.	None
maximum_sequence_length	int	Maximum encoder sequence length.	2048
maximum_decoder_length	int	Maximum decoder sequence length.	256
timestep_embedding_dimension	int	Dimension for timestep sinusoidal embedding.	256
bias	bool	Whether to use bias in linear layers.	True
normalization_epsilon	float	Epsilon for normalization layers.	1e-06
use_gating	bool	Whether to use gating in decoder AdaNorm (often referred to as AdaLNZeroNorm).	True
initializer_range	float	Standard deviation for weight initialization.	0.02

Source code in src/versatil/models/layers/diffusion_transformer/dit_block_transformer.py

def __init__(
    self,
    number_of_encoder_layers: int,
    number_of_decoder_layers: int,
    embedding_dimension: int,
    number_of_heads: int,
    number_of_key_value_heads: int | None = None,
    feedforward_dimension: int | None = None,
    dropout: float = 0.1,
    attention_dropout: float = 0.0,
    activation: str = ActivationFunction.SWIGLU.value,
    normalization_type: str = NormalizationType.RMS_NORM.value,
    attention_type: str = AttentionType.MULTI_HEAD.value,
    positional_encoding_type: str | None = None,
    maximum_sequence_length: int = 2048,
    maximum_decoder_length: int = 256,
    timestep_embedding_dimension: int = 256,
    bias: bool = True,
    normalization_epsilon: float = 1e-6,
    use_gating: bool = True,
    initializer_range: float = 0.02,
):
    """Initialize the DiffusionTransformer.

    Args:
        number_of_encoder_layers: Number of encoder layers.
        number_of_decoder_layers: Number of decoder layers.
        embedding_dimension: Hidden dimension of the transformer.
        number_of_heads: Number of attention heads.
        number_of_key_value_heads: Number of Key/Values heads (for Group Query Attention).
        feedforward_dimension: Feedforward network hidden dimension.
        dropout: Dropout rate.
        attention_dropout: Dropout rate for attention.
        activation: Activation function.
        normalization_type: Type of normalization.
        attention_type: Type of attention.
        positional_encoding_type: Type of positional encoding for encoder.
        maximum_sequence_length: Maximum encoder sequence length.
        maximum_decoder_length: Maximum decoder sequence length.
        timestep_embedding_dimension: Dimension for timestep sinusoidal embedding.
        bias: Whether to use bias in linear layers.
        normalization_epsilon: Epsilon for normalization layers.
        use_gating: Whether to use gating in decoder AdaNorm (often referred to as AdaLNZeroNorm).
        initializer_range: Standard deviation for weight initialization.
    """
    super().__init__()
    self.embedding_dimension = embedding_dimension
    self.number_of_encoder_layers = number_of_encoder_layers
    self.number_of_decoder_layers = number_of_decoder_layers
    self.initializer_range = initializer_range
    if feedforward_dimension is None:
        feedforward_dimension = 4 * embedding_dimension
    self.timestep_embedding_network = SinusoidalPositionalEncoding1D(
        embedding_dimension=timestep_embedding_dimension,
        denominator_mode=DenominatorMode.HALF_MINUS_ONE.value,
        ordering_mode=OrderingMode.CAT_COS_SIN.value,
        position_source=PositionSource.SCALAR.value,
        precompute_encodings=False,
        temperature=10000.0,
        learnable_frequencies=False,
        mlp_activation=nn.SiLU,
        mlp_hidden_dimensions=[embedding_dimension, embedding_dimension],
    )
    self.encoder = TransformerEncoder(
        number_of_layers=number_of_encoder_layers,
        embedding_dimension=embedding_dimension,
        number_of_heads=number_of_heads,
        number_of_key_value_heads=number_of_key_value_heads,
        feedforward_dimension=feedforward_dimension,
        dropout=dropout,
        attention_dropout=attention_dropout,
        activation=activation,
        normalization_type=normalization_type,
        attention_type=attention_type,
        positional_encoding_type=positional_encoding_type,
        maximum_sequence_length=maximum_sequence_length,
        bias=bias,
        normalization_epsilon=normalization_epsilon,
        initializer_range=initializer_range,
    )
    self.decoder = ConditionalBidirectionalDecoder(
        number_of_layers=number_of_decoder_layers,
        embedding_dimension=embedding_dimension,
        conditioning_dimension=embedding_dimension,
        number_of_heads=number_of_heads,
        number_of_key_value_heads=number_of_key_value_heads,
        feedforward_dimension=feedforward_dimension,
        dropout=dropout,
        attention_dropout=attention_dropout,
        activation=activation,
        normalization_type=normalization_type,
        attention_type=attention_type,
        positional_encoding_type=positional_encoding_type,
        maximum_sequence_length=maximum_decoder_length,
        bias=bias,
        normalization_epsilon=normalization_epsilon,
        use_cross_attention=False,
        use_gating=use_gating,
        initializer_range=initializer_range,
    )

forward

forward(decoder_hidden_states, timesteps, encoder_hidden_states, encoder_padding_mask=None, decoder_padding_mask=None, encoder_cache=None)

Forward pass through the transformer.

Parameters:

Name	Type	Description	Default
decoder_hidden_states	Tensor	Decoder input tokens (batch size (B), decoder sequence length (T), embedding dimension (D)).	required
timesteps	Tensor	Diffusion timesteps (B,).	required
encoder_hidden_states	Tensor	Encoder input tokens (B, encoder sequence length (S), D).	required
encoder_padding_mask	Tensor | None	Padding mask for encoder (B, S).	None
decoder_padding_mask	Tensor | None	Padding mask for decoder (B, T).	None
encoder_cache	Tensor | None	Precomputed encoder output mean (B, D) for inference.	None

Returns:

Type	Description
tuple[Tensor, Tensor, Tensor]	Tuple of (encoder_output_mean, decoder_output, conditioning): - encoder_output_mean: Mean of encoder outputs (B, D) for caching. - decoder_output: Decoder output tokens (B, T, D). - conditioning: Timestep plus encoder conditioning (B, D).

Source code in src/versatil/models/layers/diffusion_transformer/dit_block_transformer.py

def forward(
    self,
    decoder_hidden_states: torch.Tensor,
    timesteps: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    encoder_padding_mask: torch.Tensor | None = None,
    decoder_padding_mask: torch.Tensor | None = None,
    encoder_cache: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Forward pass through the transformer.

    Args:
        decoder_hidden_states: Decoder input tokens (batch size (B), decoder sequence length (T), embedding dimension (D)).
        timesteps: Diffusion timesteps (B,).
        encoder_hidden_states: Encoder input tokens (B, encoder sequence length (S), D).
        encoder_padding_mask: Padding mask for encoder (B, S).
        decoder_padding_mask: Padding mask for decoder (B, T).
        encoder_cache: Precomputed encoder output mean (B, D) for inference.

    Returns:
        Tuple of (encoder_output_mean, decoder_output, conditioning):
            - encoder_output_mean: Mean of encoder outputs (B, D) for caching.
            - decoder_output: Decoder output tokens (B, T, D).
            - conditioning: Timestep plus encoder conditioning (B, D).
    """
    return self.forward_features(
        decoder_hidden_states=decoder_hidden_states,
        timesteps=timesteps,
        encoder_hidden_states=encoder_hidden_states,
        encoder_padding_mask=encoder_padding_mask,
        decoder_padding_mask=decoder_padding_mask,
        encoder_cache=encoder_cache,
    )

forward_features

forward_features(decoder_hidden_states, timesteps, encoder_hidden_states, encoder_padding_mask=None, decoder_padding_mask=None, encoder_cache=None)

Return encoder cache, decoder hidden states, and conditioning.

Parameters:

Name	Type	Description	Default
decoder_hidden_states	Tensor	Decoder input tokens (B, T, D).	required
timesteps	Tensor	Timesteps (B,).	required
encoder_hidden_states	Tensor	Encoder input tokens (B, S, D).	required
encoder_padding_mask	Tensor | None	Padding mask (B, S).	None
decoder_padding_mask	Tensor | None	Padding mask (B, T).	None
encoder_cache	Tensor | None	Precomputed encoder output mean (B, D).	None

Returns:

Type	Description
Tensor	Tuple of encoder output mean, decoder hidden states, and
Tensor	conditioning with shapes (B, D), (B, T, D), and (B, D).

Source code in src/versatil/models/layers/diffusion_transformer/dit_block_transformer.py

def forward_features(
    self,
    decoder_hidden_states: torch.Tensor,
    timesteps: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    encoder_padding_mask: torch.Tensor | None = None,
    decoder_padding_mask: torch.Tensor | None = None,
    encoder_cache: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Return encoder cache, decoder hidden states, and conditioning.

    Args:
        decoder_hidden_states: Decoder input tokens (B, T, D).
        timesteps: Timesteps (B,).
        encoder_hidden_states: Encoder input tokens (B, S, D).
        encoder_padding_mask: Padding mask (B, S).
        decoder_padding_mask: Padding mask (B, T).
        encoder_cache: Precomputed encoder output mean (B, D).

    Returns:
        Tuple of encoder output mean, decoder hidden states, and
        conditioning with shapes ``(B, D)``, ``(B, T, D)``, and ``(B, D)``.
    """
    if encoder_cache is None:
        encoder_output_mean = self.forward_encoder(
            encoder_hidden_states,
            encoder_padding_mask,
        )
    else:
        encoder_output_mean = encoder_cache
    decoder_output, combined_conditioning = self.forward_decoder_features(
        hidden_states=decoder_hidden_states,
        timesteps=timesteps,
        encoder_output_mean=encoder_output_mean,
        padding_mask=decoder_padding_mask,
    )
    return encoder_output_mean, decoder_output, combined_conditioning

forward_encoder

forward_encoder(hidden_states, padding_mask=None)

Encode input tokens.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Input tokens (batch size (B), encoder sequence length (S), embedding dimension (D)).	required
padding_mask	Tensor | None	Padding mask (B, S) where True means masked.	None

Returns:

Type	Description
Tensor	Mean of encoder outputs (B, D) for conditioning the decoder.

Source code in src/versatil/models/layers/diffusion_transformer/dit_block_transformer.py

def forward_encoder(
    self,
    hidden_states: torch.Tensor,
    padding_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Encode input tokens.

    Args:
        hidden_states: Input tokens (batch size (B), encoder sequence length (S), embedding dimension (D)).
        padding_mask: Padding mask (B, S) where True means masked.

    Returns:
        Mean of encoder outputs (B, D) for conditioning the decoder.
    """
    encoder_output = self.encoder(
        hidden_states=hidden_states,
        padding_mask=padding_mask,
    )
    if padding_mask is None:
        encoder_output_mean = encoder_output.mean(dim=1)
    else:
        valid_mask = ~padding_mask
        valid_counts = valid_mask.sum(dim=1, keepdim=True).clamp_min(1)
        encoder_output = encoder_output * valid_mask.unsqueeze(-1).to(
            encoder_output.dtype
        )
        encoder_output_mean = encoder_output.sum(dim=1) / valid_counts.to(
            encoder_output.dtype
        )
    return encoder_output_mean

forward_decoder

forward_decoder(hidden_states, timesteps, encoder_output_mean, padding_mask=None)

Decode with timestep conditioning.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Decoder input tokens (batch size (B), decoder sequence length (T), embedding dimension (D)).	required
timesteps	Tensor	Timesteps (B,).	required
encoder_output_mean	Tensor	Mean encoder output (B, D).	required
padding_mask	Tensor | None	Padding mask (B, T).	None

Returns:

Type	Description
Tensor	Decoder output tokens and conditioning with shapes (B, T, D)
Tensor	and (B, D).

Source code in src/versatil/models/layers/diffusion_transformer/dit_block_transformer.py

def forward_decoder(
    self,
    hidden_states: torch.Tensor,
    timesteps: torch.Tensor,
    encoder_output_mean: torch.Tensor,
    padding_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Decode with timestep conditioning.

    Args:
        hidden_states: Decoder input tokens (batch size (B), decoder sequence length (T), embedding dimension (D)).
        timesteps: Timesteps (B,).
        encoder_output_mean: Mean encoder output (B, D).
        padding_mask: Padding mask (B, T).

    Returns:
        Decoder output tokens and conditioning with shapes ``(B, T, D)``
        and ``(B, D)``.
    """
    return self.forward_decoder_features(
        hidden_states=hidden_states,
        timesteps=timesteps,
        encoder_output_mean=encoder_output_mean,
        padding_mask=padding_mask,
    )

forward_decoder_features

forward_decoder_features(hidden_states, timesteps, encoder_output_mean, padding_mask=None)

Return decoder hidden states and combined conditioning.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Decoder input tokens (B, T, D).	required
timesteps	Tensor	Timesteps (B,).	required
encoder_output_mean	Tensor	Mean encoder output (B, D).	required
padding_mask	Tensor | None	Padding mask (B, T).	None

Returns:

Type	Description
Tensor	Decoder hidden states and conditioning, with shapes (B, T, D)
Tensor	and (B, D).

Source code in src/versatil/models/layers/diffusion_transformer/dit_block_transformer.py

def forward_decoder_features(
    self,
    hidden_states: torch.Tensor,
    timesteps: torch.Tensor,
    encoder_output_mean: torch.Tensor,
    padding_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return decoder hidden states and combined conditioning.

    Args:
        hidden_states: Decoder input tokens (B, T, D).
        timesteps: Timesteps (B,).
        encoder_output_mean: Mean encoder output (B, D).
        padding_mask: Padding mask (B, T).

    Returns:
        Decoder hidden states and conditioning, with shapes ``(B, T, D)``
        and ``(B, D)``.
    """
    timestep_embedding = self.timestep_embedding_network(timesteps)  # (B, D)
    combined_conditioning = timestep_embedding + encoder_output_mean
    decoder_output = self.decoder(
        hidden_states=hidden_states,
        condition=combined_conditioning,
        query_padding_mask=padding_mask,
    )
    return decoder_output, combined_conditioning