Skip to content

fix: encoder with attention - #29

Merged
jannisborn merged 9 commits into
mainfrom
fix/encoder-with-attention
Apr 14, 2026
Merged

fix: encoder with attention#29
jannisborn merged 9 commits into
mainfrom
fix/encoder-with-attention

Conversation

@jorisSchaller

Copy link
Copy Markdown
Contributor

Fix the attention issue where in case of batching, some padding is added and those tokens are not ignored in the attention.

@jorisSchaller jorisSchaller changed the title Fix/encoder with attention fix/encoder with attention Mar 23, 2026
@jorisSchaller
jorisSchaller force-pushed the fix/encoder-with-attention branch from e237d54 to 3f2c358 Compare March 23, 2026 13:12
@jorisSchaller jorisSchaller changed the title fix/encoder with attention fix: encoder with attention Mar 23, 2026
@jorisSchaller
jorisSchaller force-pushed the fix/encoder-with-attention branch from 014efa5 to d6b2695 Compare March 23, 2026 16:50

@jannisborn jannisborn left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great job overall! The forward looks good, but we are not using the mask in the inputs_embeds track. Why was this done?

Comment thread src/mblm/model/mblm.py Outdated
This allows exact replication of the input_ids path.
Expected shapes: [local_embs (B, P1', P2, ..., Pn, Dn), ..., global_embs (B, P1', D1)]
return_type: What to return - the loss, the logits, the the Hidden_state or the loss and the logits
mask: The

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

doc for the attention mask is missing. Could make sense to rename mask to attention_mask as there is also loss_mask

@jorisSchaller
jorisSchaller requested a review from jannisborn April 8, 2026 08:36
Comment thread src/mblm/model/mblm.py Outdated
if loss_mask is not None:
# potentially apply the loss mask. this does not involve
# broadcasting as after slicing above, the mask and the loss tensor
# potentially apply the loss attention_mask. this does not involve

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

revert attention_mask back to mask

Comment thread src/mblm/model/mblm.py Outdated
loss_tensor *= loss_mask

# after applying the mask, some elements might be 0 - they should not be
# after applying the attention_mask, some elements might be 0 - they should not be

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

here we mean the loss mask

@jannisborn
jannisborn merged commit 6cfdf6e into main Apr 14, 2026
7 checks passed
@jannisborn
jannisborn deleted the fix/encoder-with-attention branch April 14, 2026 10:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants