如何在PyTorch中获取BERT模型的输入梯度？

首先需要明确：如果你的输入x是离散的input_ids（整数类型），这类离散张量无法直接计算梯度——梯度是针对连续数值的数学概念，离散变量不存在导数。这种情况下，你需要针对BERT嵌入层输出的连续张量求梯度，或者构造连续的输入张量。

下面分两种场景给出实现方案：

场景1：针对连续输入张量求梯度 #

如果你的x是连续张量（比如直接传入BERT编码器的嵌入向量），可以按以下步骤操作：

• 确保输入张量开启梯度追踪（requires_grad=True）
• 前向传播得到模型输出
• 计算损失并执行反向传播
• 从输入张量的.grad属性获取梯度

示例代码（基于HuggingFace Transformers库）：

import torch
from transformers import BertModel

# 初始化BERT模型
model = BertModel.from_pretrained('bert-base-uncased')
# 若不需要更新模型参数，可关闭参数的梯度追踪以节省内存
for param in model.parameters():
    param.requires_grad = False

# 构造连续输入张量（形状：[batch_size, seq_len, hidden_size]）
batch_size, seq_len, hidden_size = 2, 10, 768
x = torch.randn(batch_size, seq_len, hidden_size, requires_grad=True)

# 前向传播（直接调用编码器，跳过嵌入层）
encoder_outputs = model.encoder(x)
y_prime = encoder_outputs.last_hidden_state

# 构造目标张量并计算损失（这里以MSE损失为例）
y = torch.randn_like(y_prime)
loss = torch.nn.functional.mse_loss(y_prime, y)

# 反向传播计算梯度
loss.backward()

# 获取输入x的梯度
input_gradient = x.grad
print(input_gradient.shape)  # 输出: torch.Size([2, 10, 768])

场景2：针对离散input_ids的嵌入向量求梯度 #

如果你的输入是离散的input_ids，需要先将其转换为BERT的嵌入向量，再对这个连续的嵌入张量求梯度：

示例代码：

import torch
from transformers import BertModel, BertTokenizer

# 初始化Tokenizer和BERT模型
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
# 关闭模型参数的梯度追踪（仅计算输入嵌入的梯度）
for param in model.parameters():
    param.requires_grad = False

# 准备文本输入并转换为input_ids
texts = ["Hello world", "PyTorch gradient"]
inputs = tokenizer(texts, return_tensors='pt', padding=True, truncation=True)
input_ids = inputs['input_ids']
attention_mask = inputs['attention_mask']

# 生成完整的BERT嵌入向量（word + position + token_type嵌入），并开启梯度追踪
position_ids = torch.arange(input_ids.size(1), device=input_ids.device).unsqueeze(0).repeat(input_ids.size(0), 1)
token_type_ids = torch.zeros_like(input_ids)

full_embeddings = model.embeddings.word_embeddings(input_ids) + \
                  model.embeddings.position_embeddings(position_ids) + \
                  model.embeddings.token_type_embeddings(token_type_ids)
full_embeddings = model.embeddings.LayerNorm(full_embeddings)
full_embeddings = model.embeddings.dropout(full_embeddings)
full_embeddings.requires_grad_(True)

# 前向传播通过编码器
encoder_outputs = model.encoder(full_embeddings, attention_mask=attention_mask)
y_prime = encoder_outputs.last_hidden_state

# 计算损失并反向传播
y = torch.randn_like(y_prime)
loss = torch.nn.functional.mse_loss(y_prime, y)
loss.backward()

# 获取嵌入向量的梯度（对应原input_ids每个位置的梯度）
embedding_gradient = full_embeddings.grad
print(embedding_gradient.shape)  # 输出: torch.Size([2, seq_len, 768])

关键注意事项 #

• 若要仅计算输入的梯度，务必关闭模型参数的梯度追踪（param.requires_grad = False），否则反向传播会同时计算参数梯度，浪费内存和计算资源。
• 离散输入（如input_ids）无法直接求梯度，必须转换为连续的嵌入张量后再计算。

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