Masked attention

Overview

Masked attention switches off some query-key pairs before the softmax. The mask holds 0 for each allowed pair and minus infinity for each blocked pair. Added to the scores, it gives blocked keys zero weight after the softmax. Cutie writes the update as a softmax over the mask plus the query-key scores, applied to the values, with a residual connection 2310.12982 §3.2.2 eq. 2. The mask is a rule about which pixels each query may see, and the rule can depend on the data. In Cutie it comes from a foreground prediction made inside the network 2310.12982 §3.2.2.

Only one card on this wiki lists this concept, so the page covers one design. It will grow as more cards are added.

Variants

Foreground/background split of object queries (Cutie)

Cutie has 16 object queries 2310.12982 §3.4.2. The first half may attend only to pixels predicted as foreground. The second half may attend only to pixels predicted as background 2310.12982 §3.2.2 eq. 3. The mask comes from a mask head on the previous pixel features, and the same mask is shared across all attention heads 2310.12982 §3.2.2. The released code refreshes the mask after each block except the last, from the updated pixels. The paper's text does not spell this out, so treat it as a code detail 2310.12982 §3.2.2. The queries still exchange information afterwards through unmasked self-attention 2310.12982 §3.2.1. Card: Putting the Object Back into Video Object Segmentation.

The paper says this extends a masked form that used only the foreground 2310.12982 §3.2.2. That earlier work has no card on this wiki, so its content is not described here.

Foreground-only masking

All queries read only predicted foreground pixels. Cutie reports this as an ablation, not as its own model: 66.7 ± 0.21 J&F against 67.3 ± 0.36 for the split 2310.12982 tab. 6.

When to use which

The evidence is from Cutie's ablations on MOSE validation, using the small model and MOSE training 2310.12982 §4.2 tab. 6. It covers video object segmentation only. No card here tests masked attention on another task.

Implementation sketch

The sketch follows the Cutie form 2310.12982 §3.2.2 eq. 2 eq. 3. Shapes: N queries, C channels, HW pixels. Projections of queries, keys and values are as in standard attention, and the positional terms are left out.

# Build the allowed-pair mask for one block  [arxiv:2310.12982 §3.2.2] [eq. 3]
# Input:  predicted foreground probability M (HW,), N queries
# Output: additive mask (N, HW), 0 where allowed, -inf where blocked

def build_mask(M, N):
    allowed = zeros(N, HW, dtype=bool)
    # First half of the queries: pixels predicted as foreground  [eq. 3]
    allowed[: N // 2] = (M >= 0.5)
    # Second half: pixels predicted as background  [eq. 3]
    allowed[N // 2 :] = (M < 0.5)
    # Rows with no allowed pixel are a special case; see the note below
    return where(allowed, 0.0, -inf)

# Masked cross-attention from queries to pixels  [arxiv:2310.12982 eq. 2]
# Input:  queries X (N, C), pixels R (HW, C), additive mask (N, HW)
# Output: X plus the attended values

def masked_cross_attention(X, R, mask):
    # Scores between queries and pixels
    scores = Q(X) @ K(R).T
    # Add the mask, so blocked pixels get zero weight after softmax
    weights = softmax(scores + mask, axis=-1)
    # Read the values and add the residual  [eq. 2]
    return weights @ V(R) + X

Notes for implementers:

• Rebuild the mask from the pixels after each block except the last, as the released code does. The paper's text does not state this schedule 2310.12982 §3.2.2.
• If a query has no allowed pixel, its softmax has only -inf inputs and is undefined. The released code lifts the mask for such rows. The paper does not mention this 2310.12982 §3.2.2.
• The paper does not describe how to mask several objects at once. Its per-object rule is the 0.5 threshold above 2310.12982 eq. 3.

Contradictions

None between cards: only one card lists this concept. The paper and its released code do differ on the rule for several objects. The paper's rule for each object is the 0.5 threshold in eq. 3. The code assigns each pixel to the object with the largest aggregated score against background. The two agree for a single object and can differ for several. The paper does not describe the multi-object rule 2310.12982 eq. 3 §3.2.2.

Papers

• Putting the Object Back into Video Object Segmentation: introduces the foreground/background masked attention between 16 object queries and pixel features, inside a three-block object transformer.