专访｜基于LSTM与TensorFlow Lite，kika输入法是如何造就的
近日，机器之心采访了 kika 的高级技术总监黄康，他向我们讲述了 kika 开发输入法 AI 引擎（项目代号：Alps）所采用的深度学习模型以及在移动端轻量化部署遇到的各种挑战。本文从输入法与语言模型开始介绍了 kika Alps 项目的理论支持与实践挑战，并重点讨论了轻量化部署方法。
深度学习模型由于强大的表征能力在很多任务上都有非常优秀的表现，但也因为模型大小和计算量很难轻量化部署到移动端。这也是目前很多研发团队都在思考如何解决的难题。
一般在我们借助 TensorFlow、MXNet、和 Caffe2 等框架构建深度学习模型后，它在服务器训练与推断往往会有非常好的效果。但如果我们将模型部署到移动端，即使只执行推断过程，也会因为硬件条件和系统环境而遇到各种各样的问题。此外，目前关注于移动端的解决方案如 TensorFlow Mobile、TensorFlow Lite 等在一定程度上并不完善（TF Mobile 的内存管理与 TF Lite 的 Operators 的缺失），在实践中可能需要更多的修正与完善。
关注于输入法的 kika 成功地将基于循环神经网络的深度学习模型应用到安卓版的手机输入法引擎中，在克服工程化问题的情况下大大提升了输入体验：不仅使基于上下文的词预测更加准确，同时还使得词纠错功能更加强大。
在构建这样的输入法引擎过程中，kika 不仅需要考虑使用 LSTM 还是 GRU 来实现高效的语言模型，同时还需要探索如何使整个方案更轻量化以及如何快速的进行部署。本文首先介绍了输入法及 kika 所采用的语言模型，并在随后展示了 Android 移动端轻量化部署所遇到的工程化挑战。最后，本文介绍了 kika 压缩模型所采用的稀疏词表征方法与 K-means 参数量化方法，它们是轻量化部署深度学习模型的重要前提。
输入法与语言模型
输入法最重要的部分就是输入法引擎，kika 很多算法和项目都围绕它展开。一般而言，输入法引擎的输入包含两部分，即已经键入的词组和当前正在输入的词汇，前者可以视为上下文，而未完成的后者可称为键码。输入法引擎的输出是给定所有上下文和当前输入键码所『预测』的词，它也包含两部分，即当前输入词汇的补全和纠错。实现这样的功能也就是输入法最为核心的模块，kika 最开始是使用谷歌半开源的 LatinIME 来实现这样的功能，但这种基于 n-gram 的方法并不能实现顶尖的用户体验，因此经过研究与开发才有了现在基于循环神经网络（RNN）的解决方案。
输入法引擎这种给定上下文和当前键码以预测下一个词的方法其实可以看成语言建模，一般来说，语言模型旨在定义自然语言中「标记」的概率分布，这种标记可以是单词、字符甚至是字节。根据 kika 介绍，LatinIME 构建语言模型的方法是 n-gram，这种模型定义了一个条件概率分布，即给定前 n-1 个单词后第 n 个词的条件概率。因为假定当前词出现的概率只与前面 n-1 个词相关，那么 n-gram 可以使用这种条件概率的乘积来定义较长序列的概率分布：
虽然 n-gram 一直以来都是统计语言模型的核心模块，但它还是有很多局限性。对于输入法而言，较小的 n（二元语法与三元语法）不足以捕捉整个上下文信息来执行预测，而增大 n 又会使计算量成指数级增加。此外，kika 还希望引擎实现其它一些智能功能，例如根据上下文自动纠错或排序等。因此，kika 选择了更强大的循环神经网络构建语言模型。
为了构建强大的语言模型，kika 选择了长短期记忆单元（LSTM）作为网络的基础。LSTM 作为标准循环神经网络的变体在语言模型上有非常好的性能，它引入自循环的巧妙构想来更新「记忆」，若通过门控控制这样的自循环，那么累积的历史记忆或时间尺度就可以动态地改变。直观来说，LSTM 会通过门控选择需要保留的上下文信息或记忆，并用于预测当前输入的词。每当输入一个词，输入门会控制当前输入对最终预测有用的信息，而遗忘门会控制前面输入词对最终预测有用的信息，最后的输出门则会直接控制前两部分最终有效的信息。
kika 表明最开始 LSTM 只是用来实现标准的语言模型，它不会将正在输入的键码作为模型输入。这一套早期方案使用 LSTM 语言模型根据上下文预测当前词，而键码部分仍然使用传统基于 n-gram 和字典等的解决方案。后来 kika 使用一种新的策略将两部分结合在一起，因此模型不仅能接受上下文的输入，同时还能接受键码的输入。这种通用解决方案的架构如下图所示，它和一般的单词级语言模型和字符级语言模型都不一样，可以说是结合了两种架构。
如上图所示，首先 LSTM 会对前面输入的词进行建模，并输出对应的隐藏状态和记忆而作为后面字符级语言模型的先验输入。后面从 Start Flag 开始对键码实现字符级的建模而最终得出预测。
根据 kika 的解释，最后这种方案统一了两种输入。它的基本思想首先考虑到前面的 LSTM 语言模型除了要根据隐藏状态预测当前时间步的输出，同时还会向后传递这个隐藏状态。且 kika 表示若网络有良好的训练，那么这个隐藏状态是可以包含足够的语意信息，因此我们可以将它作为后面字符级 LSTM 网络的初始状态。这相当给循环神经网络一个初始量，然后再接受键码的输入而作出最终的词预测和词纠错等。
其实这里还有一个非常有意思的问题，即为什么 kika 会采用 LSTM 而不是 GRU。因为我们知道若需要将深度学习模型部署到移动端，那么我们要严格控制模型与计算量的大小，kika 所采用的稀疏词表征与参数量化等方法能有效控制模型大小，这一点将在后文展开。但 LSTM 的结构比 GRU 要复杂，门控也需要得更多，因此 LSTM 的参数会比 GRU 多，那么 kika 为什么不采用 GRU 控制参数数量？
kika 就这一点对机器之心做了详细的解答。黄康说：「在层数和单元数均一致的情况下，GRU 要比 LSTM 少一些参数和矩阵运算，因此，模型体积和训练速度方面都会有一定的优势。为了严谨的进行效果对比，我们做了两组实验。其中第一组是将 LSTM 和 GRU 的超参数设置一致，结果是： GRU 的效果明显差于 LSTM，同时，由于整体模型体积的主要贡献来源于前后两个巨大的词嵌入矩阵，模型体积方面的优势也不明显。」
但在同样超参数的情况下，GRU 的实际参数数量明显少于 LSTM。因此，kika 继续做了第二组实验，在保证基本一致的参数数量而放开网络架构约束的情况下，最后得到的结论是：LSTM 与 GRU 的模型大小基本一致，效果也基本一致，实际上，在 kika 的应用场景下，LSTM 的效果略好，但也仅仅是略好一点点。此外，由于 GRU 在当时也是比较新的结构，因此在体积和效果没有优势的情况下 kika 还是倾向于选择更温和的 LSTM，从而把主要精力用于模型结构的调整与参数调优方面。其实最近 kika 也在做一些网络架构和基本单元方面的调研，因为最近在 GRU 之后又出现了非常多训练简单且高效的单元。在 kika 当前的开发过程中也出现了一些场景更为复杂的 NLP/NLU 应用，因此也在考虑采用一些训练时间上更为友好的网络结果。对于如何选择网络结构，黄康表示：「我们内部有共识，考虑新结构有一个基本原则：我们会采用类似机器翻译的复杂任务去验证此种网络是否真实有效，才会考虑在工程上采用。」
总体而言，kika 花了很大一部分时间完成参数调优，因而能基于一体化的 LSTM 实现效果非常好的输入法引擎。当然只是构建模型还远远不够，将这种深度学习模型部署到移动端还面临着非常多的挑战，例如深度学习框架的选择和模型压缩的方法等等。
轻量化部署的工程挑战
在 kika，轻量化部署包括以下四个方面：模型压缩、快速的响应时间、较低的内存占用以及 较小的 so 库（shared object，共享库）大小等。除了优化模型的设计之外，压缩模型大小的方法主要是稀疏词表征与量化，这一部分我们将在后一部分展开讨论。
响应时间与内存是去年 kika 的工作重点，它主要是需要对 TensorFlow Mobile 和 Lite 做大量的修补。最后是动态链接库文件（.so），它定义了所有需要的运算和操作。因为整个输入法的核心代码是 C++ 完成的，而它在安卓设备中是以一个库（.so 文件）的形式存在的，它的大小直接影响了安装包的大小（由于 Android 加载 so 的机制，也会影响到内存的开销）。
针对响应时间与内存，kika 最开始是基于 TensorFlow Mobile 做一些修补和改进。黄康说：「TensorFlow Mobile 有一个非常好的优势，即它在底层使用了一套很成熟很稳定的矩阵运算库。因此，我们的主要精力放在 TensorFlow Mobile 在底层矩阵运算库之上的部分。它在矩阵运算库之间采用了很多封装与调用，但是没有考虑到很多实际工业化中遇到的问题，尤其是在内存保护这一块做得相当一般。」
TF Mobile 的内存管理与内存保护设计得并不完善，存在两个主要的问题：1. 内存保护机制不完善，在实际内存不足的情况（尤其对于一部分低端机型），容易引发内存非法操作。2. 内存大小控制机制存在明显的问题，例如模型本身在计算时只有 20MB，但加载到内存之后的运行时峰值可能会达到 40 到 70MB。据 kika 的数据，基于 TF Mobile 的解决方案大概有 1% 的场景（如游戏中调起输入法）由于内存大小限制的原因会加载不了深度学习模型，只能回退到非深度的解决方案。
2017 年 11 月，谷歌正式发布了 TensorFlow Lite，这对于移动端深度学习模型来说是非常重要的框架。在 TF Lite 开源后，kika 马上就进行了测试，并重点关注内存管理模块。黄康表示：「TF Lite 的内存管理上确实有非常大的改进，加载不了深度学习模型的场景会成百倍地减少。但它最大的问题就是 Operator 的不全，它基本上只定义了最基础的运算和操作。所以 kika 为了在内存上减少 20 多 MB 的开销，我们自行编写了大量的 Operator。但目前这个还是有一定的风险，因为如果我们修改了模型结构，那还是需要手写新的运算。」
工程化挑战最后一个问题就是动态链接库的大小，这一部分还会涉及到参数量化方法的实现，我们会在参数量化方法那边讨论。其实 TF Mobile 还有一个缺点，即它会将很多冗余的操作与运算都会打包到 .so 文件中，因此也就导致了动态链接库过大。kika 为了让 .so 文件尽可能小，开发了一套全新的工具，用于自动的判断到底哪些操作与运算的定义是模型实际需要的。
稀疏词表征
深度学习模型在输入法客户端部署的一个重要问题就是模型大小，我们需要将参数数量与计算量限制绝大部分移动设备可接受的范围内。kika 发现模型体积的主要矛盾体现在词嵌入矩阵中。因此，如果能够压缩词嵌入矩阵的大小，那么就能有效地控制模型大小。kika 采用了稀疏词表征的方法以压缩词嵌入矩阵的大小，从而大幅度减少 LSTM 语言模型的参数与计算量。
其实对于语言模型，甚至是自然语言处理而言，词嵌入是非常重要的成分，我们可以使用 300 到 500 维的向量表示词表中数以万计的词汇。这种方法不会像 one-hot 编码那样使用超高维的向量表示一个词，可以说词嵌入将词的表征由|V|维减少到几百维，其中|V|表示词汇数量。但是当我们使用词嵌入作为语言模型的输入时，我们会发现尽管每个词的维度只有 n，但需要|V|个向量，而 |V| 通常要比 n 高好几个量级。因此，稀疏词表征就尝试使用少量词向量（少于|V|）而表征 |V| 个词。
这种方法的直观概念即我们的词表可以分为常见词与非常见词，而一般单个词可以由多个词定义，因此非常见词可以使用多个常见词表示。根据这样的观点，我们可以使用一组常见词的词嵌入向量作为基础，再通过组合而表示所有词的词嵌入向量。因此，我们能使用少量词嵌入向量表示大量词汇。又因为每一个词的表征都只使用少量的常见词来定义，所以这种表示方法是非常稀疏的，这也就是稀疏词表征的直观概念。
若我们将词表 V 分割为两个子集 B 和 C，第一个子集 B 为基向量集，它包含了固定数量的常见词。而 C 包含了所有不常见的词，因此现在需要使用 B 的词嵌入向量以线性组合的方式编码 C 中的词。这一过程可通过最小化由基向量集学习重构的词表征和完整表征之间的距离而学习，一般来说整个过程可表示为：
1. 使用全部词汇训练一个词嵌入矩阵。
2. 按词频选取最常见的 |B| 个词嵌入向量，并组成过完备基矩阵。
3. 非常见词的预测可表示为 B 中所有词嵌入向量的稀疏组合，即
其中 w hat 为预测的非常见词词向量、U 为常见词词向量，而 x 为稀疏矩阵。
4. 最小化预测词向量和实际词向量间的距离来学习稀疏表征，即
其中第一项表示通过稀疏表示 x 预测的词向量与完整词向量（w）间的 L2 距离。后一项为 L1 正则化，它会将矩阵 x 中的元素推向 0，从而实现稀疏表示。
在 kika 的论文 Sparse Word Representation for RNN Language Models on Cellphones 中，他们使用了以下伪代码展示了稀疏表示的学习算法：
这个算法很大的特点是实现了一个二元搜索来确定α，因为我们不能直接控制稀疏矩阵 x 的稀疏程度，所以我们根据稀疏矩阵的非零元素数来控制α的变化。整个稀疏词表征算法需要输入过完备基矩阵 U（常见词）、完整词嵌入矩阵 w、稀疏程度 s 和作为终止条件的容忍度 tol。
其中 s 是非常重要的一个参数，它控制了一个词最多需要多少个过完备基向量表征。kika 表示：「s 是一种权衡，如果参数较大，那么压缩比就会很小，模型达不到预期效果。如果参数较小，那么重构的词表征就不能有效地表示所有词。」正因为需要进行精调来确定 s 及其它超参数，kika 表明总体模型调优时间是训练时间的 4 到 5 倍，所以整个稀疏词表征的训练过程还是比较挺长的。
如上算法所示，首先我们会确定α的搜索范围，然后取α的中间值并最小化损失函数而求得稀疏表示 x*，并统计 x* 中每一个列向量的非零元素数，它们代表了一个词需要多少个常见词表示。如果 k 大于 s，那么非零的元素就过多，我们需要加大 α 以增强 L1 正则化的效果。这样的二元搜索直到α的上下界距离小于参数 tol 才会终止，且一般迭代几次就能快速收敛到合适的 α 来控制 x* 的稀疏性。在完成 x* 的学习后，我们将每一列稀疏向量抽取为对应的索引与权重，索引代表使用哪些基向量或常见词，而权重代表它们定义某个词的重要性。
又因为前面的二元搜索将 k 限制为不大于 s，所以有可能 k 是小于 s 的，因此我们需要使用零将这些向量补全。经过上面的步骤，最终我们会产生包含 s 个元素的等长向量 indices 和 weights。储存这两种向量而不直接储存稀疏矩阵 x* 能节省很多空间，这对于减小安装包大小有非常重要的作用。
论文中给出的词嵌入恢复算法以一种串行密集运算的方式进行展示，这可以令读者清晰地理解重构过程：
若给定 U、indices 和 weights，一个词的词嵌入重构可直接利用索引取对应的基向量，并与对应的权重求加权和。这种线性组合非常简单且高效，也是线性代数中非常直观的表示方法。因为任何秩为 n 的矩阵都可以由 n 个线性不相关的向量或基表示出来，完整的词嵌入矩阵也就能由过完备基的线性组合表示。算法 1.2 最后返回的 v 就是我们线性组合多个常见词词嵌入而重构出来的完整词嵌入向量。
以上是 kika 采用的稀疏词表征方法，它可以有效减少模型参数和计算量，但我们还能进一步使用参数量化来压缩模型的存储大小。
量化
一般而言，应用的安装包大小对于用户体验非常重要，这一点对于移动端尤为突出。因此，我们可以使用参数量化的方法来减小安装包大小。kika 也曾尝试使用 TensorFlow 封装的压缩方法，但仍发现一些难以解决的问题，因此他们最终使用 k-means 方法重新构建参数量化而解决包体增大的问题。
kika 最开始尝试使用官方的 tf.quantize 执行参数量化，并用 tf.dequantize 恢复参数。这个方法非常迅速，基本上几十兆的模型只需要分钟级的时间就能完成压缩。但 kika 发现这种方法有一个非常大的问题，即如果当我们希望读取量化后的模型时，TensorFlow 会引入大量的 Operator，这势必会造成动态链接库（.so）的体积增大，因而会加大安装包的大小。因为动态链接库包含了所有 TF 定义的加法、减法、卷积和归一化等模型需要使用的运算，因此调用 TF 的量化方法同样会将相关的运算添加到动态链接库中。
根据 kika 的实验，使用 TF 官方的量化方法大概会使动态链接库增加 1 到 2 MB 的体积，对应的安装包大小也会增加这么多。由于这样的原因，kika 最后选择基于 k-means 的方法实现参数量化。简单而言，这个方法会先使用 k-means 将相似的向量聚类在一起，然后储存聚类中心，原参数矩阵就只需要存储聚类中心的索引就行了。kika 表明这种方法的有点在于不会额外增加动态链接库和安装包的大小。因此下面将简要介绍这种基于 k-means 的参数量化方法。
量化即从权重中归纳一些特征，这些特征会储存在码表（codebook）并以具体数值表示某一类权重，而原来的权重矩阵只需要存储索引来表示它们属于哪一类特征就行了，这种方法能很大程度上降低存储成本。
kika 使用的标量量化算法基本思路是，对于每一个 m×n 维的权重矩阵 W，首先将其转化为包含 m×n 个元素的向量 w。然后再对该权重向量的元素聚类为 k 个集群，这可借助经典的 k 均值聚类算法快速完成：
现在，我们只需储存 k 个聚类中心 c_j，而原权重矩阵只需要记录各自聚类中心的索引就行。在韩松 ICLR 2016 的最佳论文中，他用如下一张图非常形象地展示了量化的概念与过程。
如上所示权重矩阵的所有参数可以聚类为 4 个类别，不同的类别使用不同的颜色表示。上半部分的权重矩阵可以取聚类中心，并储存在 centroids 向量中，随后原来的权重矩阵只需要很少的空间储存对应的索引。下半部是韩松等研究者利用反向传播的梯度对当前 centroids 向量进行修正的过程。
稀疏词表征与参数量化是 kika 控制参数大小的主要方法，黄康表示：「实际上模型的大小可以分为两阶段，首先如果原模型是 40MB 的话，稀疏词表征可以将模型减少到 20MB 左右，这个大小是实际在内存中的大小。而进一步采用参数量化可以将大小压缩到 4MB 左右，它解决的问题是 APK 安装包大小，APK 大小也是非常重要的，毕竟作为输入法这样的应用，APK 的大小是非常重要的。不论使不使用参数量化，模型最终在计算上需要的内存就是稀疏词向量表征后的大小。」
最后两部分基本上就是 kika 解决模型大小的方案，它们令深度学习模型在实践中有了应用的可能。当然，要将深度学习模型嵌入输入法和移动端会有很多的挑战，仅仅控制模型大小是不够的，因此也就有了上文 kika 在内存大小、响应时间和动态链接库等方面的努力。
整个模型效果和工程化实践都是 kika 在过去 2 年来对输入法引擎的探索，未来还有很多优化与提升的方向，例如使用新型循环单元或新型强化学习来根据用户习惯调整输入法等。这些新功能与新方向将赋予输入法引擎更多的特性，也能适应性地为不同的用户提供最好的体验。