次世代生成モデル、Diffusion Modelの紹介

こんにちは。次世代システム研究室のA.Z.です。

今回は最近話題になっている、確率モデル＋深層ニューラルネットワーク（統合モデル）の一つのdiffusionモデルについて紹介したいと思います。

はじめに

• 生成プロセスは学習した分布から潜在変数をsamplingし、ニューラルネットワーク上に変化させ、という生成プロセス
• モデルのloss関数の計算方法はにいています（Variational Inference, Reparameterization Tricksなど)。

Diffusion modelについて

引用：https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

引用: https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

Forward Process

def sinusoidal_embedding(x):
    embedding_min_frequency = 1.0
    frequencies = tf.exp(
        tf.linspace(
            tf.math.log(embedding_min_frequency),
            tf.math.log(embedding_max_frequency),
            embedding_dims // 2,
            )
    )
    angular_speeds = 2.0 * math.pi * frequencies
    embeddings = tf.concat(
        [tf.sin(angular_speeds * x), tf.cos(angular_speeds * x)], axis=3
    )
    return embeddings

def get_network(image_size, widths, block_depth):
    noisy_images = keras.Input(shape=(image_size, image_size, 3))
    noise_variances = keras.Input(shape=(1, 1, 1))

    e = layers.Lambda(sinusoidal_embedding)(noise_variances)
    .....

#sample uniform random diffusion times
diffusion_times = tf.random.uniform(
    shape=(batch_size, 1, 1, 1), minval=0.0, maxval=1.0
)
noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
# mix the images with noises accordingly
noisy_images = signal_rates * images + noise_rates * noises

Reverse Process

with tf.GradientTape() as tape:
            # train the network to separate noisy images to their components
            pred_noises, pred_images = self.denoise(
                noisy_images, noise_rates, signal_rates, training=True
            )

            noise_loss = self.loss(noises, pred_noises)  # used for training
            image_loss = self.loss(images, pred_images)  # only used as metric

        gradients = tape.gradient(noise_loss, self.network.trainable_weights)
        self.optimizer.apply_gradients(zip(gradients, self.network.trainable_weights))

        self.noise_loss_tracker.update_state(noise_loss)
        self.image_loss_tracker.update_state(image_loss)

def denoise(self, noisy_images, noise_rates, signal_rates, training):
       # the exponential moving average weights are used at evaluation
       if training:
           network = self.network
       else:
           network = self.ema_network

       # predict noise component and calculate the image component using it
       pred_noises = network([noisy_images, noise_rates**2], training=training)
       pred_images = (noisy_images - noise_rates * pred_noises) / signal_rates

       return pred_noises, pred_images

画像生成の例

def diffusion_schedule(diffusion_times, min_signal_rate, max_signal_rate):
    start_angle = tf.acos(max_signal_rate)
    end_angle = tf.acos(min_signal_rate)

    diffusion_angles = start_angle + diffusion_times * (end_angle - start_angle)

    signal_rates = tf.cos(diffusion_angles)
    noise_rates = tf.sin(diffusion_angles)

    return noise_rates, signal_rates


def generate_images(diffusion_steps, stochasticity, min_signal_rate, max_signal_rate):
    step_size = 1.0 / diffusion_steps
    initial_noise = tf.random.normal(shape=(num_images, image_size, image_size, 3))

    # reverse diffusion
    noisy_images = initial_noise
    for step in range(diffusion_steps):
        diffusion_times = tf.ones((num_images, 1, 1, 1)) - step * step_size
        next_diffusion_times = diffusion_times - step_size

        noise_rates, signal_rates = diffusion_schedule(diffusion_times, min_signal_rate, max_signal_rate)
        next_noise_rates, next_signal_rates = diffusion_schedule(next_diffusion_times, min_signal_rate, max_signal_rate)

        sample_noises = tf.random.normal(shape=(num_images, image_size, image_size, 3))
        sample_noise_rates = stochasticity * (1.0 - (signal_rates / next_signal_rates) ** 2) ** 0.5 * (
                next_noise_rates / noise_rates)

        pred_noises, pred_images = model([noisy_images, noise_rates, signal_rates])
        noisy_images = (
                next_signal_rates * pred_images
                + (next_noise_rates ** 2 - sample_noise_rates ** 2) ** 0.5 * pred_noises
                + sample_noise_rates * sample_noises
        )

    # denormalize
    data_mean = tf.constant([[[[0.4705, 0.3943, 0.3033]]]])
    data_std_dev = tf.constant([[[[0.2892, 0.2364, 0.2680]]]])
    generated_images = data_mean + pred_images * data_std_dev
    generated_images = tf.clip_by_value(generated_images, 0.0, 1.0)

    # make grid
    generated_images = tf.image.resize(generated_images, (plot_image_size, plot_image_size), method="nearest")
    generated_images = tf.reshape(generated_images, (num_rows, num_cols, plot_image_size, plot_image_size, 3))
    generated_images = tf.transpose(generated_images, (0, 2, 1, 3, 4))
    generated_images = tf.reshape(generated_images, (num_rows * plot_image_size, num_cols * plot_image_size, 3))
    return generated_images.numpy()

<code>initial_noise = tf.random.normal(shape=(num_images, image_size, image_size, <span class="hljs-number">3</span>))</code>

まとめ

最後に