能否在HuggingFace预训练语言模型中实现类GNN的2D注意力掩码？

可以实现，绝大多数基于Transformer的HuggingFace预训练语言模型都支持传入自定义二维注意力掩码，以此限制每个token仅关注相邻token，模拟GNN中的邻接注意力行为。

核心原理 #

默认的注意力掩码（一维）仅用于屏蔽padding token，而二维注意力掩码是形状为[batch_size, seq_len, seq_len]的张量，每个元素直接控制第i个token是否能对第j个token产生注意力——通常用1表示允许注意力，0（或负无穷）表示禁止。

构造邻接注意力掩码 #

以“每个token仅关注自身及左右各1个相邻token”为例，构造掩码的逻辑如下：

1. 假设序列长度为seq_len，初始化全0的seq_len×seq_len矩阵
2. 对每个token位置i，标记其窗口范围内（max(0,i-1)到min(seq_len,i+2)）的位置为1

示例代码（PyTorch）：

import torch
from transformers import BertModel, BertTokenizer

# 加载模型和分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 处理输入文本
text = "This is a sample sentence for testing adjacency attention"
inputs = tokenizer(text, return_tensors='pt')
seq_len = inputs['input_ids'].size(1)

# 构造窗口大小为3（自身+左右各1）的二维掩码
window_size = 1
attention_mask = torch.zeros((seq_len, seq_len), dtype=torch.float32)
for i in range(seq_len):
    start_idx = max(0, i - window_size)
    end_idx = min(seq_len, i + window_size + 1)
    attention_mask[i, start_idx:end_idx] = 1.0

# 扩展batch维度（适配单样本输入）
attention_mask = attention_mask.unsqueeze(0)

# 传入模型执行前向计算
outputs = model(**inputs, attention_mask=attention_mask)

注意事项 #

• 掩码格式适配：不同模型对掩码的数值类型要求不同——部分模型接受布尔型掩码，部分需要用-1e9代替0表示禁止（避免softmax后产生无效权重），建议查看对应模型的官方文档确认格式要求。
• 微调必要性：预训练模型是在全局注意力机制下训练的，直接改用邻接注意力会导致性能下降，需要在目标任务上重新微调。
• 自回归模型适配：对于GPT这类自回归模型，本身自带因果掩码（仅允许关注前文token），需要将自定义邻接掩码与因果掩码做逻辑与操作，确保同时满足“邻接”和“不能关注后文”的约束。

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