scikit-learnを用いた古典的検索アルゴリズム

本記事はAI声づくり技術研究会 Advent Calendar 2024 16日目の記事です。

Parakeet株式会社でResearcherをしている金子(nadare)です。普段はCPUで動く軽量AIボイスチェンジャーParavoの研究開発をしております。

概要

本記事ではscikit-learnというオープンソースの機械学習ライブラリに実装されている手法と活用のコツについて、近傍探索の点に絞って一部紹介したいと思います。

近傍探索の技術は音声分野でも幅広く活用されています。直接的な例ではRVCやkNN-VCにおいては事前に自己教師あり学習等で学習したモデルの特徴量に対し近傍探索を行い音声変換に活用しています。また、HuBERTやWavLMはモデルの特徴量に対してKMeansで作成したラベルのうち最も近いものを検索して正解ラベルとすることで自己教師あり学習を行っています。Neural Encodingやそれを用いたTTS、例えばEnCodecやVALL-Eも近傍探索技術が貢献しています。

このように様々な場面で活躍する技術ですが、近年深層学習からAIに入門する傾向から背景にある古典的な機械学習の技術をキャッチアップするのが難しくなっています。本記事ではscikit-learnを例に、今でも現役で使われる古典的な機械学習の手法について、近傍探索・クラスタリング・行列分解の三つの章に分けて紹介します。

対象読者

• 機械学習のアルゴリズムに興味のある方
• AIを自身で研究開発する方

近傍探索

まずは近傍探索について紹介します。近傍探索は二種類に分けられます。

• 厳密に近傍を求めるもの
• 精度を犠牲に高速に解を求めるもの(近似近傍探索)

scikit-learnには前者の近傍探索としてNearestNeighborsが実装されており、様々な距離指標と探索手法を用いて検索することができます。

K-D TreeとBall Treeを用いた高速な近傍探索

scikit-learnのNearest Neighborsではalgorithmの引数から探索に用いるアルゴリズムを選択可能です。
D次元のN個のデータについて全点に対する最近傍をみつけるには下記の計算量で可能です。

• 全探索bruteforceではO[DN^2]
• K-D TreeやBall Treeでは平均でO[DNlog(N)]

適切なアルゴリズムはデータの数やスパース性、用いる距離指標によって異なりますがalgorithm=autoで自動で選択してくれます。

これらのアルゴリズムの詳細はscikit-learnのユーザーガイドのNeighborsの章にあり、最近傍アルゴリズムの選択基準や各メトリックで使える手法が紹介されています。CPU上においてユークリッド距離により近傍探索がしたいのであれば、まずはscikit-learnのNearestNeighborsを用いて速度感を見てみるのもよいでしょう。

クラスタリング

次に、クラスタリングの中でも深層学習とともによく使われるKMeansについてscikit-learnにおける実装を紹介します。KMeansは事前に指定したクラスタ数に応じて、各クラスターのデータとクラスタ中心との距離が最小になるように中心を求めます。KMeansを活用することで、疑似ラベルの作成や近似近傍探索のIVFにおけるインデックス作成に使えます。

初期化

kmeansのクラスタの初期化は、k-means++とrandomの2つから選べます。どちらもデータの中から適当な点を選んで初期値としますが、k-means++は選択済みの初期値との距離に応じた確率でクラスタの初期値を貪欲に決定していくことで高速に収束するような初期値を選択できます。クラスタ数xデータ数x特徴量の次元だけ初期化に計算量がかかりますが、計算時間に余裕があればデフォルトのk-means++を選んで問題ありません。また、n_initの引数で初期化の回数を選べますが、1回で十分なことが多いです。

学習

scikit-learnには2つのKMeansの実装があります。

• 更新のたびにfitで渡した全データをもちいて更新を行うKMeans
• ミニバッチに分割したデータで更新を行うMiniBatchKMeans

MiniBatchKMeansはやや精度に劣る可能性があるものの、例えばaugmentしたデータに対して、SSL出力を計算しながらKMeansを更新したいという場合はMiniBatchKMeansを利用することもあります。HuBERTの論文ではIV.Bの章にscikit-learnのMiniBatchKMeansを用いた疑似ラベル作成のパラメータが紹介されており、10000frame毎にクラスタを更新しています。

パラメータ決定

KMeansでクラスタリングを行う際の適切なクラスタ数は、シルエット分析を参考にすることができます。scikit-learnの公式ドキュメントにはSelecting the number of clusters with silhouette analysis on KMeans clusteringというドキュメントがあり、scikit-learnに実装されているシルエット係数を求めるモジュールを用いた分析の例が紹介されています。

行列分解

sklearn.decompositionにはSVDやPCA、NMFといった様々な行列分解のアルゴリズムが実装されています。行列分解を活用することでベクトル検索の定数倍高速化や精度向上を狙うことが可能です。その中でSVDとPCAに絞って紹介します。

SVDとPCAの使い分け

scikit-learnに実装されているPCAは入力の行列に対して中央化を行った上でSVDを実行します。どちらを用いても大体同じですが、スパース行列を扱う場合は注意が必要です。scikit-learnのTruncatedSVDもPCAもどちらもscipyのsparse matrixを受け付けますが、TruncatedSVDは疎行列で結果を保持できるのに対し、PCAは密行列で結果を持ちます。そのため、自然言語の文章に対して文と単語の共起行列を計算し行列分解するようなケースにおいてはTruncatedSVDを選択しないとメモリがあふれる可能性があります。

Incremental PCA

次元数がそれほど大きくない密行列でも、PCAを計算するにはデータ量が多すぎてメモリに乗り切らない場合があります。この場合はIncrementalPCAを活用することで、バッチ入力に対応させることができます。計算効率や精度は落ちるものの、深層学習のモデル出力を行列分解したいときなどに便利です。

PCA + random rotation

PCAを活用したベクトル検索のテクニックを紹介します。PCAやSVDで変換したベクトルは重要度が特定の次元に偏ります。そのため行列分解で次元を落としたベクトルに対してユークリッド距離等で比較した場合、重要な次元とそうでない次元の差が等しく扱われ、ベクトル検索の手法によっては検索精度が低くなるケースがあります。
これに対し、PCAで変換後のベクトルに対しランダム回転行列を適用して次元ごとの重要度をならすテクニックがあります。例えばLocality Sensitive Hashingの応用のITQに関する論文(Iterative Quantization:A Procrustean Approach to Learning Binary Codes for Large-scale Image Retrieval)を見ても、凝った手法に比べシンプルなPCA-RRが十分な精度を出せることがわかります。

まとめ

ベクトル検索を中心とした、現在でもよく使われる古典的な機械学習アルゴリズムについてscikit-learn実装を紹介しました。scikit-learnのコードは読みやすいので、MiniBatch KMeansやIterative PCAぐらいの手法でもソースコードを確認すればpytorch実装することも容易です。また、scikit-learnの使用感をそのままにGPUで様々なアルゴリズムを使えるcuMLがNVIDIAのRAPIDSで提供されているため、そちらのコードを使用することも可能です。古典的なテクニックも抑えた上で最先端のAIアルゴリズムと組み合わせると新たな発見があるかもしれません。