如何获取GPT（BioGPT）模型词汇表内任意Token的Embedding？

核心方法：访问模型的输入嵌入层 #

BioGPT的词汇表Embedding直接存储在模型的输入嵌入层权重中，无需依赖输入batch即可获取：

# 获取输入嵌入层
embedding_layer = model.get_input_embeddings()
# 获取整个词汇表的Embedding，形状为 [42386, 1024]（词汇表大小×隐藏层维度）
full_vocab_embeddings = embedding_layer.weight

提取指定特殊Token的Embedding #

针对你关注的索引1、2、6、112的特殊Token，直接通过索引提取即可：

# 提取目标特殊Token的Embedding，形状为 [4, 1024]
special_token_embeds = full_vocab_embeddings[[1, 2, 6, 112], :]

生成你所需的[2, 1024, 42386, 1024]形状张量 #

如果需要将每个序列位置的隐藏状态与整个词汇表的Embedding结合，得到[batch_size, seq_len, vocab_size, hidden_size]的张量，可通过张量广播实现：

# 假设已获取模型输出的最后一层隐藏状态，形状为 [2, 1024, 1024]
last_hidden_states = out.hidden_states[-1]

# 将词汇表Embedding扩展维度以匹配广播要求，形状变为 [1, 1, 42386, 1024]
expanded_vocab_embeds = full_vocab_embeddings.unsqueeze(0).unsqueeze(0)

# 广播后得到目标形状的张量 [2, 1024, 42386, 1024]
# 这里以加法为例，可根据实际需求替换为其他运算
combined_embeddings = last_hidden_states.unsqueeze(2) + expanded_vocab_embeds

针对特殊Token生成对应形状张量 #

如果只需要结合指定特殊Token的Embedding，同样用广播实现：

# 将特殊Token的Embedding扩展维度，形状变为 [1, 1, 4, 1024]
expanded_special_embeds = special_token_embeds.unsqueeze(0).unsqueeze(0)

# 广播后得到 [2, 1024, 4, 1024] 的张量
combined_special_embeddings = last_hidden_states.unsqueeze(2) + expanded_special_embeds

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