Skip to content

rslearn.models.attention_pooling

attention_pooling

An attention pooling layer.

SimpleAttentionPool

Bases: IntermediateComponent

Simple Attention Pooling.

Given a token feature map of shape BCHWN, learn an attention layer which aggregates over the N dimension.

This is done simply by learning a mapping D->1 which is the weight which should be assigned to each token during averaging:

output = sum [feat_token * W(feat_token) for feat_token in feat_tokens]

Source code in rslearn/models/attention_pooling.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class SimpleAttentionPool(IntermediateComponent):
    """Simple Attention Pooling.

    Given a token feature map of shape BCHWN,
    learn an attention layer which aggregates over
    the N dimension.

    This is done simply by learning a mapping D->1 which is the weight
    which should be assigned to each token during averaging:

    output = sum [feat_token * W(feat_token) for feat_token in feat_tokens]
    """

    def __init__(self, in_dim: int, hidden_linear: bool = False) -> None:
        """Initialize the simple attention pooling layer.

        Args:
            in_dim: the encoding dimension D
            hidden_linear: whether to apply an additional linear transformation D -> D
                to the feat tokens. If this is True, a ReLU activation is applied
                after the first linear transformation.
        """
        super().__init__()
        if hidden_linear:
            self.hidden_linear = nn.Linear(in_features=in_dim, out_features=in_dim)
        else:
            self.hidden_linear = None
        self.linear = nn.Linear(in_features=in_dim, out_features=1)

    def forward_for_map(self, feat_tokens: torch.Tensor) -> torch.Tensor:
        """Attention pooling for a single feature map (BCHWN tensor)."""
        B, D, H, W, N = feat_tokens.shape
        feat_tokens = rearrange(feat_tokens, "b d h w n -> (b h w) n d")
        if self.hidden_linear is not None:
            feat_tokens = torch.nn.functional.relu(self.hidden_linear(feat_tokens))
        attention_scores = torch.nn.functional.softmax(self.linear(feat_tokens), dim=1)
        feat_tokens = (attention_scores * feat_tokens).sum(dim=1)
        return rearrange(feat_tokens, "(b h w) d -> b d h w", b=B, h=H, w=W)

    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
        """Forward pass for attention pooling linear probe.

        Args:
            intermediates: the output from the previous component, which must be a TokenFeatureMaps.
                We pool over the final dimension in the TokenFeatureMaps. If multiple maps
                are passed, we apply the same linear layers to all of them.
            context: the model context.

        Returns:
            torch.Tensor:
                - output, attentioned pool over the last dimension (B, C, H, W)
        """
        if not isinstance(intermediates, TokenFeatureMaps):
            raise ValueError("input to Attention Pool must be a TokenFeatureMaps")

        features = []
        for feat in intermediates.feature_maps:
            features.append(self.forward_for_map(feat))
        return FeatureMaps(features)

forward_for_map 

forward_for_map(feat_tokens: Tensor) -> Tensor

Attention pooling for a single feature map (BCHWN tensor).

Source code in rslearn/models/attention_pooling.py 
48
49
50
51
52
53
54
55
56
def forward_for_map(self, feat_tokens: torch.Tensor) -> torch.Tensor:
    """Attention pooling for a single feature map (BCHWN tensor)."""
    B, D, H, W, N = feat_tokens.shape
    feat_tokens = rearrange(feat_tokens, "b d h w n -> (b h w) n d")
    if self.hidden_linear is not None:
        feat_tokens = torch.nn.functional.relu(self.hidden_linear(feat_tokens))
    attention_scores = torch.nn.functional.softmax(self.linear(feat_tokens), dim=1)
    feat_tokens = (attention_scores * feat_tokens).sum(dim=1)
    return rearrange(feat_tokens, "(b h w) d -> b d h w", b=B, h=H, w=W)

forward 

forward(intermediates: Any, context: ModelContext) -> FeatureMaps

Forward pass for attention pooling linear probe.

Parameters:

Name Type Description Default
intermediates Any

the output from the previous component, which must be a TokenFeatureMaps. We pool over the final dimension in the TokenFeatureMaps. If multiple maps are passed, we apply the same linear layers to all of them.

 required
context ModelContext

the model context.

 required

Returns:

Type Description
FeatureMaps

torch.Tensor: - output, attentioned pool over the last dimension (B, C, H, W)
Source code in rslearn/models/attention_pooling.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
    """Forward pass for attention pooling linear probe.

    Args:
        intermediates: the output from the previous component, which must be a TokenFeatureMaps.
            We pool over the final dimension in the TokenFeatureMaps. If multiple maps
            are passed, we apply the same linear layers to all of them.
        context: the model context.

    Returns:
        torch.Tensor:
            - output, attentioned pool over the last dimension (B, C, H, W)
    """
    if not isinstance(intermediates, TokenFeatureMaps):
        raise ValueError("input to Attention Pool must be a TokenFeatureMaps")

    features = []
    for feat in intermediates.feature_maps:
        features.append(self.forward_for_map(feat))
    return FeatureMaps(features)

AttentionPool

Bases: IntermediateComponent

Attention Pooling.

Given a feature map of shape BCHWN, learn an attention layer which aggregates over the N dimension.

We do this by learning a query token, and applying a standard attention mechanism against this learned query token.

Source code in rslearn/models/attention_pooling.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
class AttentionPool(IntermediateComponent):
    """Attention Pooling.

    Given a feature map of shape BCHWN,
    learn an attention layer which aggregates over
    the N dimension.

    We do this by learning a query token, and applying a standard
    attention mechanism against this learned query token.
    """

    def __init__(self, in_dim: int, num_heads: int, linear_on_kv: bool = True) -> None:
        """Initialize the attention pooling layer.

        Args:
            in_dim: the encoding dimension D
            num_heads: the number of heads to use
            linear_on_kv: Whether to apply a linear layer on the input tokens
            to create the key and value tokens.
        """
        super().__init__()
        self.query_token: nn.Parameter = nn.Parameter(torch.empty(in_dim))
        if linear_on_kv:
            self.k_linear = nn.Linear(in_dim, in_dim)
            self.v_linear = nn.Linear(in_dim, in_dim)
        else:
            self.k_linear = None
            self.v_linear = None
        if in_dim % num_heads != 0:
            raise ValueError(
                f"in_dim must be divisible by num_heads. Got {in_dim} and {num_heads}."
            )
        self.num_heads = num_heads
        self.init_weights()

    def init_weights(self) -> None:
        """Initialize weights for the probe."""
        nn.init.trunc_normal_(self.query_token, std=0.02)

    def forward_for_map(self, feat_tokens: torch.Tensor) -> torch.Tensor:
        """Attention pooling for a single feature map (BCHWN tensor)."""
        B, D, H, W, N = feat_tokens.shape
        feat_tokens = rearrange(feat_tokens, "b d h w n -> (b h w) n d")
        collapsed_dim = B * H * W
        q = self.query_token.expand(collapsed_dim, 1, -1)
        q = q.reshape(
            collapsed_dim, 1, self.num_heads, D // self.num_heads
        )  # [B, 1, head, D_head]
        q = rearrange(q, "b h n d -> b n h d")
        if self.k_linear is not None:
            assert self.v_linear is not None
            k = self.k_linear(feat_tokens).reshape(
                collapsed_dim, N, self.num_heads, D // self.num_heads
            )
            v = self.v_linear(feat_tokens).reshape(
                collapsed_dim, N, self.num_heads, D // self.num_heads
            )
        else:
            k = feat_tokens.reshape(
                collapsed_dim, N, self.num_heads, D // self.num_heads
            )
            v = feat_tokens.reshape(
                collapsed_dim, N, self.num_heads, D // self.num_heads
            )
        k = rearrange(k, "b n h d -> b h n d")
        v = rearrange(v, "b n h d -> b h n d")

        # Compute attention scores
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(
            D // self.num_heads
        )
        attn_weights = F.softmax(attn_scores, dim=-1)
        x = torch.matmul(attn_weights, v)  # [B*H*W, num_heads, 1, D_head]
        x = x.squeeze(-2)  # [B*H*W, num_heads, D_head]
        return rearrange(
            x, "(b h w) nh dh -> b (nh dh) h w", b=B, h=H, w=W
        )  # [B, D, H, W]

    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
        """Forward pass for attention pooling linear probe.

        Args:
            intermediates: the output from the previous component, which must be a TokenFeatureMaps.
                We pool over the final dimension in the TokenFeatureMaps. If multiple feature
                maps are passed, we apply the same attention weights (query token and linear k, v layers)
                to all the maps.
            context: the model context.

        Returns:
            torch.Tensor:
                - output, attentioned pool over the last dimension (B, C, H, W)
        """
        if not isinstance(intermediates, TokenFeatureMaps):
            raise ValueError("input to Attention Pool must be a TokenFeatureMaps")

        features = []
        for feat in intermediates.feature_maps:
            features.append(self.forward_for_map(feat))
        return FeatureMaps(features)

init_weights 

init_weights() -> None

Initialize weights for the probe.

Source code in rslearn/models/attention_pooling.py 
115
116
117
def init_weights(self) -> None:
    """Initialize weights for the probe."""
    nn.init.trunc_normal_(self.query_token, std=0.02)

forward_for_map 

forward_for_map(feat_tokens: Tensor) -> Tensor

Attention pooling for a single feature map (BCHWN tensor).

Source code in rslearn/models/attention_pooling.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def forward_for_map(self, feat_tokens: torch.Tensor) -> torch.Tensor:
    """Attention pooling for a single feature map (BCHWN tensor)."""
    B, D, H, W, N = feat_tokens.shape
    feat_tokens = rearrange(feat_tokens, "b d h w n -> (b h w) n d")
    collapsed_dim = B * H * W
    q = self.query_token.expand(collapsed_dim, 1, -1)
    q = q.reshape(
        collapsed_dim, 1, self.num_heads, D // self.num_heads
    )  # [B, 1, head, D_head]
    q = rearrange(q, "b h n d -> b n h d")
    if self.k_linear is not None:
        assert self.v_linear is not None
        k = self.k_linear(feat_tokens).reshape(
            collapsed_dim, N, self.num_heads, D // self.num_heads
        )
        v = self.v_linear(feat_tokens).reshape(
            collapsed_dim, N, self.num_heads, D // self.num_heads
        )
    else:
        k = feat_tokens.reshape(
            collapsed_dim, N, self.num_heads, D // self.num_heads
        )
        v = feat_tokens.reshape(
            collapsed_dim, N, self.num_heads, D // self.num_heads
        )
    k = rearrange(k, "b n h d -> b h n d")
    v = rearrange(v, "b n h d -> b h n d")

    # Compute attention scores
    attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(
        D // self.num_heads
    )
    attn_weights = F.softmax(attn_scores, dim=-1)
    x = torch.matmul(attn_weights, v)  # [B*H*W, num_heads, 1, D_head]
    x = x.squeeze(-2)  # [B*H*W, num_heads, D_head]
    return rearrange(
        x, "(b h w) nh dh -> b (nh dh) h w", b=B, h=H, w=W
    )  # [B, D, H, W]

forward 

forward(intermediates: Any, context: ModelContext) -> FeatureMaps

Forward pass for attention pooling linear probe.

Parameters:

Name Type Description Default
intermediates Any

the output from the previous component, which must be a TokenFeatureMaps. We pool over the final dimension in the TokenFeatureMaps. If multiple feature maps are passed, we apply the same attention weights (query token and linear k, v layers) to all the maps.

 required
context ModelContext

the model context.

 required

Returns:

Type Description
FeatureMaps

torch.Tensor: - output, attentioned pool over the last dimension (B, C, H, W)
Source code in rslearn/models/attention_pooling.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
    """Forward pass for attention pooling linear probe.

    Args:
        intermediates: the output from the previous component, which must be a TokenFeatureMaps.
            We pool over the final dimension in the TokenFeatureMaps. If multiple feature
            maps are passed, we apply the same attention weights (query token and linear k, v layers)
            to all the maps.
        context: the model context.

    Returns:
        torch.Tensor:
            - output, attentioned pool over the last dimension (B, C, H, W)
    """
    if not isinstance(intermediates, TokenFeatureMaps):
        raise ValueError("input to Attention Pool must be a TokenFeatureMaps")

    features = []
    for feat in intermediates.feature_maps:
        features.append(self.forward_for_map(feat))
    return FeatureMaps(features)