spatial_decay

spatial_decay

SpatialDecayMask

SpatialDecayMask(number_of_heads, initial_decay=5.0, decay_range=3.0)

Bases: Module

Generates attention decay based on spatial distance between positions.

Farther positions receive exponentially decaying attention weights, with per-head decay rates allowing different receptive fields.

Initializes spatial decay mask generator.

Parameters:

Name	Type	Description	Default
number_of_heads	int	Number of attention heads.	required
initial_decay	float	Initial decay rate.	5.0
decay_range	float	Range of decay rates across heads.	3.0

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

def __init__(
    self, number_of_heads: int, initial_decay: float = 5.0, decay_range: float = 3.0
):
    """Initializes spatial decay mask generator.

    Args:
        number_of_heads: Number of attention heads.
        initial_decay: Initial decay rate.
        decay_range: Range of decay rates across heads.
    """
    super().__init__()
    self.number_of_heads = number_of_heads

    decay_rates = self._compute_per_head_decay(
        number_of_heads=number_of_heads,
        initial_decay=initial_decay,
        decay_range=decay_range,
    )
    self.register_buffer("decay_rates", decay_rates)

compute_2d_distance_matrix

compute_2d_distance_matrix(height, width)

Computes pairwise Manhattan distances for 2D grid.

Parameters:

Name	Type	Description	Default
height	int	Grid height.	required
width	int	Grid width.	required

Returns:

Type	Description
Tensor	Distance matrix of shape (HW, HW).

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

def compute_2d_distance_matrix(self, height: int, width: int) -> torch.Tensor:
    """Computes pairwise Manhattan distances for 2D grid.

    Args:
        height: Grid height.
        width: Grid width.

    Returns:
        Distance matrix of shape (H*W, H*W).
    """
    height_indices = torch.arange(height, device=self.decay_rates.device)
    width_indices = torch.arange(width, device=self.decay_rates.device)

    grid = torch.stack(
        torch.meshgrid(height_indices, width_indices, indexing="ij"), dim=-1
    )
    grid_flat = grid.reshape(height * width, 2)

    distance_matrix = grid_flat[:, None, :] - grid_flat[None, :, :]
    distance_matrix = distance_matrix.abs().sum(dim=-1)

    return distance_matrix

compute_1d_distance_matrix

compute_1d_distance_matrix(length)

Computes pairwise distances for 1D sequence.

Parameters:

Name	Type	Description	Default
length	int	Sequence length.	required

Returns:

Type	Description
Tensor	Distance matrix of shape (length, length).

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

def compute_1d_distance_matrix(self, length: int) -> torch.Tensor:
    """Computes pairwise distances for 1D sequence.

    Args:
        length: Sequence length.

    Returns:
        Distance matrix of shape (length, length).
    """
    indices = torch.arange(length, device=self.decay_rates.device)
    distance_matrix = (indices[:, None] - indices[None, :]).abs()
    return distance_matrix

forward

forward(height, width, decomposition_mode=value)

Generates spatial decay mask(s).

Parameters:

Name	Type	Description	Default
height	int	Grid height.	required
width	int	Grid width.	required
decomposition_mode	str	Whether to generate full or separable masks.	value

Returns:

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

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

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

    Args:
        height: Grid height.
        width: Grid width.
        decomposition_mode: Whether to generate full or separable masks.

    Returns:
        If FULL: Single mask of shape (number_of_heads, H*W, H*W).
        If SEPARABLE: Tuple of (height_mask, width_mask) each (number_of_heads, H, H) and (number_of_heads, W, W).
    """
    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_distances = self.compute_1d_distance_matrix(height)
        width_distances = self.compute_1d_distance_matrix(width)

        height_mask = height_distances * self.decay_rates[:, None, None]
        width_mask = width_distances * self.decay_rates[:, None, None]

        return height_mask, width_mask
    distances = self.compute_2d_distance_matrix(height, width)
    mask = distances * self.decay_rates[:, None, None]
    return (mask,)