💨

オートエンコーダーとトランスフォーマーを組み合わせてGenie 2。

2024/12/05に公開

https://deepmind.google/discover/blog/genie-2-a-large-scale-foundation-world-model/

タイトル: 未来を紡ぐプログラム

東京・渋谷のカフェ。主人公である若手プログラマー、拓海はノートパソコンを前にコーヒーを一口飲みながら考え込んでいた。彼が取り組んでいるのは、新しい映像生成アルゴリズムの開発だった。

「オートエンコーダーとトランスフォーマーを組み合わせて、映像に連続性を持たせる…。言葉では簡単だけど、実現するのは難しいな。」

彼の目の前には、先ほど動かしたプログラムのエラーメッセージが並んでいる。実験は失敗続きだったが、それでも彼は手を止めなかった。この技術が完成すれば、単なる画像生成を超え、ストーリー性のある動画がAIによって生み出されるという夢があった。

その夜、渋谷の高層ビル群の中で灯るオフィスの一角に彼は座っていた。机には無数のデータシートと数学的なメモが散らばり、パソコンの画面にはオートエンコーダーで圧縮された潜在ベクトルと、トランスフォーマーの生成結果が映し出されている。

「よし、今回は連続性を意識したフレーム生成にフォーカスしよう。」
拓海は再びキーボードを叩き始めた。トランスフォーマーを改良し、過去のフレーム情報を考慮するようモデルを調整する。

数時間後、生成された映像を見た瞬間、拓海の目が輝いた。ぼんやりとしたモノクロの世界が、一つの連続した物語として動き出していたのだ。最初は抽象的な模様だったが、徐々に形が整い、流れるような映像となっていった。

「これだ…!これが俺が目指していた映像の連続性だ!」

実行結果。

オートエンコーダーとデコーダーを通じて、画像データを潜在ベクトル空間に変換することで、各画像(フレーム)の特徴が圧縮され、その違いが潜在ベクトルの差異として明確に表現されます。この潜在空間の変化をうまくモデル化することで、連続的なデータを生成できます。

トランスフォーマーモデルの役割は、この潜在ベクトル間の時間的または連続的な関係性を学習することです。特に以下の点が重要です:

潜在ベクトルの特徴抽出:

潜在ベクトルは元の画像データよりも次元が低いため、計算効率が高く、モデルが学習しやすい特徴が含まれています。

連続性の学習:

トランスフォーマーは、自己注意(Self-Attention)メカニズムを使い、時間的に連続するフレーム間の関係性や変化パターンを捉えます。

フレームの生成:

トランスフォーマーの出力は、新しい潜在ベクトルを生成する形でモデル化され、それをデコーダーに通して次のフレーム(画像)を復元します。
これにより、単なる画像生成ではなく、連続性を考慮した映像的な生成が可能になります。

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.datasets import mnist
import numpy as np
import matplotlib.pyplot as plt

# --- 1. データ準備 ---
(x_train, _), (_, _) = mnist.load_data()
x_train = x_train.astype("float32") / 255.0
x_train = np.expand_dims(x_train, axis=-1)  # (28, 28, 1)に変換

# --- 2. オートエンコーダーの定義 ---
latent_dim = 64  # 潜在空間の次元数

# エンコーダー
encoder = tf.keras.Sequential([
    layers.Input(shape=(28, 28, 1)),
    layers.Conv2D(32, kernel_size=3, activation="relu", strides=2, padding="same"),
    layers.Conv2D(64, kernel_size=3, activation="relu", strides=2, padding="same"),
    layers.Flatten(),
    layers.Dense(latent_dim),
])

# デコーダー
decoder = tf.keras.Sequential([
    layers.Input(shape=(latent_dim,)),
    layers.Dense(7 * 7 * 64, activation="relu"),
    layers.Reshape((7, 7, 64)),
    layers.Conv2DTranspose(64, kernel_size=3, activation="relu", strides=2, padding="same"),
    layers.Conv2DTranspose(32, kernel_size=3, activation="relu", strides=2, padding="same"),
    layers.Conv2DTranspose(1, kernel_size=3, activation="sigmoid", padding="same"),
])

# オートエンコーダーモデル
input_tensor = layers.Input(shape=(28, 28, 1))
encoded = encoder(input_tensor)
decoded = decoder(encoded)
autoencoder = tf.keras.Model(inputs=input_tensor, outputs=decoded)
autoencoder.compile(optimizer="adam", loss="mse")
autoencoder.fit(x_train, x_train, epochs=5, batch_size=128)

# --- 3. トランスフォーマーによるオートリグレッシブ自己回帰生成 ---
from tensorflow.keras.layers import MultiHeadAttention, Dense, LayerNormalization

class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super().__init__()
        self.att = MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.ffn = tf.keras.Sequential([
            Dense(ff_dim, activation="relu"),
            Dense(embed_dim),
        ])
        self.layernorm1 = LayerNormalization(epsilon=1e-6)
        self.layernorm2 = LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)

    def call(self, inputs, training=False):
        attn_output = self.att(inputs, inputs)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(out1 + ffn_output)

# Latent inputs preparation
latent_inputs = encoder.predict(x_train[:10])  # 10枚の画像を潜在表現に変換
latent_inputs = latent_inputs[:, np.newaxis, :]  # (10, 1, 64)

# Transformer-based autoregressive generation
transformer_block = TransformerBlock(embed_dim=latent_dim, num_heads=4, ff_dim=128)
generated_frames = []
current_frame = latent_inputs

for _ in range(10):  # Generate 10 frames
    next_frame = transformer_block(current_frame, training=False)
    generated_frames.append(next_frame)
    current_frame = next_frame


# --- 4. 再構築と可視化 ---
# Decode the generated frames
generated_frames = np.array(generated_frames)
print("Shape of generated_frames before squeeze:", generated_frames.shape)

# Squeeze only if the second axis has size 1
if generated_frames.shape[1] == 1:
    generated_frames = np.squeeze(generated_frames, axis=1)
else:
    generated_frames = generated_frames.reshape(-1, latent_dim)  # Reshape to match decoder input

# Decode the generated frames
reconstructed_images = decoder.predict(generated_frames)

# Visualization
plt.figure(figsize=(10, 2))
for i, img in enumerate(reconstructed_images[:10]):
    plt.subplot(1, 10, i + 1)
    plt.imshow(img.squeeze(), cmap="gray")
    plt.axis("off")
plt.show()



このコードは、MNISTデータセットを使ってオートエンコーダー(Autoencoder)を構築し、トランスフォーマー(Transformer)を使用して潜在空間(latent space)からオートリグレッシブ(autoregressive)な生成を行うプロセスを示しています。

以下は、各ステップの計算プロセスの解説です:

  1. データ準備
    MNISTデータセットのロードと前処理:
    MNISTデータ(28x28の手書き数字画像)を0〜1の範囲に正規化します。
    np.expand_dims を使用してチャンネル次元を追加し、データを (28, 28, 1) の形に整形します。
  2. オートエンコーダーの構築
    エンコーダー:

入力画像(28x28x1)を受け取り、畳み込み層(Conv2D)で特徴を抽出します。
ダウンサンプリング(strides=2)で画像のサイズを縮小しながら特徴マップを学習します。
Flatten層で1次元ベクトルに変換し、全結合層(Dense)で潜在ベクトル(latent_dim=64)を生成します。
デコーダー:

潜在ベクトルを入力として受け取り、全結合層で形状を (7, 7, 64) に拡大します。
転置畳み込み層(Conv2DTranspose)を使って元の画像サイズ(28x28x1)に復元します。
出力層のシグモイド(sigmoid)関数でピクセル値を [0, 1] の範囲に制限します。
モデル学習:

入力画像と再構築された画像のMSE(平均二乗誤差)を最小化するように学習します。

  1. トランスフォーマーによるオートリグレッシブ生成
    TransformerBlockの設計:

MultiHeadAttention: 入力間の関係性を学習します。
Feed-Forward Network (FFN): 特徴を強化します。
Layer Normalization: 各層の出力を正規化して安定化します。
Residual Connection: 各層の出力に元の入力を加算して勾配消失問題を軽減します。
潜在空間の初期化:

MNIST画像をエンコーダーで潜在空間にマッピングし、形状を (10, 1, 64) に整形します。
オートリグレッシブ生成:

潜在空間をトランスフォーマーに入力し、次のフレーム(潜在ベクトル)を生成します。
新しいフレームを現在の入力として更新し、10フレームを生成します。

  1. 再構築と可視化
    生成されたフレームの整形:
    トランスフォーマーの出力(generated_frames)の形状を確認し、必要に応じて軸を削除(squeeze)または再整形(reshape)します。

デコーダーによる再構築:
各フレームをデコーダーに通して画像を再構築します。

可視化:
再構築された画像をプロットし、生成結果を視覚的に確認します。

計算の要点
オートエンコーダーは、画像データの高次元特徴を潜在空間に圧縮し、情報を効率的に表現します。
トランスフォーマーは、生成された潜在ベクトル間の関係を学習して、次の潜在ベクトルを生成します。
最後に、潜在ベクトルをデコーダーで再構築し、生成された画像を確認します。
この設計により、エンコーダーが抽出した潜在空間をもとに新しい画像を生成できるモデルが完成します。

Discussion