# The way to implement attention-mask/uni-direction attention in TransformerDecoder

**URL:** <https://discuss.pytorch.org/t/the-way-to-implement-attention-mask-uni-direction-attention-in-transformerdecoder/73124>\
**Category:** nlp\
**Created:** [March 13, 2020, 12:39pm UTC](https://discuss.pytorch.org/t/the-way-to-implement-attention-mask-uni-direction-attention-in-transformerdecoder/73124 "2020-03-13T12:39:37Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![AOZMH](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aozmh/32/33497_2.png) [@AOZMH](https://discuss.pytorch.org/u/AOZMH)\
**Post date:** [March 13, 2020, 12:39pm UTC](https://discuss.pytorch.org/t/the-way-to-implement-attention-mask-uni-direction-attention-in-transformerdecoder/73124/1 "2020-03-13T12:39:37Z")

</div>

Hi guys,

I’m learning about nn.Transformer in pytorch these days and I’m a bit confused about the implementation of the attention mask in decoder. Say we’re doing a machine translation task using Transformer, when inferencing, the output of each time step can only “see” the tokens before it. However, when training, we simply feed the correct sequence to the decoder and thus a token in the decoded sequence can see the tokens both before and after it, which I guess is not suitable for a robust model.

I know that GPT adopted single directional attention when decoding (and I guess it’s suitable for a decoder), but I’m wondering if the APIs in nn.Transformer can realize such feature (e.g. using attention masks for different tokens to avoid it from attending on later tokens)? If so, what should I do?

Any help would be appreciated, thanks in advance!

---

<div class="post-metadata">

**Author:** ![zhangguanheng66](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/z/8edcca/32.png) [@zhangguanheng66](https://discuss.pytorch.org/u/zhangguanheng66)\
**Post date:** [March 13, 2020, 2:46pm UTC](https://discuss.pytorch.org/t/the-way-to-implement-attention-mask-uni-direction-attention-in-transformerdecoder/73124/2 "2020-03-13T14:46:03Z")

</div>

You should set up the attention mask for the decoder to mask the token from see the one before it.

---

<div class="post-metadata">

**Author:** ![AOZMH](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aozmh/32/33497_2.png) [@AOZMH](https://discuss.pytorch.org/u/AOZMH)\
**Post date:** [March 13, 2020, 3:10pm UTC](https://discuss.pytorch.org/t/the-way-to-implement-attention-mask-uni-direction-attention-in-transformerdecoder/73124/3 "2020-03-13T15:10:19Z")

</div>

Thanks for the reply! Could you please offer more concrete hints, e.g. ways to setup attention masks, sample codes or a link that provide these? Many thanks!

---

<div class="post-metadata">

**Author:** ![zhangguanheng66](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/z/8edcca/32.png) [@zhangguanheng66](https://discuss.pytorch.org/u/zhangguanheng66)\
**Post date:** [March 13, 2020, 4:18pm UTC](https://discuss.pytorch.org/t/the-way-to-implement-attention-mask-uni-direction-attention-in-transformerdecoder/73124/4 "2020-03-13T16:18:37Z")

</div>

> <https://github.com/pytorch/examples/blob/5551061414d3bcf202de520d20e8163f58eb664a/word_language_model/model.py#L126>

---

<div class="post-metadata">

**Author:** ![AOZMH](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aozmh/32/33497_2.png) [@AOZMH](https://discuss.pytorch.org/u/AOZMH)\
**Post date:** [March 14, 2020, 2:44pm UTC](https://discuss.pytorch.org/t/the-way-to-implement-attention-mask-uni-direction-attention-in-transformerdecoder/73124/5 "2020-03-14T14:44:49Z")

</div>

Many thanks for the code, it’s much clearer right now!