JoeyNMT

📌

到達目標:LLMの内部を理解し操作することができる


1. JoeyNMTの構築&訓練と評価
1. JoeyNMTの訓練&評価, small_parallel_enja翻訳(ja→en)

small_parallel_enjaコーパスのダウンロード

※orig→元のコーパス, tok→トークナイズされたコーパス

git clone https://github.com/odashi/small_parallel_enja.git

JoeyNMTのインストール

git clone https://github.com/joeynmt/joeynmt.git
uv add -r joeynmt/requirements.txt
uv add sacrebleu # sacreBLEUのインストール
uv add sacrebleu[ja]
cd joeynmt

コードテスト

python3 -m unittest

# エラーが出た時
uv add importlib_metadata
uv add --editable . --dev

small_parallel_enja翻訳用のディレクトリを作成&学習用のファイルを用意

※./joeynmt/scripts/にbuild_vocab.pyがあったため、そちらを使った方が良いかもしれない

mkdir -p small_parallel_jaen
cp /home/nishida/b4/joeynmt/small_parallel_enja/train.ja small_parallel_jaen/train.ja
cp /home/nishida/b4/joeynmt/small_parallel_enja/train.en small_parallel_jaen/train.en

SentencePieceモデルの学習

spm_train --input=small_parallel_jaen/train.ja --model_prefix=small_parallel_jaen/spm.ja --vocab_size=9000 --character_coverage=0.9995 --model_type=bpe
spm_train --input=small_parallel_jaen/train.en --model_prefix=small_parallel_jaen/spm.en --vocab_size=5000 --character_coverage=1.0 --model_type=bpe

SentencePieceの語彙ファイルをJoeyNMT用に変換

cut -f1 small_parallel_jaen/spm.ja.vocab | head -n 9000 > small_parallel_jaen/src_vocab.txt
cut -f1 small_parallel_jaen/spm.en.vocab | head -n 5000 > small_parallel_jaen/trg_vocab.txt

yamlファイルの設定(./joeynmt/configsファイル内で自分で作成)

※early_stoppingが効かない不具合が発生しているため、自分で作った方が良いかも。おそらくschedulingのせいで学習率が低くなっているので、"plateau"にすると解決するかも。

name: "small_parallel_jaen"

data:
  train: "/home/nishida/b4/joeynmt/small_parallel_enja/train"
  dev: "/home/nishida/b4/joeynmt/small_parallel_enja/dev"
  test: "/home/nishida/b4/joeynmt/small_parallel_enja/test"
  dataset_type: "plain"
  src:
    lang: "ja"
    level: "bpe"
    voc_file: "small_parallel_jaen/src_vocab.txt"
    tokenizer_type: "sentencepiece"
    tokenizer_cfg:
      model_file: "small_parallel_jaen/spm.ja.model"
  trg:
    lang: "en"
    level: "bpe"
    voc_file: "small_parallel_jaen/trg_vocab.txt"
    tokenizer_type: "sentencepiece"
    tokenizer_cfg:
      model_file: "small_parallel_jaen/spm.en.model"

testing:
  n_best: 1
  beam_size: 6
  beam_alpha: 1.0
  batch_size: 1024
  batch_type: "token"
  max_output_length: 100
  eval_metrics: ["bleu"]
  return_prob: "none"
  return_attention: False
  generate_unk: False
  no_repeat_ngram_size: -1
  repetition_penalty: -1
  sacrebleu_cfg:
    tokenize: "intl" # ja-mecab

training:
	# load_model: "small_parallel_enja/model/???.ckpt"
	reset_best_ckpt: False
	reset_scheduler: False
	reset_optimizer: False
	reset_iter_state: False
	random_seed: 42
	optimizer: "adamw"
	normalization: "tokens"
	adam_betas: [0.9, 0.98]
	scheduling: "warmupinversesquareroot"
	loss: "crossentropy"
	learning_rate: 0.001
	learning_rate_min: 1.0e-06
	learning_rate_warmup: 4000
	clip_grad_norm: 1.0
	weight_decay: 0.0
	label_smoothing: 0.1
	batch_multiplier: 8
	batch_size: 512 # 2048 per device
	batch_type: "token"
	early_stopping: True
	early_stopping_metric: "bleu"
	patience: 5
	minimize_valid_metric: True
	epochs: 10000
	updates: 100000
	validation_freq: 200
	logging_freq: 200
  model_dir: "small_parallel_jaen/model_1"
  overwrite: False # small_parallel_jaen内の全てのファイルが消される
  shuffle: True
  use_cuda: True
  fp16: True
  print_valid_sents: [0, 1, 2, 3, 4]
  keep_best_ckpts: 5
  num_workers: 0

model:
  initializer: "xavier_uniform"
	embed_initializer: "xavier_uniform"
  embed_init_gain: 1.0
  init_gain: 1.0
  bias_initializer: "zeros"
  tied_embeddings: False
  tied_softmax: False
  encoder:
    type: "transformer"
    num_layers: 8
    num_heads: 16
    embeddings:
      embedding_dim: 1024
      scale: True
      dropout: 0.
    hidden_size: 1024
    ff_size: 4096
    dropout: 0.3
    layer_norm: "pre"
  decoder:
    type: "transformer"
    num_layers: 6
    num_heads: 16
    embeddings:
      embedding_dim: 1024
      scale: True
      dropout: 0.
    hidden_size: 1024
    ff_size: 4096
    dropout: 0.3
    layer_norm: "pre"

訓練の実行

screen -S train_joeynmt_1 bash -c 'CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt train configs/small_parallel_jaen.yaml'

評価の実行

CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt test small_parallel_jaen/model_1/config.yaml

sacreBLEUの結果

Evaluation result (beam search): bleu:  36.76, 0.0195[sec]

2. 事前学習済みJoeyNMTの評価, KFTT翻訳(ja→en)

KFTTコーパスのダウンロード

※orig→元のコーパス, tok→トークナイズされたコーパス

wget https://phontron.com/kftt/download/kftt-data-1.0.tar.gz
tar -zxvf kftt-data-1.0.tar.gz
rm kftt-data-1.0.tar.gz

事前学習済みJoeyNMTのインストール(./joeynmtファイル内)

wget https://cl.uni-heidelberg.de/statnlpgroup/joeynmt2/jparacrawl_jaen.tar.gz
tar -zxvf jparacrawl_jaen.tar.gz
rm jparacrawl_jaen.tar.gz

yamlファイルの設定(検証と評価で使うファイルをkfttコーパスに変更, ファイルは元からある)

data:
  dev: "/home/nishida/b4/joeynmt/kftt-data-1.0/data/orig/kyoto-dev"
  test: "/home/nishida/b4/joeynmt/kftt-data-1.0/data/orig/kyoto-test"

評価の実行

CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt test jparacrawl_jaen/config.yaml

sacreBLEUの結果

 Evaluation result (beam search): bleu:  11.49, 0.1077[sec]

追. model_1のいらないファイルを削除

.hypsといったファイルを削除するスクリプトの作成(./joeynmt/scriptsファイル内で自分で作成)

import os
import glob

def clean_files(directory):
    extensions = [".hyps", ".refs", ".loss", ".bleu"]
    if not os.path.exists(directory):
        print(f"Directory does not exist: {directory}")
        return

    removed = 0
    for ext in extensions:
        pattern = os.path.join(directory, f"*{ext}")
        for file_path in glob.glob(pattern):
            os.remove(file_path)
            print(f"Deleted: {file_path}")
            removed += 1

    if removed == 0:
        print("No files matched for deletion.")
    else:
        print(f"{removed} files deleted.")

clean.pyをコマンドラインから実行するためのスクリプトの作成(./joeynmt/scriptsファイル内で自分で作成)

import sys
from .clean import clean_files

def main():
    if len(sys.argv) < 3:
        print("Usage: python3 -m scripts <command> <target_directory>")
        sys.exit(1)

    command = sys.argv[1]
    target_dir = sys.argv[2]

    if command == "clean":
        clean_files(target_dir)
    else:
        print(f"Unknown command: {command}")
        sys.exit(1)

if __name__ == "__main__":
    main()

必要がないファイルを削除

python3 -m scripts clean small_parallel_jaen/model_1

没. JoeyNMTの訓練&評価, KFTT翻訳(ja→en)

JoeyNMTのインストール

git clone https://github.com/joeynmt/joeynmt.git
uv add -r joeynmt/requirements.txt
uv add sacrebleu # sacreBLEUのインストール
cd joeynmt

コードテスト

python3 -m unittest

# エラーが出た時
uv add importlib_metadata
uv add --editable . --dev

KFTT翻訳用のディレクトリを作成&学習用のファイルを用意

mkdir -p kftt_jaen
cp /home/nishida/b4/joeynmt/kftt-data-1.0/data/orig/kyoto-train.ja kftt_jaen/train.ja
cp /home/nishida/b4/joeynmt/kftt-data-1.0/data/orig/kyoto-train.en kftt_jaen/train.en

SentencePieceモデルの学習

spm_train --input=kftt_jaen/train.ja --model_prefix=kftt_jaen/spm.ja --vocab_size=9000 --character_coverage=0.9995 --model_type=bpe
spm_train --input=kftt_jaen/train.en --model_prefix=kftt_jaen/spm.en --vocab_size=5000 --character_coverage=1.0 --model_type=bpe

SentencePieceの語彙ファイルをJoeyNMT用に変換

cut -f1 kftt_jaen/spm.ja.vocab | head -n 9000 > kftt_jaen/src_vocab.txt
cut -f1 kftt_jaen/spm.en.vocab | head -n 5000 > kftt_jaen/trg_vocab.txt

yamlファイルの設定(./joeynmt/configsファイル内で自分で作成)

name: "kftt_jaen"

data:
  train: "kftt_jaen/train"
  dev: "/home/nishida/b4/joeynmt/kftt-data-1.0/data/orig/kyoto-dev"
  test: "/home/nishida/b4/joeynmt/kftt-data-1.0/data/orig/kyoto-test"
  dataset_type: "plain"
  src:
    lang: "ja"
    level: "bpe"
    voc_file: "kftt_jaen/src_vocab.txt"
    tokenizer_type: "sentencepiece"
    tokenizer_cfg:
      model_file: "kftt_jaen/spm.ja.model"
  trg:
    lang: "en"
    level: "bpe"
    voc_file: "kftt_jaen/trg_vocab.txt"
    tokenizer_type: "sentencepiece"
    tokenizer_cfg:
      model_file: "kftt_jaen/spm.en.model"

testing:
  n_best: 1
  beam_size: 6
  beam_alpha: 1.0
  batch_size: 1024
  batch_type: "token"
  max_output_length: 100
  eval_metrics: ["bleu"]
  return_prob: "none"
  return_attention: False
  generate_unk: False
  no_repeat_ngram_size: -1
  repetition_penalty: -1
  sacrebleu_cfg:
    tokenize: "intl"

training:
  # load_model: "kftt_jaen/model/???.ckpt"
  reset_best_ckpt: False
  reset_scheduler: False
  reset_optimizer: False
  reset_iter_state: False
  random_seed: 42
  optimizer: "adam"
  normalization: "tokens"
  adam_betas: [0.9, 0.98]
  scheduling: "warmupinversesquareroot"
  loss: "crossentropy"
  learning_rate: 0.001
  learning_rate_min: 1.0e-09
  learning_rate_warmup: 4000
  clip_grad_norm: 1.0
  weight_decay: 0.0
  label_smoothing: 0.1
  batch_multiplier: 8
  batch_size: 512 # 2048 per device
  batch_type: "token"
  early_stopping_metric: "bleu"
  epochs: 5
  updates: 100000
  validation_freq: 1000
  logging_freq: 200
  model_dir: "kftt_jaen/model"
  overwrite: False # kftt_jaen内の全てのファイルが消される
  shuffle: True
  use_cuda: True
  fp16: False
  print_valid_sents: [2000, 2001, 2002, 2003, 2004]
  keep_best_ckpts: 5
  num_workers: 0

model:
  initializer: "xavier"
  embed_initializer: "xavier"
  embed_init_gain: 1.0
  init_gain: 1.0
  bias_initializer: "zeros"
  tied_embeddings: False
  tied_softmax: False
  encoder:
    type: "transformer"
    num_layers: 8
    num_heads: 16
    embeddings:
      embedding_dim: 1024
      scale: True
      dropout: 0.
    hidden_size: 1024
    ff_size: 4096
    dropout: 0.3
    layer_norm: "pre"
  decoder:
    type: "transformer"
    num_layers: 6
    num_heads: 16
    embeddings:
      embedding_dim: 1024
      scale: True
      dropout: 0.
    hidden_size: 1024
    ff_size: 4096
    dropout: 0.3
    layer_norm: "pre"

訓練の実行(1エポック3時間くらいかかる)

screen -S train_joeynmt bash -c 'CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt train configs/kftt_jaen.yaml'

評価の実行

CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt test kftt_jaen/model/config.yaml

2. 相対位置エンコーディングの実装
1. 相対位置エンコーディングの理解

相対位置エンコーディングの重要性[1][2]

機械翻訳: 翻訳においては、文法構造や意味の正確な伝達は単語間の相対的な位置関係に強く依存する。文の構造は言語間で異なるため、相対位置が重要な指標となる。

相対位置エンコーディングの計算方法[3][4][5]

  • Self-Attentionにおいて、QWKの計算中で実装を行うものである。
  • 通常のAttentionScore=QK^TにS_relを加算
    AttentionScore=QK+Srel\text{AttentionScore} = QK^\top + S_{\text{rel}}

    S_rel:トークンiとそこからの距離に基づいて決定されるQK^Tの補正値

  • S_relを、embeddingとQueryとの内積で計算

    効率的に求めるため、計算にSkewアルゴリズムを用いる

参考資料

[1] Positional Encoding徹底解説:Sinusoidal(絶対位置)から相対位置エンコーディング - nomulog

[2] Transformerとは?世界を変えた深層学習モデルの仕組みをわかりやすく徹底解説 - nomulog

[3] Transformerにおける相対位置エンコーディングを理解する。 #機械学習 - Qiita

[4] [1803.02155] Self-Attention with Relative Position Representations

[5] MusicTransformer-pytorch (GitHub)

2. 相対位置エンコーディングの実装

transformer_layers.pyにあるMultiHeadedAttentionで相対位置エンコーディングを実装

self.len_k = None
self.max_seq = 2048
self.E = nn.Parameter(torch.randn(self.max_seq, self.head_size))
# [batch_size, num_heads, query_len, key_len]
scores = torch.matmul(q, k.transpose(2, 3))

# 相対位置エンコーディングの計算
if torch.equal(q, k): # Self-Attentionのみ
    self.len_k = k.size(2)
    self.len_q = q.size(2)
    E = self._get_left_embedding(self.len_q, self.len_k).to(q.device)
    QE = torch.einsum('bhld,md->bhlm', [q, E])
    QE = self._qe_masking(QE)
    Srel = self._skewing(QE)
    scores += Srel

# compute scores
scores = scores / math.sqrt(self.head_size)

それに伴い、MultiHeadedAttentionに関数を追加

def _get_left_embedding(self, len_q, len_k):
    starting_point = max(0,self.max_seq-len_q)
    e = self.E[starting_point:,:]
    return e

def _skewing(self, tensor: torch.Tensor):
    padded = F.pad(tensor, [1, 0, 0, 0, 0, 0, 0, 0])
    reshaped = torch.reshape(padded, shape=[padded.size(0), padded.size(1), padded.size(-1), padded.size(-2)])
    Srel = reshaped[:, :, 1:, :]
    if self.len_k > self.len_q:
        Srel = F.pad(Srel, [0, 0, 0, 0, 0, 0, 0, self.len_k-self.len_q])
    elif self.len_k < self.len_q:
        Srel = Srel[:, :, :, :self.len_k]
  return Srel

@staticmethod
def _qe_masking(qe):
    query_len = qe.size(-2)
    key_len = qe.size(-1)
    mask = torch.arange(key_len, device=qe.device).unsqueeze(0) <= \
        torch.arange(query_len, device=qe.device).unsqueeze(1)
    mask = mask.to(qe.dtype)
    return qe * mask.unsqueeze(0).unsqueeze(0) 

small_parallel_enja翻訳(ja→en)での評価 ※追をやるならそちらを先に

任意:yamlファイルの設定を書き換える

training:
  model_dir: "small_parallel_jaen/model_2"

訓練の実行

screen -S train_joeynmt_2 bash -c 'CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt train configs/small_parallel_jaen.yaml'

評価の実行

CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt test small_parallel_jaen/model_2/config.yaml

sacreBLEUの結果

Evaluation result (beam search): bleu:  38.22, 0.0192[sec]

比較:実装前のsacreBLEUの結果

Evaluation result (beam search): bleu:  36.76, 0.0195[sec]

追. yamlファイルから実装の有無を設定

model.pyにあるTransformerEncoder, TransformerDecoderの変数に以下を追加

use_relative_pos_enc=cfg.get("relative_position_encoding", False)

encoders.pyにあるTransformerEncoderの__init__()を以下のように設定

def __init__(self, ..., use_relative_pos_enc=False, **kwargs):
    self.use_relative_pos_enc = use_relative_pos_enc

また、TransformerEncoderLayerの変数に以下を追加

use_relative_pos_enc = use_relative_pos_enc

decoders.pyにあるTransformerDecoderの__init__()を以下のように設定

def __init__(self, ..., use_relative_pos_enc = False, **kwargs):
    self.use_relative_pos_enc = use_relative_pos_enc

また、TransformerDecoderLayerの変数に以下を追加

use_relative_pos_enc = use_relative_pos_enc

tranformer_layers.pyにあるTransformerEncoderLayerとTransformerDecoderLayerの__init__()を以下のように設定

def __init__(self, ..., use_relative_pos_enc = False):
    self.use_relative_pos_enc = use_relative_pos_enc

また、MultiHeadedAttentionの変数に以下を追加

※src_trgはCross_Attentionのため入力しない

use_relative_pos_enc=use_relative_pos_enc

MultiHeadedAttentionの__init__()を以下のように設定

def __init__(self, ..., use_relative_pos_enc = False):
    self.use_relative_pos_enc = use_relative_pos_enc

MultiHeadedAttentionの相対位置エンコーディングの条件を以下のように変更

if self.use_relative_pos_enc:

yamlファイルの設定を書き換える

※次からはここを変えるだけで相対位置エンコーディングの有無を決められる

model:
  relative_position_encoding: True


3. LLMとKLダイバージェンスを使ったMTの高精度化
1. LLMとKLダイバージェンスについての理解

MTにLMを使用する事の重要性[1]

モノリンガルデータの活用: MTには、大量の並列データが必要であり、低リソースの言語には性能が低下する。単言語のデータは比較的多く存在するため、そのデータで学習されたLMを使用し、低リソースの言語に対しても高精度を保つようにする。

LLMとKLダイバージェンスの適用方法[1]

  • MTの出力分布をLMに近づけるために、目的関数に正則化項を追加する。
    • 通常の目的関数
      LNMT=t=1NlogpTM(yty<t,x)\begin{equation} \mathcal{L}_{\text{NMT}} = \sum_{t=1}^{N} - \log p_{\text{TM}}(\bm{y}_t | \bm{y}_{<t}, \bm{x}) \end{equation}
    • 提案手法の目的関数
      L=LNMT+λτ2DKL(pTM(yty<t,x;τ)pLM(yty<t;τ))\begin{equation} \mathcal{L} = \mathcal{L}_{\text{NMT}} + \lambda \bm{\tau}^2 D_\text{{KL}}\left( p_{\text{TM}}(\bm{y}_t | \bm{y}_{<t}, \bm{x}; \bm{\tau}) \parallel p_{\text{LM}}(\bm{y}_t | \bm{y}_{<t};\bm{\tau}) \right) \end{equation}

      D_KL:MTとLMの2つの出力分布間のKLダイバージェンス

  • LMはターゲット言語を学習したモデルで、LMの代わりにLLMを使用する。

参考資料

[1] https://aclanthology.org/2020.emnlp-main.615v2.pdf

2. トークナイザをLLM(Llama-3.2-1B-Instruct)に変更

LLM(Llama-3.2-1B-Instruct)のvocabを作成

transformersをインストールしておく

uv add transformers

.bashrcに以下を追加

export HUGGING_FACE_TOKEN={トークン名}

llamaのvocabを保存するスクリプトの作成(./joeynmt/scriptsファイル内で自分で作成)

import os
from transformers import AutoTokenizer

def build_llama_vocab(output_dir):
    token = os.environ['HUGGING_FACE_TOKEN']
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", token=token)

    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, "llama_vocab.txt")
    vocab_items = sorted(tokenizer.get_vocab().items(), key=lambda x: x[1])

    with open(output_path, "w", encoding="utf-8") as f:
        for token, _ in vocab_items:
            f.write(token + "\n")

    print(f"Vocabulary saved to {output_path}")

__main__.pyファイルに条件分岐を追加。ないならそのままスクリプトを実行。

from .build_llama_vocab import build_llama_vocab

elif command == "build_llama_vocab":
    build_llama_vocab(target_dir)

llamaのvocabを追加

python3 -m scripts build_llama_vocab small_parallel_jaen

JoeyNMTのTokenizerをLLM(Llama-3.2-1B-Instruct)に設定

tokenizers.pyにHuggingFaceTokenizerクラスを追加

import os
from transformers import AutoTokenizer

class HuggingFaceTokenizer(BasicTokenizer):
    def __init__(
        self,
        access_token_name: str,
        level: str = "bpe",
        lowercase: bool = False,
        normalize: bool = False,
        max_length: int = -1,
        min_length: int = -1,
        **kwargs,
    ):
        super().__init__(level, lowercase, normalize, max_length, min_length, **kwargs)
        token = os.environ[access_token_name]
        self.model_file: Path = Path(kwargs["model_file"])
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_file, token=token, use_fast=True)
        special_tokens_dict = {"pad_token": "<|finetune_right_pad_id|>", "unk_token": "<|reserved_special_token_0|>"}
        self.tokenizer.add_special_tokens(special_tokens_dict)

    def __call__(self, raw_input: str, is_train: bool = False) -> List[str]:
        if raw_input is None:
            return None
        tokens = self.tokenizer.tokenize(raw_input)
        if is_train and self._filter_by_length(len(tokens)):
            return None
        return tokens

    def post_process(
        self,
        sequence: Union[List[str], str],
        generate_unk: bool = True,
        cut_at_sep: bool = True,
    ) -> str:
        if isinstance(sequence, list):
            if cut_at_sep and self.sep_token in sequence:
                try:
                    sep_pos = sequence.index(self.sep_token)
                    sequence = sequence[sep_pos + 1:]
                except ValueError:
                    pass
            sequence = self._remove_special(sequence, generate_unk=generate_unk)
            sequence = self.tokenizer.convert_tokens_to_string(sequence)

        if self.normalize:
            sequence = remove_extra_spaces(sequence)
        return sequence

    def set_vocab(self, vocab) -> None:
        super().set_vocab(vocab)

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(tokenizer={self.tokenizer.name_or_path}, "
            f"lowercase={self.lowercase}, normalize={self.normalize}, "
            f"filter_by_length=({self.min_length}, {self.max_length}))"
        )

_build_tokenizer()関数の中に以下を追加

elif tokenizer_type == "huggingface":
    assert "model_file" in tokenizer_cfg
    tokenizer = HuggingFaceTokenizer(
		    access_token_name = cfg["access_token_name"],
        level=cfg["level"],
        lowercase=cfg.get("lowercase", False),
        normalize=cfg.get("normalize", False),
        max_length=cfg.get("max_length", -1),
        min_length=cfg.get("min_length", -1),
        **tokenizer_cfg,
    )

set_vocab()関数のunk_tokenとeos_tokenを以下のようにする

※Llamaの特殊トークンの位置が異なるため

self.unk_token = vocab.specials[0]
self.eos_token = vocab.specials[3]

Vocabularyクラス__init__()関数の中の関数定義を以下に変更

if not (cfg.unk_token and cfg.pad_token and cfg.bos_token and cfg.eos_token):
		self.add_tokens(tokens=self.specials + self.lang_tags + tokens)
else:
    self.add_tokens(tokens=self.lang_tags + tokens)

small_parallel_jaen.yamlのdata:のsrc, trgを以下の設定にしておく

voc_file: "small_parallel_jaen/llama_vocab.txt"
tokenizer_type: "huggingface"
access_token_name: "HUGGING_FACE_TOKEN"
tokenizer_cfg:
  model_file: "meta-llama/Llama-3.2-1B-Instruct"

また、small_parallel_jaen.yamlのdata:に以下を追加しておく

special_symbols:
  pad_token: "<|finetune_right_pad_id|>"   # Llamaの語彙にはあるが認識されていない
  unk_token: "<|reserved_special_token_0|>"# Llamaの語彙にはあるが認識されていない
  bos_token: "<|begin_of_text|>" # JoeyNMTがLlamaの特殊トークン位置を認識していない
  eos_token: "<|eot_id|>"        # JoeyNMTがLlamaの特殊トークン位置を認識していない
  pad_id: 128004
  unk_id: 128002
  bos_id: 128000
  eos_id: 128009

課題1をsmall_parallel_enja翻訳(ja→en)で評価

任意:yamlファイルの設定を書き換える

※relative_position_encodingは課題2の追加をやっていないと設定できません

name: "small_parallel_jaen"
use_cuda: True
fp16: True

data:
  train: "/home/nishida/b4/joeynmt/small_parallel_enja/train"
  dev: "/home/nishida/b4/joeynmt/small_parallel_enja/dev"
  test: "/home/nishida/b4/joeynmt/small_parallel_enja/test"
  src:
    lang: "ja"
    level: "bpe"
    remove_space: True
    voc_file: "small_parallel_jaen/llama_vocab.txt"
    tokenizer_type: "huggingface"
    access_token_name: "HUGGING_FACE_TOKEN"
    tokenizer_cfg:
      model_file: "meta-llama/Llama-3.2-1B-Instruct"
    lowercase: True
    max_sent_length: 50
  trg:
    lang: "en"
    level: "bpe"
    voc_file: "small_parallel_jaen/llama_vocab.txt"
    tokenizer_type: "huggingface"
    access_token_name: "HUGGING_FACE_TOKEN"
    tokenizer_cfg:
      model_file: "meta-llama/Llama-3.2-1B-Instruct"
    lowercase: True
    max_sent_length: 50
  special_symbols:
    pad_token: "<|finetune_right_pad_id|>"
    unk_token: "<|reserved_special_token_0|>"
    bos_token: "<|begin_of_text|>"
    eos_token: "<|eot_id|>"
    pad_id: 128004
    unk_id: 128002
    bos_id: 128000
    eos_id: 128009

testing:
  beam_size: 10
  alpha: 1.0
  eval_metrics: ["bleu"]

training:
  random_seed: 42
  label_smoothing: 0.1
  optimizer: "adamw"
  normalization: "tokens"
  adam_betas: [0.9, 0.999]
  learning_rate: 0.0001
  learning_rate_min: 0.00005
  batch_size: 64
  scheduling: "plateau"
  patience: 5
  decrease_factor: 0.5
  early_stopping_metric: "loss"
  epochs: 100000
  validation_freq: 600
  logging_freq: 100
  eval_metric: ["bleu"]
  model_dir: "small_parallel_jaen/model_3-1"
  overwrite: False # small_parallel_jaen内の全てのファイルが消される
  shuffle: True
  use_cuda: True
  max_output_length: 100
  print_valid_sents: [0, 1, 2, 3, 4]

model:
  initializer: "xavier_uniform"
  init_gain: 1.0
  bias_initializer: "zeros"
  embed_initializer: "xavier_uniform"
  embed_init_gain: 1.0
  tied_embeddings: False
  tied_softmax: False
  relative_position_encoding: False
  encoder:
    type: "transformer"
    num_layers: 6
    num_heads: 8
    embeddings:
      embedding_dim: 512
      scale: True
      freeze: False
    hidden_size: 512
    ff_size: 128
    dropout: 0.3
    layer_norm: "pre"
    activation: "relu"
  decoder:
    type: "transformer"
    num_layers: 6
    num_heads: 8
    embeddings:
      embedding_dim: 512
      scale: True
      freeze: False
    hidden_size: 512
    ff_size: 128
    dropout: 0.3
    freeze: False
    layer_norm: "pre"
    activation: "relu"

訓練の実行

screen -S train_joeynmt_3-1 bash -c 'CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt train configs/small_parallel_jaen.yaml'

評価の実行

CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt test small_parallel_jaen/model_3-1/config.yaml

sacreBLEUの結果

Evaluation result (beam search): bleu:  36.90, 0.0363[sec]

課題2をsmall_parallel_enja翻訳(ja→en)で評価

任意:yamlファイルの設定を書き換える

※relative_position_encodingは課題2の追加をやっていないと設定できません

training:
  model_dir: "small_parallel_jaen/model_3-2"

model:
  relative_position_encoding: True

訓練の実行

screen -S train_joeynmt_3-2 bash -c 'CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt train configs/small_parallel_jaen.yaml'

評価の実行

CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt test small_parallel_jaen/model_3-2/config.yaml

sacreBLEUの結果

Evaluation result (beam search): bleu:  38.05, 0.0363[sec]

比較:課題1のsacreBLEUの結果

Evaluation result (beam search): bleu:  36.90, 0.0363[sec]

3. LLM(Llama-3.2-1B-Instruct)とKLダイバージェンスの実装

MTとLLMでのKLダイバージェンスの実装

model.pyにあるbuild_model関数に以下を追加

lm_prior = cfg.get("lm_prior", {})

同じ関数のModelに引数を渡しているところに以下を追加

lm_prior=lm_prior

model.pyにあるModelクラスの__init__()に以下を追加

def __init__(self, ..., lm_prior: dict = None) -> None:
	self.lm_prior = lm_prior

model.pyにあるModelクラスの@loss_function.setterを以下のように設定

@loss_function.setter
def loss_function(self, cfg: Tuple):
    loss_type, label_smoothing = cfg
    assert loss_type == "crossentropy"
    self._loss_function = XentLoss(
        pad_index=self.pad_index, smoothing=label_smoothing, lm_prior=self.lm_prior
    )

model.pyにあるModelクラスの__forward__()を以下のように変更

# compute log pro
if all([self.loss_function.kl_lambda, self.loss_function.kl_tau, self.loss_function.lm_model]):
    log_probs = F.log_softmax(out, dim=-1)
    kl_log_probs = F.log_softmax(out / self.loss_function.kl_tau, dim=-1)
else:
    log_probs = F.log_softmax(out, dim=-1)
    kl_log_probs = torch.zeros_like(log_probs)

# compute batch loss
batch_loss = self.loss_function(log_probs, kl_log_probs, **kwargs)

loss.pyにあるXentLossの__init__()に以下を追加

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch.nn.functional as F

def __init__(self,..., lm_prior: dict = None): # lm_priorを追加
	  self.kl_lambda = lm_prior.get("kl_lambda", 0.0)
	  self.kl_tau = lm_prior.get("kl_tau", 0.0)
	  self.token = lm_prior.get("access_token_name", None)
	  self.model_name = lm_prior.get("model_file", None)
    if all([self.kl_lambda, self.kl_tau, self.model_name, self.token]):
        self.lm_model = AutoModelForCausalLM.from_pretrained(self.model_name, token=self.token)
        self.lm_tokenizer = AutoTokenizer.from_pretrained(self.model_name, token=self.token)
        special_tokens_dict = {"pad_token": "<|finetune_right_pad_id|>", "unk_token": "<|reserved_special_token_0|>"}
        self.lm_tokenizer.add_special_tokens(special_tokens_dict)
    else:
        self.lm_model = None

XentLossの__forward__()の引数にkl_log_probs: Tensorを追加する(28行目)

また、XentLossの__forward__()を以下に変更

assert "trg" in kwargs
log_probs, targets = self._reshape(log_probs, kwargs["trg"])

# compute loss
logits = self.criterion(log_probs, targets)
if all([self.kl_lambda, self.kl_tau, self.lm_model]):
    kl_log_probs, _ = self._reshape(kl_log_probs, kwargs["trg"])
    with torch.no_grad():
        lm_input_ids = insert_eos_before_padding(
            kwargs["trg_input"], 
            eos_token_id=self.lm_tokenizer.eos_token_id, 
            pad_token_id=self.lm_tokenizer.pad_token_id
            )
        lm_inputs = {
            "input_ids": lm_input_ids.to(log_probs.device),
            "attention_mask": (lm_input_ids != self.lm_tokenizer.pad_token_id).to(log_probs.device)
            }
        lm_logits = self.lm_model(**lm_inputs).logits[:, 1:, :]
        lm_probs = F.softmax(lm_logits / self.kl_tau, dim=-1)
        lm_probs_flat = lm_probs.reshape(-1, lm_probs.size(-1))
        non_pad_mask = (kwargs["trg"].contiguous().view(-1) != self.pad_index)
    return logits + self.kl_lambda * self.kl_tau * self.kl_tau * F.kl_div(kl_log_probs[non_pad_mask], lm_probs_flat[non_pad_mask], reduction='batchmean')
else:
    return logits

関数の追加

def insert_eos_before_padding(input_ids: torch.Tensor, eos_token_id: int, pad_token_id: int) -> torch.Tensor:
    B, L = input_ids.size()
    output_ids = torch.full((B, L + 1), pad_token_id, dtype=input_ids.dtype, device=input_ids.device)

    for i in range(B):
        seq = input_ids[i]
        pad_pos = (seq == pad_token_id).nonzero(as_tuple=True)[0]
        if len(pad_pos) > 0:
            insert_pos = pad_pos[0].item()
        else:
            insert_pos = L

        output_ids[i, :insert_pos] = seq[:insert_pos]
        output_ids[i, insert_pos] = eos_token_id
        output_ids[i, insert_pos + 1:L + 1] = seq[insert_pos:]

    return output_ids

yamlファイルの設定を書き換える

※次からはここを変えるだけでKLダイバージェンスの設定を決められる

model:
  lm_prior:
    kl_lambda: 0.5
    kl_tau: 2.0
    access_token_name: "HUGGING_FACE_TOKEN"
    model_file: "meta-llama/Llama-3.2-1B-Instruct"

small_parallel_enja翻訳(ja→en)での評価

任意:yamlファイルの設定を書き換える

training:
  model_dir: "small_parallel_jaen/model_3-3"

訓練の実行

screen -S train_joeynmt_3-3 bash -c 'CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt train configs/small_parallel_jaen.yaml'

評価の実行

CUDA_VISIBLE_DEVICES=0 python3 -m joeynmt test small_parallel_jaen/model_3-3/config.yaml

課題3のsacreBLEUの結果

Evaluation result (beam search): bleu:  39.09, 0.0391[sec]

比較:課題1のsacreBLEUの結果

Evaluation result (beam search): bleu:  36.90, 0.0363[sec]

比較:課題2のsacreBLEUの結果

Evaluation result (beam search): bleu:  38.05, 0.0363[sec]