如何同时获取预训练RoBERTa的pooled_output与Splade风格嵌入？

你不需要分别加载两个独立模型，通过复用基础RoBERTa模型并手动构建MLM预测头，就能一次前向传播同时拿到两种输出，既省资源又高效。

方案一：自定义双输出模型（推荐） #

直接基于基础RoBERTa模型搭建带MLM头的自定义模型，同时输出pooled_output和MLM logits：

import tensorflow as tf
from transformers import TFAutoModel, AutoTokenizer

model_name = "roberta-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
base_model = TFAutoModel.from_pretrained(model_name)
config = base_model.config

# 实现RoBERTa的MLM预测头，与官方TFAutoModelForMaskedLM结构一致
class TFRobertaMLMHead(tf.keras.layers.Layer):
    def __init__(self, config):
        super().__init__()
        self.dense = tf.keras.layers.Dense(config.hidden_size, activation=config.hidden_act, name="dense")
        self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")
        self.decoder = tf.keras.layers.Dense(config.vocab_size, name="decoder")

    def call(self, features):
        x = self.dense(features)
        x = self.layer_norm(x)
        return self.decoder(x)

# 构建双输出模型
input_ids = tf.keras.layers.Input(shape=(None,), dtype=tf.int32, name="input_ids")
attention_mask = tf.keras.layers.Input(shape=(None,), dtype=tf.int32, name="attention_mask")

# 基础模型输出
base_outputs = base_model(input_ids, attention_mask=attention_mask)
pooled_output = base_outputs.pooler_output
sequence_output = base_outputs.last_hidden_state

# MLM头输出logits
mlm_head = TFRobertaMLMHead(config)
mlm_logits = mlm_head(sequence_output)

# 定义模型，同时返回两种结果
dual_model = tf.keras.Model(inputs=[input_ids, attention_mask], outputs=[pooled_output, mlm_logits])

使用示例 #

# 处理输入文本
text = "Your input text here"
inputs = tokenizer(text, return_tensors="tf", padding=True, truncation=True)

# 一次前向传播获取双输出
pooled_embedding, mlm_logits = dual_model(inputs)

# 常规嵌入：直接用pooled_embedding
print("Pooled Embedding Shape:", pooled_embedding.shape)

# Splade风格嵌入：对logits取绝对值后按token维度取最大值
splade_embedding = tf.reduce_max(tf.abs(mlm_logits), axis=1)
print("Splade Embedding Shape:", splade_embedding.shape)

方案二：复用预训练MLM模型（简单但稍低效） #

如果你不想手动构建MLM头，也可以直接加载TFAutoModelForMaskedLM，然后从它的base_model属性获取基础模型的pooled_output：

from transformers import TFAutoModelForMaskedLM

mlm_model = TFAutoModelForMaskedLM.from_pretrained(model_name)
base_model = mlm_model.base_model

# 前向传播获取MLM logits
mlm_outputs = mlm_model(inputs)
mlm_logits = mlm_outputs.logits

# 前向传播获取pooled_output
base_outputs = base_model(inputs)
pooled_output = base_outputs.pooler_output

这种方式需要两次前向传播，资源消耗比方案一高，适合快速验证场景。

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