geometric_attention_encoder

geometric_attention_encoder

Taken from DFormerV2 paper: https://arxiv.org/pdf/2504.04701

GeometricFeedForwardNetwork

GeometricFeedForwardNetwork(embedding_dimension, ffn_dimension)

Bases: Module

DFormerv2-style feed-forward network with an inner depthwise convolution.

Mirrors the reference FeedForwardNetwork: fc1 -> GELU -> 3x3 depthwise convolution with an inner residual -> fc2. The convolution injects local spatial mixing that a plain MLP lacks, and pretrained DFormerv2 checkpoints carry its weights.

Initialize the feed-forward network.

Parameters:

Name	Type	Description	Default
embedding_dimension	int	Input and output feature dimension.	required
ffn_dimension	int	Hidden dimension between the two linears.	required

Source code in src/versatil/models/layers/geometric_attention/geometric_attention_encoder.py

def __init__(self, embedding_dimension: int, ffn_dimension: int):
    """Initialize the feed-forward network.

    Args:
        embedding_dimension: Input and output feature dimension.
        ffn_dimension: Hidden dimension between the two linears.
    """
    super().__init__()
    self.fc1 = nn.Linear(embedding_dimension, ffn_dimension)
    self.fc2 = nn.Linear(ffn_dimension, embedding_dimension)
    self.dwconv = DepthwiseConv2D(
        dimension=ffn_dimension, kernel_size=3, stride=1, padding=1
    )
    self.activation = nn.GELU()

forward

forward(features)

Apply the feed-forward network to (B, H, W, C) features.

Source code in src/versatil/models/layers/geometric_attention/geometric_attention_encoder.py

def forward(self, features: torch.Tensor) -> torch.Tensor:
    """Apply the feed-forward network to (B, H, W, C) features."""
    features = self.activation(self.fc1(features))
    residual = features
    features = self.dwconv(features) + residual
    return self.fc2(features)

GeometricAttentionEncoderBlock

GeometricAttentionEncoderBlock(decomposition_mode, embedding_dimension, number_of_heads, ffn_dimension, drop_path_rate=0.0, use_layer_scale=False, layer_scale_init_value=1e-05, initial_decay=2.0, decay_range=4.0, value_dimension_factor=1, depthwise_kernel_size=5, depthwise_padding=2, input_positional_kernel_size=3, input_positional_padding=1, use_raster_positions=False, use_feedforward_convolution=False)

Bases: Module

Geometric attention encoder block for conditioning with depth maps on RGB images.

Integrates depth-conditioned attention, feed-forward network, and optional layer scaling for residual connections.

Initializes the geometric attention encoder block.

Parameters:

Name	Type	Description	Default
decomposition_mode	AttentionDecompositionMode	Attention mode (full or separable).	required
embedding_dimension	int	Feature dimension.	required
number_of_heads	int	Number of attention heads.	required
ffn_dimension	int	Hidden dimension for the fully-connected layer that follows the self-attention layer.	required
drop_path_rate	float	Stochastic depth rate.	0.0
use_layer_scale	bool	Whether to use layer scaling.	False
layer_scale_init_value	float	Initial value for layer scale parameters.	1e-05
initial_decay	float	Initial decay rate for spatial biases.	2.0
decay_range	float	Range of decay rates across heads.	4.0
value_dimension_factor	int	Expansion factor for value dimension.	1
depthwise_kernel_size	int	Kernel size for value positional encoding.	5
depthwise_padding	int	Padding for value positional encoding.	2
input_positional_kernel_size	int	Kernel size for input positional encoding.	3
input_positional_padding	int	Padding for input positional encoding.	1
use_raster_positions	bool	Whether rotary encoding uses flattened raster grid positions (the DFormerv2 reference convention).	False
use_feedforward_convolution	bool	Whether the feed-forward network uses the DFormerv2 inner depthwise convolution instead of a plain MLP. Required for pretrained DFormerv2 checkpoints.	False

Source code in src/versatil/models/layers/geometric_attention/geometric_attention_encoder.py

def __init__(
    self,
    decomposition_mode: AttentionDecompositionMode,
    embedding_dimension: int,
    number_of_heads: int,
    ffn_dimension: int,
    drop_path_rate: float = 0.0,
    use_layer_scale: bool = False,
    layer_scale_init_value: float = 1e-5,
    initial_decay: float = 2.0,
    decay_range: float = 4.0,
    value_dimension_factor: int = 1,
    depthwise_kernel_size: int = 5,
    depthwise_padding: int = 2,
    input_positional_kernel_size: int = 3,
    input_positional_padding: int = 1,
    use_raster_positions: bool = False,
    use_feedforward_convolution: bool = False,
):
    """Initializes the geometric attention encoder block.

    Args:
        decomposition_mode: Attention mode (full or separable).
        embedding_dimension: Feature dimension.
        number_of_heads: Number of attention heads.
        ffn_dimension: Hidden dimension for the fully-connected layer that follows the self-attention layer.
        drop_path_rate: Stochastic depth rate.
        use_layer_scale: Whether to use layer scaling.
        layer_scale_init_value: Initial value for layer scale parameters.
        initial_decay: Initial decay rate for spatial biases.
        decay_range: Range of decay rates across heads.
        value_dimension_factor: Expansion factor for value dimension.
        depthwise_kernel_size: Kernel size for value positional encoding.
        depthwise_padding: Padding for value positional encoding.
        input_positional_kernel_size: Kernel size for input positional encoding.
        input_positional_padding: Padding for input positional encoding.
        use_raster_positions: Whether rotary encoding uses flattened raster
            grid positions (the DFormerv2 reference convention).
        use_feedforward_convolution: Whether the feed-forward network uses
            the DFormerv2 inner depthwise convolution instead of a plain
            MLP. Required for pretrained DFormerv2 checkpoints.
    """
    super().__init__()
    self.use_layer_scale = use_layer_scale
    self.embedding_dimension = embedding_dimension

    self.norm1 = nn.LayerNorm(embedding_dimension, eps=1e-6)
    self.norm2 = nn.LayerNorm(embedding_dimension, eps=1e-6)

    self.attention = GeometricSelfAttention(
        embedding_dimension=embedding_dimension,
        number_of_heads=number_of_heads,
        value_dimension_factor=value_dimension_factor,
        decomposition_mode=decomposition_mode.value,
        initial_decay=initial_decay,
        decay_range=decay_range,
        depthwise_convolution_kernel_size=depthwise_kernel_size,
        depthwise_convolution_padding=depthwise_padding,
        use_raster_positions=use_raster_positions,
    )

    self.drop_path = DropPath(drop_path_rate)
    self.mlp: nn.Module
    if use_feedforward_convolution:
        self.mlp = GeometricFeedForwardNetwork(
            embedding_dimension=embedding_dimension,
            ffn_dimension=ffn_dimension,
        )
    else:
        self.mlp = MLP(
            input_dimension=embedding_dimension,
            hidden_dimensions=[ffn_dimension],
            output_dim=embedding_dimension,
            activation_function=nn.GELU,
            dropout=0.0,
        )

    self.input_positional_encoding = DepthwiseConv2D(
        dimension=embedding_dimension,
        kernel_size=input_positional_kernel_size,
        stride=1,
        padding=input_positional_padding,
    )

    if use_layer_scale:
        self.gamma1 = nn.Parameter(
            layer_scale_init_value * torch.ones(1, 1, 1, embedding_dimension),
            requires_grad=True,
        )
        self.gamma2 = nn.Parameter(
            layer_scale_init_value * torch.ones(1, 1, 1, embedding_dimension),
            requires_grad=True,
        )

forward

forward(rgb_tensor, depth_map)

Applies the geometric attention encoder block.

Parameters:

Name	Type	Description	Default
rgb_tensor	Tensor	Input images of shape (B, H, W, C).	required
depth_map	Tensor	Depth map of shape (B, 1, H, W).	required

Returns:

Type	Description
Tensor	Updated features of shape (B, H, W, C).

Source code in src/versatil/models/layers/geometric_attention/geometric_attention_encoder.py

def forward(
    self, rgb_tensor: torch.Tensor, depth_map: torch.Tensor
) -> torch.Tensor:
    """Applies the geometric attention encoder block.

    Args:
        rgb_tensor: Input images of shape (B, H, W, C).
        depth_map: Depth map of shape (B, 1, H, W).

    Returns:
        Updated features of shape (B, H, W, C).
    """
    features = rgb_tensor + self.input_positional_encoding(rgb_tensor)

    residual = features
    features = self.norm1(features)
    attention_output = self.attention(features, depth_map)

    if self.use_layer_scale:
        features = residual + self.drop_path(self.gamma1 * attention_output)
    else:
        features = residual + self.drop_path(attention_output)

    residual = features
    features = self.norm2(features)
    mlp_output = self.mlp(features)

    if self.use_layer_scale:
        features = residual + self.drop_path(self.gamma2 * mlp_output)
    else:
        features = residual + self.drop_path(mlp_output)

    return features