使用TensorFlow Lite进行基于移动设备的对话建模

文 / Google Expander 团队研究员 Sujith Ravi

今年早些时候，我们发布了 Android Wear 2.0，其中包含第一个用于智能消息的“基于移动设备的”机器学习技术。这让智能回复（之前在 Gmail、Inbox 和 Allo 中提供）一类的云技术首次可以直接用于任何应用（包括第三方消息应用），无需连接到云。因此，您可以随时随地直接从您的智能手表响应传入的聊天消息。

我们宣布推出 TensorFlow Lite，这个版本是 TensorFlow 面向移动和嵌入式设备的轻量级解决方案。此框架针对机器学习模型的低延迟推理进行优化，占用内存小，并具有快速性能。

作为库的一部分，我们还发布了一个基于移动设备的对话模型和一个演示应用：

http://download.tensorflow.org/deps/tflite/SmartReplyDemo.apk

这个应用提供了 TensorFlow Lite 实现的一个自然语言应用示例，这些举措旨在让开发者和研究人员更轻松地构建由机器上推理驱动的新机器智能功能。此模型可以生成对输入对话聊天消息的回复建议，其有效推理可以轻松集成到您的聊天应用中来增强设备上对话智能。

我们发布的基于移动设备的对话模型使用一个全新的机器学习架构根据一个联合优化框架训练紧凑型神经网络（以及其他机器学习模型），这个框架最初在 ProjectionNet：利用神经投影学习有效的基于移动设备的深度网络中提供。此架构可以在具有有限计算能力和内存的移动设备上有效地运行，它利用高效的“投影”运算将任何输入转换成紧凑的位矢量表示 - 相似的输入投影到附近密集或稀疏的矢量上，具体取决于投影的类型。例如，“hey, how's it going?”和“How's it going buddy?”消息可能会投影到相同的矢量表示上。

根据这个理念，对话模型将这些占用较少计算和内存资源的有效运算结合起来。我们使用一个联合训练两种模型的机器学习框架对此设备上模型进行端到端训练 - 这两个模型分别是紧凑的投影模型（如上所述）和教练模型。两个模型联合接受训练，投影模型从教练模型学习 - 教练具备专家特质，并使用更大、更复杂的机器学习架构建模，而投影模型则类似于从专家那里学习的学生。在训练期间，我们还可以堆叠其他技术（例如量化或提炼）来实现更进一步的压缩或选择性地优化目标函数的特定部分。完成训练后，较小的投影模型能够直接用于在移动设备上进行推理。

对于推理，训练的投影模型将编译到一组 TensorFlow Lite 运算中，这些运算已经针对在移动平台上快速执行和直接在设备上执行进行了优化。下面是基于移动设备的对话模型的 TensorFlow Lite 推理图。

设备上对话模型的 TensorFlow Lite 执行

最新发布的开放源代码对话模型（以及代码）是使用上面介绍的联合机器学习架构以端到端方式训练的。今天的发布还包括一个演示应用，您可以轻松下载这个应用并在您的移动设备上试用一键智能回复。利用此架构，可以根据应用需求针对模型大小和预测质量进行轻松配置。您可以在这里找到此模型非常适合的示例消息列表。系统还可以回退，从在聊天对话中观察到的热门响应意图学习和编译的一个固定集合建议回复。基础模型与 Google 在其应用中为“智能回复”响应使用的模型不同1。

超越对话模型

有趣的是，上面介绍的机器学习架构允许灵活选择基础模型。我们还设计了与不同的机器学习方式兼容的架构 - 例如，与 TensorFlow 深度学习结合，我们可以让基础模型学习轻量级神经网络 (ProjectionNet)，而另一个架构 (ProjectionGraph) 则使用图框架而不是神经网络表示模型。

联合框架也可用于为使用不同机器学习建模架构的其他任务训练轻量级基于移动设备的模型。作为示例，我们衍生了 ProjectionNet 架构，它为教练模型使用复杂的前馈或循环架构（例如 LSTM），并搭配一个由动态投影运算和一些完全连接的窄层组成的简单投影架构。整个模型在 TensorFlow 中使用反向传播进行端到端训练，紧凑型 ProjectionNet 直接用于推理。利用此方法，我们成功训练了微小的 ProjectionNet 模型，这些模型的大小显著减小（高达七个数量级），并在多个视觉和语言分类任务上取得了较高的准确率。类似地，我们使用了我们的图学习框架在半监督设置中训练了其他轻量级模型。

用于训练基于移动设备模型的机器学习架构：使用深度学习训练的 ProjectionNet（左），使用图学习训练的 ProjectionGraph（右）

我们将在开放源代码项目中继续改进和发布更新的 TensorFlow Lite 模型。我们认为，使用这些机器学习架构学习的产品模型（以及未来模型）可重用于许多自然语言和计算机视觉应用，或者集成到现有应用中来实现机器智能。我们希望机器学习和自然语言处理社区能够基于这些成果开展构建工作，解决我们还未想到的新问题和案例。

致谢

Yicheng Fan 和 Gaurav Nemade 对这项工作贡献巨大。特别感谢 TensorFlow 团队的 Rajat Monga、Andre Hentz、Andrew Selle、Sarah Sirajuddin 和 Anitha Vijayakumar，以及 Robin Dua、Patrick McGregor、Andrei Broder、Andrew Tomkins 和 Google Expander 团队。

1 发布的基于移动设备的模型经过专门训练，并针对移动电话和穿戴式设备上的小型和低延迟应用进行优化。不过，Google 应用中的“智能回复”预测使用更大、更复杂的模型生成。在生产系统中，我们还使用多个经过训练的分类器检测不当内容和应用进一步的筛选和调整来优化用户体验和质量的层次。我们建议使用开放源代码 TensorFlow Lite 版本的开发者也为他们的终端应用遵循此类做法。