混合密度ネットワークをPythonコードを交えて紹介

はじめに

PRML(Pattern Recognition and Machine Learning)でニューラルネットワークの、混合密度ネットワークについて学んだ内容をまとめて、実際のデータを使って学習しました。主に5章の内容です。学習アルゴリズムの全体像を示すことを意識したので、導出は端折っている部分が多いです。詳しい解の導出は、PRMLの本文や演習問題にあります。

混合密度ネットワークのアルゴリズム

今回は目的変数が1次元の回帰問題を考えます。目標は、入力に対する出力の条件付き分布$p(t | \mathbf{x})$をモデル化することです。まず、2層ニューラルネットワーク（重みのある層が2層）で解くことを考えます。活性化関数にはシグモイド関数を用います。
$x_i(i = 1, \cdots, D)$を入力とする基本的な2層ニューラルネットワークは以下のようになります。

$$\begin{aligned} a_j &= \sum^D_{i=1}w_{ji}^{(1)} x_i + w_{j0}^{(1)} \\ z_j &= \frac{1}{1 + exp(-a_j)} \\ y &= \sum^M_{j=1} w_{j}^{(2)} z_j + w_{0}^{(2)} \end{aligned}$$

このモデルでは、$y$を平均としたガウス分布を条件付き分布としますが、次のようなデータ（横軸が入力、縦軸が出力）をうまく表現することができません。入力に対して出力が多峰性になっているからです。なお、このデータは後ほど使います。

このような時にも、適切な条件付き分布を見つけたい時、次のように先ほどのモデルを修正します。

$$\begin{aligned} a_j &= \sum^D_{i=1}w_{ji}^{(1)} x_i + w_{j0}^{(1)} \\ z_j &= \frac{1}{1 + exp(-a_j)} \\ a_k &= \sum^M_{j=1} w_{kj}^{(2)} z_j + w_{k0}^{(2)} \\ \end{aligned}$$

ここで、$a_k$を$a_k^{\pi}, a_k^{\mu}, a_k^{\sigma}$のいずれかに対応させます。

$$\begin{aligned} \pi_k &= \frac{exp(a_k^{\pi})}{\sum^K_{l=1}exp(a_l^{\pi})} \\ \mu_k &= a_k^{\mu} \\ \sigma_k &= exp({a_k^{\sigma}}) \\ y &= \sum^K_{k=1} \pi_k \mathcal{N}(y | \mu_k, \sigma_k^2) \end{aligned}$$

このようにすると、平均を$\mu_k$、分散を$\sigma_k$とした$K$個のガウス分布を混合係数$\pi_k$で足し合わせた条件付き分布によって、多峰性のモデルを表現することができます。
$p(t | \mathbf{x}) = y$となるので、誤差関数$E(\mathbf{w})$を負の条件付き分布（事後確率）とすると、以下のようになります。

$$\begin{aligned} E(\mathbf{w}) = \sum^N_{n=1}E_n = -\sum^N_{n=1}ln\left\{ \sum^K_{k=1}\pi_k\mathcal{N}(t_n | \mu_k, \sigma_k^2) \right\} \end{aligned}$$

導出は省略しますが、$a_k^{\pi}, a_k^{\mu}, a_k^{\sigma}$に関しての微分は以下のようになります。あとは通常通り、誤差逆伝播法によってそれぞれの重みパラメータに関する微分を計算できます。

$$\begin{aligned} \frac{\partial E_n}{\partial a_k^{\pi}} &= \pi_k - \gamma_{nk} \\ \frac{\partial E_n}{\partial a_k^{\mu}} &= \gamma_{nk} \left\{ \frac{\mu_k - t_n}{\sigma_k^2} \right\} \\ \frac{\partial E_n}{\partial a_k^{\mu}} &= \gamma_{nk} \left\{ 1 - \frac{(t_n-\mu_k)^2}{\sigma_k^2} \right\} \\ where \quad \gamma_{nk} &= \frac{\pi_k\mathcal{N}(t_n | \mu_k, \sigma_k^2)}{\sum^K_{k=1}\pi_l\mathcal{N}(t_n | \mu_l, \sigma_l^2)} \end{aligned}$$

混合密度ネットワークの実装

実装は、書籍『ゼロから作る Deep Learning』(O'Reilly Japan, 2016)のコードを一部改変して行いました。githubリポジトリは以下です。

https://github.com/oreilly-japan/deep-learning-from-scratch

import文

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import norm
import seaborn as sns

np.random.seed(28)

レイヤ

全結合レイヤです。ほとんど変更していません。先ほどの図で言うと、$x$と$a_j$の間と、$z$と$a_k$の間にあります。

class Affine:
    def __init__(self, W, b):
        self.W =W
        self.b = b
        self.x = None
        # 重み・バイアスパラメータの微分
        self.dW = None
        self.db = None

    def forward(self, x):
        self.x = x
        out = np.dot(self.x, self.W) + self.b

        return out

    def backward(self, dout):
        dx = np.dot(dout, self.W.T)
        self.dW = np.dot(self.x.T, dout)
        self.db = np.sum(dout, axis=0)
        
        return dx

Sigmoidレイヤです。ほとんど変更していません。先ほどの図で言うと、$a_j$と$z$の間にあります。

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

class Sigmoid:
    def __init__(self):
        self.out = None

    def forward(self, x):
        out = sigmoid(x)
        self.out = out
        return out

    def backward(self, dout):
        dx = dout * (1.0 - self.out) * self.out

        return dx

MixtureDensityWithLossレイヤは、先ほどの図で言うと、$a_k$以降の部分です。このレイヤは新たに実装しました。

def softmax(x):
    x = x - np.max(x, axis=-1, keepdims=True)   # オーバーフロー対策
    return np.exp(x) / np.sum(np.exp(x), axis=-1, keepdims=True)


class MixtureDensityWithLoss:
    def __init__(self, num_components):
        self.loss = None
        self.pi = None
        self.mu = None
        self.sigma = None
        self.K = num_components

    def forward(self, x, t):
        self.t = t
        self.pi = softmax(x[:, :self.K])
        self.mu = x[:, self.K:2*self.K]
        self.sigma = np.exp(x[:, 2*self.K:])
        self.calc_loss()

        return self.loss

    def backward(self, dout=1):
        batch_size = self.t.shape[0]
        dx = np.empty((batch_size, 3*self.K))
        gamma = self.calc_gamma()

        for n in range(batch_size):
            for k in range(self.K):
                # piの微分
                dx[n][k] = self.pi[n][k] - gamma[n][k]
                # muの微分
                dx[n][k + self.K] = gamma[n][k] * (self.mu[n][k] - self.t[n]) / self.sigma[n][k]
                # sigmaの微分
                dx[n][k + 2*self.K] = gamma[n][k] * (1 - (self.t[n] - self.mu[n][k])**2 / self.sigma[n][k]**2)

        return dx / batch_size

    def calc_gamma(self):
        batch_size = self.t.shape[0]

        numerator = np.empty((batch_size, self.K))
        for n in range(batch_size):
            for k in range(self.K):
                numerator[n][k] = self.pi[n][k] * norm.pdf(self.t[n], self.mu[n][k], self.sigma[n][k])

        return numerator / np.sum(numerator, axis=1).reshape((batch_size, 1))

    def calc_loss(self):
        batch_size = self.t.shape[0]

        self.loss = 0
        for n in range(batch_size):
            tmp_loss = 0
            for k in range(self.K):
                tmp_loss += self.pi[n][k] * norm.pdf(self.t[n], self.mu[n][k], self.sigma[n][k])
            self.loss -= np.log(tmp_loss)
        self.loss /= batch_size

ネットワーク

以下のような2層の混合密度ニューラルネットワークになっています。

from collections import OrderedDict

class TwoLayerNet:

    def __init__(self, input_size, hidden_size, output_size, K, weight_init_std = 0.01):
        # 重みの初期化
        self.params = {}
        self.params['W1'] = weight_init_std * np.random.randn(input_size, hidden_size)
        self.params['b1'] = np.zeros(hidden_size)
        self.params['W2'] = weight_init_std * np.random.randn(hidden_size, output_size)
        self.params['b2'] = np.zeros(output_size)

        # レイヤの生成
        self.layers = OrderedDict()
        self.layers['Affine1'] = Affine(self.params['W1'], self.params['b1'])
        self.layers['Sigmoid'] = Sigmoid()
        self.layers['Affine2'] = Affine(self.params['W2'], self.params['b2'])

        self.lastLayer = MixtureDensityWithLoss(K)
        
    def predict(self, x):
        for layer in self.layers.values():
            x = layer.forward(x)
        
        return x
        
    # x:入力データ, t:教師データ
    def loss(self, x, t):
        y = self.predict(x)
        return self.lastLayer.forward(y, t)
        
    def gradient(self, x, t):
        # forward
        self.loss(x, t)

        # backward
        dout = 1
        dout = self.lastLayer.backward()
        
        layers = list(self.layers.values())
        layers.reverse()
        for layer in layers:
            dout = layer.backward(dout)

        # 設定
        grads = {}
        grads['W1'], grads['b1'] = self.layers['Affine1'].dW, self.layers['Affine1'].db
        grads['W2'], grads['b2'] = self.layers['Affine2'].dW, self.layers['Affine2'].db

        return grads

データ

scikit-learnのアヤメのデータセットを用いて、学習を行います。sepal_width（がく片の幅）を入力、petal_width（花びらの幅）を出力とするようなモデルを作ります。つまり、条件付き分布$p(petal\_width | sepal\_width)$を求めます。訓練データは100個、テストデータは50個です。

iris_df = sns.load_dataset('iris').sample(frac=1)

t_vanilla_df = iris_df.loc[:, 'petal_width']
t_df = (t_vanilla_df - t_vanilla_df.mean()) / t_vanilla_df.std()
x_vanilla_df = iris_df.loc[:, 'sepal_width']
x_df = (x_vanilla_df - x_vanilla_df.mean()) / x_vanilla_df.std()

N = 100 # 訓練データの数
x_train = x_df[:N].to_numpy(copy=True).reshape((-1, 1))
t_train = t_df[:N].to_numpy(copy=True)
x_test = x_df[N:].to_numpy(copy=True).reshape((-1, 1))
t_test = t_df[N:].to_numpy(copy=True)

このデータセットには複数の品種のアヤメのデータが入っているので、以下のように二峰性の分布になっています。なので、今回は2つのガウス分布を混合させることにします。

作図コード

plt.scatter(x_train, t_train, label='train data')
plt.scatter(x_test, t_test, label='test data')
plt.xlabel('sepal_width')
plt.ylabel('petal_width')
plt.legend()

学習

あまり優れた方法ではないですが、簡単のために確率的勾配降下法を用いて、重りパラメータを更新します。また、種々の正則化手法はここでは使いません。

K = 2 # 混合させるガウス分布の数

network = TwoLayerNet(input_size=1, hidden_size=4, output_size=K*3, K=K)

iters_num = 3000
train_size = x_train.shape[0]
batch_size = 10
learning_rate = 0.1

iter_per_epoch = max(train_size / batch_size, 1)

for i in range(iters_num):
    batch_mask = np.random.choice(train_size, batch_size)
    x_batch = x_train[batch_mask]
    t_batch = t_train[batch_mask]
    
    # 勾配
    grad = network.gradient(x_batch, t_batch)
    
    # 更新
    for key in ('W1', 'b1', 'W2', 'b2'):
        network.params[key] -= learning_rate * grad[key]
    
    if (i+1) % (iter_per_epoch*30) == 0:
        train_acc = network.loss(x_train, t_train)
        test_acc = network.loss(x_test, t_test)
        print("iteration, train_acc, test_acc | {0:4d}, {1:.5f}, {2:.5f}".format(i+1, train_acc, test_acc))

iteration, train_acc, test_acc |  300, 1.39369, 1.44268
iteration, train_acc, test_acc |  600, 1.30168, 1.35834
iteration, train_acc, test_acc |  900, 1.16585, 1.20975
iteration, train_acc, test_acc | 1200, 0.88597, 0.94173
iteration, train_acc, test_acc | 1500, 0.68086, 0.67129
iteration, train_acc, test_acc | 1800, 0.67816, 0.66658
iteration, train_acc, test_acc | 2100, 0.64981, 0.67003
iteration, train_acc, test_acc | 2400, 0.65283, 0.64939
iteration, train_acc, test_acc | 2700, 0.63715, 0.64105
iteration, train_acc, test_acc | 3000, 0.66861, 0.64719

結果

学習した重りパラメータで条件付き分布を作図しました。

作図コード

x1_line = np.linspace(-3, 3, 61).reshape(-1, 1)
y = network.predict(x1_line)
pi = softmax(y[:, :K])
mu = y[:, K:2*K]
sigma = np.exp(y[:, 2*K:])

x2_line = np.linspace(-3, 3, 61).reshape(-1, 1)
x1_grid, x2_grid = np.meshgrid(x1_line, x2_line)
ps = np.zeros((61, 61))

for x1 in range(61):
    for x2 in range(61):
        for k in range(2):
            ps[x2][x1] += pi[x1][k] * norm.pdf(x2_line[x2], mu[x1][k], sigma[x1][k])

plt.scatter(x_train, t_train, label='train data')
plt.scatter(x_test, t_test, label='test data')
plt.xlabel('sepal_width')
plt.ylabel('petal_width')
plt.contour(x1_grid, x2_grid, ps, levels=[0.025, 0.1, 0.4, 1.6])
plt.colorbar()
plt.legend()

しっかりとデータの多いところの確率密度が大きくなっていることが見てとれます。
次に、テストデータの入力を用いて得られた条件付き分布から、ランダムにデータをサンプリングして作図しました。左がテストデータ、右が新たにサンプリングしたデータです。

作図コード

y = network.predict(x_test)
pi = softmax(y[:, :K])
mu = y[:, 2:2*K]
sigma = np.exp(y[:, 2*K:])

ts = np.empty(len(x_test))

for i in range(len(x_test)):
    if np.random.random() < pi[i][0]:
        ts[i] = np.random.normal(mu[i][0], sigma[i][0])
    else:
        ts[i] = np.random.normal(mu[i][1], sigma[i][1])

fig = plt.figure(figsize=(4 * 2 * np.sqrt(2), 4))

ax1 = fig.add_subplot(1, 2, 1)
ax2 = fig.add_subplot(1, 2, 2)

ax1.scatter(x_test, t_test, label='test data')
ax2.scatter(x_test, ts, label='sampling data')

y_min = min(np.min(t_test), np.min(ts))
y_max = max(np.max(t_test), np.max(ts))
ax1.set_ylim(y_min - 0.2, y_max + 0.2)
ax2.set_ylim(y_min - 0.2, y_max + 0.2)

ax1.set_xlabel('sepal_width')
ax1.set_ylabel('petal_width')
ax2.set_xlabel('sepal_width')
ax2.set_ylabel('petal_width')
ax2.legend()
ax1.legend()

plt.show()

似たような図になっていることから、うまく二峰性のデータを表現するモデルを学習できたことがわかります。