MADDi attention implementation: single-token sequences in code vs. attention over feature matrices in paper

[abidabrar]

Hi, and thank you for releasing the MADDi code and the paper. It has been very helpful for my work on multimodal modelling of our data.

While trying to reproduce the attention-based fusion described in the paper, I found a possible mismatch between the attention mechanism in the paper and the attention mechanism in the code (training/train_all_modalities_maddi.py).

In the paper, the Multimodal framework and Neural network attention sections describe self-attention and cross-modal attention over latent feature matrices (I, G, C) for imaging, genetic, and clinical data, using the usual formulation:

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T} {\sqrt{d}}\right) V$

This suggests that each modality is represented as a sequence of features, where attention can highlight interactions across tokens.

In the released Keras code, each modality encoder produces a single 50-dimensional embedding vector:

dense_clinical = create_model_clinical()(in_clinical)   # (batch, 50)
dense_snp      = create_model_snp()(in_snp)             # (batch, 50)
dense_img      = create_model_img()(in_img)             # (batch, 50)

The attention helpers then expand each embedding to a sequence length of 1:

def cross_modal_attention(x, y, num_heads=4, key_dim=50):
    x = tf.expand_dims(x, axis=1)  # (batch, 1, 50)
    y = tf.expand_dims(y, axis=1)  # (batch, 1, 50)

    a1 = MultiHeadAttention(num_heads=num_heads, key_dim=key_dim)(x, y)
    a2 = MultiHeadAttention(num_heads=num_heads, key_dim=key_dim)(y, x)

    a1 = a1[:, 0, :]
    a2 = a2[:, 0, :]
    return concatenate([a1, a2])

self_attention(x) uses the same pattern.

According to the keras.layers.MultiHeadAttention documentation (Keras 2), for query shape (B, T, dim) and value shape (B, S, dim), the attention scores have shape (B, num_heads, T, S). In this implementation, T = 1 and S = 1. So each head's attention score is a single scalar (1, 1). The softmax over a single scalar is always [1.0], which means the attention weights are always 1, and the output for each head reduces to the projected value.

So cross_modal_attention(x, y) becomes a learned projection of y only. The query input x does not change which positions in y are weighted, because there is only one position to attend over.

So, from my understanding, the current implementation behaves like a stack of additional linear projections on top of the 50-dim embeddings rather than attention over multiple latent features/tokens, whereas the paper text and Figure 2 suggest a richer token-wise interaction and feature importance.

My question is: is this single-token attention behavior (sequence length = 1 per modality) intentional and identical to what was used for the results reported in the MADDi paper? If I have misunderstood some parts, I would really appreciate clarification.