Large Concept ModelsのSONAR論文とコードを読む

Metaが公開した拡散モデルを使った言語生成モデルであるLarge Concept Models(LCM)では、
従来のトークン単位での予測ではなく、より大きなまとまりである「概念」を用います。

https://ai.meta.com/research/publications/large-concept-models-language-modeling-in-a-sentence-representation-space/

LCMでは概念として文章単位の埋め込み表現を用います。文章の埋め込み表現はSONARと呼ばれるネットワークで行います。

どのように文章の埋め込み表現を行っているか気になったので、SONARの論文とコードを読んでいきます。LCMに関しては触れません。

SONARとは

https://arxiv.org/abs/2308.11466
https://github.com/facebookresearch/SONAR

SONARとは、翻訳タスクをベースにTransformerのEncoder/Decoder構成でマルチリンガル/マルチモーダル対応の固定長の文章の埋め込み表現を行うネットワークとなります。

特に、TransformerのEncoder/Decoder構成、マルチモーダル、マルチリンガルで埋め込み表現を行うところが目新しいポイントとなっていそうです。

SONARの構成

参照: 論文中の図

図にもあるようにSONARは、SpeachEncoders、TextEncoder、TextDecoderを持ちます。SpeachEncodersとTextEncoderの出力が近づくように学習することでマルチモーダルに対応してます。

embeeding時にはTextEncoder(SpeachEncoder)を使い、学習時にはTextEncoderによるembeedingをTextDecoderで復元するようなイメージです。

TextEncoderとTextDecoderには、Metaの開発した翻訳モデルであるNNLB 1Bのweightを初期値として用います。
NLLB: https://ai.meta.com/blog/nllb-200-high-quality-machine-translation/ja/

この記事では主に、TextEncoder、TextDecoderについて見ていきます。(マルチモーダルに関してはあまり触れません)

TextEncoder

TextEncoderは、Encoder部分とPooler部分で構成されます。

Encoder部分

Encoder部分には、Transformer構造におけるEncoderが使われます。

Encoder Layerは以下のコードで作られ、attentionやffnを持つことがわかります。
https://github.com/facebookresearch/SONAR/blob/549d287466443bd8720f938047882630c1c5c3f7/sonar/models/sonar_text/builder.py#L261-L268

例えば、Attentionについて見てみるとその実態はbuild_attentionにより作成されるStandardMultiheadAttentionを使っていることがわかります。
https://github.com/facebookresearch/SONAR/blob/549d287466443bd8720f938047882630c1c5c3f7/sonar/models/sonar_text/builder.py#L270-L282

TextEncoderのEncoder部分の出力は、shape=(batch_size, seq_len, hidden_dim)です。

Pooler部分

Poolerとは、TextEncoderのEncoder部分の出力を固定長ベクトルとするため、Poolingを行う部分です。論文によると、Poolingにはいくつか比較の結果、Mean-Poolingを選択しています。

実装的には、LAST, MAX, MEAN, ATTENTIONの4種類が設定可能で、LAST,　MAX, MEANは以下のように実装されています。
https://github.com/facebookresearch/SONAR/blob/549d287466443bd8720f938047882630c1c5c3f7/sonar/models/sonar_text/model.py#L100-L124

EOS pooling

PoolingにAttentionを使う場合です。

実装を見ていて一番気になった部分でしたが、EOS poolingを使うと学習が不安定になったということで不採用なっています

EOS Poolingの場合、TransformerにおけるDecoderを用います。つまり、self_attentionとcross_attentionを持ちます。

https://github.com/facebookresearch/SONAR/blob/549d287466443bd8720f938047882630c1c5c3f7/sonar/models/sonar_text/builder.py#L340-L357

入力は

• seqとしてbos_token列(shape=(batch_size, 1))
• encoder_outputとしてencoder_outputを入力

出力はshape=(batch_size, 1, hidden_dim)となる。
https://github.com/facebookresearch/SONAR/blob/549d287466443bd8720f938047882630c1c5c3f7/sonar/nn/encoder_pooler.py#L70-L81

bosのみembeddingする層を追加
https://github.com/facebookresearch/SONAR/blob/549d287466443bd8720f938047882630c1c5c3f7/sonar/models/sonar_text/builder.py#L307-L323

self_attentionは、BOSトークン列のself_attentionなので実質機能しておらず、cross_attetionについてもBOSトークンに対して入力文章の埋込表現が割り当てられるイメージ？　BOSに対してattentionを取っているので不安定にはなりそう。

TextDecoder

TextDecoderは、Transformer構造におけるDecoderが使われます。

https://github.com/facebookresearch/SONAR/blob/549d287466443bd8720f938047882630c1c5c3f7/sonar/models/sonar_text/builder.py#L579-L593

入力はSequence列とTextEncoderの出力するsentence embeddingになります。

https://github.com/facebookresearch/SONAR/blob/549d287466443bd8720f938047882630c1c5c3f7/sonar/nn/conditional_decoder_model.py#L67-L88

学習

マルチリンガルとマルチモーダルに対応するために2段階の学習を行います。今回は、第１段階のマルチリンガルに対応する学習のみ見ていきます。

学習は、翻訳タスクをベースとして、Auto-Encodeing的な学習も行います。

Translation Objective

入力文を多言語に翻訳することで、異なる言語の文が同じ意味を持つように埋め込み空間を調整
例: 「Hello, how are you?」と「Bonjour, comment ça va?」が近い埋め込みになるように学習

オートエンコーディング (Auto-Encoding)

文を一度埋め込み空間に変換し、再び元の言語にデコードすることで情報を保持する能力を向上

デノイジング・オートエンコーディング (Denoising Auto-Encoding)

ランダムなノイズ（単語の削除・順序の変更）を加えた文を元の文に復元するように学習

Lossについて

Lossは次のものを使います。

$$L=LMT+α⋅LMSE+β⋅LAE/DAE L = L_{MT} + \alpha · L_{MSE} + \beta · L_{AE/DAE}$$

$LMTL_{MT}$、$LMSEL_{MSE}$、$LAE/DAEL_{AE/DAE}$はそれぞれ、翻訳タスクのLoss、MSE、Auto-EncodingのLossとなっています。各Lossの計算は、翻訳用のデータセットを使い「source: 英語の文章、 target: フランス語」の組み合わせで$LMTL_{MT}$と$LMSEL_{MSE}$を算出、「source: 英語の文章、 target: 英語」の組み合わせで、$LAEL_{AE}$を算出しているようです。

$LMTL_{MT}$と$LMSEL_{MSE}$については、特徴量の生成に使われるAuto-Encoderですが、Auto-Encoderだけでは文章の特徴量を獲得するのは難しく、さらに翻訳タスクに比べて学習が簡単すぎることが報告されており、Auto-encodingとdenoising auto-encodingを翻訳タスクの事前学習とするのが良いそうです。

結果

興味のあった部分のみ

既存のモデルとの比較として、LASER3(LSTMベースの双方向エンコーダ)、LaBSE(BERTベース)、NLLB 1B(翻訳タスク向けTransformer)と比較しています。

• 翻訳タスク: NLLB 1Bに匹敵する翻訳精度
• 翻訳ペアが埋め込み空間で近くなるか(xsim-クロス言語類似度検索): LASER3やLaBSEよりも優れた結果となり、LASER3と比較してxsim++のエラー率を72%削減し
• オートエンコーディング: 固定サイズの文埋め込みを行いつつ、NLLB 1Bと同等の性能

所感

翻訳タスクベースでの埋め込み表現の獲得は面白いポイントでした。Auto-Encodingタスクも行なっていましたが、翻訳ペアが翻訳ペアが埋め込み空間で近くなるような学習で文章の埋め込み表現がうまく獲得できるのかは疑問。例えば、語順は異なるが同じ意味の文章などの評価はなかったので、どのような結果となるのか気になります。