Translation model Cross Attention problem, inconsistent with paper !

[42694647426]

When I'm inspecting the cross-attention layers from the pretrained transformer translation model (MarianMT model), It is very strange that the cross attention from layer 0 and 1 provide best alignment between input and output. I used bertviz to visualize all heads from all 6 layers, and tried different language, english to german and english to chinese, it all gives the same results, which does not make sense because the last layers should be more accurate according to the paper Jointly Learning to Align and Translate with Transformer Models https://arxiv.org/pdf/1909.02074.pdf

But when I'm looking at the cross attention of model Helsinki-NLP/opus-mt-en-de and Helsinki-NLP/opus-mt-en-zh , the layer 1 gives the best alignment. the code is below:

from transformers import AutoTokenizer, AutoModel
import os
os.environ['TRANSFORMERS_CACHE'] = '/data2/hanyings/.cache'

tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-de")
model = AutoModel.from_pretrained("Helsinki-NLP/opus-mt-en-de", output_attentions=True)

encoder_input_ids = tokenizer("She sees the small elephant.", return_tensors="pt", add_special_tokens=True).input_ids
with tokenizer.as_target_tokenizer():
    decoder_input_ids = tokenizer("Sie sieht den kleinen Elefanten.", return_tensors="pt", add_special_tokens=True).input_ids

outputs = model(input_ids=encoder_input_ids, decoder_input_ids=decoder_input_ids)

encoder_text = tokenizer.convert_ids_to_tokens(encoder_input_ids[0])
decoder_text = tokenizer.convert_ids_to_tokens(decoder_input_ids[0])

from bertviz import model_view
model_view(
    encoder_attention=outputs.encoder_attentions,
    decoder_attention=outputs.decoder_attentions,
    cross_attention=outputs.cross_attentions,
    encoder_tokens= encoder_text,
    decoder_tokens = decoder_text
)

And the results are:


From the above pictures, I observed that the first 2 layers give the best alignment whereas the last layers do not align the input and output tokens properly. Can you please help me to explain why this happens? and If the alignment of the last layer is not accurate, how does the model provide correct predictions? Please! It is very important for my research project!
@patrickvonplaten @sshleifer @Tiedeman @joaogante @bobosui

[Kai-Piontek]

Did you solve this eventually?