depth_decay

depth_decay

DepthAwareDecayMask

DepthAwareDecayMask(number_of_heads)

Bases: Module

Generates attention decay based on depth discontinuities.

Reduces attention across object boundaries by penalizing depth differences between positions.

Initializes depth-aware decay mask.

Parameters:

Name	Type	Description	Default
number_of_heads	int	Number of attention heads.	required

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

def __init__(self, number_of_heads: int):
    """Initializes depth-aware decay mask.

    Args:
        number_of_heads: Number of attention heads.
    """
    super().__init__()
    self.number_of_heads = number_of_heads

compute_depth_difference_matrix staticmethod

compute_depth_difference_matrix(depth_map, height, width)

Computes pairwise depth differences.

Parameters:

Name	Type	Description	Default
depth_map	Tensor	Depth values of shape (B, 1, H, W).	required
height	int	Target height.	required
width	int	Target width.	required

Returns:

Type	Description
Tensor	Depth difference matrix of shape (B, HW, HW).

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

@staticmethod
def compute_depth_difference_matrix(
    depth_map: torch.Tensor, height: int, width: int
) -> torch.Tensor:
    """Computes pairwise depth differences.

    Args:
        depth_map: Depth values of shape (B, 1, H, W).
        height: Target height.
        width: Target width.

    Returns:
        Depth difference matrix of shape (B, H*W, H*W).
    """
    depth_map = F.interpolate(
        depth_map, size=(height, width), mode="bilinear", align_corners=False
    )

    batch_size = depth_map.shape[0]
    depth_flat = depth_map.reshape(batch_size, height * width, 1)
    depth_differences = depth_flat[:, :, None, :] - depth_flat[:, None, :, :]
    depth_differences = depth_differences.abs().sum(dim=-1)
    return depth_differences

compute_1d_depth_difference_matrix staticmethod

compute_1d_depth_difference_matrix(depth_map, axis, height, width)

Computes depth differences along one axis.

Parameters:

Name	Type	Description	Default
depth_map	Tensor	Depth values of shape (B, 1, H, W).	required
axis	str	Either 'height' or 'width'.	required
height	int	Target height.	required
width	int	Target width.	required

Returns:

Type	Description
Tensor	Depth differences of shape (B, secondary_length, primary_length, primary_length).

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

@staticmethod
def compute_1d_depth_difference_matrix(
    depth_map: torch.Tensor, axis: str, height: int, width: int
) -> torch.Tensor:
    """Computes depth differences along one axis.

    Args:
        depth_map: Depth values of shape (B, 1, H, W).
        axis: Either 'height' or 'width'.
        height: Target height.
        width: Target width.

    Returns:
        Depth differences of shape (B, secondary_length, primary_length, primary_length).
    """
    depth_map = F.interpolate(
        depth_map, size=(height, width), mode="bilinear", align_corners=False
    )
    if axis == Axis.HEIGHT.value:
        depth_map = depth_map.transpose(-2, -1)
    depth_differences = depth_map[:, :, :, :, None] - depth_map[:, :, :, None, :]
    depth_differences = depth_differences.abs()
    return depth_differences.squeeze(1)

forward

forward(depth_map, height, width, decay_rates, decomposition_mode=value)

Generates depth-aware decay mask(s).

Parameters:

Name	Type	Description	Default
depth_map	Tensor	Depth map of shape (B, 1, H, W).	required
height	int	Target height.	required
width	int	Target width.	required
decay_rates	Tensor	Per-head decay rates of shape (number_of_heads,).	required
decomposition_mode	str	Whether to generate full or separable masks.	value

Returns:

Type	Description
Tensor	If FULL: Single mask of shape (B, number_of_heads, HW, HW).
...	If SEPARABLE: Tuple of (height_mask, width_mask).

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

def forward(
    self,
    depth_map: torch.Tensor,
    height: int,
    width: int,
    decay_rates: torch.Tensor,
    decomposition_mode: str = AttentionDecompositionMode.FULL.value,
) -> tuple[torch.Tensor, ...]:
    """Generates depth-aware decay mask(s).

    Args:
        depth_map: Depth map of shape (B, 1, H, W).
        height: Target height.
        width: Target width.
        decay_rates: Per-head decay rates of shape (number_of_heads,).
        decomposition_mode: Whether to generate full or separable masks.

    Returns:
        If FULL: Single mask of shape (B, number_of_heads, H*W, H*W).
        If SEPARABLE: Tuple of (height_mask, width_mask).
    """
    valid_modes = [mode.value for mode in AttentionDecompositionMode]
    if decomposition_mode not in valid_modes:
        raise ValueError(
            f"Unknown decomposition_mode '{decomposition_mode}'; "
            f"expected one of {valid_modes}."
        )
    if decomposition_mode == AttentionDecompositionMode.SEPARABLE.value:
        height_depth_diffs = self.compute_1d_depth_difference_matrix(
            depth_map, axis=Axis.HEIGHT.value, height=height, width=width
        )
        width_depth_diffs = self.compute_1d_depth_difference_matrix(
            depth_map, axis=Axis.WIDTH.value, height=height, width=width
        )

        height_mask = (
            height_depth_diffs.unsqueeze(1) * decay_rates[None, :, None, None, None]
        )
        width_mask = (
            width_depth_diffs.unsqueeze(1) * decay_rates[None, :, None, None, None]
        )

        return height_mask, width_mask
    else:
        depth_diffs = self.compute_depth_difference_matrix(depth_map, height, width)
        mask = depth_diffs.unsqueeze(1) * decay_rates[None, :, None, None]
        return (mask,)