# C4\_W2 ungraded lab masking, regarding addition of an extra dimension while creating padding mask

**URL:** <https://community.deeplearning.ai/t/c4-w2-ungraded-lab-masking-regarding-addition-of-an-extra-dimension-while-creating-padding-mask/620223>\
**Category:** NLP with Attention Models\
**Tags:** course-related, week-module-2\
**Created:** [May 2, 2024, 6:29am UTC](https://community.deeplearning.ai/t/c4-w2-ungraded-lab-masking-regarding-addition-of-an-extra-dimension-while-creating-padding-mask/620223 "2024-05-02T06:29:23Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![God\_of\_Calamity](https://sea1.discourse-cdn.com/dlai/user_avatar/community.deeplearning.ai/god_of_calamity/32/85939_2.png) [@God\_of\_Calamity](https://community.deeplearning.ai/u/God_of_Calamity)\
**Post date:** [May 2, 2024, 6:29am UTC](https://community.deeplearning.ai/t/c4-w2-ungraded-lab-masking-regarding-addition-of-an-extra-dimension-while-creating-padding-mask/620223/1 "2024-05-02T06:29:23Z")

</div>

![Screenshot 2024-05-02 at 11.49.21 AM](https://global.discourse-cdn.com/dlai/original/3X/b/8/b8a3cf1142a2cdd03c1b76e62d9ad750b798ba2b.png)  
As we can see in the above image we add an extra dimension to padding mask like [:, tf.newaxis, :]

When we feed data to attention block it will be like:  
input: (batch\_size, m, n) m is number of token, n is embedding dimensions  
Q: input \* Wq which is of shape (batch, m, dk)  
K: input \* Wk which is of shape (batch, m, dk)  
V: input \* Wv which is of shape (batch, m, dv)  
Here we assumen dk = dv = dm / h

Then we do (Q \* K.T) which is of shape (batch, m, m).  
And then we do (Q \* K.T + M), so mask should be of shape (batch, m, m)  
but as we can see in the above image we return a mask of shape  
(n, 1, m), should we add the new dimension to in the front like (1, n, m) so that it can be broadcasted, also should the mask be like (1, m, m) why is there an ‘n’ in it.

---

<div class="post-metadata">

**Author:** ![arvyzukai](https://sea1.discourse-cdn.com/dlai/user_avatar/community.deeplearning.ai/arvyzukai/32/165165_2.png) [@arvyzukai](https://community.deeplearning.ai/u/arvyzukai)\
**Post date:** [May 2, 2024, 7:27am UTC](https://community.deeplearning.ai/t/c4-w2-ungraded-lab-masking-regarding-addition-of-an-extra-dimension-while-creating-padding-mask/620223/2 "2024-05-02T07:27:40Z")

</div>

Hi @God_of_Calamity

> [@God\_of\_Calamity](#):
>
> Then we do (Q \* K.T) which is of shape (batch, m, m).  
> And then we do (Q \* K.T + M), so mask should be of shape (batch, m, m)

Yes, that is correct.

> [@God\_of\_Calamity](#):
>
> but as we can see in the above image we return a mask of shape  
> (n, 1, m),

I’m no expert on keras but I believe this is how the “1” dimension is broadcasted m times to become (batch, m, m). From the [docs](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention):

> Broadcasting can happen for the missing batch dimensions and the head dimension.

The behavior is not precisely explained, but I assume the 1 would represent the missing head dimension…  
You should check the tensorflow implementation to confirm or dispute that since I don’t want to go deeper. (Btw, the link “View source on GitHub” is “404 - page not found” as is often the case with TensorFlow 🙄)

> [@God\_of\_Calamity](#):
>
> should we add the new dimension to in the front like (1, n, m) so that it can be broadcasted

I don’t think that makes sense. What do you mean by that? What would be the shape after broadcasting?

> [@God\_of\_Calamity](#):
>
> also should the mask be like (1, m, m) why is there an ‘n’ in it.

The mask would be (1, m, m) if the batch size was 1 (for example, for inference). But since we have n sequences, we want the (n, m, m) shape. In other words, n here represent the batch size.

Cheers