Add configurable bottleneck blocks and aggregation strategies

[mbsantiago]

Summary

• Add VerticalMean as a lightweight frequency aggregation block alongside the existing VerticalConv.
• Add EfficientSelfAttention, using fused QKV projection and PyTorch scaled dot-product attention, as an alternative bottleneck attention layer.
• Make bottleneck frequency aggregation configurable through BottleneckConfig.frequency_aggregation.
• Update bottleneck construction so frequency aggregation is built before bottleneck layers, matching the forward execution order.
• Keep minimal bottleneck configs consistent by defaulting frequency aggregation channels from BottleneckConfig.channels.
• Bump Torch/Torchaudio minimum versions to >=2.0.0 for scaled dot-product attention support.
• Add regression coverage for new blocks, configurable aggregation behavior, and bundled checkpoint config-schema loading.

Rationale

This PR adds more flexibility to how bottlenecks are constructed, with the goal of making the architecture easier to experiment with and potentially leaner and more efficient.

Two bottleneck components are especially worth exploring:

• The self-attention layer is compute-heavy, so this adds an EfficientSelfAttention variant that uses PyTorch’s optimized scaled dot-product attention implementation.
• The frequency aggregation step currently uses a full-column VerticalConv before self-attention. That preserves the desired shape, but it contributes a large share of the model parameters, around half of the total parameter count. This PR adds VerticalMean as a more parameter-efficient alternative that still collapses the frequency axis to the same shape before attention.

Benchmark

In a CPU run, the efficient self-attention layer was about 40% faster than the original implementation.
Keep in mind that this layer is only one part of the full model forward pass, so the overall model speedup is much more modest.

Validation

• uv run pytest tests/test_models/test_blocks.py tests/test_models/test_bottleneck.py tests/test_model.py::test_bundled_checkpoint_loads_with_current_config_schema
• uv run ty check src tests/test_models/test_blocks.py tests/test_models/test_bottleneck.py tests/test_model.py
• just check

[Copilot]

Pull request overview

This PR expands the model bottleneck to support configurable frequency aggregation (adding VerticalMean alongside VerticalConv) and introduces an EfficientSelfAttention alternative that uses PyTorch’s scaled dot-product attention, with accompanying config/schema updates and regression tests.

Changes:

• Add VerticalMean and EfficientSelfAttention blocks (plus config models) and extend the layer factory unions accordingly.
• Make bottleneck frequency aggregation configurable via BottleneckConfig.frequency_aggregation and adjust bottleneck construction to match forward execution order.
• Update dependency minimums and add regression tests for new blocks, configurable behavior, and bundled checkpoint config-schema loading.

Reviewed changes

Copilot reviewed 6 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
uv.lock	Updates locked metadata/deps consistent with the dependency bumps and new resolution markers.
pyproject.toml	Bumps minimum torch/torchaudio versions for SDPA usage.
src/batdetect2/models/blocks.py	Adds EfficientSelfAttention and VerticalMean, updates exported symbols and layer config unions, and adjusts LayerGroup to be a Block.
src/batdetect2/models/bottleneck.py	Adds configurable frequency_aggregation, introduces a frequency-aggregation builder, and updates bottleneck construction/channel wiring.
tests/test_models/test_blocks.py	Adds coverage for VerticalMean and EfficientSelfAttention (including temporary output parity via copied weights).
tests/test_models/test_bottleneck.py	Adds tests ensuring post-aggregation layers are built for aggregated channels and that minimal configs remain consistent.
tests/test_model.py	Adds a regression test ensuring the bundled checkpoint remains loadable with the current config schema.

Suppressed comments (3)

src/batdetect2/models/blocks.py:395

• torch.nn.functional.scaled_dot_product_attention in PyTorch 2.0 does not accept the scale= keyword argument. With the current project minimum of torch>=2.0.0, this will raise a TypeError at runtime on Torch 2.0.x.

        att = F.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=None,

src/batdetect2/models/bottleneck.py:236

• build_frequency_aggregation docstring is out of sync with the signature/behavior: config is not optional/nullable here, and the return value is not always a VerticalConv (it can be VerticalMean).

    config : FrequencyAggregationLayerConfig, optional
        Configuration specifying the output channel count and any
        additional layers. Uses ``VerticalConvConfig`` if ``None``.

    Returns

src/batdetect2/models/bottleneck.py:203

• The BottleneckConfig docstring is now outdated: it doesn't mention frequency_aggregation, and it still says only SelfAttentionConfig is supported in layers even though the type union includes EfficientSelfAttentionConfig. This can mislead config authors and schema consumers.

    channels: int
    frequency_aggregation: FrequencyAggregationLayerConfig = Field(
        default_factory=lambda data: VerticalConvConfig(
            channels=data["channels"]
        )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.