banner
Nagi-ovo

Nagi-ovo

Breezing
github

LLM演進史(四):WaveNet——シーケンスモデルの畳み込み革新

本節内容のソースコードリポジトリ

前の部分で多層パーセプトロンの文字レベルの言語モデルを構築しましたが、今はその構造をより複雑にする時です。現在の目標は、入力シーケンスが現在の 3 文字ではなく、より多くの文字を入力できるようにすることです。それに加えて、すべてを 1 つの隠れ層に入れることは避け、情報を圧縮しすぎないようにします。これにより、WaveNetに似たより深いモデルが得られます。

WaveNet#

2016 年に発表され、基本的には言語モデルの一種ですが、予測対象は文字レベルや単語レベルのシーケンスではなく、音声シーケンスです。しかし、基本的なモデリング設定は同じです —— どちらも自己回帰モデル (Autoregressive Model) で、シーケンス内の次の文字を予測しようとしています。

Screenshot 2024-03-08 at 15.00.26

論文では、このツリー状の階層構造を使用して予測を行い、本節ではこのモデルを実装します。

nn.Module#

前のセクションの内容をクラスにカプセル化し、PyTorch の nn.Module の API を模倣します。これにより、「Linear」、「1 次元 Batch Norm」、「Tanh」などのモジュールをレゴブロックのように考え、それらを使って神経ネットワークを構築できます:

class Linear:
  
  def __init__(self, fan_in, fan_out, bias=True):
    self.weight = torch.randn((fan_in, fan_out), generator=g) / fan_in**0.5
    self.bias = torch.zeros(fan_out) if bias else None
  
  def __call__(self, x):
    self.out = x @ self.weight
    if self.bias is not None:
      self.out += self.bias
    return self.out
  
  def parameters(self):
    return [self.weight] + ([] if self.bias is None else [self.bias])

Linear、線形層のこのモジュールの役割は、フォワードパスの過程で行う行列の乗算です。

class BatchNorm1d:
  
  def __init__(self, dim, eps=1e-5, momentum=0.1):
    self.eps = eps
    self.momentum = momentum
    self.training = True
    # 逆伝播で訓練されたパラメータ
    self.gamma = torch.ones(dim)
    self.beta = torch.zeros(dim)
    # モーメンタム更新を使用して訓練するバッファ
    self.running_mean = torch.zeros(dim)
    self.running_var = torch.ones(dim)
  
  def __call__(self, x):
    # フォワードパスを計算
    if self.training:
      xmean = x.mean(0, keepdim=True) # バッチ平均
      xvar = x.var(0, keepdim=True) # バッチ分散
    else:
      xmean = self.running_mean
      xvar = self.running_var
    xhat = (x - xmean) / torch.sqrt(xvar + self.eps) # データを単位分散に標準化
    self.out = self.gamma * xhat + self.beta
    # バッファを更新
    if self.training:
      with torch.no_grad():
        self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean
        self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar
    return self.out
  
  def parameters(self):
    return [self.gamma, self.beta]

Batch-Norm:

  1. バックプロップ外で訓練された running mean & variance を持つ
  2. self.training = True、これはバッチノルムが訓練と評価の 2 つの段階で異なる動作をするため、バッチノルムの状態を追跡するためのトレーニングフラグが必要です
  3. バッチ内の要素の結合計算を処理し、活性化の統計的特性を制御し、内部共変量シフト (Internal Covariate Shift) を減少させます
class Tanh:
  def __call__(self, x):
    self.out = torch.tanh(x)
    return self.out
  def parameters(self):
    return []

以前の局所設定で g のtorch.Generatorを設定するのに対し、後でグローバルなランダムシードを設定します。

torch.manual_seed(42);

以下の内容は非常に馴染みがあるはずで、embedding table C と私たちのレイヤー構造を含んでいます:

n_embd = 10 # 文字埋め込みベクトルの次元
n_hidden = 200 # MLPの隠れ層におけるニューロンの数

C = torch.randn((vocab_size, n_embd))
layers = [
	Linear(n_embd * block_size, n_hidden, bias=False), 
	BatchNorm1d(n_hidden), 
	Tanh(),
	Linear(n_hidden, vocab_size),
]

# パラメータの初期化
with torch.no_grad():
	layers[-1].weight *= 0.1 # 最後の層(ここでは出力層)をスケールダウンし、初期モデルの予測に対する自信を減少させる

parameters = [C] + [p for layer in layers for p in layer.parameters()]
'''
リスト内包表記は、次のように相当します:
for layer in layers:
	for p in layer.parameters():
		p...
'''

print(sum(p.nelement() for p in parameters)) # 総パラメータ数
for p in parameters:
  p.requires_grad = True

最適化訓練部分はまだ変更せず、次に進んで私たちの損失関数の曲線が大きく変動しているのを見ます。これは 32 のバッチサイズが小さすぎるためで、各バッチ内での予測が非常に幸運または不幸である可能性がある(ノイズが大きい)からです。

Screenshot 2024-03-08 at 17.05.44

評価段階では、すべての層のトレーニングフラグを False に設定する必要があります(現在はバッチノルム層にのみ影響します):

# レイヤーを評価状態に設定
for layer in layers:
	layer.training = False

まず、損失関数のグラフの問題を解決します:

lossi はすべての損失を含むリストで、今私たちがやるべき基本的なことは、その中の値を平均して、より代表的な値を得ることです。

torch.view()の使用を復習します:

Screenshot 2024-03-08 at 17.25.53

view(5, -1)と同等です

これにより、リスト内の値を簡単に展開できます。

torch.tensor(lossi).view(-1, 1000).mean(1)

Screenshot 2024-03-08 at 20.09.18

今はずっと良く見えます。グラフから、学習率の減少が局所最小値に達したことが観察できます。

次に、以下に示す元の Embedding と Flattening 操作もモジュールに変えます:

emb = C[Xb]
x = emb.view(emb.shape[0], -1)
class Embedding:
  
  def __init__(self, num_embeddings, embedding_dim):
    self.weight = torch.randn((num_embeddings, embedding_dim))
    # 現在Cは埋め込みの重みになりました
    
  def __call__(self, IX):
    self.out = self.weight[IX]
    return self.out
  
  def parameters(self):
    return [self.weight]


class FlattenConsecutive:
    
  def __call__(self, x):
    self.out = x.view(x.shape[0], -1)
    return self.out
  
  def parameters(self):
    return []

PyTorch にはコンテナの概念もあり、基本的にはレイヤーをリストや辞書などに整理する方法です。その中にSequentialというものがあり、基本的な役割は与えられた入力をすべての層に順番に渡すことです:

class Sequential:
  
  def __init__(self, layers):
    self.layers = layers
  
  def __call__(self, x):
    for layer in self.layers:
      x = layer(x)
    self.out = x
    return self.out
  
  def parameters(self):
    # すべてのレイヤーのパラメータを取得し、それらをリストに展開します。
    return [p for layer in self.layers for p in layer.parameters()]

これでモデルの概念が得られました:

model = Sequential([
  Embedding(vocab_size, n_embd),
  Flatten(),
  Linear(n_embd * block_size, n_hidden, bias=False),
  BatchNorm1d(n_hidden), Tanh(),
  Linear(n_hidden, vocab_size),
])

parameters = model.parameters()
print(sum(p.nelement() for p in parameters)) # 総パラメータ数
for p in parameters:
  p.requires_grad = True

これにより、さらに簡素化されました:

# フォワードパス
  logits = model(Xb)
  loss = F.cross_entropy(logits, Yb) # 損失関数

# 損失を評価
  logits = model(x)
  loss = F.cross_entropy(logits, y)

# モデルからサンプリング
  # ニューラルネットをフォワードパス 
  logits = model(torch.tensor([context]))
  probs = F.softmax(logits, dim=1)

層状構造の実装#

現在のモデルのように、1 つのステップで情報をすべて 1 つの層に圧縮したくありません。WaveNet のように、次の文字を予測する際に 2 つの文字を統合して二文字表現を作り、それから 4 つの文字レベルの小さなブロックに合成し、このようなツリー状の階層構造で徐々に情報をネットワークに融合させたいと考えています。

Screenshot 2024-03-08 at 15.00.26

WaveNet の例では、この図は「Dilated causal convolution layer」(拡張因果畳み込み層)の可視化であり、具体的に何であるかは気にせず、私たちはその核心思想「Progressive fusion(漸進的融合)」を学べば良いです。

コンテキスト入力を増やし、この 8 つの入力文字をツリー構造で処理します。

# block_size = 3
# train 2.0677597522735596; val 2.1055991649627686
block_size = 8

単にコンテキストの長さを拡大するだけで性能が向上しました:

Screenshot 2024-03-08 at 20.49.15

私たちが何をしているのかを理解するために、各レイヤーを通過する際のテンソルの形状を観察します:

Screenshot 2024-03-08 at 21.02.13

4 つのランダムな数を入力すると、モデル内の形状は 4x8(block_size=8)になります。

  1. 最初の層(embedding)を通過すると、4x8x10 の出力が得られます。これは、私たちの埋め込みテーブルが各文字に対して学習する 10 次元ベクトルを持っていることを意味します。
  2. 2 番目の層(flatten)を通過すると、前述のように 4x80 に変わります。この層の効果は、8 つの文字の 10 次元埋め込みを 1 つの長い行に引き伸ばすことです。これは接続操作のようなものです。
  3. 3 番目の層(linear)は、この 80 を行列の乗算を通じて 200 のチャネル (channel) を作成します。

再度まとめると、Embedding 層が最終的に行う作業

この回答に非常に良い説明があります:
1. スパース行列を線形変換(ルックアップ)して密な行列に変換します。
2. この密な行列は、N 個の特徴を使用してすべての単語を表現します。密な行列の表現は、単語と特徴の関係係数ですが、実際には単語間の内在的な関係を多く含んでいます。
3. それらの間の重みパラメータは、埋め込み層で学習されたパラメータを使用して表現されます。神経ネットワークの逆伝播最適化の過程で、このパラメータも継続的に更新され最適化されます。

線形層は、フォワードパス中に入力 X を受け取り、それを重みと乗算し、オプションでバイアスを追加します:

def __init__(self, fan_in, fan_out, bias=True):
    self.weight = torch.randn((fan_in, fan_out)) / fan_in**0.5 # 注意: kaiming init
    self.bias = torch.zeros(fan_out) if bias else None

ここでの重みは 2 次元で、バイアスは 1 次元です。

入力出力の形状に基づいて、この線形層の内部は次のようになります:

(torch.randn(4, 80) @ torch.randn(80, 200) + torch.randn(200)).shape

出力は 4x200 で、最後に加えられるバイアスはここでブロードキャストの意味が発生します。

補足として、PyTorch の行列乗算演算子は非常に強力で、高次元テンソルを受け入れることができ、行列乗算は最後の次元でのみ機能し、他のすべての次元はバッチ処理次元(batch dimensions)として扱われます。

Screenshot 2024-03-09 at 16.53.25

これは、後で行うことに非常に役立ちます:並列のバッチ処理次元。私たちは一度に 80 個の数字を入力したくないので、最初の層で 2 つの統合された文字を持ちたいのです。つまり、20 個の数字の入力だけが必要です。以下のように:

# (1 2) (3 4) (5 6) (7 8)

(torch.randn(4, 4, 20) @ torch.randn(20, 200) + torch.randn(200)).shape

こうすることで、4 つのビグラムが得られ、ビグラムグループの各々は 10 次元ベクトルになります。

このような構造を実現するために、Python にはリスト内の偶数と奇数部分を取得する便利な方法があります:

Screenshot 2024-03-09 at 17.04.08

e = torch.randn(4, 8, 10)
torch.cat([e[:, ::2, :], e[:, 1::2, :]], dim=2)
# torch.Size([4, 4, 20])

これにより、偶数と奇数部分が明確に抽出され、4x4x10 の部分が結合されます。

Screenshot 2024-03-09 at 17.10.43

強力なview()も同等の作業を行うことができます。

次に、Flatten 層を完成させ、コンストラクタを作成し、出力の最後の次元で接続したい連続要素の数を取得します。基本的には n 個の連続要素を平坦化し、それらを最後の次元に配置することです。

class FlattenConsecutive:
  
  def __init__(self, n):
    self.n = n
    
  def __call__(self, x):
    B, T, C = x.shape
    x = x.view(B, T//self.n, C*self.n)
    if x.shape[1] == 1:
      x = x.squeeze(1)
    self.out = x
    return self.out
  
  def parameters(self):
    return []
  • B: バッチサイズ(批サイズ)、バッチ内に含まれるサンプルの数を表します。
  • T: タイムステップ(時間ステップ)、シーケンス内の要素の数、つまりシーケンスの長さを表します。
  • C: チャネルまたは特徴(通道または特徴)、各タイムステップ内のデータの特徴の数を表します。
  1. 入力テンソル: 入力xは三次元テンソルで、形状は(B, T, C)です。

  2. 平坦化操作: x.view(B, T//self.n, C*self.n)を呼び出すことで、このクラスは元のデータ内の連続したタイムステップを統合します。ここでself.nは統合するタイムステップの数を示します。操作の結果、タイム次元Tn倍減少し、特徴次元Cn倍増加します。新しい形状は(B, T//n, C*n)になり、各新しいタイムステップには元のn個のタイムステップの情報が含まれます。

  3. 単一のタイムステップ次元の削除: 統合後のタイムステップ長が 1 の場合、つまりx.shape[1] == 1であれば、x.squeeze(1)操作を通じてこの次元を削除します。これは、以前直面した二次元ベクトルの状況です。

code

修正後、中間各層の形状を確認します:

Screenshot 2024-03-09 at 18.09.42

バッチノルムでは、68 個のチャネルの平均と分散のみを維持し、32x4 次元ではないように、既存の BatchNorm1D の実装を変更します:

class BatchNorm1d:
  
  def __call__(self, x):
    # フォワードパスを計算
    if self.training:
      if x.ndim == 2:
        dim = 0
      elif x.ndim == 3:
        dim = (0,1) # torch.mean()はタプルを受け入れ、複数の次元のdimを指定できます
        
      xmean = x.mean(dim, keepdim=True) # バッチ平均
      xvar = x.var(dim, keepdim=True) # バッチ分散

現在、running_mean.shape は [1, 1, 68] です。

ニューラルネットワークの拡張#

上記の改善を完了し、ネットワークのサイズを増やすことで性能をさらに向上させます。

n_embd = 24    # 埋め込みベクトルの次元
n_hidden = 128 # MLP隠れ層のニューロン数 

現在のパラメータ数は 76579 に達し、性能も 2.0 の壁を突破しました:

Screenshot 2024-03-09 at 21.45.20

これまでのところ、神経ネットワークの訓練に必要な時間は大幅に増加しましたが、性能が向上したにもかかわらず、学習率などのハイパーパラメータの正しい設定には困惑しています。ただ訓練の損失を見つめながら、デバッグと修正を繰り返しています。

畳み込み#

本節では、WaveNet の主要なアーキテクチャを実装しましたが、その中で特定のフォワードパス、つまりより複雑な線形層:ゲート付き線形層 (gated linear layer)、残差接続 (Residual connection)、およびスキップ接続 (Skip connection) は実装していません。

Screenshot 2024-03-09 at 21.52.42

ここでは、私たちが実装したツリー状構造と WaveNet 論文で使用されている畳み込みニューラルネットワークとの関連を簡単に理解します。

基本的に、ここで畳み込み (Convolution) を使用するのは効率を高めるためです。畳み込みは、入力シーケンス上でモデルをスライドさせることを可能にし、この部分の for ループ(畳み込みカーネルのスライドと計算)が CUDA カーネル内で完了します。

Screenshot 2024-03-08 at 15.00.26

私たちは、図に示された黒い構造を単一で実装し、出力を得ただけですが、畳み込みはこの黒い構造を通じて入力シーケンスに適用し、線形フィルターのようにすべてのオレンジ色の出力を同時に計算することを可能にします。

効率が向上する理由は次のとおりです:

  1. for ループが CUDA コア内で完了します。
  2. 変数を再利用します。たとえば、2 番目の層の 1 つの白点は、3 番目の層の白点の左子ノードであり、別の白点の右子ノードでもあります。このノードとその値は 2 回使用されます。

まとめ#

本節を経て、torch.nn モジュールが解放されました。今後はモデルの実装をそれを使用するように変更します。

本節の作業を振り返ると、多くの時間が各レイヤーの形状を正しくするために費やされました。そのため、Andrej は常に Jupyter Notebook で形状のデバッグを行い、満足した後にそれを vscode にコピーしていました。

読み込み中...
文章は、創作者によって署名され、ブロックチェーンに安全に保存されています。