Parametric UMAP とは
高次元のデータ列の可視化手法として、低次元埋め込みアルゴリズムである t-SNE や UMAP よく用いられます。これらのアルゴリズムはデータ点どうしの距離関係を2次元などの低次元でも維持するようにすることで可視化を実現しています。
特にUMAPは巧みなアルゴリズムと実装により高速な動作を実現しており、様々な分野で利用されています。このUMAPに、2021年の1月に 埋め込み部分にニューラルネットワークを用いた亜種として Parametric UMAP が登場しました。
ParametricUMAP uses a neural network to learn a UMAP embedding. This allows for a number of significant advantages. 2/14 pic.twitter.com/oDCcpW5ORb
— Leland McInnes (@leland_mcinnes) January 12, 2021
何が新しくなったのか?
UMAP アルゴリズムは大雑把に次の2ステップから成ります。
- データ点どうしの距離を元に $k$-近傍グラフを作成。密集している点同士が連結されているようなグラフ構造を構築
- ステップ1 で構築した$k$-近傍グラフを低次元空間(2次元など)に埋め込み
Parametric UMAP では、このステップ2にニューラルネットワークが用いられるようになりました。これにより、データの情報をニューラルネットのパラメータとして保持できるようになったので、次のようなことが可能になりました。
- 増分データに対する埋め込みが瞬時にできるようになった。
- 半教師あり学習時の品質が向上した。
- 低次元埋め込みの逆写像も学習可能に。つまり、二次元の埋め込みから、高次元データを復元(推定)する関数を学習可能になった。
なお、UMAPの詳細な仕組みをサクッと理解したい方にはこの記事が簡潔です: t-SNEより強いUMAPを(工学的に)理解したい
使い方
Parametric UMAP の主な機能は UMAPリポジトリ (v0.5+) の umap.parametric_umap
で利用できます。最新版のインストールは
git clone https://github.com/lmcinnes/umap.git
cd umap
python setup.py develop
従来ながらの埋め込み
従来通りの埋め込みアルゴリズム的な使い方をする場合は、 umap.umap.UMAP
と同様のメソッドで埋め込みが可能です。
MNIST データセットに対する使用例は次の通り。
from umap.parametric_umap import ParametricUMAP
# データセットの準備
from tensorflow.keras.datasets import mnist
(train_images, Y_train), (test_images, Y_test) = mnist.load_data()
train_images = train_images.reshape((train_images.shape[0], -1))/255.
test_images = test_images.reshape((test_images.shape[0], -1))/255.
# 埋め込み実行
embedder = ParametricUMAP(verbose=True)
embedding = embedder.fit_transform(train_images)
matplotlibで可視化する場合は次のように可視化できます。
import matplotlib.pyplot as plt
fig, ax = plt.subplots( figsize=(8, 8))
sc = ax.scatter(
embedding[:, 0],
embedding[:, 1],
c=Y_train.astype(int),
cmap="tab10",
s=0.1,
alpha=0.5,
rasterized=True,
)
ax.axis('equal')
ax.set_title("Parametric UMAP embedding", fontsize=20)
plt.colorbar(sc, ax=ax)
可視化はこんな感じに。
学習に使うニューラルネットを独自で定義する場合
デフォルトでは、3層100ニューロンの全結合ネットワークが用いられます。
tf.Keras.Sequential
によりネットワークを定義し、ParametricUMAP
の生成時にencoder
引数として渡すことにより、畳み込みニューラルネットワークなどの独自定義したネットワークを用いることができます。
上記の MNIST データに対する畳込みニューラルネットで encoder
を定義する例は以下の通りです。
import tensorflow as tf
dims = (28, 28, 1)
n_components = 2 # 出力次元
encoder = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=dims),
tf.keras.layers.Conv2D(
filters=32, kernel_size=3, strides=(2, 2), activation="relu", padding="same"
),
tf.keras.layers.Conv2D(
filters=64, kernel_size=3, strides=(2, 2), activation="relu", padding="same"
),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(units=256, activation="relu"),
tf.keras.layers.Dense(units=256, activation="relu"),
tf.keras.layers.Dense(units=n_components),
])
# encoder と入力次元をモデルに渡す
embedder = ParametricUMAP(encoder=encoder, dims=dims)
embedding = embedder.fit_transform(train_images)
上記の例で雰囲気は掴めると思うのですが、要点は以下の通りです。
-
encoder
のニューラルネットワークは、入力がデータの次元、出力が埋め込み次元になるようにします。
MNISTの例では、入力層はinput_shape=(28, 28, 1)
で、出力層は2次元になるように設計します。 - 埋め込みを学習する際には各学習データを平滑化した1次元のベクトルにして
fit_transform
関数に渡す必要があります。MNISTの場合だと、train_images
の 形状は(60000, 784)
です。
モデルの保存と読み込み
内部にKerasのモジュールを用いているため、build in function で保存や読み込みをする必要があります。
保存
embedder.save('<path to save>')
読み込み
from umap.parametric_umap import load_ParametricUMAP
embedder = load_ParametricUMAP('<path to save>')
増分データに対する埋め込み
多くの距離学習ベースの低次元可視化アルゴリズムでは、マップ構築に使用するデータ間の近傍グラフの計算がボトルネックとなっていました。UMAPもその例外ではなかったのですが、Parametric UMAPでは写像をニューラルネットで保持しているため、増分データの埋め込みを計算する際は順方向にネットワークにデータを流すだけで済みます。
埋め込みを学習せずに写像のみを計算する場合には transform
メソッドを使います。
例として、test_images
を増分データと見立てて新たな埋め込みをする場合次のようになります。
additional_embedding = embedder.transform(test_images)
この際、ニューラルネットによって入力サンプルに対する予測が行われているので、下記のコードと基本的に等価です。
additional_embedding = embedder.encoder.predict(test_images)
ハードウェアにも依りますが、埋め込み学習時には通常のUMAPより時間がかかるのに対し、増分データの推測は瞬時に行われます。事前に学習をしておき、リアルタイムで増加するデータを可視化するシステムの実装などに応用ができそうです。
半教師あり学習をする場合
従来のUMAPは教師なしの低次元可視化アルゴリズムとして用いる他に、ラベル付きのデータに対してラベルを参考にすることで、より高品質な埋め込みを作成することが可能でした。
Parametric UMAP ではUMAPの損失関数と、ラベル誤差を同時に学習することで(半)教師あり学習を実現しており、従来のUMAPより品質の高い埋め込みが可能になったようです。
ラベルデータを利用する場合は、sklearn のインターフェイスを踏襲し、ParametricUMAP クラス作成時にy
をラベルデータの引数として渡します。半教師あり学習をする場合、すなわちラベルの無いデータがある場合は、そのラベルを -1
としておきます。
以下は、Fashon-MNIST データセットに対して ParametricUMAPで半教師あり学習をする例です。
import numpy as np
from keras.datasets import fashion_mnist
from umap import ParametricUMAP
(train_images, Y_train), (test_images, Y_test) = fashion_mnist.load_data()
train_images = train_images.reshape((train_images.shape[0], -1))/255.
test_images = test_images.reshape((test_images.shape[0], -1))/255.
masked_y_train = Y_train.copy().astype(np.int8)
# ランダムに10000個のデータをマスキング
masked_y_train[np.random.choice(60000, size=10000, replace=False)] = -1
classes = [
'T-shirt/top',
'Trouser',
'Pullover',
'Dress',
'Coat',
'Sandal',
'Shirt',
'Sneaker',
'Bag',
'Ankle boot']
embedding = ParametricUMAP(verbose=True).fit_transform(train_images, y=masked_y_train)
逆写像を学習する場合
低次元埋め込みアルゴリズムは、高次元データから低次元(2次元)への写像を学習するアルゴリズムと言えます。
Parametric UMAPでは、埋め込み構成時に 低次元から高次元へのニューラルネットを同時に学習することで、低次元埋め込みの逆写像も学習することができます。逆写像も学習する際はparametric_reconstruction
引数を True
にします。
また、学習に使うニューラルネットを独自で定義する場合 で encoder を定義したときのように、decoder も独自に定義することができます。また、学習に用いないデータセットを、逆写像のバリデーションデータとして用いることも可能です。
import tensorflow as tf
dims = (28, 28, 1)
n_components = 2 # 出力次元
encoder = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=dims),
tf.keras.layers.Conv2D(
filters=32, kernel_size=3, strides=(2, 2), activation="relu", padding="same"
),
tf.keras.layers.Conv2D(
filters=64, kernel_size=3, strides=(2, 2), activation="relu", padding="same"
),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(units=256, activation="relu"),
tf.keras.layers.Dense(units=256, activation="relu"),
tf.keras.layers.Dense(units=n_components),
])
decoder = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(n_components)),
tf.keras.layers.Dense(units=256, activation="relu"),
tf.keras.layers.Dense(units=7 * 7 * 256, activation="relu"),
tf.keras.layers.Reshape(target_shape=(7, 7, 256)),
tf.keras.layers.UpSampling2D((2)),
tf.keras.layers.Conv2D(
filters=64, kernel_size=3, padding="same", activation="relu"
),
tf.keras.layers.UpSampling2D((2)),
tf.keras.layers.Conv2D(
filters=32, kernel_size=3, padding="same", activation="relu"
),
])
# テストデータを逆写像のバリデーションとして用いる
validation_images = test_images.reshape((test_images.shape[0], -1))/255.
embedder = ParametricUMAP(
encoder=encoder,
decoder=decoder,
dims=dims,
parametric_reconstruction= True,
reconstruction_validation=validation_images,
verbose=True,
)
mapper = embedder.fit(train_images)
逆写像を利用する場合は
inv_transformed_points = mapper.inverse_transform(two_dim_points)
で可能です。これは以下の操作と同じです。
inv_transformed_points = mapper.decoder.predict(two_dim_points)
その他の機能につて
Autoencoding UMAP
上記の逆写像を学習する例では、encoder と decoder で異なる損失を学習していましたが、これらを共通化して学習することで、UMAPによる埋め込みをオートエンコーダとして考えることができます。オートエンコーダ自体は強力な次元削減アルゴリズムですが、UMAPと組み合わせることでデータセット上のより大局的な構造を捉えることができるとされています。
Autoencoding UMAP を用いる際には、ParametricUMAPの引数にautoencoder_loss=True
を追加します。
embedder = ParametricUMAP(
encoder=encoder,
decoder=decoder,
dims=dims,
parametric_reconstruction= True,
reconstruction_validation=validation_images,
autoencoder_loss = True,
verbose=True,
)
学習時の損失のプロット
学習後にembedder._history['loss']
でアクセスできる。
fig, ax = plt.subplots()
ax.plot(embedder._history['loss'])
ax.set_ylabel('Cross Entropy')
ax.set_xlabel('Epoch')
ParametricUMAP の内部の話
Parametric UMAPが、どのようにしてUMAPをニューラルネットに落とし込んでいるか
学習するencoderニューラルネットは任意の 高次元 → 低次元のモデルです。ネットワークの構造は全結合でもCNNでもResNetでもなんでも構いません。
ParametricUMAP は encoder ネットワークに、UMAP の損失関数を課すことで実現しています。
UMAPは以下の損失を最適化することで実現していました。
$$
C_{\textrm{UMAP}}= \sum_{i \neq j} p_{i j} \log \left(\frac{p_{i j}}{q_{i j}}\right)+\left(1-p_{i j}\right) \log \left(\frac{1-p_{i j}}{1-q_{i j}}\right)
$$
数式の気持ちを大まかに説明すると、
$p_{ij}$は 高次元空間で、データ点$x_i$, $x_j$ 間の距離に基づいて計算される値です。
$q_{ij}$は 埋め込み空間で、データ点$x_i$, $x_j$を埋め込んだ点$z_i$, $z_j$ 間の距離に基づいて計算される値です。
これらの値は(0,1)の値を取り、各距離が近ければ近いほど大きい値を取ります。t-SNE などでは、データ点間の確率分布として考えています。
また、$i$, $j$は全ての点に対して計算をしているわけではなく、左の項ではデータ点$i$ の第$k$近傍までの点を$j$ としており、右の項ではデータ点$i$に対し、近傍でない点を複数とってきて(ネガティブ・サンプリング)、$j$ としています。
左の項の損失を最適化することで、データ空間で近傍同士にある点が低次元空間でも近くなるように学習されます。そのためこの項を引力項と呼びます。右の項の最適化は、関係の無いデータ同士は低次元空間で離れることに寄与します。そのため斥力項と呼びます。
ちなみに、$p_{ij}$, $q_{ij}$は次の式で計算されます。
$$
p_{i j}=\left(p_{j \mid i}+p_{i \mid j}\right)-p_{j \mid i} p_{i \mid j}, ただし p_{j\mid i} = \exp\left(-\frac{||x_i - x_j|| - \rho_i}{\sigma_i}\right)
$$
$$
q_{i j}=\left(1+a\ ||z_{i}-z_{j}||^{2 b}\right)^{-1}
$$
$\rho$, $\sigma$, $a$, $b$はハイパーパラメータです。
ネガティブサンプリングは実装時にはデータからランダムにサンプルするので問題ありません。そのためこの損失関数は、バッチ化することが可能で、GPUによるニューラルネットの並列処理が適用できます。
t-SNE のパラメトリック版である Parametric t-SNE では損失関数の性質のため、大規模なデータに対するオンライン学習が困難だっという欠点がありましたが、 Parametric UMAP ではその点においてニューラルネット化を克服しています。
近傍グラフの構築について
UMAPの内部では、データ間の近傍探索をしています。一般に高次元のデータで厳密な近傍探索を高速に行うことは困難とされており、高速な探索を行うには近似的な手法を用いる必要があります。
Parametric UMAP の実装以降、(近似)k-近傍探索のライブラリとしてUMAPのチームが開発しているpynndescent が用いられています。以前は spotify の Annoy が用いられていましたが、高速に動作する近似近傍探索を自分たちで実装してしまったようです。
アルゴリズムとしては、探索空間をランダムに分割し、木構造を生成して近傍を探索するタイプの、 Random Projection Forests というアルゴリズムを用いています。
Dong, Wei, Charikar Moses, and Kai Li. "Efficient k-nearest neighbor graph construction for generic similarity measures." Proceedings of the 20th international conference on World wide web. ACM, 2011
参考文献
Parametric UMAP の提案論文
PARAMETRIC UMAP: LEARNING EMBEDDINGS WITH DEEP NEURAL NETWORKS FOR REPRESENTATION AND SEMI-SUPERVISED LEARNING, arxiv
公式ドキュメント
UMAP Documentation: Parametric Embedding
UMAP の解説記事
t-SNEより強いUMAPを(工学的に)理解したい
謝辞
この記事は研究室インターンで取り組みました:https://kojima-r.github.io/kojima/