置换BERT位置编码未达预期效果，是否存在其他位置相关机制？

我正在开发一个Transformer相关的Jupyter Notebook，目的是演示Transformer完全依赖位置编码来理解序列顺序。之前了解到这个概念仅适用于无掩码注意力的模型（比如GPT-2），但用BERT模型（带交叉注意力）预测[MASK] token时，得到了意外结果。

预期结果 #

• 不做任何置换时，模型预测的token分布（分布A）保持一致
• 仅置换输入ID得到分布B
• 仅置换位置嵌入得到分布B
• 同时置换输入ID和位置嵌入得到分布A

实际结果 #

有时结果符合预期，但有时仅置换输入ID或位置嵌入会产生不同结果，偶尔二者结果又相同。

核心问题 #

Hugging Face的BERT模型中，除位置编码外，是否存在其他受位置影响的机制？

复现代码 #

import torch
import ipywidgets as widgets
from IPython.display import display
from transformers import BertForMaskedLM, AutoTokenizer
import matplotlib.pyplot as plt
import torch.nn.functional as F
import logging
import warnings

# surpress renaming warnings
logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
warnings.simplefilter("ignore", FutureWarning)

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

input_ids = torch.Tensor([[]])
tokens = []
permutation = []

output = widgets.Output()

def permute_columns(matrix, permutation=None):
    n = len(permutation)
    first_n_columns = matrix[:, :n]
    permuted_columns = first_n_columns[:, permutation]
    remaining_columns = matrix[:, n:]
    new_matrix = torch.hstack((permuted_columns, remaining_columns))
    return new_matrix

def update_permutation(ordered_tags):
    global permutation
    fixed_tokens = [tokens[0]] + ordered_tags + [tokens[-1]]
    
    permutation = [tokens.index(tag) for tag in fixed_tokens]
    

def tokenize(text):
    global input_ids, tokens
    input_ids = tokenizer(text, return_tensors="pt").input_ids
    tokens = [tokenizer.decode([token_id]).strip() for token_id in input_ids[0]]
    
    if len(tokens) > 2:
        reorderable_tokens = tokens[1:-1]
    else:
        reorderable_tokens = []
    
    with output:
        output.clear_output(wait=True)
        tags_input.allowed_tags = reorderable_tokens
        tags_input.value = reorderable_tokens
        update_permutation(tags_input.value)

def on_tags_change(change):
    if len(change['new']) != len(tags_input.allowed_tags):
        tags_input.value = tags_input.allowed_tags  # Restore original value


def masked_prediction(input_ids, permutation, permute_input, permute_encoding):
    
    with output:
        output.clear_output(wait=True)  # Clear previous outputs
        
        if input_ids.numel() == 0:
            print("You can't use an empty sequence for prediction")
            return
        
        model = BertForMaskedLM.from_pretrained("bert-base-uncased")
        
        if permute_encoding:
            model.bert.embeddings.position_embeddings.weight.data = permute_columns(model.bert.embeddings.position_embeddings.weight.T, permutation).T
        if permute_input:
            input_ids = permute_columns(input_ids, permutation)
            
        decoded_text = tokenizer.decode(input_ids[0], skip_special_tokens=False)
            
        with torch.no_grad():
            outputs = model(input_ids)
            
        logits = outputs.logits

        top_k = 5

        mask_token_indices = torch.where(input_ids == tokenizer.mask_token_id)[1]
        print(decoded_text, mask_token_indices, permutation)
        num_masks = len(mask_token_indices)
        if num_masks == 0:
            print("You need to include a [MASK] token for prediction")
            return

        fig, axs = plt.subplots(1, num_masks, figsize=(15, 6))
        
        if num_masks == 1:
            axs = [axs]

        for i, idx in enumerate(mask_token_indices):
            mask_token_logits = logits[0, idx, :]

            softmax_probs = F.softmax(mask_token_logits, dim=0)

            top_token_probs, top_token_ids = torch.topk(softmax_probs, top_k, dim=0)

            predicted_tokens = [tokenizer.decode([token_id]).strip() for token_id in top_token_ids]
            predicted_confidences = top_token_probs.tolist()

            axs[i].bar(predicted_tokens, predicted_confidences, color='blue')
            axs[i].set_xlabel('Predicted Tokens')
            axs[i].set_ylabel('Confidence')
            axs[i].set_title(f'Masked Token at Position {idx.item()}')
            axs[i].set_ylim(0, 1)

        plt.show()

def on_predict_button_click(b):
    masked_prediction(input_ids, permutation, permute_input_checkbox.value, permute_encoding_checkbox.value)

text_input = widgets.Text(placeholder='Write text here to encode.', description='Input:')
text_input.observe(lambda change: tokenize(change['new']), names='value')
tags_input = widgets.TagsInput(value=[], allowed_tags=[], allow_duplicates=False)

# Observe changes in tags order to update the permutation and prevent deletion
tags_input.observe(on_tags_change, names='value')
tags_input.observe(lambda change: update_permutation(change['new']), names='value')

# Create checkboxes for permute_input and permute_encoding
permute_input_checkbox = widgets.Checkbox(value=False, description='Permute Inputs')
permute_encoding_checkbox = widgets.Checkbox(value=False, description='Permute Encodings')

# Create a button to trigger the prediction
predict_button = widgets.Button(description="Run Prediction")
predict_button.on_click(on_predict_button_click)

# Display the widgets
display(text_input)
display(tags_input)
display(permute_input_checkbox)
display(permute_encoding_checkbox)
display(predict_button)
display(output)

内容的提问来源于stack exchange，提问作者Lukas