0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

1M級の言語モデルで埋め込み行列の34.9%を消す — 語の同一視の実装と測定

0
Posted at

個人で、1M級(パラメータ100万規模)の日本語言語モデルを「ノートPCに収まる大きさ」まで圧縮する研究をしています。CPUのみ、学習も推論も手元です。

この記事は、3か月で試した100を超える構成のうち、最も明確に成功に近づいたパターンの実装と測定の記録です。

TL;DR

  • 1M級の言語モデルでは、**埋め込み行列が全パラメータの71%**を占める。削るならここ。
  • 埋め込み行をコサイン類似度で同一視して代表行にまとめると、16,000行 → 10,423行(34.9%減)。
  • 語彙→代表行の写像は 28,000 B。それを払っても同じ品質でファイルが8.3%小さくなった。
  • 落とし穴:最適な閾値は、量子化だけ(PTQ)と再学習込み(QAT)で違う(0.94〜0.96 → 0.98)。
  • 語彙を32,000に倍増すると削減率は**54.9%**まで上がる。語彙が2倍でも代表行は1.385倍にしかならない。

ノートPCのCPUだけで学習・評価しています。GPUは使っていません。

なぜ埋め込みなのか

対象は自作の小型言語モデル(d_model=144、4層、語彙16,000、tied embeddings)です。内訳はこうなります。

部位 パラメータ 割合
本体(Attention + FFN + 正規化) 959,760 29.4%
埋め込み(出力層と共有) 2,304,000 70.6%
合計 3,263,760 fp32 で 13,055,040 B

モデルを小さくしても語彙サイズは減らせない(減らすとトークン列が伸びて別の損をする)ので、
小さいモデルほど埋め込みの比重が上がります。

私は最初の2か月、本体側(29%)を削る手法ばかり試して全部失敗しました。
低ランク近似は同じパラメータ削減率で量子化より誤差が4.25倍大きく、
Tensor-Train 分解は許容誤差を20%にしてもランクが一度も打ち切られませんでした。
残りの71%を見たのは、打つ手が尽きてからです。

手法:語の同一視

埋め込み行列の行どうしを比べると、非常に近い対が大量にあります。
表記ゆれ、活用形、同じ文脈でしか現れない低頻度語などです。
学習が進むほど、同じ役割の語は同じ方向を向きます(後述)。

やることは単純です。

  1. 行の最近傍を取り、cos ≥ 閾値 の対を同じとみなす
  2. union-find で連結成分にまとめる
  3. 各グループの出現頻度で重み付けした平均を代表行にする
  4. 語彙→代表行の写像を保存する

1〜2:併合

全ペアの類似度行列は $V^2$ で重いので、ブロックに切って最近傍だけ見ます。

import numpy as np

def merge_groups(E, cos_th, block=2048):
    """E: (V, d) の埋め込み。返り値 inv: (V,) 各語のグループ番号"""
    N = E / np.maximum(np.linalg.norm(E, axis=1, keepdims=True), 1e-12)
    parent = np.arange(len(E))

    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]   # 経路圧縮
            x = parent[x]
        return x

    for s in range(0, len(E), block):
        S = N[s:s + block] @ N.T                       # (block, V)
        rows = np.arange(len(S))
        S[rows, np.arange(s, s + len(S))] = -1.0       # 自分自身を外す
        j, v = S.argmax(axis=1), S.max(axis=1)
        for i in np.flatnonzero(v >= cos_th):
            a, b = find(s + i), find(int(j[i]))
            if a != b:
                parent[max(a, b)] = min(a, b)

    root = np.array([find(i) for i in range(len(E))])
    _, inv = np.unique(root, return_inverse=True)
    return inv

最近傍だけを見ても、union-find が推移的に連結するので3語以上のグループもできます。

3:頻度で重み付けした代表行

ここが効きます。グループには高頻度語と低頻度語が混ざっていることが多く、
単純平均だと低頻度語に引きずられて高頻度語の表現がぼやけます。

inv = merge_groups(E, 0.98)
R = int(inv.max()) + 1
w = np.maximum(counts, 1e-3)              # counts: 学習コーパスでの出現頻度

num = np.zeros((R, E.shape[1]))
den = np.zeros(R)
np.add.at(num, inv, E * w[:, None])
np.add.at(den, inv, w)
E_rep = num / den[:, None]                # (R, d) これだけを量子化する

# 復元は E_rep[inv] で元の (V, d) に戻る

以後のパイプライン(回転 → ビット配分 → 非一様代表値 → 符号化)は、
代表行 E_rep だけに掛けます。量子化対象が34.9%減るわけです。

4:写像のコスト

写像 inv は語彙ぶんの整数列なので、保存が要ります。

値
語彙 16,000
代表行 10,423
索引ビット ceil(log2(10423)) = 14 bit
写像の容量 16,000 × 14 / 8 = 28,000 B

ここは正直に書いておきます。 この写像にも範囲符号(静的エントロピー符号)を掛けていますが、
取り分はゼロでした。 28,000 B は 14 bit 固定長で詰めた値とぴったり一致します。

理由は分かっていて、k-means やグループ番号のような索引は設計上ほぼ一様分布だからです。
一方、非一様代表値(Lloyd–Max)で作ったスカラー量子化の符号はヒストグラムが尖るので、
同じ範囲符号が実効ビット幅を 4.000 → 3.921 まで削ります。
同じエントロピー符号でも、符号の作り方で取り分が5%変わります。

結果

評価は dev を10万トークンの塊に分け、偶数塊で設定を決め、奇数塊で確かめるという分け方にしています。
100を超える条件を試したので、分けていないと評価データに過剰適合します。以下はすべて奇数側(40バッチ)です。

容量(数え) ファイル実測 倍率 dev ppl 基準比
同一視なし(QAT) 1,215,071 B — 10.74× 31.307 0.999
同一視 cos 0.98(QAT) 1,030,852 B — 12.66× 32.237 1.029
+下層VQを厚く 1,007,070 B 1,022,875 B 12.96× 32.290 1.030

写像 28,000 B を払ってなお、同じ品質で8.3%小さくなりました。

落とし穴1:PTQ と QAT で最適な閾値が違う

閾値の掃き方で結論が変わります。

最適な cos
量子化しただけ(PTQ)で掃く 0.94〜0.96
量子化したまま再学習(QAT)まで通して掃く 0.98

PTQ の結果だけを信じると、畳みすぎた設定を選びます。

解釈はこうです。

QAT が取り返せるのは「ズレ」であって、「消えた区別」ではない。

丸め誤差は重みを動かせば補償できますが、2行を1行にまとめたら、その2語の違いは
再学習しても戻りません。情報そのものが消えているからです。
したがって QAT を前提にするなら、あとで直せる側(量子化)を攻めて、直せない側(同一視)は控えるのが正解になります。

同じ理由で、PTQ 段階で出た改善を最終結果に持ち込んではいけません。
私は回転を高ビット層で切る案で PTQ 1.37% の改善を出して採用しかけ、QAT を通したら差が消えました。

落とし穴2:学習を進めるほど、行は「分化」ではなく「統合」する

直感に反したので測りました。同じ run の途中 checkpoint で削減率を並べます。

語彙 ステップ 代表行 削減率
16,000 12,000 12,240 23.5%
16,000 18,000 11,596 27.5%
16,000 24,000 10,418 34.9%
32,000 2,000 18,669 41.7%
32,000 4,000 15,523 51.5%
32,000 6,000 14,441 54.9%

両系列とも単調に減り続け、終端でもまだ落ちています。

学習が浅いうちは初期値が残っていて行が似たまま → 学習が進むと分化する、と予想していましたが逆でした。
学習は同じ役割の語を同じ方向に集めます。収束とは分化ではなく統合です。

実務上の含意:学習が終わっていないモデルで同一視の効きを測ると、過小評価になります。

規模を上げるとどうなるか

語彙だけを 16,000 → 32,000 に倍増して同じパイプラインを通しました(本体は同一)。

語彙16,000 語彙32,000
埋め込み 2,304,000 4,608,000
埋め込みの占有率 70.6% 82.8%
代表行 10,423 14,441
削減率 34.9% 54.9%
ファイル実測 1,029,748 B 1,250,508 B
倍率 12.678× 17.810×
基準比(品質) 1.1147 1.1118

語彙が2.00倍でも、代表行は1.385倍にしかなりません。
増えた語彙の多くは既存の意味の重複で、同じ cos 0.98 で機械的に回収できています。

結果として品質の落ち方をほぼ変えないまま、倍率が 12.68 → 17.81 に伸びました。
この手法は規模に対して有利です。

(注:ppl 50.30 と 34.93 を直接比べてはいけません。トークナイザが違うと1トークンあたりの難しさが変わります。
判定は各モデル内部の基準比だけで行っています。)

既存研究との関係

正直な自己採点です。

  • 回転で量子化しやすい表現にする(QuaRot / SpinQuant 系)→ 新規性なし
  • 機能的距離でニューロンを統合する(Neuron Merging, NeurIPS 2020)→ 新規性は小さい
  • 実バイト数を目的関数にする(DeepCABAC / ISO-IEC 15938-17)→ 目的関数自体は既存

「語の同一視 + 写像の保存 + その状態で量子化込み再学習」の組み合わせは、探した範囲では見つかりませんでした。
埋め込み行を減らす研究も語彙を刈り込む研究もありますが、
「畳んだ上で写像を保存し、その状態で QAT を掛ける」という形は見当たりません。
ご存じの方がいたら教えてください。

まとめ

  • 小さいモデルは埋め込みが支配的。削るなら語彙側。
  • cos 0.98 の同一視で 34.9% の行が消え、同品質で 8.3% 小さくなる。写像は 28,000 B。
  • 閾値は QAT まで通して決める。PTQ の最適点は信用できない。
  • 学習は語を統合する。未収束のモデルで測ると効きを過小評価する。
  • 索引のような一様分布の符号にはエントロピー符号の取り分がない。

この連載について

書いている人

個人開発者です。1M級の言語モデルを、ノートPCの中で完結する大きさまで圧縮する研究を3か月続けています。現時点で 13,055,040 バイト → 1,022,875 バイト(約13倍)、品質の低下は3%ほど。

何を書くか

記録しているのは主に、失敗したパターンのほうです。

試した構成は100を超えますが、成功に近づいたのはこの記事の内容を含めてごく一部でした。期待していた手法ほど効果が出ず、特に幾何学的な構造を使う方法(格子量子化、低ランク近似、テンソル分解)はいずれも改善に至りませんでした。

成功例の記事は多くありますが、同じことをやろうとしている人には 「どこで失敗するか」のほうが役に立つはずです。自分がそれを探して見つからず、何度も時間を無駄にしました。

今後扱う予定:

  • 量子化(PTQ / QAT)で実際に何が起きたか
  • ベクトル量子化と残差VQ を打ち切った判断の根拠
  • 測定方法そのものの設計 — 比較が成立する条件を揃える話

頻度

月1本程度。測定が終わった分だけ書きます。 定期更新はしません。手元に数字がない状態で記事を出さない、というのがこの研究の方針です。

訂正について

出した数字が誤っていた場合は、訂正記事を書きます。

直近でも、同じ主張を5回訂正しました。「語彙を大きくすると有利」という結論が、条件を揃え直すたびに +44% → +22% → +6% と縮み、最終的には符号が反転して「目標とする領域では不利」になりました。原因は毎回こちら側(学習量を揃えていない、品質を揃えていない、自分で禁じたはずの圧縮率で比較していた)でした。

訂正の過程も含めて記録します。


同じ研究を、専門知識なしで読める読み物として note でも連載しています → https://note.com/ai_engineer_k

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?