#

線形バイアス (AliBi) 実験による注意

これは、AliBi モデルをトレーニングするための注釈付きの PyTorch 実験です。

14import torch
15from torch.utils.data import DataLoader
16
17from labml import experiment, tracker
18from labml.configs import option, calculate
19from labml_helpers.datasets.text import SequentialUnBatchedDataset
20from labml_nn.transformers.alibi import AlibiMultiHeadAttention
21from labml_nn.experiments.nlp_autoregression import transpose_batch
22from labml_nn.transformers import TransformerConfigs
23from labml_nn.transformers.gpt import Configs as GPTConfigs

#

コンフィギュレーション

GPT構成を拡張し、アテンションメカニズムを変更します。

26class Configs(GPTConfigs):

#

ALiBi ベースのトランスフォーマー (以下に定義)

34    transformer: TransformerConfigs = 'GPT_ALiBi'

#

より長い検証セット

36    valid_seq_len: int = 128
37    valid_loader = 'shuffled_longer_valid_loader'

#

最初のトークンと最後のトークンで損失を記録する

39    def other_metrics(self, output: torch.Tensor, target: torch.Tensor):

#

トレーニングシーケンスの長さよりも多くのトークンがある場合 (検証中)、

44        if self.seq_len < output.shape[0]:

#

トレーニングシーケンスの長さで損失を記録します

46            tracker.add(f'loss.{self.seq_len - 1}.', self.loss_func(output[self.seq_len - 1], target[self.seq_len - 1]))

#

最初のトークンで損失を記録する

48            tracker.add(f'loss.0.', self.loss_func(output[0], target[0]))

#

最後のトークンで損失を記録する

50        tracker.add(f'loss.{int(output.shape[0]) - 1}.', self.loss_func(output[-1], target[-1]))

#

AliBi アテンションモジュールを作成する

53def _alibi_mha(c: TransformerConfigs):

#

57    return AlibiMultiHeadAttention(c.n_heads, c.d_model, dropout_prob=c.dropout)

#

すべての注意メカニズムをAliBiに設定

61calculate(TransformerConfigs.encoder_attn, 'alibi_mha', _alibi_mha)
62calculate(TransformerConfigs.decoder_attn, 'alibi_mha', _alibi_mha)
63calculate(TransformerConfigs.decoder_mem_attn, 'alibi_mha', _alibi_mha)

#

シーケンス長のシャッフル検証データローダー valid_seq_len

66@option(Configs.valid_loader)
67def shuffled_longer_valid_loader(c: Configs):

#

71    return DataLoader(SequentialUnBatchedDataset(text=c.text.valid,
72                                                 dataset=c.text,
73                                                 seq_len=c.valid_seq_len),
74                      batch_size=c.batch_size,
75                      collate_fn=transpose_batch,
76                      shuffle=True)

#

AliBi ベースのトランスフォーマー構成

79@option(Configs.transformer, 'GPT_ALiBi')
80def _transformer_configs(c: Configs):

#

設定可能なトランス実装を使用しています

87    conf = TransformerConfigs()

#

埋め込みやロジットの生成に使用するボキャブラリーサイズを設定

89    conf.n_src_vocab = c.n_tokens
90    conf.n_tgt_vocab = c.n_tokens

#

GPT は GELU アクティベーションを使用して位置ごとのフィードフォワードを行います

92    conf.ffn.activation = 'GELU'

#

AliBi は位置埋め込みを使用しません

95    conf.src_embed = 'no_pos'
96    conf.tgt_embed = 'no_pos'

#

すべての注意メカニズムをAliBiに設定

99    conf.encoder_attn = 'alibi_mha'
100    conf.decoder_attn = 'alibi_mha'
101    conf.decoder_mem_attn = 'alibi_mha'

#

104    return conf

#

107def main():

#

実験を作成

109    experiment.create(name="gpt_alibi")

#

コンフィグの作成

111    conf = Configs()

#

オーバーライド設定

113    experiment.configs(conf, {

#

キャラクターレベルのトークナイザーを使う

115        'tokenizer': 'character',

#

プロンプトセパレータが空白

117        'prompt_separator': '',

#

サンプリングの開始プロンプト

119        'prompt': 'It is ',

#

タイニー・シェイクスピア・データセットを使う

121        'text': 'tiny_shakespeare',

#

'テキスト': 'tiny_shakespeare_no_split'、

#

コンテキストサイズを次の値にしてください $128$

125        'seq_len': 64,

#

コンテキストサイズを次の値にしてください $128$

127        'valid_seq_len': 80,

#

$32$ 時代に合わせた列車

129        'epochs': 128,

#

バッチサイズ $128$

131        'batch_size': 128,

#

エポックごとにトレーニングと検証を切り替える $10$

134        'inner_iterations': 10,

#

変圧器構成

137        'transformer.d_model': 128,
138        'transformer.ffn.d_ff': 512,
139        'transformer.n_heads': 8,
140        'transformer.n_layers': 4,
141        'transformer.dropout': 0.1,
142    })

#

保存および読み込み用のモデルを設定する

145    experiment.add_pytorch_models({'model': conf.model})

#

実験を始める

148    with experiment.start():

#

トレーニングを実行

150        conf.run()

#

154if __name__ == '__main__':
155    main()