專訪｜基於LSTM與TensorFlow Lite，kika輸入法是如何造就的

摘要：近日，機器之心採訪了kika
的高級技術總監黃康，他向我們講述了kika開發輸入法AI
引擎（項目代號：Alps）所採用的深度學習模型以及在移動端輕量化部署遇到的各種挑戰。 本文從輸入法與語言模型開始介

近日，機器之心採訪了 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 年來對輸入法引擎的探索，未來還有很多優化與提升的方向，例如使用新型循環單元或新型強化學習來根據用戶習慣調整輸入法等。 這些新功能與新方向將賦予輸入法引擎更多的特性，也能適應性地為不同的用戶提供最好的體驗。