MLP Head的输入输出是什么

### MLP Head 的输入输出细节

在多层感知器（MLP）结构中，特别是在涉及自注意力机制的架构里，MLP Head 起着至关重要的作用。具体到分类任务中的 MLP Head 部分，在接收来自前一层特征表示之后，会进一步处理这些信息以生成最终类别预测。

对于 MLP Head 来说：

- **输入**: 接收到的是经过一系列变换后的向量序列，通常是从 Transformer 编码器最后一层提取出来的隐藏状态。每个 token 对应一个 d-dimensional 向量，其中 d 表示模型设定好的隐含单元数。特别地，在某些实现方式下，可能会有一个特殊的 [CLS] 标记作为整个句子或片段的代表[^1]。

- **内部操作**: 这些向量会被送入一个多层全连接网络 (Fully Connected Layers)，每一层都可能应用激活函数来引入非线性特性。此过程可以看作是对原始特征空间的一种映射转换，旨在捕捉更复杂的模式并增强表达能力。

- **输出**: 如果用于分类，则最后一个 FC 层将把上述得到的新特征映射至指定数量的目标类别的概率分布上；如果是回归或其他类型的下游任务，则相应调整该层的设计目标。例如，在图像分类场景中，如果存在 C 类标签，则输出形状为 (batch_size, num_classes)。

import torch.nn as nn

class MLPClsHead(nn.Module):
    def __init__(self, hidden_dim=768, output_dim=10):
        super().__init__()
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):  # 假设x.shape=(batch_size, seq_len, hidden_dim)
        cls_token_output = x[:, 0, :]  # 只取[CLS]对应的hidden state
        logits = self.fc(cls_token_output)  # shape变为(batch_size, output_dim)
        return logits