Mask in TokenAttentionPooling

[lhqing]

Report

Hi, very nice work!

I'm trying to understand the code of TokenAttentionPooling, it seems to me the class token will always attend to every payload token. And since there is only a single attention call, the class token is return as the result, the input provide mask should not have any effect? I used some different masks and the result of this class is identical. Or maybe I missed something here?

CellFlow/src/cellflow/networks/_utils.py

Line 420 in 4970ad4

class TokenAttentionPooling(BaseModule):

Besides, I wonder in practice, how do you choose between TokenAttentionPooling and SeedAttentionPooling?

Thanks!

Version information

No response

[MUCDK]

Hi @lhqing ,

Thanks a lot for your interest!
I'm not entirely sure whether I understand what you mean, but let me try to explain the way we intended to implement this:

The class token indeed attends to all input tokens, but the attention is ignored whenever we padded the input, which is given by the mask argument.

Could you please provide a code snippet which reproduces your result? This would also help me to understand better what you mean!

[lhqing]

@MUCDK thank you for the reply.

What I meant is the attn_mask doesn't seem to impact the result of TokenAttentionPooling, a minimum example is in below.

I think this is due to this two lines, where the first token will always attend to entire sequence regardless of what the mask is:

CellFlow/src/cellflow/networks/_utils.py

Lines 468 to 469 in 970acd8

	token_mask = token_mask.at[:, :, 0, :].set(1)
	token_mask = token_mask.at[:, :, :, 0].set(1)

# to test mask impact on TokenAttentionPooling:

import abc
from collections.abc import Callable, Sequence

import jax
import jax.numpy as jnp
from cellflow._types import Layers_t
from cellflow.networks._utils import BaseModule, TokenAttentionPooling
from flax import linen as nn
from flax.linen import initializers


def _get_masks(conditions):
    """Get mask for padded conditions tensor."""
    # mask of shape (batch_size, set_size)
    mask = 1 - jnp.all(
        jnp.array(
            [jnp.all(c == 0.0, axis=-1) for c in conditions.values()],
        ),
        axis=0,
    )
    mask = jnp.expand_dims(mask, -1)

    # attention mask of shape (batch_size, 1, set_size, set_size)
    attention_mask = mask & jnp.matrix_transpose(mask)
    attention_mask = jnp.expand_dims(attention_mask, 1)

    return mask, attention_mask


rng = jax.random.PRNGKey(0)

# create some toy condition of (1, 28, 6)
cond = jnp.array([[[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]] * 28])
cond = cond.at[0, 20, 4].set(3)
cond = cond.at[0, 15, 3].set(2)
mask, attn_mask = _get_masks({"conditions": cond})
random_mask = jax.random.bernoulli(rng, 0.5, attn_mask.shape).astype(jnp.int32)

agg = TokenAttentionPooling()

# random mask or attn_mask produce the same result
variables = agg.init(rng, cond, random_mask, training=True)
out = agg.apply(variables, cond, attn_mask, training=True)
out_rand = agg.apply(variables, cond, random_mask, training=True)
(out == out_rand).all()  # Array(True, dtype=bool)

[MUCDK]

Thank you very much @lhqing for spotting this. I will open a PR shortly with a fix!

[MUCDK]

Fixed in #241