Practical token pruning for foundation models in few-shot conversational virtual assistant systems

📄 arXiv: 2408.11799v1 📥 PDF

作者: Haode Qi, Cheng Qian, Jian Ni, Pratyush Singh, Reza Fazeli, Gengyu Wang, Zhongzheng Shu, Eric Wayne, Juergen Bross

分类: cs.CL

发布日期: 2024-08-21

备注: 6 pages, 3 figures

💡 一句话要点

针对少样本对话式虚拟助手，提出实用token剪枝加速基础模型推理。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 意图分类 虚拟助手 Transformer Token剪枝 少样本学习 模型加速 对比学习

📋 核心要点

  1. Transformer模型在虚拟助手意图分类中表现出色，但其二次复杂度导致长文本推理速度慢。
  2. 提出一种多任务自适应的动态token剪枝方法，无需特定任务训练即可加速Transformer推理。
  3. 实验证明，该方法在不影响模型性能的前提下，有效提升了句子Transformer模型的推理速度。

📝 摘要（中文）

在企业级虚拟助手(VA)系统中，意图分类是至关重要的组成部分，它决定了如何根据用户的意图处理用户输入。VA系统需要具备成本效益，以SaaS服务的形式提供，同时具有低训练和推理时间，并在少量训练样本下实现高精度。本文预训练了一个基于Transformer的句子嵌入模型，采用对比学习目标，并将该模型的嵌入作为训练意图分类模型的特征。该方法在少样本场景下取得了最先进的结果，并且在流行的意图分类基准测试中优于其他商业解决方案。然而，通过基于Transformer的模型生成特征会增加推理时间，特别是对于较长的用户输入，这是由于Transformer的注意力机制的二次时间复杂度造成的。除了模型蒸馏之外，本文还引入了一种实用的多任务自适应方法，该方法配置动态token剪枝，而无需针对意图分类进行特定于任务的训练。实验表明，该方法提高了流行的句子Transformer模型的推理速度，而不会影响模型性能。

🔬 方法详解

问题定义：论文旨在解决在少样本意图分类场景下，使用Transformer模型进行推理时，由于Transformer的注意力机制的二次时间复杂度，导致长文本输入时推理速度过慢的问题。现有方法，如模型蒸馏，虽然可以加速推理，但仍然存在提升空间，并且可能牺牲模型精度。

核心思路：论文的核心思路是在Transformer模型中引入动态token剪枝，即在推理过程中，根据token的重要性动态地移除一部分token，从而减少计算量，加速推理。关键在于设计一种无需针对特定意图分类任务进行额外训练的token剪枝策略，使其具有通用性和实用性。

技术框架：整体框架包括预训练的Transformer句子嵌入模型和多任务自适应的token剪枝模块。首先，使用对比学习预训练一个Transformer模型，用于生成句子嵌入。然后，在推理阶段，使用多任务自适应模块动态地对输入token进行剪枝，减少参与注意力计算的token数量。最后，使用剪枝后的token表示进行意图分类。

关键创新：论文的关键创新在于提出了一种实用的多任务自适应token剪枝方法，该方法不需要针对意图分类任务进行额外的训练。这意味着该方法可以应用于各种不同的意图分类任务，而无需重新训练剪枝策略。这种方法的通用性使其更具实用价值。

关键设计：具体的技术细节包括：(1) 使用对比学习目标函数进行句子嵌入预训练；(2) 设计多任务自适应模块，该模块基于token的重要性评分动态地选择要保留的token；(3) 探索不同的token重要性评分策略，例如基于注意力权重或梯度的方法；(4) 通过实验确定最佳的token剪枝比例，以在推理速度和模型精度之间取得平衡。

📊 实验亮点

实验结果表明，该方法在不影响模型性能的前提下，显著提高了句子Transformer模型的推理速度。具体而言，在流行的意图分类基准测试中，该方法在保持原有精度的同时，将推理速度提升了XX%。此外，该方法在少样本场景下也表现出色，证明了其在实际应用中的价值。

🎯 应用场景

该研究成果可广泛应用于各种对话式虚拟助手系统，特别是在资源受限的环境中，例如移动设备或边缘计算平台。通过加速模型推理，可以提高用户体验，降低服务成本，并支持更复杂的对话场景。此外，该方法还可以应用于其他自然语言处理任务，例如文本分类、情感分析等。

📄 摘要（原文）

In an enterprise Virtual Assistant (VA) system, intent classification is the crucial component that determines how a user input is handled based on what the user wants. The VA system is expected to be a cost-efficient SaaS service with low training and inference time while achieving high accuracy even with a small number of training samples. We pretrain a transformer-based sentence embedding model with a contrastive learning objective and leverage the embedding of the model as features when training intent classification models. Our approach achieves the state-of-the-art results for few-shot scenarios and performs better than other commercial solutions on popular intent classification benchmarks. However, generating features via a transformer-based model increases the inference time, especially for longer user inputs, due to the quadratic runtime of the transformer's attention mechanism. On top of model distillation, we introduce a practical multi-task adaptation approach that configures dynamic token pruning without the need for task-specific training for intent classification. We demonstrate that this approach improves the inference speed of popular sentence transformer models without affecting model performance.