pooling_head

pooling_head

Pooling strategies for spatial feature maps and token sequences.

PoolingHead

PoolingHead(input_dimension)

Bases: Module, ABC

Abstract base class for pooling operations on spatial feature maps or token sequences.

Parameters:

Name	Type	Description	Default
input_dimension	int	Feature vector size, i.e. channel count for spatial feature maps (B, C, H, W), or hidden dimension for token sequences (B, S, D).	required

Source code in src/versatil/models/layers/pooling/pooling_head.py

def __init__(self, input_dimension: int):
    super().__init__()
    self.input_dimension = input_dimension

output_dim abstractmethod property

output_dim

Output dimension after pooling.

forward abstractmethod

forward(features)

Apply pooling to input features.

Source code in src/versatil/models/layers/pooling/pooling_head.py

@abstractmethod
def forward(self, features: torch.Tensor) -> torch.Tensor:
    """Apply pooling to input features."""
    raise NotImplementedError

SpatialSoftmaxPooling

SpatialSoftmaxPooling(input_dimension, spatial_height, spatial_width)

Bases: PoolingHead

Spatial softmax pooling on feature maps.

Initialize spatial softmax pooling.

Parameters:

Name	Type	Description	Default
input_dimension	int	Number of feature channels.	required
spatial_height	int	Height of the feature map.	required
spatial_width	int	Width of the feature map.	required

Source code in src/versatil/models/layers/pooling/pooling_head.py

def __init__(self, input_dimension: int, spatial_height: int, spatial_width: int):
    """Initialize spatial softmax pooling.

    Args:
        input_dimension: Number of feature channels.
        spatial_height: Height of the feature map.
        spatial_width: Width of the feature map.
    """
    super().__init__(input_dimension=input_dimension)
    self.spatial_softmax = SpatialSoftmax(
        spatial_height, spatial_width, input_dimension
    )

output_dim property

output_dim

Two coordinates (x, y) per input channel.

forward

forward(features)

Compute per-channel spatial softmax keypoints from (B, C, H, W).

Source code in src/versatil/models/layers/pooling/pooling_head.py

def forward(self, features: torch.Tensor) -> torch.Tensor:
    """Compute per-channel spatial softmax keypoints from (B, C, H, W)."""
    result: torch.Tensor = self.spatial_softmax(features)
    return result

GlobalAveragePooling

GlobalAveragePooling(input_dimension)

Bases: PoolingHead

Global average pooling over spatial dimensions.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def __init__(self, input_dimension: int):
    super().__init__()
    self.input_dimension = input_dimension

output_dim property

output_dim

Channel count of the pooled vector.

forward

forward(features)

Average (B, C, H, W) over the spatial dimensions.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def forward(self, features: torch.Tensor) -> torch.Tensor:
    """Average (B, C, H, W) over the spatial dimensions."""
    return features.mean(dim=[2, 3])

MaxPooling

MaxPooling(input_dimension)

Bases: PoolingHead

Global max pooling over spatial dimensions.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def __init__(self, input_dimension: int):
    super().__init__()
    self.input_dimension = input_dimension

output_dim property

output_dim

Channel count of the pooled vector.

forward

forward(features)

Max-reduce (B, C, H, W) over the spatial dimensions.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def forward(self, features: torch.Tensor) -> torch.Tensor:
    """Max-reduce (B, C, H, W) over the spatial dimensions."""
    return torch.amax(features, dim=[2, 3])

SpatialIdentityPooling

SpatialIdentityPooling(input_dimension)

Bases: PoolingHead

No pooling — returns spatial feature maps unchanged.

output_dim returns (C, -1, -1) where -1 indicates dynamic spatial dimensions resolved at forward time from the actual feature map.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def __init__(self, input_dimension: int):
    super().__init__()
    self.input_dimension = input_dimension

output_dim property

output_dim

Channel count with dynamic spatial dimensions.

forward

forward(features)

Return the spatial feature maps unchanged.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def forward(self, features: torch.Tensor) -> torch.Tensor:
    """Return the spatial feature maps unchanged."""
    return features

SpatialLearnedAggregationPooling

SpatialLearnedAggregationPooling(input_dimension)

Bases: PoolingHead

Learned aggregation of spatial feature maps through attention.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def __init__(self, input_dimension: int):
    super().__init__(input_dimension=input_dimension)
    self.learned_aggregation = LearnedAggregation(feature_dimension=input_dimension)

output_dim property

output_dim

Channel count of the aggregated vector.

forward

forward(features)

Aggregate (B, C, H, W) through learned attention.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def forward(self, features: torch.Tensor) -> torch.Tensor:
    """Aggregate (B, C, H, W) through learned attention."""
    return self.learned_aggregation(features)

TokenPoolingHead

TokenPoolingHead(input_dimension, pooling_method, sequence_length=-1, num_prefix_tokens=0)

Bases: PoolingHead

Pooling head for token sequences (B, S, D).

Reduces a sequence of token embeddings to a single vector (B, D) via CLS token selection, mean pooling, or learned aggregation. With pooling_method=NONE, returns the sequence with prefix tokens stripped.

Parameters:

Name	Type	Description	Default
input_dimension	int	Hidden dimension of the token embeddings.	required
pooling_method	str	Pooling strategy from PoolingMethod enum.	required
sequence_length	int	Fixed sequence length for NONE output dim (-1 for variable).	-1
num_prefix_tokens	int	Number of prefix tokens (CLS, registers) to exclude from AVERAGE, LEARNED_AGGREGATION, and NONE pooling. The first prefix token is still used for DEFAULT (CLS) pooling.	0

Source code in src/versatil/models/layers/pooling/pooling_head.py

def __init__(
    self,
    input_dimension: int,
    pooling_method: str,
    sequence_length: int = -1,
    num_prefix_tokens: int = 0,
):
    super().__init__(input_dimension=input_dimension)
    self.pooling_method = pooling_method
    self.sequence_length = sequence_length
    self.num_prefix_tokens = num_prefix_tokens
    self.learned_aggregation: LearnedAggregation | None = None
    if pooling_method == PoolingMethod.LEARNED_AGGREGATION.value:
        self.learned_aggregation = LearnedAggregation(
            feature_dimension=input_dimension
        )

output_dim property

output_dim

Feature dimension, with sequence length when pooling is disabled.

forward

forward(hidden_states, padding_mask=None)

Pool token sequence.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Token embeddings of shape (B, S, D).	required
padding_mask	Tensor | None	Optional padding mask of shape (B, S), where True means padded. Used by pooling methods that aggregate tokens.	None

Returns:

Type	Description
Tensor	Pooled features of shape (B, D) or (B, S', D) for NONE.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def forward(
    self,
    hidden_states: torch.Tensor,
    padding_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Pool token sequence.

    Args:
        hidden_states: Token embeddings of shape (B, S, D).
        padding_mask: Optional padding mask of shape (B, S), where True
            means padded. Used by pooling methods that aggregate tokens.

    Returns:
        Pooled features of shape (B, D) or (B, S', D) for NONE.
    """
    start = self.num_prefix_tokens
    match self.pooling_method:
        case PoolingMethod.DEFAULT.value:
            return hidden_states[:, 0]  # CLS token
        case PoolingMethod.AVERAGE.value:
            token_padding_mask = self._slice_padding_mask(
                padding_mask=padding_mask,
                hidden_states=hidden_states,
            )
            return self._masked_average(
                tokens=hidden_states[:, start:],
                padding_mask=token_padding_mask,
            )
        case PoolingMethod.LEARNED_AGGREGATION.value:
            token_padding_mask = self._slice_padding_mask(
                padding_mask=padding_mask,
                hidden_states=hidden_states,
            )
            return self.learned_aggregation(
                hidden_states[:, start:],
                padding_mask=token_padding_mask,
            )
        case PoolingMethod.NONE.value:
            return hidden_states[:, start:]
        case _:
            raise ValueError(
                f"Unsupported token pooling method: {self.pooling_method}. "
                f"Supported: {[e.value for e in PoolingMethod]}"
            )

create_spatial_pooling_head

create_spatial_pooling_head(pooling_method, input_dimension, spatial_height, spatial_width)

Create a pooling head for spatial feature maps (B, C, H, W).

Parameters:

Name	Type	Description	Default
pooling_method	str	Pooling strategy from PoolingMethod enum.	required
input_dimension	int	Number of feature channels.	required
spatial_height	int	Height of the feature map.	required
spatial_width	int	Width of the feature map.	required

Returns:

Type	Description
PoolingHead	Configured spatial pooling head.

Raises:

Type	Description
ValueError	If pooling_method is not supported for spatial features.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def create_spatial_pooling_head(
    pooling_method: str,
    input_dimension: int,
    spatial_height: int,
    spatial_width: int,
) -> PoolingHead:
    """Create a pooling head for spatial feature maps (B, C, H, W).

    Args:
        pooling_method: Pooling strategy from PoolingMethod enum.
        input_dimension: Number of feature channels.
        spatial_height: Height of the feature map.
        spatial_width: Width of the feature map.

    Returns:
        Configured spatial pooling head.

    Raises:
        ValueError: If pooling_method is not supported for spatial features.
    """
    match pooling_method:
        case PoolingMethod.SPATIAL_SOFTMAX.value:
            return SpatialSoftmaxPooling(
                input_dimension=input_dimension,
                spatial_height=spatial_height,
                spatial_width=spatial_width,
            )
        case PoolingMethod.AVERAGE.value:
            return GlobalAveragePooling(input_dimension=input_dimension)
        case PoolingMethod.MAX.value | PoolingMethod.DEFAULT.value:
            return MaxPooling(input_dimension=input_dimension)
        case PoolingMethod.NONE.value:
            return SpatialIdentityPooling(input_dimension=input_dimension)
        case PoolingMethod.LEARNED_AGGREGATION.value:
            return SpatialLearnedAggregationPooling(input_dimension=input_dimension)
        case _:
            raise ValueError(
                f"Unsupported spatial pooling method: {pooling_method}. "
                f"Supported: {[e.value for e in PoolingMethod]}"
            )

create_token_pooling_head

create_token_pooling_head(pooling_method, input_dimension, sequence_length=-1, num_prefix_tokens=0)

Create a pooling head for token sequences (B, S, D).

Parameters:

Name	Type	Description	Default
pooling_method	str	Pooling strategy from PoolingMethod enum.	required
input_dimension	int	Hidden dimension of the token embeddings.	required
sequence_length	int	Fixed sequence length for NONE output dim (-1 for variable).	-1
num_prefix_tokens	int	Number of prefix tokens (CLS, registers) to strip.	0

Returns:

Type	Description
TokenPoolingHead	Configured token pooling head.

Source code in src/versatil/models/layers/pooling/pooling_head.py

def create_token_pooling_head(
    pooling_method: str,
    input_dimension: int,
    sequence_length: int = -1,
    num_prefix_tokens: int = 0,
) -> TokenPoolingHead:
    """Create a pooling head for token sequences (B, S, D).

    Args:
        pooling_method: Pooling strategy from PoolingMethod enum.
        input_dimension: Hidden dimension of the token embeddings.
        sequence_length: Fixed sequence length for NONE output dim (-1 for variable).
        num_prefix_tokens: Number of prefix tokens (CLS, registers) to strip.

    Returns:
        Configured token pooling head.
    """
    return TokenPoolingHead(
        input_dimension=input_dimension,
        pooling_method=pooling_method,
        sequence_length=sequence_length,
        num_prefix_tokens=num_prefix_tokens,
    )