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]

If the input TokenFeatureMaps has masks, invalid tokens receive zero attention weight.

Source code in rslearn/models/attention_pooling.py 
 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
 78
 79
 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
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]

    If the input TokenFeatureMaps has masks, invalid tokens receive zero attention
    weight.
    """

    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, mask: torch.Tensor | None = None
    ) -> torch.Tensor:
        """Attention pooling for a single feature map (BCHWN tensor).

        Args:
            feat_tokens: the BCHWN token feature map.
            mask: optional (B*H*W) x N bool mask of valid tokens. Invalid tokens get
                zero attention weight. If every token at a location is invalid, the
                weights degrade to uniform.
        """
        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))
        logits = self.linear(feat_tokens)  # (B*H*W, N, 1)
        if mask is not None:
            logits = logits.masked_fill(
                ~mask.unsqueeze(-1), torch.finfo(logits.dtype).min
            )
        attention_scores = torch.nn.functional.softmax(logits, 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, mask in zip(
            intermediates.feature_maps, _get_flat_masks(intermediates)
        ):
            features.append(self.forward_for_map(feat, mask))
        return FeatureMaps(features)

forward_for_map 

forward_for_map(feat_tokens: Tensor, mask: Tensor | None = None) -> Tensor

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

Parameters:

Name Type Description Default
feat_tokens Tensor

the BCHWN token feature map.

 required
mask Tensor | None

optional (BHW) x N bool mask of valid tokens. Invalid tokens get zero attention weight. If every token at a location is invalid, the weights degrade to uniform.

 None
Source code in rslearn/models/attention_pooling.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def forward_for_map(
    self, feat_tokens: torch.Tensor, mask: torch.Tensor | None = None
) -> torch.Tensor:
    """Attention pooling for a single feature map (BCHWN tensor).

    Args:
        feat_tokens: the BCHWN token feature map.
        mask: optional (B*H*W) x N bool mask of valid tokens. Invalid tokens get
            zero attention weight. If every token at a location is invalid, the
            weights degrade to uniform.
    """
    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))
    logits = self.linear(feat_tokens)  # (B*H*W, N, 1)
    if mask is not None:
        logits = logits.masked_fill(
            ~mask.unsqueeze(-1), torch.finfo(logits.dtype).min
        )
    attention_scores = torch.nn.functional.softmax(logits, 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 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
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, mask in zip(
        intermediates.feature_maps, _get_flat_masks(intermediates)
    ):
        features.append(self.forward_for_map(feat, mask))
    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.

If the input TokenFeatureMaps has masks, invalid tokens receive zero attention weight.

Source code in rslearn/models/attention_pooling.py 
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
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
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.

    If the input TokenFeatureMaps has masks, invalid tokens receive zero attention
    weight.
    """

    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, mask: torch.Tensor | None = None
    ) -> torch.Tensor:
        """Attention pooling for a single feature map (BCHWN tensor).

        Args:
            feat_tokens: the BCHWN token feature map.
            mask: optional (B*H*W) x N bool mask of valid tokens. Invalid tokens get
                zero attention weight. If every token at a location is invalid, the
                weights degrade to uniform.
        """
        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
        )  # [B*H*W, num_heads, 1, N]
        if mask is not None:
            attn_scores = attn_scores.masked_fill(
                ~mask[:, None, None, :], torch.finfo(attn_scores.dtype).min
            )
        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, mask in zip(
            intermediates.feature_maps, _get_flat_masks(intermediates)
        ):
            features.append(self.forward_for_map(feat, mask))
        return FeatureMaps(features)

init_weights 

init_weights() -> None

Initialize weights for the probe.

Source code in rslearn/models/attention_pooling.py 
147
148
149
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, mask: Tensor | None = None) -> Tensor

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

Parameters:

Name Type Description Default
feat_tokens Tensor

the BCHWN token feature map.

 required
mask Tensor | None

optional (BHW) x N bool mask of valid tokens. Invalid tokens get zero attention weight. If every token at a location is invalid, the weights degrade to uniform.

 None
Source code in rslearn/models/attention_pooling.py 
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
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
def forward_for_map(
    self, feat_tokens: torch.Tensor, mask: torch.Tensor | None = None
) -> torch.Tensor:
    """Attention pooling for a single feature map (BCHWN tensor).

    Args:
        feat_tokens: the BCHWN token feature map.
        mask: optional (B*H*W) x N bool mask of valid tokens. Invalid tokens get
            zero attention weight. If every token at a location is invalid, the
            weights degrade to uniform.
    """
    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
    )  # [B*H*W, num_heads, 1, N]
    if mask is not None:
        attn_scores = attn_scores.masked_fill(
            ~mask[:, None, None, :], torch.finfo(attn_scores.dtype).min
        )
    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 
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
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, mask in zip(
        intermediates.feature_maps, _get_flat_masks(intermediates)
    ):
        features.append(self.forward_for_map(feat, mask))
    return FeatureMaps(features)