PEFT与Sentence Transformers结合时的张量维度不匹配问题分析

问题背景

在使用PEFT(Parameter-Efficient Fine-Tuning)对基于BERT的Sentence Transformer模型进行提示调优(Prompt Tuning)时，开发人员遇到了一个关键的技术问题：在前向传播过程中，词嵌入张量与注意力掩码在池化操作时出现维度不匹配的情况。

问题现象

具体表现为当使用PromptTuningConfig配置时，模型在前向传播时会抛出RuntimeError，错误信息显示张量扩展尺寸不匹配。有趣的是，当使用LoRA或Prefix Tuning等其他PEFT方法时，模型能够正常工作。

技术分析

1. 问题根源

通过深入分析，发现问题出在模型架构的两个部分之间的交互上：

1. Transformer部分：经过PEFT的提示调优改造，这部分会插入虚拟标记(prompt embeddings)，形状为[1, num_virtual_tokens, hidden_size]，同时会相应扩展注意力掩码的维度。

2. 池化(Pooler)部分：属于Sentence Transformer的组件，这部分没有经过PEFT改造，不知道提示调优的存在，因此仍使用原始的注意力掩码维度。

2. 维度变化细节

以示例中的参数为例：

• 批大小：64
• 序列长度：53
• 虚拟标记数：10

经过Transformer部分处理后：

• 词嵌入维度变为[64, 63, 768]（原始53 + 虚拟10）
• 但池化部分仍期望[64, 53, 1]的注意力掩码

3. 解决方案探讨

目前可行的解决方案方向包括：

1. 整体模型封装：将整个模型封装，确保在调用池化层之前，注意力掩码已经根据虚拟标记数进行了适当扩展。

2. 修改池化层实现：调整Sentence Transformer的池化层实现，使其能够接收和处理来自Transformer部分的扩展后的注意力掩码。

3. 配置协调：在PEFT配置和Sentence Transformer配置之间建立协调机制，确保维度变化信息能够传递到所有相关组件。

技术影响

这个问题揭示了深度学习模型组件化开发中的一个常见挑战：当对模型的某一部分进行修改时，可能会破坏与其他组件的兼容性。特别是在参数高效微调场景下，原始模型和改造后的模型在输入输出维度上可能存在差异，需要特别注意。

最佳实践建议

对于需要在Sentence Transformer上使用PEFT提示调优的开发人员，目前建议：

1. 暂时使用其他PEFT方法（如LoRA或Prefix Tuning）作为替代方案
2. 如果必须使用提示调优，可以考虑以下临时解决方案：
   • 手动扩展注意力掩码的维度以匹配词嵌入
   • 创建自定义池化层处理扩展后的维度
3. 关注官方更新，等待更完善的集成解决方案

总结

这个问题体现了深度学习框架集成中的复杂性，特别是在参数高效微调这种相对前沿的技术领域。开发者在结合使用不同技术栈时需要特别注意组件间的兼容性，特别是在张量维度变化方面。随着PEFT和Sentence Transformer等技术的不断发展，这类集成问题有望得到更系统化的解决方案。