TLM（Task-driven Language Modeling）の手法を使って、効率的に事前学習済BERTモデルの精度を超えられるか検証してみる

# はじめに
TLM（Task-drive language Modeling）という手法を提唱する論文を読んで、内容が面白そうだったので、実際に実装してみて、TLMの効果を検証してみました。
実装する上で色々と勉強になったので、誰かのためになるのでは、と思い記事としてまとめてみようと思います。

# TLM: Task-drive language Modelingとは？

[NLP From Scratch Without Large-Scale Pretraining:
A Simple and Efficient Framework](https://arxiv.org/pdf/2111.04130.pdf)という論文で紹介されている手法です。

TLMの課題背景は以下のような感じです。

- 近年のNLPはとんでもなく巨大なデータセットで何十億や何百億のパラメータを持つTransformerベースのモデルを扱うのが主流の1つに感じますが、そんな巨大なデータを扱えるのはごく一部の巨大企業だったりして、個人レベルでは事前学習モデルの学習段階から改良しようと思ってもコスト面でなかなか難しい
- このような課題感はNLPの研究において大きな障壁となってて、NLPの長期的な発展を考えると言語モデルの事前学習の効率化を研究・改善する取り組みが重要である


アブストを簡単にまとめると

- 昨今のNLPは大規模なリソースで事前学習されたモデルを扱うことが主流だけど、事前学習モデルは学習コストが大きく誰でも簡単に作れるわけじゃない
- TLMは大規模な事前学習モデルを必要とせず、特定のタスクを効率的に解くフレームワークである
- 4つのドメインの8つの分類データセットにおいて、TLMは学習FLOPを2桁削減しながら、事前に学習した言語モデル（例：RoBERTa-Large）と同等以上の結果を達成した。

→高い精度と効率性を持つTLMが、自然言語処理の民主化と開発の促進に貢献することを期待する、だそうです。

![スクリーンショット 2022-01-29 18.34.34.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/240999/775979e0-a5c1-0b5c-fbc3-d2251ee0008d.png)


## TLMの手法

![スクリーンショット 2022-01-29 18.27.35.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/240999/7397b8e7-dfb9-d247-9175-3824bede7df2.png)

ここで用語の定義をしときます。

- 一般コーパス: Wikipediaやcc100など事前学習に使われる巨大なデータセットを想定
- 固有タスクデータ: Fine Tuningで使うデータ（livedoorニュースコーパスのデータなど）を想定

TLMの手法はとてもシンプルで以下の3stepで構成されます。

1. **Data Selection**  
固有タスクデータ($\cal{T}$)を検索クエリーとして一般コーパスのデータを検索(※1)し、一般コーパスの部分集合$\cal{S}$を取得する
1. **Joint Learning**  
一般コーパスデータ$\cal{S}$と、固有タスクデータ$\cal{T}$を使って、事前学習タスク(MLMなど)と固有タスク(分類タスクなど)を同時に学習させる。その際、最適化させる損失関数は以下の通りで、$\rho_1$と$\rho_2$はハイパーパラメータ(※2)である
<img src=https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/240999/b95d1298-0782-16ff-5c20-abf3a19a59c0.png width=70%>
1. **Fine Tuning**  
最後に固有タスクデータ$\cal{T}$でFine Tuningする。

上記の3ステップを得ることで、**大規模なデータで事前学習されたモデルに頼ることなく**、下流タスクを効率的に解けるよ、ってのかTLMの主張です。

手法に関して2点ほど補足があります。
※1: Data Selectionのステップの検索アルゴリズムは**BM25**を用いています。もちろんここでの検索の目的は$\cal T$に似ているデータセットを用意することなわけで、BM25よりも何かしらの埋め込み表現を使った類似度検索のほうが検索精度は良いでしょう。ただTLMは効率性に重きを置いているため、学習とかが必要なく素早く精度良く検索するためにBM25を用いているようです。

※2: $\rho_1, \rho_2$の決め方は固有タスクによって最適値が全然違うようです。詳しくは論文を参照していただきたいですが、論文ではグリッドサーチで決定しているようです。この$\rho_1, \rho_2$を最適化するのにかなりの計算コストかかってるんじゃないかと予想してて、全体的に本当に効率化されているのかやや懐疑的に思います。

# 実装
TLMの主張をだいたい理解したところで実際にTLMの手法を試してみて、効率的に下流タスクが解けるのか検証してみようと思います。
検証の対象は以下とします。

- アーキテクチャー  
日本語BERT-base. 東北大の事前学習済BERTモデルを比較対象とします。
- 一般コーパス  
Wikipedia
- 固有タスクデータと下流タスク  
Livedoorニュースコーパスで、本文のカテゴリー分類問題を解くことを目的とします。

実装は論文の著者のGithubを参考にしています。本記事と合わせてご参照ください。

- https://github.com/yaoxingcheng/TLM

上で紹介した3ステップを1つ1つ実装で確認していきます。

**注意**
本記事を検証した際の`transformers`のバージョンは以下の通りです。今後の`transformers`のバージョンアップにより、本記事で紹介している実装が動作しなくなる可能性があります。

```
transformers              4.15.0
datasets                  1.18.0
```


## Step1. Data Selection

### 固有タスクデータを準備する
まず、事前にカテゴリー、タイトル、本文に分けたLivedoorニュースコーパスを手元に用意します。

```python
import pickle
import pandas as pd

with open('./livedoor_data.pickle', 'rb') as r:
    df = pickle.load(r)

display(df.sample(3))
```
![スクリーンショット 2022-01-30 0.15.01.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/240999/2d8a1f30-5581-925d-20cf-7fd1e89dc9b1.png)

Livedoorニュースコーパスを学習データとテストデータに分けます。

```python
from sklearn.model_selection import train_test_split

train_df, test_df = train_test_split(df, train_size=0.8)
print('train size', train_df.shape)
print('test size', test_df.shape)
# train size (5900, 3)
# test size (1476, 3)

# カテゴリーのID列を付与しておく
categories = df['category'].unique().tolist()
category2id = {cat: categories.index(cat) for cat in categories}
train_df['category_id'] = train_df['category'].map(lambda x: category2id[x])
test_df['category_id'] = test_df['category'].map(lambda x: category2id[x])
```
上で分けた学習データを$\cal T$とし、Wikipediaを検索する上での検索クエリーに使います。

### 一般コーパスを準備する
huggingfaceが提供している`datasets`ライブラリを使って簡単にWikipediaデータを取得することができます。2022年1月20日時点のデータを以下のようにダウンロードすると、キャッシュとして、18GBのデータが保存されます。（これ、結構時間かかりました。。。パース？する処理がかなり時間かかります。）
キャッシュは`$HOME/.cache/huggingface/datasets/wikipedia`に保存されてました。2回目以降はキャッシュをロードするだけなので、すぐにwikipediaデータを呼び出すことができます。

```python
from datasets import load_dataset

wiki_dataset = load_dataset("wikipedia",
                            language="ja",
                            date="20220120",
                            beam_runner="DirectRunner"
                           )
wiki_dataset = wiki_dataset['train']
wiki_dataset
# Dataset({
#     features: ['title', 'text'],
#     num_rows: 1624858
# })
```

### Elasticsearchを使ってBM25で一般データを検索する
参考にしているGithubのソースコードでもそうしているように、BM25による検索はElasticsearchを使います。
こういうニーズのためなのか、`datasets`ライブラリにはElasticsearchと連携する機能が備わっています。

- https://huggingface.co/docs/datasets/faiss_es.html#id2

まずはElasticsearch本体を裏で起動しておく必要があるのですが、公式サイトを参考にDockerで起動しました。

- https://www.elastic.co/guide/en/elasticsearch/reference/7.16/docker.html#docker

Elasticsearchを起動させた状態にして、python側では以下のように実行します。

```python
# 事前にpipでElasticsearchクライアントをインストールしておきましょう
# pip install elasticsearch
from elasticsearch import Elasticsearch

es = Elasticsearch(
    [{"host": "localhost", "port": "9200"}],
    timeout=300,
)
es_config = {
    "settings": {
        "number_of_shards": 1,
        "analysis": {"analyzer": "kuromoji"},
    },
    "mappings": {
        "properties": {
            "text": {
                "type": "text",
                "analyzer": "standard",
                "similarity": "BM25"
            },
        }
    },
}

# Elasticsearchにwikipediaデータのindexをはる
wiki_dataset.add_elasticsearch_index(
    'text',
    es_client=es,
    es_index_name='wiki_es_bm25',
    es_index_config=es_config,
)
```

ElasticsearchにWikipediaデータのindexがはれたところで、固有タスクデータで検索をします。TLMの論文では固有タスクデータ1件に対して類似度TOP50件のデータを一般コーパスから検索したりしてますが、Top50件はさすがに結構な量になってしまいますので、今回はTOP1件だけにしときましょう。

```python
from tqdm import tqdm

nearest_corpus = []
train_corpus = train_df['body'].tolist()

# elasticsearchの検索クエリーは1024文字以内にしないといけないので、冒頭1024文字でカットしています。
train_corpus = [t[:1024] for t in train_corpus]

query_batch = 100
for i in tqdm(range(0, len(train_corpus), query_batch)):
    query_corpus = train_corpus[i:i+query_batch]

    # バッチ検索が早くて便利ですが、メモリに注意！
    nearest_batch = wiki_dataset.get_nearest_examples_batch('text', query_corpus, k=1)
    for batch in nearest_batch.total_examples:
        nearest_corpus += batch['text']

# もちろん異なる検索クエリーに対して同一データがTop1でヒットする可能性もあるので、重複削除します。
print(len(nearest_corpus)) # 5900
nearest_corpus = list(set(nearest_corpus))
print(len(nearest_corpus)) # 3226
```

上で作成された`nearest_corpus`が論文でいうところの$\cal S$になるんじゃないかと思います。$\cal T$は`train_df`ですね。
Data Selectionのステップは以上になります。


## Step2. Joint Learning

Joint Learningってつまりはマルチタスク事前学習なわけですが、BERTの事前学習モデルはhuggingfaceの`Trainer`クラスを使い倒すことで実現が容易になります。`Trainer`クラスを使うために、TLM用のDataset、DataCollator、モデルの3つを準備します。

以降のソースコードで変数名の接頭詞に`internal`とか`external`とかを使いますが、`internal`→$\cal T$、`external`→$\cal S$に関する変数だと理解していただければと思います。

今回の実装でここが一番難しいところでした。というのも、今回のデータは$\cal S$と$\cal T$で損失を計算する上でのラベルが違います。$\cal S$はMLMのラベルだけで良いですが、$\cal T$はMLM＋クラス分類ラベルが必要です。そのように目的変数が異なるデータをまとめて1つのDatasetとして学習させたいです。論文著者のGithubのソースコードでは$\cal S$と$\cal T$のDataLoaderを分けて実装していますが、`Trainer`クラスはDatasetを1つしか受け取ることができないので、論文著者はTrainerクラスをスクラッチで実装しています。さすがにそれはかなり大変なので、Datasetを1つでまとめる代わりにバッチ内に目的変数が異なるデータを混在させながら、Joint Learningを実現することを試みます。

### Datasetを準備
MLMの正解ラベルの付与は次に出てくる`DataCollatorForLanguageModeling`に丸投げしたいです。（自分でMLMを実装しようとするとちょっと大変だと思ったので）。
`DataCollatorForLanguageModeling`はtoken id列に変換したデータを受け取ることを想定しているっぽいので、Datasetは以下のようなデータ構造をもつものとして用意したいです。

```python
print(tlm_dataset[0])
{'input_ids': [2, 70, 657, 12, 5,・・・
 'class_labels': 2,
 'data_type': 1}
```
ここで、`data_type`はこのデータが$\cal S$のデータなのか$\cal T$のデータなのかを区別する識別子を想定しています。後続のモデルを作成する際の`forward`関数内で使う予定です。（がなくても実装できる気がする。。。）
`class_labels`はLivedoorニュースコーパスのカテゴリーIDです。
実際のDatasetクラスは以下のように実装しました。

```python
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer

class TLMDataset(Dataset):
    def __init__(self, tokenizer, external_corpus=None, internal_df=None, max_length=512):
        self.sentences = []
        
        # external data
        if external_corpus is not None:
            external_sentences = []
            for text in external_corpus:
                # wikipediaデータを軽く前処理しとく
                text = text.replace('\u3000', ' ') # 全角スペースを半角スペースに変換
                text = text.replace('\n\n', '\n').replace('\n', ' ')
                sents = text.split('。')
                external_sentences += [sent+'。' for sent in sents]
            # 重複削除
            external_sentences = list(set(external_sentences))
            # 外部データはカテゴリーのラベルはないので、-100をセットする
            # -100はCrossEntropyLossのignore_indexで、-100のデータの損失は計算されません。
            external_sentences = [{
                'input_ids': tokenizer(text, truncation=True, max_length=max_length)['input_ids'],
                'class_labels': -100,
                'data_type': 0
                } for text in tqdm(external_sentences)]
            self.sentences += external_sentences

        # internal data
        if internal_df is not None:
            internal_sentences = [{
                'input_ids': tokenizer(row.body, truncation=True, max_length=512)['input_ids'],
                'class_labels': row.category_id,
                'data_type': 1
            } for row in tqdm(internal_df.itertuples(), total=internal_df.shape[0])]
            self.sentences += internal_sentences
    
    def __len__(self):
        return len(self.sentences)
    
    def __getitem__(self, idx):
        return self.sentences[idx]

# Tokenizerは東北大BERTのtokenizerを拝借します。
tokenizer = AutoTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')

tlm_dataset = TLMDataset(tokenizer, external_corpus=nearest_corpus, internal_df=train_df)
```

### DataCollatorを用意

用意も何も、`transformers`の`DataCollatorForLanguageModeling`を使うだけです。

```python
from transformers import DataCollatorForLanguageModeling

tlm_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=True,
    mlm_probability= 0.15
)
```

ここまでで用意したDatasetとDataCollatorが想定する動きをするか確認しましょう。

```python
tlm_loader = DataLoader(tlm_dataset, collate_fn=tlm_collator, batch_size=8, shuffle=True)
batch = next(iter(tlm_loader))
print(batch.keys())
# dict_keys(['input_ids', 'class_labels', 'data_type', 'attention_mask', 'labels'])
print(batch)
# {'input_ids': tensor([[    2, 26992,     9,    36,     4, 17358, 28454,    38,・・・
# 'class_labels': tensor([-100, -100, -100, -100, -100, -100, -100,・・・
# 'data_type': tensor([0, 0, 0, 0, 0, 0, 0, 0]),
# 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1・・・
# 'labels': tensor([[ -100,  -100,  -100,  -100,   630,  -100,  -100,  -100,・・・}
```

問題なく実装できてそうです。


### モデルを用意

上の動作確認で見た`batch`をモデルが受け取ることになります。

#### MLMのweightについて
`data_type`によってMLMの損失のweight($\rho_1, \rho_2$)を調整できるように実装したいです。論文では解きたいタスクによって$\rho_1, \rho_2$の値は様々なようですが、論文内のAGNewsデータによる検証が今回のLivedoorニュースコーパスに近いデータかなと思われるので、AGNewsデータの検証時に使われたパラメータ$\rho_1 = 4$、$\rho_2 = 20$を今回は採用しようと思います。

#### ネットワークの実装について
著者のGithubのソースコードを参考にしています。BERT本体は`BertConfig`でアーキテクチャーを設定して、それを`BertModel`に渡すだけです。`BertModel`の最終層の出力のうち、CLSベクトルはカテゴリー分類問題を解く用のネットワークに、各トークンベクトルはMLMを解く用のネットワーク(`BertOnlyMLMHead`)に渡してやります。`Trainer`クラスはモデルのアウトプットに`loss`という名前のkeyに最適化したい損失が格納されていることを前提としているので、モデルのアウトプットはその要件を満たすように実装します（Githubの実装と同様に`MaskedLMOutput`を使います）。


```python
import torch
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers import BertConfig, BertModel
from transformers import PreTrainedModel
from transformers.models.bert.modeling_bert import BertOnlyMLMHead
from transformers.modeling_outputs import MaskedLMOutput


class BertForTLM(PreTrainedModel):
    def __init__(self, config, num_labels, external_mlm_weight=4, internal_mlm_weight=20):
        super().__init__(config)
        self.config = config
        self.bert = BertModel(config, add_pooling_layer=False)
        self.mlm_head = BertOnlyMLMHead(config)

        # Githubの実装を参考にしていますが、普通にnn.Linearだけだとどうなるかも気になる
        self.class_head = nn.Sequential(
            nn.Linear(config.hidden_size, config.hidden_size),
            nn.GELU(),
            nn.Linear(config.hidden_size, num_labels)
        )
        
        # batch内の各tensorのlossをそのまま保持するためにreduction='noone'を設定
        # defaultはreduction='mean'、つまりbatchのlossの平均を返す
        self.loss_function = nn.CrossEntropyLoss(reduction='none')
        
        # TLMの論文のρ1とρ2に相当
        # タスクによって最適な重みは全然違うようです
        self.external_mlm_weight = external_mlm_weight
        self.internal_mlm_weight = internal_mlm_weight
    
    def _get_mean_loss(self, loss):
        return loss.sum() / (loss != 0.).sum().clamp(min=1)

    def forward(self, input_ids, attention_mask, class_labels=None, labels=None, data_type=None):
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        outputs = outputs['last_hidden_state']
        cls_vec = outputs[:, 0, :]
        class_outputs = self.class_head(cls_vec)

        class_loss = 0.0
        # このモデルを学習後、推論時など正解のラベルがないときでもこのモデルを使えるように
        # class_labelsが与えられたときだけlossを計算するようにしています。
        if class_labels is not None:
            class_loss = self.loss_function(class_outputs, class_labels)
            class_loss = self._get_mean_loss(class_loss)

        if data_type is not None:
            mlm_weight_tensor = (data_type == 0).int() * self.external_mlm_weight + (data_type == 1).int() * self.internal_mlm_weight
            mlm_weight_tensor = torch.stack([mlm_weight_tensor] * input_ids.shape[1]).T.flatten()
        
        mlm_loss = 0.0
        # このモデルを学習後、推論時など正解のラベルがないときでもこのモデルを使えるように
        # labelsが与えられたときだけlossを計算するようにしています。
        if labels is not None:
            mlm_outputs = self.mlm_head(outputs)
            mlm_loss = self.loss_function(mlm_outputs.view(-1, self.config.vocab_size), labels.view(-1))
            mlm_loss *= mlm_weight_tensor
            mlm_loss = self._get_mean_loss(mlm_loss)
    

        total_loss = class_loss + mlm_loss

        return MaskedLMOutput(
            loss=total_loss,
            logits=class_outputs
        )


num_labels = len(categories) # Livedoorニュースコーパスのカテゴリー数（9つ）
config = BertConfig(vocab_size=len(tokenizer))
model = BertForTLM(config, num_labels)
```


### 事前学習

後は上で実装した`tlm_dataset`, `tlm_collator`, `model`をTrainerクラスに渡して学習するだけです。バッチサイズがやや大きめで設定していますが、colabなどで学習する場合は8くらいで12GBほどGPUメモリを喰いました。エポック数はもう少し多くてもいいかもしれませんが、エポック数3でもバッチサイズ32で3時間30分ほど学習に時間かかります。（バッチサイズ8でcolab上で実行すると30時間ほどかかる試算でした。なんだかんだ結構時間かかっちゃいます。。。）

```python
from transformers import TrainingArguments
from transformers import Trainer

training_args = TrainingArguments(
    output_dir= './pretrained_TLM',
    overwrite_output_dir=True,
    num_train_epochs=3,
    per_device_train_batch_size=32,
    save_steps=10000,
    save_total_limit=2,
    prediction_loss_only=True
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=tlm_collator,
    train_dataset=tlm_dataset
)

trainer.train()
```

学習が終わったらパラメータを保存します。以下では2つの方法でモデルを保存しています。`trainer.save_model`で保存すると、その後`AutoModel.from_pretrained`でいつもの呼び出し方で学習済モデルをロードできますが、その際`class_head`のパラメータのロードの仕方がよくわかりませんでした。`class_head`の学習済パラメータもロードしたいときは`torch.save`で`model.state_dict()`を保存し、呼び出すときは、`BertForTLM`のインスタンス変数に対して`load_state_dict`でパラメータをロードすれば良いです。
TLMの手法は事前学習の段階で下流タスクも学習させているので、下流タスク用のhead(`class_head`)のパラメータをFine Tuningに使うべきかなと思っているのですが、実際に検証してみたところ、学習済`class_head`を使っても、Fine Tuning時に新しくhead(例えば`nn.Linear`)を追加して学習してもさほど結果は変わらなかったです。この辺はどうするべきかは模索中です。

```python
trainer.save_model('./pretrained_TLM')
torch.save(model.state_dict(), './pretrained_TLM/tlm_params.pth')
```

## Step3. Fine Tuning

TLMの方法で事前学習したモデルと東北大が提供しているBERTモデル（`cl-hotoku/bert-base-japanese-whole-word-masking`）
のFine Tuningにおける精度を比較してみます。
Fine Tuningの方法は上の学習のところでも触れたように`trainer.save_model`で保存したモデルに対して、Fine Tuning用のhead（`nn.Linear`）を追加して学習する方法を取りました。（つまり、TLMの方法で学習されたBERTモデルの部分のパラメータの比較をしているイメージです。）
論文でも行っているように、$\cal S$だけで事前学習した場合と一般コーパスからランダムにデータをピックアップしてJoint Learningしたときの精度も合わせて確認してみました。

まず、Fine Tuningで使うネットワークはこんな感じで実装します。

```python
import torch
import torch.nn as nn

class LivedoorNet(nn.Module):
    def __init__(self, pretrained_model, num_labels):
        super().__init__()
        self.hidden_size = pretrained_model.config.hidden_size
        self.bert = pretrained_model
        self.linear = nn.Linear(self.hidden_size, num_labels)
    
    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        outputs = outputs['last_hidden_state']
        outputs = outputs[:, 0, :]
        outputs = self.linear(outputs)
        return outputs
```

Fine Tuningで学習させるところはこんな感じで実装しました。とりあえず10エポック学習させてみました。

```python
from tqdm import tqdm
from sklearn.metrics import f1_score
import torch.optim as optim
from transformers import AutoModel, AutoTokenizer

# GPU使う
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')
loss_function = nn.CrossEntropyLoss()

pretrained_model = AutoModel.from_pretrained('./pretrained_TLM/')
# pretrained_model = AutoModel.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')

# Sだけで事前学習した場合
# pretrained_model = AutoModel.from_pretrained('./pretrained_only_random_external')

# ランダムに抽出したデータセットだけで事前学習した場合
# pretrained_model = AutoModel.from_pretrained('./pretrained_only_external')

# DatasetとDataLoaderを定義
# collatorはTLMの事前学習時に使ったtlm_collatorを使ってます。これ使うともちろんMLMのlabelもbatch内に含まれてしまいますが、まぁいいでしょう。
train_dataset = TLMDataset(tokenizer, internal_df=train_df)
test_dataset = TLMDataset(tokenizer, internal_df=test_df)
train_loader = DataLoader(train_dataset, collate_fn=tlm_collator, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, collate_fn=tlm_collator, batch_size=64, shuffle=False)

net = LivedoorNet(pretrained_model, len(categories))
# GPU2枚並列で使います。
net = nn.DataParallel(net, device_ids=[0, 1])
net.to(device)

# nn.DataParallel使うときはnet.”module”が必要ですが、nn.DataParallel使わないときは不要です。
# 事前学習済のパラメータ学習率を小さめに設定
optimizer = optim.Adam([
            {'params': net.module.bert.parameters(), 'lr': 5e-5},
            {'params': net.module.linear.parameters(), 'lr': 1e-4},
        ])


test_fscores = []
for epoch in range(10):

    # 学習ステップ
    all_loss = 0
    net.train()
    for batch in tqdm(train_loader):
        optimizer.zero_grad()
        input_ids = batch['input_ids'].cuda()
        attention_mask = batch['attention_mask'].cuda()
        y = batch['class_labels'].cuda()
        pred = net(input_ids, attention_mask=attention_mask)
        loss = loss_function(pred, y)
        loss.backward()
        optimizer.step()
        all_loss += loss.item()

    # 推論ステップ
    answer = []
    prediction = []
    net.eval()
    with torch.no_grad():
        for batch in tqdm(test_loader):
            input_ids = batch['input_ids'].cuda()
            attention_mask = batch['attention_mask'].cuda()
            y = batch['class_labels'].cuda()
            pred = net(input_ids, attention_mask=attention_mask)
            _, pred = torch.max(pred, 1)
            prediction += list(pred.cpu().numpy())
            answer += list(y.cpu().numpy())
        test_fscore = f1_score(answer, prediction, average='macro')
        test_fscores.append(test_fscore)

    print("epoch", epoch, "\ttrain loss", all_loss, "\ttest fscore", test_fscore)
```

最初にテストデータとして分けたデータのエポック毎のF1-scoreを比較してみます。

![download.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/240999/846a0a38-39a7-e1ac-eea9-70e1fbebe4c8.png)

- 東北大BERTモデルが圧倒的に強い
- TLMと$\cal S$のみで事前学習したモデルを比較すると、ややTLMが優勢？
- ランダムにピックアップして事前学習したモデルは立ち上がりの精度が悪く、Fine Tuningでだいぶ精度が上がったけど、他のモデルに劣る

東北大BERTには勝てなかったけど、TLMの手法であるJoint Learningの効果も少し実感できたかなという印象です。（自分の実装方法や検証方法が微妙なのかもですが、もう少し良い結果を期待してましたが。。。）
TLMのモデルの精度を向上させるためには、BM25による検索をTOP1だけじゃなく、論文でもやってるようにTOP50とかにして$\cal S$のデータ量をもっと増やすとか、$\rho_1, \rho_2$を調整する、とかが必要かも。事前学習時のエポック数も合わせて増やしたほうがいいかも。

## Attention weightの分析
最後に論文でも紹介されているAttention weightの分析を紹介します。
![スクリーンショット 2022-01-31 16.45.53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/240999/54ce7254-f34c-2135-04d2-1d6c7979b686.png)

詳しくは論文をご参照してほしいのですが、Multi-head Attentionの各weightの傾向として以下の2点が知られているようです。

**positional head**
最大のAttention weightの90％以上が隣接するトークンに割り当てられるheadであり、これらのheadはモデルの最終予測に重要な貢献をしている。（上図の赤枠）

**vertical head**
ほとんどの最大のAttention weightが[CLS],[SEP]またはピリオドトークン(".")に割り当てられているもので、これらは潜在的にあまり意味的または構文的な情報をエンコードしない。（上図のグレーアウトしている箇所）

これらの観点で上図を見てみるとTLMのほうが比較対象のBERT-baseやRoBERTa-baseと比べて、positional headの数が多く、vertical headの数が少ないです。つまり、TLMのほうがより情報量の多いAttentionを学習していることを示唆する、ようです。他の文章で実験してみても同様の傾向があるようで、この分析の仕方面白いなぁと思ったので、今回のケースでも各Attention headのweightを可視化してみました。

インプットに使った文章はこちらです。
> [CLS]音楽とスポーツをファッショナブルに楽しむためのヘッドホンが誕生[SEP]


![図1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/240999/caca2d0a-27dd-71c5-99bd-2190dbe31b09.png)

軸の説明がなくて恐縮ですが、左から右にかけてhead1->head12、上から下にかけてLayer1->Layer12のAttention weightを可視化しています。
vertical headは今回はスルーしちゃってます。パット見これはpositional headだろうな、と思った箇所を青枠で囲って見ました。
精度評価でそもそも東北大BERTのほうが全然良かったので、ある意味その結果の裏付けみたいになっちゃいますが、論文で紹介されているような結果にはならなかったです。東北大BERTのほうがpositional headの数が明らかに多いし、特に層が深くなるにつれて、positional headの数が増えているのに対し、TLMは層が深くなるとpostional headが出てこなくなる、という結果になってしまいました。

# おわりに

大規模データによる事前学習の強さを改めて実感しました。実際のところBERTやRoBERTaなんかは日本語の事前学習済モデルが公開されているわけですし、それをありがたく使えばいいわけですが、もっと新しい独自のネットワークによる事前学習モデルを作りたい、けど大規模データを用意できない、とか学習コストが高すぎでそこまで事前学習してられない、ってときにはこのTLMの手法で事前学習モデルを作ってみるのもいいかもしれませんね。

ただTLMの論文でも書かれていますが、事前学習モデルの汎用性（Fine Tuningするだけでいろんなタスクを高精度で解ける）はどうしてもTLMは劣るようなので、いろんなタスクにモデルを転用したいという場合はやはり大規模データによる事前学習モデルに頼らざるを得ないですかねぇ。

おわり