データナード

機械学習と自然言語処理についての備忘録 (旧ナード戦隊データマン)

keras-laser: LASERをkerasで実装

LASERはマルチリンガルのSentence Embeddingを獲得するためのモデルで、facebook researchによって開発されています。

LASERのオリジナルの事前訓練済みモデルはfairseqを用いて訓練されているようです。事前訓練済みモデルをそのまま使うこともできますが、事前訓練をするためのオリジナルのコードは公開されていません。事前訓練を自前で行うために、kerasでイチからLASERの実装を試みます。

https://github.com/sugiyamath/keras-laser

※ 現在、このリポジトリはprivateに設定しています。

概要

LASERのモデルの概要は以下記事から見ることができます。

また、LASERの実装の詳細は以下の論文に書かれています:

ただ、肝心の「訓練のためのコード」がありません。issueには以下のように書かれています:

Hello, We are aware that there's a lot of interest in the training code. The original LASER training code was based on the version of fairseq which now dates back almost 1 year. We are working on a substantially improved version of LASER training which will use the current fairseq and scales much better to many languages. Please be patient :-) https://github.com/facebookresearch/LASER/issues/70#issuecomment-540647189

現時点ではまだ訓練コードがないため、kerasで実装することにしました。

事前準備

  • 訓練データを用意してください。
  • 訓練データは、パラレルコーパスです。
  • 出力言語として英語とスペイン語を使うため、英語・スペイン語とペアになったパラレルコーパスを複数言語で用意する必要があります。
  • 用意した複数言語のパラレルコーパスから、BPEを訓練しておく必要があります。

主要なコード

現状のモデル

import tensorflow as tf
import tensorflow.keras.layers as layers

from data_generator2 import data_generate_phase1, data_generate_phase2

MAX_OUT_LANG = 3
BS = 3
STEPS = 100
EPOCHS = 10000

CALLBACKS = [
    tf.keras.callbacks.ModelCheckpoint(
        filepath="/root/work/keras-laser/models/model_last_euro.h5",
        save_weights_only=False,
        monitor='val_acc',
        mode='max',
        save_best_only=False),
    tf.keras.callbacks.CSVLogger(
        "/root/work/keras-laser/models/training_euro.log")
]


def _validation_data(gen):
    return next(gen)


def build_model(max_features=50001,
                max_out_features=10001,
                emb_dim=320,
                max_len=256,
                lstm_dim=512,
                lstm_depth=4,
                lidemb_dim=32,
                drate=0.1,
                lr=0.001):
    bpe1 = layers.Input(batch_shape=(None, max_len),
                        name="bpe_enc",
                        dtype=tf.int32)
    bpe2 = layers.Input(batch_shape=(None, max_len),
                        name="bpe_dec",
                        dtype=tf.int16)
    lid = layers.Input(batch_shape=(None, max_len), name="lid", dtype=tf.uint8)
    #shared_emb = layers.Embedding(max_features,
    #                              emb_dim,
    #                              input_length=max_len,
    #                              mask_zero=True)
    #encoder = shared_emb(bpe1)
    #decoder = shared_emb(bpe2)

    encoder = layers.Embedding(max_features,
                               emb_dim,
                               input_length=max_len,
                               mask_zero=True)(bpe1)
    decoder = layers.Embedding(max_features,
                               emb_dim,
                               input_length=max_len,
                               mask_zero=True)(bpe2)
    for i in range(lstm_depth):
        encoder = layers.Bidirectional(
            layers.LSTM(lstm_dim, return_sequences=True))(encoder)
    sentemb = layers.GlobalMaxPooling1D(name="sentemb")(encoder)
    sentemb_fx = layers.RepeatVector(max_len)(sentemb)
    est_h = layers.Dense(lstm_dim * 4,
                         activation="linear")(sentemb)
    est_c = layers.Dense(lstm_dim * 4,
                         activation="linear")(sentemb)
    lid_emb = layers.Embedding(MAX_OUT_LANG,
                               lidemb_dim,
                               input_length=max_len,
                               mask_zero=True)(lid)
    cat_emb = layers.Concatenate()([decoder, sentemb_fx, lid_emb])
    decoder = layers.LSTM(lstm_dim * 4,
                          return_sequences=True)(cat_emb,
                                                 initial_state=[est_h, est_c])
    decoder = layers.Dropout(drate)(decoder)
    out = layers.TimeDistributed(
        layers.Dense(max_out_features, activation="softmax"))(decoder)
    model = tf.keras.Model([bpe1, bpe2, lid], out)
    adam = tf.keras.optimizers.Adam(learning_rate=lr)
    model.compile(optimizer=adam,
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    return model


def main():
    gen = data_generate_phase2(data_generate_phase1(), BS)
    X_val, Y_val = _validation_data(gen)
    model = build_model()
    model.fit_generator(gen,
                        steps_per_epoch=STEPS,
                        epochs=EPOCHS,
                        validation_data=(X_val, Y_val),
                        callbacks=CALLBACKS)


if __name__ == "__main__":
    main()

現時点での考察

  • 訓練に非常に時間がかかります。バッチサイズを100、ステップ数を10としたとき、1 epochにかかる時間はおよそ390sです (私の環境での時間)。
  • モデルがこれで正しいのか、個人的に確信が持てません (訓練中であるため)。例えば、LSTMレイヤーのパラメータは適切でしょうか?
  • 基本的な部分はseq2seqモデルに似ています。
  • このモデルはSentence Embeddingを獲得することが目的なので、Sentence Embeddingのレイヤーの出力を使って文類似度を取れるか検証する必要があります。
  • Sentence Embedding以外の情報がデコーダに漏れないようにしたほうが良い気がします。
  • 基本的に、ある言語をエンコーダに入力としたときのデコーダの出力は英語またはスペイン語にします。英語を入力とした表現を獲得するためには、出力が他の言語(スペイン語)である必要があると考えられます。そして、ある言語との対として最も多くのデータが存在するのは英語で、二番目がスペイン語です。おそらく、これがこの2つの言語を出力にする理由です。
  • エンコーダとデコーダのbpe embeddingは共有しています。ただ、オリジナルの実装がこの部分を共有しているのか不明です。
  • オリジナルのモデルではBidirectional LSTMの層は5層になっていますが、ここが多いほうが多くの情報を保持できる可能性があります。
  • Sentence Embeddingの出力を特徴量として様々なタスクを多言語で解けるのがLASERの利点であるため、凍結せずにファインチューニングをすると多言語への対応力が失われる可能性があります。
  • TimeDistributedを必ず使ったほうが良いようです。他の時間の情報が漏れると、leakageと同じようなことになっていると思います。
  • 可変入出力にする場合としない場合にsentembに違いが現れるでしょうか。
  • leakageを防ぐことを目的にする場合、デコーダ側の結合EmbeddingをSpatialDropout1Dなどでdropoutしたほうがよいかもしれません。
  • 通常、事前訓練済みのオリジナルのLASERを使ったほうが精度も高くて楽です。
  • kerasの熱狂的信者はFacebookという組織を嫌う傾向があり、kerasで実装したLASERを作る気がありません(要出典)。

参考

  1. [1812.10464] Massively Multilingual Sentence Embeddings for Zero-Shot Cross-Lingual Transfer and Beyond
  2. LASER natural language processing toolkit - Facebook Engineering
  3. Sequence to sequence - training - Keras Documentation