Word2Vec メモその2 - @nojima's blog

Word2Vec を Chainer で実装していく。完全なコードは以下の URL にある。

workspace/learning-chainer/word2vec at master · nojima/workspace · GitHub

model

誤差関数をネットワークとして表現すると下図のようになる。

f:id:nojima718:20170821234058p:plain

$W$ と $W'$ が入力である単語の hot vector を低次元のベクトル表現に変換するための行列で、Chainer 本の図だと $W = W'$ になっている。今回の実験では、 $W$ と $W'$ が同じ行列の場合と違う行列の場合を両方試してみた。

$W = W'$ の方のモデルを Chainer のコードに落とすとこんな感じになる。

class Word2Vec(Chain):
    def __init__(self, n_vocabulary: int, n_units: int) -> None:
        super().__init__()
        with self.init_scope():
            self._embed = L.EmbedID(n_vocabulary, n_units)

    def __call__(self, x1: Variable, x2: Variable, t: Variable) -> Variable:
        output = self.forward(x1, x2)
        return F.sigmoid_cross_entropy(output, t)

    def forward(self, x1: Variable, x2: Variable) -> Variable:
        h1 = self._embed(x1)
        h2 = self._embed(x2)
        return F.sum(h1 * h2, axis=1)

（変数名をこれまでの説明と揃えるべきなんだけど、時間がなくて全然揃えれてないので、すみませんが雰囲気で読み替えてください）

L.EmbedID は one-hot ベクトルに最適化された線形レイヤーで、L.Linear よりも forward も backward も効率的に実装されている。L.EmbedID.__call__ は単語ベクトルじゃなくて単語IDを取ることに注意。

$W$ と $W'$ を異なる行列にする場合は L.EmbedID を２つ用意して使い分ければいい。

train

モデルを学習するためのコードを書いていく。

まずハイパーパラメータを保持しておくための型を定義する。

HyperParameters = namedtuple('HyperParameters', [
    'vocabulary_size',      # 語彙数
    'vector_dimension',     # 分散表現ベクトルの次元
    'window_size',          # ウィンドウの半径
    'n_negative_samples',   # １単語につきネガティブサンプルする単語の数
    'batch_size',           # １回のミニバッチで処理する単語数
    'n_epoch',              # 総エポック数
])

次に negative sampling で使うためのサンプラを作る。

Chainer には WalkerAlias という便利なサンプラがある。 WalkerAlias のコンストラクタに頻度の列を渡すと、その分布に従うサンプラが作れる。

def _make_sampler(dataset: DataSet) -> walker_alias.WalkerAlias:
    _, counts = np.unique(dataset.data, return_counts=True)
    counts = np.power(counts, 0.75)
    return walker_alias.WalkerAlias(counts)

（DataSet は、単語IDの列 data と語彙集合 vocabulary を保持する型）

次に、train 関数を定義する。普通にミニバッチで学習する。ここで出てくる _make_batch_set は訓練データを作る関数で、後で定義する。

def train(dataset: DataSet, params: HyperParameters) -> Word2Vec:

    model = Word2Vec(params.vocabulary_size, params.vector_dimension)

    optimizer = optimizers.Adam()
    optimizer.setup(model)

    sampler = _make_sampler(dataset)

    batch_size = params.batch_size
    n_epoch = params.n_epoch

    for epoch in range(n_epoch):
        indices = np.random.permutation(dataset.size)

        for i in range(0, dataset.size, batch_size):
            model.cleargrads()

            x1, x2, t = _make_batch_set(
                indices[i:i+batch_size], sampler, dataset, params)
            loss = model(x1, x2, t)
            loss.backward()
            optimizer.update()

    return model

後回しにしていた _make_batch_set の定義はこれ。昨日の記事で説明した方法にしたがって正例と負例を作って返す。

def _make_batch_set(
        indices: np.ndarray,
        sampler: walker_alias.WalkerAlias,
        dataset: DataSet,
        params: HyperParameters) -> Tuple[Variable, Variable, Variable]:

    window_size = params.window_size
    n_negative_samples = params.n_negative_samples

    x1, x2, t = [], [], []

    for index in indices:
        id1 = dataset.data[index]

        for i in range(-window_size, window_size+1):
            p = index + i
            if i == 0 or p < 0 or p >= dataset.size:
                continue
            id2 = dataset.data[p]
            x1.append(id1)
            x2.append(id2)
            t.append(1)
            for nid in sampler.sample(n_negative_samples):
                x1.append(id1)
                x2.append(nid)
                t.append(0)

    return (Variable(np.array(x1, dtype=np.int32)),
            Variable(np.array(x2, dtype=np.int32)),
            Variable(np.array(t,  dtype=np.int32)))

serialize/deserialize

枝葉だけど、モデルの保存と読み込みに結構嵌ったのでここにメモしておく。

モデルの保存は save_npz を使えば簡単にできる。しかし save_npz ではモデルパラメータのみが保存され、語彙数、ベクトルの次元、ウィンドウサイズなどのハイパーパラメータを保存することができない。これらのパラメータも一緒に保存したい場合は serializers.DictionarySerializer を使ってハイパーパラメータを追加でシリアライズすればいい。

def save_model(filename: str, model: Word2Vec, params: HyperParameters) -> None:
    serializer = serializers.DictionarySerializer()

    pickled_params = np.frombuffer(pickle.dumps(params), dtype=np.uint8)
    serializer("hyper_parameters", pickled_params)

    serializer["model"].save(model)

    np.savez_compressed(filename, **serializer.target)

モデルの読み込みも、保存のときと同様、 load_npz で一発でできるが、これだとモデルパラメータしかロードできない。よって、serializers.NpzDeserializer を使ってまずハイパーパラメータを復元し、それを使ってモデルを作り、そのモデルを使ってモデルパラメータを読み込む。

def load_model(filename: str) -> Tuple[Word2Vec, HyperParameters]:
    with np.load(filename) as f:
        deserializer = serializers.NpzDeserializer(f)

        pickled_params = deserializer("hyper_parameters", None)
        params = pickle.loads(pickled_params.tobytes())  # type: HyperParameters

        model = Word2Vec(params.vocabulary_size, params.vector_dimension)
        deserializer["model"].load(model)

        return model, params

実験

Chainer 本で使っていたデータセットである ptb.train.txt を使ってモデルを学習してみた。

ハイパーパラメータがいくつかあるので、以下のすべての組み合わせ(12通り)を使って学習を行った。

モデル: SingleMatrix( $W$ と $W'$ が等しいモデル), DoubleMatrix( $W$ と $W'$ が等しくないモデル)
ベクトルの次元: 50, 100, 200
ウィンドウサイズ: 3, 5

類義語

ベクトルのコサイン類似度を使って"意味"が近い単語をアドホックに調べてみる。（Search はコサイン類似度を使って似ている単語を探してくるヘルパークラス。名前が適当なのでもっといい名前がほしい。）

In [7]: model, params = load_model("./word2vec/results/word2vec_singlematrix_vd200_w5_ns5_bs100_epoch29.npz")

In [8]: s = Search(dataset.vocabulary, model)

In [9]: s.find_similar_words("ibm")
Out[9]: 
['ibm',
 'mainframes',
 'computer',
 'digital',
 'machine',
 'mainframe',
 'machines',
 'armonk',
 'software',
 'chips']

In [10]: s.find_similar_words("monday")
Out[10]: 
['monday',
 'late',
 'friday',
 'tuesday',
 'thursday',
 'wednesday',
 'sell-off',
 'opened',
 'close',
 'points']

In [11]: model, params = load_model("./word2vec/results/word2vec_doublematrix_vd200_w5_ns5_bs100_epoch29.npz")

In [12]: s = Search(dataset.vocabulary, model)

In [13]: s.find_similar_words("ibm")
Out[13]: 
['ibm',
 'digital',
 'mainframe',
 'armonk',
 'customers',
 'computer',
 'machine',
 'machines',
 'hewlett-packard',
 'computers']

In [14]: s.find_similar_words("monday")
Out[14]: 
['monday',
 'friday',
 'tuesday',
 'wednesday',
 'thursday',
 'advancers',
 'october',
 'quoted',
 'week',
 'trading']

一応それっぽい結果になったけど、これを見ただけだと上手く学習できているのかはいまいちわからない。