feat(map_encoder): add deformable cross-attention map BEV fusion mode

[intisar1020]

feat(map_encoder): add deformable cross-attention as a production-viable map BEV fusion mode

Motivation & Architectural Context

Map BEV fusion currently offers two modes (introduced in #55):

• residual: image_bev + alpha * map_bev with a per-channel gate initialized
  to zero. Cheap and stable at any BEV resolution, but spatially rigid -- every
  pixel only ever sees its own map cell.
• cross_attn: dense spatial cross-attention where each image BEV query attends
  to every map BEV token. Spatially adaptive, but O(N^2): at the production BEV
  grid (450x300 = 135K tokens) the score matrix is ~10^11 elements and OOMs,
  which is why the module currently raises a guard above 4096 tokens.

This PR adds a third mode, deformable, that gives attention-based map fusion
the spatial adaptivity of cross-attention at linear cost: each query attends to
only K learned-offset sample points in the map BEV (K=4 by default), dropping
the complexity from O(N^2) to O(N*K) and making it viable at production
resolution.

The design follows the deformable spatial cross-attention from BEVFormer
(Li et al., 2022): a query predicts sampling offsets relative to its own
reference position, features are sampled at those locations, and the K samples
are aggregated with per-head softmax weights. Since the reference plane here is
the map BEV itself, each query's reference point is simply its own pixel
position, and sampling uses F.grid_sample (bilinear, border padding) -- no
custom CUDA kernels required.

Summary of Changes

1. New fusion module: MapDeformableCrossAttentionFusion

model_components/map_encoder/map_bev_fusion/deformable_cross_attention_fusion.py

class MapDeformableCrossAttentionFusion(nn.Module):
    def __init__(self, embed_dim=256, num_sample_points=4, num_heads=8, dropout=0.1): ...

• offset_proj predicts K 2D pixel displacements from the query feature;
  attn_proj predicts per-head softmax weights over the K sample points.
• Reference grid is each query's own position in grid_sample coordinates;
  sampling positions are reference + offset, with out-of-map positions clamped
  by padding_mode="border".
• F.grid_sample yields (B, C, N, K); the channel dim is split per head and
  the K samples are aggregated by weighted sum, followed by a pre-norm output
  projection with residual, then the same FFN + residual block used by the
  cross-attention fusion.
• Interface is identical to the other modes: forward(image_bev, map_bev)
  with (B, C, H, W) tensors in and out.

Why this avoids OOM

Standard attention materializes an attention score matrix of shape
(B, num_heads, N, N) -> (1, 8, 135000, 135000) at the production grid,
which is ~145.8B elements -- roughly 583 GB in float32 -- and exactly why the
dense cross_attn mode carries the >4096-token guard.

Deformable attention never builds that matrix. Its largest intermediate tensor
is the sampled per-head features, (B, N, num_heads, K, head_dim) ->
(1, 135000, 8, 4, 32) (~138M elements), which is only ~553 MB in float32
(~276 MB under fp16 autocast) -- a ~1000x reduction that fits comfortably on a
12 GB GPU at production resolution.

2. Registry wiring

model_components/map_encoder/map_bev_fusion/__init__.py registers the new
mode so it can be selected end-to-end via map_fusion_mode="deformable"
(AutoE2E -> ReactiveE2E -> build_map_bev_fusion):

MAP_FUSION_REGISTRY = {
    "residual": ResidualMapFusion,
    "cross_attn": MapCrossAttentionFusion,
    "deformable": MapDeformableCrossAttentionFusion,
}

Constructor options such as num_sample_points are forwarded through the
existing map_fusion_kwargs plumbing, so no model changes are required.

3. Unit and integration tests

tests/test_map_encoder.py adds:

• TestMapDeformableCrossAttentionFusion: output shape, map influence on
  output, gradient flow, all-parameters-receive-gradients, NaN safety with zero
  inputs, non-square grids, and configurable K.
• Registry tests for the "deformable" key via build_map_bev_fusion.
• AutoE2E integration tests: map_fusion_mode="deformable" forward pass and
  gradient flow through MapBEVFusion (mock-backbone harness).

---

Verification

• All 50 tests in tests/test_map_encoder.py pass, including the new
  deformable unit and integration cases.
• ruff check and mypy are clean for the touched code.
• Full CI unit suite (pytest Model/tests) passes.
• Forward pass at the production BEV grid 450x300 (135K tokens) with K=4
  completes without OOM -- the exact failure mode that motivates this mode.

Training sanity check on a small subset of KITScenes tar files (loss goes
down over the first ~6-12 steps):

Metric	Start (step 0)	Best (step ~6-12)
Val loss	1.012	0.083
ADE	36.3 m	11.2 m
FDE	70.9 m	28.2 m

[riita10069]

@intisar1020

Regarding the changes in this PR, I believe the approach of using Deformable Cross Attention to reduce computational cost is appropriate for addressing the issue that dense cross-attention is impractical in production BEV. Therefore, I plan to resolve the conflicts separately and merge #184 itself.

That said, I would like to request additional improvements regarding the handling of Route information.

Currently, map_context and route_mask are concatenated channel-wise and fed into a single NavigationEncoder.

navigation_input = torch.cat([gated_map, gated_route], dim=1)
navigation_bev = self.NavigationEncoder(navigation_input)

With this configuration, the sparse and thin Route features are mixed with the larger-area Map features at an early stage, which may result in the model not utilizing them sufficiently. Additionally, since the Deformable Cross Attention being added in this PR improves the fusion of image_bev and navigation_bev, if the Route features become weakened within the NavigationEncoder, that information cannot be recovered in later stages.

As a follow-up change, could you consider a configuration that separates the encoding of Map and Route, and performs Route-gated fusion at a later stage?

The concept would be a configuration like the following:

Map Raster
    ↓
MapEncoder
    ↓
map_bev ──────────────────┐
                          ├─ Route-gated fusion
Route Mask                │
    ↓                     │
RouteEncoder              │
    ↓                     │
route_bev ────────────────┘
    ↓
navigation_bev
    ↓
Deformable Cross Attention with image_bev

For example, I think a per-channel gated residual like the following would be sufficient as an initial implementation:

map_bev = self.MapEncoder(gated_map)
route_bev = self.RouteEncoder(gated_route)

route_gate = torch.sigmoid(self.route_gate).view(1, -1, 1, 1)
navigation_bev = map_bev + route_gate * route_bev

The RouteEncoder doesn't need to be a heavy backbone equivalent to the MapEncoder; a lightweight CNN would be fine.

With this configuration, the following can be handled separately:

• MapEncoder: Static structures such as road structure, lanes, boundaries, intersections
• RouteEncoder: Selected travel route and branch directions
• Route-gated fusion: Emphasize Map features related to the selected Route
• Deformable Cross Attention: Absorb positional misalignment between image BEV and navigation BEV

Additionally, I would like you to add a test that verifies that when only the Route is changed for the same camera/map input, the final trajectory changes appropriately. I want to validate that the Route is actually being used in the final output, not just that gradients flow through.

This doesn't need to be included in the scope of #184; it's fine to handle it as a follow-up PR after merging.

[intisar1020]

@riita10069
Thank you for the detailed proposal. I agree that concatenating the route channel with the map at the input level is not an effective way to carry route information, because the two are stacked before any learnable weighting exists, so the sparse route signal can be easily diluted/dominated by the large-area map features inside the shared encoder. I think the idea to encode them separately and modulate the map with a learnable route gate is intuitive and can be quickly implemented and test. I will proceed with it.

A clarifying note on my current baseline: I have been training the reactive branch with the route mask gated to zero on recently on purpose to test it, because I wanted to first establish a model conditioned only on the map input. so next logical step woudld be I will enable the route when testing the route-gated fusion.

Planned configuration:

Map Raster
    ↓
MapEncoder
    ↓
map_bev ──────────────────┐
                          ├─ Route-gated fusion ── navigation_bev
Route Mask                │                            ↓
    ↓                     │                   Deformable Cross Attention
RouteEncoder              │                    with image_bev
    ↓                     │
route_bev ────────────────┘

with a per-channel learnable gate:

map_bev   = self.MapEncoder(gated_map) # swin
route_bev = self.RouteEncoder(gated_route)          # lightweight CNN

route_gate = torch.sigmoid(self.route_gate).view(1, -1, 1, 1)   # learnable, per-channel
navigation_bev = map_bev + route_gate * route_bev

For the RouteEncoder I plan a lightweight strided-CNN stack. Each stage is Conv2d -> BatchNorm -> GELU, downsampling by 2, then a 1x1 projection to embed_dim and an adaptive pool to the BEV size:

class RouteEncoder(nn.Module):
    def __init__(self, in_channels=2, embed_dim=256, out_h=450, out_w=300):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels, 32, 3, stride=2, padding=1),   # 256 -> 128
            nn.BatchNorm2d(32), nn.GELU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),            # 128 -> 64
            nn.BatchNorm2d(64), nn.GELU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),           # 64 -> 32
            nn.BatchNorm2d(128), nn.GELU(),
        )
        self.proj = nn.Conv2d(128, embed_dim, 1)
        self.pool = nn.AdaptiveAvgPool2d((out_h, out_w))
    def forward(self, route_mask):
        return self.pool(self.proj(self.encoder(route_mask)))

This keeps the encoder to a few hundred thousand parameters.

For validation I will run an ablation of with-route vs without-route on a subset of the dataset, targeting ADE / FDE (and a route-corridor compliance metric such as mean path delta when the route is swapped). In addition to the numbers, I will visually verify by overlaying the predicted trajectory against the route map for a few samples, which makes it concrete whether the route changes the output path. This covers the point you wanted to vrify --> holding the camera and map fixed while changing only the route, then checking that the predicted trajectory changes appropriately, confirming the route influences the final output rather than only providing a gradient path.
I will handle this as a follow-up PR. thanks.

[riita10069]

@intisar1020 / cc. @m-zain-khawaja

Thanks again for the suggestion. Working at #186, the route reconstruction loss is now there, but Map and Route are still concatenated before the shared NavigationEncoder, so the original concern remains: the route signal can still get buried by the map features.

Could you try youridea on top of #186 as a follow-up?

Map → MapEncoder ───────┐
                        ├─ route-gated fusion → navigation_bev
Route → RouteEncoder ───┘
                                                 ↓
                                       route reconstruction head

I think the combination of separate encoders + learnable route gate + the reconstruction auxiliary from #186 would be a cleaner way to make sure the route is preserved and actually used.

It would be great if you could also validate on your side whether this improves the results compared with the current shared-encoder baseline, especially ADE/FDE and route-sensitive cases such as route swap / junction branch behavior.