如何在德语GPT-2模型中单独使用每层进行下词预测

核心思路 #

要让12层各自单独预测下词，核心是让仅目标层参与输入变换，其余层完全不处理输入（直接传递原始值），再将目标层的输出接入最后的语言模型头（lm_head）生成预测结果。无需修改模型权重，只需手动控制每一层的前向计算逻辑即可。

实现步骤与代码 #

1. 基础准备（加载模型与处理输入） #

首先加载预训练模型、分词器，并处理输入文本：

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 加载德语文本训练的GPT-2模型与分词器
tokenizer = AutoTokenizer.from_pretrained("dbmdz/german-gpt2")
model = AutoModelForCausalLM.from_pretrained("dbmdz/german-gpt2", pad_token_id=tokenizer.eos_token_id)

# 设置模型为评估模式（关闭训练相关的dropout等操作）
model.eval()
# 选择设备（GPU优先）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 处理输入文本
input_text = "Orlando liebte von Natur aus einsame Orte, weite Ausblicke und das Gefühl, für immer und ewig"
inputs = tokenizer(input_text, return_tensors="pt").to(device)

2. 计算初始嵌入 #

GPT-2的输入先经过词嵌入（wte）和位置嵌入（wpe）的叠加，得到初始的隐藏状态：

# 词嵌入 + 位置嵌入
seq_len = inputs["input_ids"].size(1)
position_ids = torch.arange(seq_len, dtype=torch.long, device=device).unsqueeze(0)
embeds = model.transformer.wte(inputs["input_ids"]) + model.transformer.wpe(position_ids)

3. 遍历每一层，单独生成预测 #

对每一层（0到11），仅让目标层处理输入，其余层直接传递输入，最后通过lm_head预测下一个词：

layer_predictions = []

with torch.no_grad():  # 关闭梯度计算，提升效率
    for target_layer_idx in range(12):
        # 重置当前隐藏状态为初始嵌入
        current_hidden = embeds.clone()
        
        # 遍历所有Transformer层
        for layer_idx, layer in enumerate(model.transformer.h):
            if layer_idx == target_layer_idx:
                # 仅目标层正常执行前向计算
                current_hidden = layer(current_hidden)[0]  # GPT2Block返回元组，第一个元素为隐藏状态
        
        # 取最后一个token的隐藏状态预测下词
        logits = model.lm_head(current_hidden[:, -1, :])
        # 获取概率最高的词
        predicted_token_id = torch.argmax(logits, dim=-1).item()
        predicted_word = tokenizer.decode(predicted_token_id).strip()
        
        layer_predictions.append({
            "层索引": target_layer_idx,
            "预测下词": predicted_word
        })

# 打印所有层的预测结果
for pred in layer_predictions:
    print(f"层 {pred['层索引']}: 预测下词 -> {pred['预测下词']}")

关键说明 #

• 为什么不修改权重？：直接控制层的执行逻辑比修改权重更高效且可逆，不会破坏原模型的权重参数，无需重新加载模型即可恢复正常使用。
• 关于“失效”其他层：非目标层直接跳过计算，输入值完全不被修改，确保仅目标层对初始嵌入进行特征变换，符合“单独预测”的需求。
• 评估模式的作用：model.eval()会关闭模型中的dropout、层归一化的训练模式，保证预测结果稳定。
• 关于“设置注意力权重为0”的思路：该方法只能让注意力层失效，但层内的前馈网络仍会处理输入，无法实现“完全不参与”的效果。直接跳过层的计算是更彻底的方案。

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