Custom attention mask

BramVanroy · August 1, 2020, 11:00pm

Is it possible in either the py or tf version to add a custom attention mask to the input? So that in addition to the padding tokens also other tokens are masked?

If it is not currently supported, where should I have a look to implement this myself? I imagine a separate file with 0s and 1s for all tokens, or some sort of word feature appended to each token to indicate which words should get attention and which can be ignored. Perhaps an alternative is swapping out the respective tokens with padding tokens. Would that work? If so, what’s onmt’s padding token?

Thanks

guillaumekln · August 4, 2020, 1:55pm

This sounds possible but you’ll still need to edit the code to generate masks based on the padding id and not the sequence length. See for example the Transformer encoder in OpenNMT-py (the code is similar in the decoder):

github.com

OpenNMT/OpenNMT-py/blob/master/onmt/encoders/transformer.py#L124


        embeddings,
        opt.max_relative_positions)

def forward(self, src, lengths=None):
    """See :func:`EncoderBase.forward()`"""
    self._check_args(src, lengths)

    emb = self.embeddings(src)

    out = emb.transpose(0, 1).contiguous()
    mask = ~sequence_mask(lengths).unsqueeze(1)
    # Run the forward pass of every layer of the tranformer.
    for layer in self.transformer:
        out = layer(out, mask)
    out = self.layer_norm(out)

    return emb, out.transpose(0, 1).contiguous(), lengths

def update_dropout(self, dropout, attention_dropout):
    self.embeddings.update_dropout(dropout)
    for layer in self.transformer:

The padding token is <blank>.

Adding a new input file is generally more painful. If you go in that direction, I believe it is easier to do in OpenNMT-tf as there is no preprocessing step.

BramVanroy · August 4, 2020, 2:20pm

Thanks for the reply.

Modifying source and our input data seems easier, indeed. So if I am correct the steps are the following:

replace the respective tokens by <blank>;
modify the forward of the encoder so that the mask is determined by the padding tokens rather than lengths alone. Probably want to return the mask here rather than just the lengths;
in the model, pass the returned src_mask (previously lengths) to the decoder;
in the decoder, directly use the passed src_mask rather than (again) calculating the mask based on sequence length.

In onmt, should the masked tokens (e.g. the padding token) be 0 or 1? From reading the code I would guess 0.

guillaumekln · August 4, 2020, 2:55pm

This sounds about right.

The padding index is 1 in OpenNMT-py.

github.com

OpenNMT/OpenNMT-py/blob/master/onmt/inputters/inputter.py#L342-L344


all_specials = [
    field.unk_token, field.pad_token, field.init_token, field.eos_token
]

BramVanroy · August 4, 2020, 3:07pm

I think you misunderstood. I don’t mean the padding index but I wonder what the value for masked tokens in the mask tensor is. There are some differences out there: sometimes 1 is used for the masked tokens, other times 0.

guillaumekln · August 4, 2020, 3:12pm

Looks like padding positions should be set to 1:

github.com

OpenNMT/OpenNMT-py/blob/master/onmt/modules/multi_headed_attn.py#L199


query_key = torch.matmul(query, key.transpose(2, 3))

if self.max_relative_positions > 0 and attn_type == "self":
    scores = query_key + relative_matmul(query, relations_keys, True)
else:
    scores = query_key
scores = scores.float()

if mask is not None:
    mask = mask.unsqueeze(1)  # [B, 1, 1, T_values]
    scores = scores.masked_fill(mask, -1e18)

# 3) Apply attention dropout and compute context vectors.
attn = self.softmax(scores).to(query.dtype)
drop_attn = self.dropout(attn)

context_original = torch.matmul(drop_attn, value)

if self.max_relative_positions > 0 and attn_type == "self":
    context = unshape(context_original
                      + relative_matmul(drop_attn,

BramVanroy · August 4, 2020, 6:01pm

Great, I’m having a quick look now. Could you explain to me the sizes of the inputs in the encoder/decoder forward? For batch size three (as a test) I am getting

src torch.Size([6, 3, 2])
tgt, torch.Size([11, 3, 1])

What is that last dimension, and why is there a difference between the encoder (src) and decoder (tgt)?

BramVanroy · August 4, 2020, 6:44pm

I found out this last dimension are the features, including the tokens themselves. (So in my case two source features (token, and another) and one target (token))

BramVanroy · August 6, 2020, 1:36pm

I am running into some issues during translation. During training everything works fine, but during translation the multiheadattention in the decoder throws an error about the mask size.

Traceback (most recent call last):
  File "/home/bram/.local/share/virtualenvs/nfr-experiments-3R5lX5O6/bin/onmt_translate", line 11, in <module>
    load_entry_point('OpenNMT-py', 'console_scripts', 'onmt_translate')()
  File "/home/bram/Python/projects/nfr-experiments/OpenNMT-py/onmt/bin/translate.py", line 48, in main
    translate(opt)
  File "/home/bram/Python/projects/nfr-experiments/OpenNMT-py/onmt/bin/translate.py", line 25, in translate
    translator.translate(
  File "/home/bram/Python/projects/nfr-experiments/OpenNMT-py/onmt/translate/translator.py", line 361, in translate
    batch_data = self.translate_batch(
  File "/home/bram/Python/projects/nfr-experiments/OpenNMT-py/onmt/translate/translator.py", line 550, in translate_batch
    return self._translate_batch_with_strategy(batch, src_vocabs,
  File "/home/bram/Python/projects/nfr-experiments/OpenNMT-py/onmt/translate/translator.py", line 674, in _translate_batch_with_strategy
    log_probs, attn = self._decode_and_generate(
  File "/home/bram/Python/projects/nfr-experiments/OpenNMT-py/onmt/translate/translator.py", line 589, in _decode_and_generate
    dec_out, dec_attn = self.model.decoder(
  File "/home/bram/.local/share/virtualenvs/nfr-experiments-3R5lX5O6/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/bram/Python/projects/nfr-experiments/OpenNMT-py/onmt/decoders/transformer.py", line 319, in forward
    output, attn, attn_align = layer(
  File "/home/bram/.local/share/virtualenvs/nfr-experiments-3R5lX5O6/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/bram/Python/projects/nfr-experiments/OpenNMT-py/onmt/decoders/transformer.py", line 93, in forward
    output, attns = self._forward(*args, **kwargs)
  File "/home/bram/Python/projects/nfr-experiments/OpenNMT-py/onmt/decoders/transformer.py", line 165, in _forward
    mid, attns = self.context_attn(memory_bank, memory_bank, query_norm,
  File "/home/bram/.local/share/virtualenvs/nfr-experiments-3R5lX5O6/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/bram/Python/projects/nfr-experiments/OpenNMT-py/onmt/modules/multi_headed_attn.py", line 202, in forward
    scores = scores.masked_fill(mask, -1e18)
RuntimeError: The size of tensor a (30) must match the size of tensor b (150) at non-singleton dimension 0

After doing some logging I found that the encoder does not have any issues, but that the decoder cannot simply use the src_mask that was created in the encoder. I am not sure why. Here are the logs, and you can see that in the first decoder layer multi-head attention will fail because there is a shape mismatch in the code snippet that you linked. So the mask should be filled for the padded tokens, but the dimensions do not match.

Entering encoder layer 0
score torch.Size([30, 8, 223, 223])
mask torch.Size([30, 1, 1, 223])

Entering encoder layer 1
score torch.Size([30, 8, 223, 223])
mask torch.Size([30, 1, 1, 223])

Entering encoder layer 2
score torch.Size([30, 8, 223, 223])
mask torch.Size([30, 1, 1, 223])

Entering encoder layer 3
score torch.Size([30, 8, 223, 223])
mask torch.Size([30, 1, 1, 223])

Entering encoder layer 4
score torch.Size([30, 8, 223, 223])
mask torch.Size([30, 1, 1, 223])

Entering encoder layer 5
score torch.Size([30, 8, 223, 223])
mask torch.Size([30, 1, 1, 223])

decoder_input torch.Size([1, 150, 1])
src_pad_mask torch.Size([30, 1, 223]) # here
tgt_pad_mask torch.Size([150, 1, 1])

Entering decoder layer 0
score torch.Size([150, 8, 1, 223])
mask torch.Size([30, 1, 1, 223]) # here

So it seems that the decoder expects a much larger attention mask then the encoder returns, but I do not understand why. Do you have any thoughts on this? Especially odd to me since this does work in training.

I only made a few changes to the encoder/decoder, which you can find here: https://github.com/BramVanroy/OpenNMT-py/tree/masked_tokens

guillaumekln · August 6, 2020, 1:41pm

Most likely you are running beam search. You should repeat each batch in the tensor beam_size times.

For example, you can see the memory outputs and state being repeated in this function:

github.com

OpenNMT/OpenNMT-py/blob/master/onmt/translate/beam_search.py#L96


    # "global state" of the old beam
    self._prev_penalty = None
    self._coverage = None

    self._stepwise_cov_pen = (
        stepwise_penalty and self.global_scorer.has_cov_pen)
    self._vanilla_cov_pen = (
        not stepwise_penalty and self.global_scorer.has_cov_pen)
    self._cov_pen = self.global_scorer.has_cov_pen

def initialize(self, memory_bank, src_lengths, src_map=None, device=None,
               target_prefix=None):
    """Initialize for decoding.
    Repeat src objects `beam_size` times.
    """

    def fn_map_state(state, dim):
        return tile(state, self.beam_size, dim=dim)

    if isinstance(memory_bank, tuple):
        memory_bank = tuple(tile(x, self.beam_size, dim=1)

BramVanroy · August 6, 2020, 2:07pm

Ah that makes sense. I never realised that beam search was enabled by default in translation mode.

So in the Strategy initializer I now tile the mask, like so:

tile(src_mask, self.beam_size)

And return the value, too so that it can be used in the decoding steps loop. Updated in this commit: https://github.com/BramVanroy/OpenNMT-py/commit/9271fc1483b2ac4c2e786ed9e63195b078782d3e

Funnily enough, I am still getting a size mismatch:

decoder_input torch.Size([1, 145, 1])

src_pad_mask torch.Size([150, 1, 223])
tgt_pad_mask torch.Size([145, 1, 1])

Entering decoder layer 0
score torch.Size([145, 8, 1, 223])
mask torch.Size([150, 1, 1, 223])

Should the target mask be padded up to the max src length? I don’t think so, or at least I don’t see that happening in the original forward:

github.com

OpenNMT/OpenNMT-py/blob/592a2f62f098c69beeb5487f93d6c74321363630/onmt/decoders/transformer.py#L306


emb = self.embeddings(tgt, step=step)
assert emb.dim() == 3  # len x batch x embedding_dim

output = emb.transpose(0, 1).contiguous()
src_memory_bank = memory_bank.transpose(0, 1).contiguous()

pad_idx = self.embeddings.word_padding_idx
src_lens = kwargs["memory_lengths"]
src_max_len = self.state["src"].shape[0]
src_pad_mask = ~sequence_mask(src_lens, src_max_len).unsqueeze(1)
tgt_pad_mask = tgt_words.data.eq(pad_idx).unsqueeze(1)  # [B, 1, T_tgt]

with_align = kwargs.pop('with_align', False)
attn_aligns = []

for i, layer in enumerate(self.transformer_layers):
    layer_cache = self.state["cache"]["layer_{}".format(i)] \
        if step is not None else None
    output, attn, attn_align = layer(
        output,
        src_memory_bank,

guillaumekln · August 6, 2020, 2:11pm

Another thing that is happening is that finished translations are removed from the batch. So you should update your tensor accordingly:

github.com

OpenNMT/OpenNMT-py/blob/master/onmt/translate/translator.py#L692



decode_strategy.advance(log_probs, attn)
any_finished = decode_strategy.is_finished.any()
if any_finished:
    decode_strategy.update_finished()
    if decode_strategy.done:
        break

select_indices = decode_strategy.select_indices

if any_finished:
    # Reorder states.
    if isinstance(memory_bank, tuple):
        memory_bank = tuple(x.index_select(1, select_indices)
                            for x in memory_bank)
    else:
        memory_bank = memory_bank.index_select(1, select_indices)

    memory_lengths = memory_lengths.index_select(0, select_indices)

    if src_map is not None:

BramVanroy · August 6, 2020, 2:24pm

Ah yes, noted. It seems to work well now.

Is this change useful for upstream? If so, I can do a PR which basically allows tokens to be inserted anywhere in source and not receiving attention.

guillaumekln · August 6, 2020, 2:25pm

@vince62s @francoishernandez Would you accept a PR for that?

francoishernandez · August 17, 2020, 1:08pm

This looks like an interesting idea indeed. I thought of implementing such a feature at some point but never got around to doing it.
How would this be specified in the args though? Some additional file with the masks on each line? Or would it be a list of specific tokens to ignore?

BramVanroy · August 17, 2020, 1:46pm

I would prefer either

an additional file that should contain 0 and 1 to mask/unmask tokens. That means that every line should the same number of tokens as there are 0s and 1s; or, preferably:
a flag in args such as --custom_attention_mask, which uses the features pipe format. If the flag is enabled, the first “feature” would actually indicate whether or not a token is masked or not. So e.g. i|0|c like|1|l cookies|1|l

francoishernandez · August 17, 2020, 1:48pm

Your first proposition is what I meant by my first one. Seems good to me.
(I fear the feature one might be a bit misleading.)
Edit: sorry, everything was clear in your initial post but I read too quickly.

francoishernandez · August 17, 2020, 7:09pm

Also, @BramVanroy may I ask what would be your use case for such a feature?What I had in mind may not exactly be what you imply here in fact.

Would the rationale be something like MASS or BART ? Probably not though, since you would use a fixed mask here.

Would be nice if you could elaborate a bit. Thanks!

BramVanroy · August 17, 2020, 7:54pm

We are building on previous work considering neural fuzzy repair, which builds on integrating fuzzy match targets on the source side to direct the system to a correct translation (so your source input consists of the source text + noisy (fuzzy) target). For a new paper we are trying out a number of different strategies. One of the things that we wanted to try was to see how we can tell the system which parts of the fuzzy match are relevant and which aren’t.

We tried a number of approaches (look forward to the paper, which is in the works), and one of them is to add a feature to fuzzy tokens, telling the system which ones are actually matched and which are just noise. This binary feature works quite well. I personally also wanted to train what would happen if we just block all attention that these irrelevant tokens receive, hence this use case. I found that this does not work well and that a matched/not_matched feature works much better. My explanation would be that the context that these “irrelevant tokens” provide is still quite useful to the system, even though their direct, concrete meaning is not relevant.

francoishernandez · August 18, 2020, 7:01am

Interesting.
But if I understand correctly, the mask does not work well compared to the features in your experiments, so it might not be that useful as a feature after all.
Anyways, if it’s not too much of a hassle you can always sent a PR and we might merge it. It could be good if some people want to further experiment down this path.