sage.architecture.zoo.cross_attention

Filename : cross_attention.py Description : Short description of the file

Created on 2026-03-19 23:36:04

__author__ = Narenraju Nagarajan __copyright__ = Copyright 2026, ProjectName __license__ = MIT Licence __version__ = 0.0.1 __maintainer__ = Narenraju Nagarajan __affiliation__ = N/A __email__ = N/A __status__ = [‘inProgress’, ‘Archived’, ‘inUsage’, ‘Debugging’]

GitHub Repository: NULL

Documentation: NULL

Classes

CrossAttention2D	Full cross-attention between two 2D feature maps.
AxialCrossAttention2D	Axial cross-attention for 2-detector 2D features.

Module Contents

class CrossAttention2D(embed_dim=128, num_heads=4)

Bases: torch.nn.Module

Full cross-attention between two 2D feature maps.

Projects each detector’s 2D feature map (B, 1, H, W) to embed_dim channels, flattens the spatial grid into a sequence, and lets each detector attend to the other using torch.nn.MultiheadAttention. Both detectors receive mutually informed representations: H1 attends to L1 and vice versa.

Memory scales as O(H²W²) — use AxialCrossAttention2D for large feature maps.

Parameters:

• embed_dim (int) – Projection and attention dimension (default 128).

• num_heads (int) – Number of parallel attention heads (default 4).

• state (Initialize internal Module)

• ScriptModule. (shared by both nn.Module and)

embed_dim = 128

num_heads = 4

proj1

proj2

attn

forward(f1, f2)

Apply cross-attention between two detector feature maps.

Parameters:

• f1 (torch.Tensor, shape (B, 1, H, W)) – Feature map for detector 1 (e.g. H1).

• f2 (torch.Tensor, shape (B, 1, H, W)) – Feature map for detector 2 (e.g. L1).

Returns:

• f1_attn (torch.Tensor, shape (B, embed_dim, H, W))

• f2_attn (torch.Tensor, shape (B, embed_dim, H, W))

class AxialCrossAttention2D(embed_dim=128, num_heads=4)

Bases: torch.nn.Module

Axial cross-attention for 2-detector 2D features. Attends along features (H) and time (W) separately to save memory. Input: f1, f2 -> (B, 1, H, W) Output: f1_attn, f2_attn -> same shape

Initialize internal Module state, shared by both nn.Module and ScriptModule.

embed_dim = 128

num_heads = 4

proj1

proj2

attn_H

attn_W

forward(f1, f2)