TensorFlow Similarityを読み解いてみる

この記事はフューチャー Advent Calendar 2021の18日目の記事です。遅刻しました。

はじめに

　TensorFlow Similarityは2021年9月14日に初公開されたライブラリです。本記事ではこのライブラリについて実装の内容や意図を読み解いていきたいと思います。流れとしてはまず初めにSimilarityの全体の流れについて説明し、次に実装について説明や考察を行い、最後にどのような用途でTensorFlow Similarityを使うべきかや未来について思いをはせます。

Similarityについて

距離学習とTensorflow Similarity

　TensorFlow Similarityは距離学習を行い、モデルを提供するためのパイプラインを揃えたライブラリです。距離学習は古典的なクラス分類の学習に近いですが、目的が少し異なります。クラス分類の学習ではクラスの分類誤差の最小化を目指しますが、距離学習では同じラベルのサンプル同士は距離が近く・異なるラベルのサンプル同士は距離が遠くなるように学習を行います。より具体的には、前者が各クラスの分類確率を出力するのに対し、後者ではサンプルを適切な空間に埋め込んだベクトル(embedding)を出力します。
　距離学習のモデルで画像やテキストをembeddingに変換すると、embedding同士で定量的に類似度を計算することが可能です。これにより、サンプルに近いサンプルを持ってくる類似画像・類似文書検索や、近いサンプルのラベルや数値からzero-shotでの予測や未知のラベルに対する予測、近いサンプルがどれくらい存在するかから異常検知等を行うことができます。

距離学習のためのlossとデータセットの条件

　TensorFlow Similarityは現状教師有りでの距離学習のみに対応しているので、これに沿って説明をします。距離学習を教師有りで行う方法は大きく分けて二つあります。一つがTripletLossのようなlossを用いサンプル同士のembeddingの距離を学習する方法、もう一つがArcFaceのような従来のクラス分類を拡張し各クラスのサンプルと各クラスの代表点同士の距離を学習する方法です。
　TensorFlow Similarityは前者のサンプル同士のembeddingを学習するためのlossが実装されています。それぞれのlossはおおよそ考えは共通していて、以下のようにlossを計算します。

Loss = (バッチ内のラベルが違うサンプル同士の距離) - (バッチ内のラベルが同じサンプル同士の距離) + マージン

　このlossはサンプルごとではなく、ミニバッチ単位で計算を行います。そのため、ミニバッチ内にはクラスごとに同じクラスのサンプルが2つ以上必要です。TensorFlow Similarityではこのような条件を持つミニバッチを作成するために、samplerを何種類か提供しています。samplerの中身については実装の章で説明を行います。

モデルの提供と近傍探索

　TensorFlow Similarityではtf.keras.Modelを拡張したSimilarityModelによって学習からモデルの提供をスムーズに行えるようにしています。特徴的なのが、Indexerを用いた近似近傍探索の実行で、これにより大規模なデータセットに対しても効率よくデータを探すことができます。他にも、SimilarityModelは検索の評価を行うための関数や、Similarity用に訓練したmodelをクラス分類用にcalibrateする関数も用意されています。このSimilarityModelを使うことで手軽に距離学習を活用できます。

embeddingの可視化

　Tensorflow Similarityではembeddingの可視化を行うための、bokehを活用したビジュアライザ関数も整備されています。これによりデータの性質やどのクラス同士が近いのかを見ることも可能です。

実装の読解と考察

　この章ではコードからTensorFlow Similarityのできることできないこと、なぜその実装なのかについて考察を行います。今回はデータを取り出すSampler、SimilarityModelのコアになるIndexer、可視化を行うvisualizerの順に説明します。

Sampler

Samplerの種類

　前述のとおり、TensorFlow Similarityで実装されている損失関数はバッチ内にそれぞれのラベルのサンプルが2つ以上あることを前提にしていて、そのために以下のような関数がsamplersに実装されています。

• memory_samplers.pyにはメモリに画像が乗り切るとき用に、
  • ラベル毎に画像を取り出すMultishotMemorySampler
  • 一つの画像をaugmentして正例の組み合わせを作るSingleShotMemorySampler
• tfdataset_samplers.pyにはTensorFlow datasets catalogueからMultishotMemorySamplerを作るTFDatasetMultiShotMemorySampler
• tfrecords_sampler.pyはメモリに画像が乗り切らないときにtfrecordsを用いてストレージから画像を読み込むTFRecordDatasetSampler

　それぞれの使い方は、Sampler IO Cookbookに書いてありますが、そのうちのMultiShotMemorySamplerの内部の動きについて追っていきます。

MultiShotMemorySamplerの動き

　MultiShotSamplerの引数は次のようになっています。

引数名	機能
x	入力画像のFloatTensor
y	入力ラベルのIntTensor
class per batch	バッチ当たりのラベルの種類
examples_per_class_per_batch	バッチ当たりラベル当たりのサンプルの数
class_list	データセット中のラベルの種類(デフォルトだとyから自動で作成)
total_examples_per_class	抽出対象の各ラベルの画像の枚数を指定するか
augmentor	image augment用の関数
warmup	augmentorにwarmupであるかを伝えるための引数

　MultiShotMemorySamplerではxにFloatTensorの画像を、yにIntTensorラベルを入力すると、各ラベルが出現する確率が均等になるようにデータセットから復元抽出を繰り返します。一つのバッチに出現するラベルの種類はclasses_per_batchで指定し、ラベル毎のサンプル数はexamples_per_class_per_batchで指定します。バッチ当たりのサンプル数はclasses_per_batch * examples_per_class_per_batchになり、この抽出をsteps_per_epochで指定した回数繰り返します。バッチ毎ラベル毎にデータを抽出する際は、各ラベル毎のデータの数がexamples_per_class_per_batchより多ければ重複の無い抽出が、少なければ重複ありの抽出が行われます。

注意点

　上記のような実装で容易にsamplerは作れますが、各ラベルは均等に選ばれるためデータセットのラベルが不均衡な場合、データの少ないラベルでは同じサンプルが繰り返し抽出されてしまいます。また、入力のx, yに型ヒントはついていないですが、出力はxがFloatTensor、yはIntTensorの型ヒントがついているため、自然言語やデータセットのパスをxに入れるとeditorに怒られる可能性があります。

Indexer

IndexerはSimilarityModelのコアとなる部分で、nmslibをラップして近似近傍探索を行います。Indexerには以下のようなメソッドが実装されています。

• データ追加のaddとbatch_add
• 近傍を探索するsingle_lookupとbatch_lookup
• クラス分類を行うためのcalibrateとmatch
• モデル評価のためのevaluate_retrivalやevaluate_classification
• モデル保存のためのsaveとload
• indexの統計を返すstatsとprint_stats

使用例は公式のtutorialを見ていただくとして、NMSLIBが使われている理由について考察していきます。

近似近傍探索ライブラリNMSLIBについて

　NMSLIBはNon-Metric Space Libraryの略で、様々な距離の指標で高パフォーマンスな近似近傍探索を行うライブラリです。近似近傍探索は正確な距離計算を行わない代わりにある程度近いデータを検索するための技術で、日本語では東大の松井先生の資料に分かりやすく解説されています。
　TensorFlow SimilarityではデフォルトでNMSLIBに実装されたHNSWアルゴリズムによるコサイン類似度を用いた検索を行います。それぞれの特徴は以下の通りです。

• HNSWアルゴリズムは様々な距離指標で優れたパフォーマンスを出すグラフ系の手法です。生データを保存するという性質上、データがメモリに乗り切る範囲で有効なアルゴリズムで、およそミリオンスケールのデータ探索に向いています。
• コサイン類似度はl2ノルムで正規化したベクトルの同士の内積で、word2vecや距離学習で学習したembeddingに対してはユークリッド距離での比較に比べてより意味的に近いベクトルを取得することができます。

　ちなみにTensorFlow SimilarityのHNSWはデフォルトパラメーターですが、NMSLIBのpythonバインディングではグラフ構築のハイパーパラメータを設定することができます。ハイパーパラメータについてはnmslibのマニュアルにコツを含め書かれているので興味があれば読んでみてください。

NMSLIBの利点の考察

　近似近傍探索の手法のパフォーマンスはANN-Benchmarksで比較されています。NMSLIB実装のHNSWはScaNNを除いてはコサイン類似度(angluar)での検索で高いパフォーマンスを出していることが分かります。このScaNNはgoogleの開発したコサイン類似度での検索に特化した量子化を用いたライブラリで、Vertex Matching Engineにも採用されているのですが、なぜTensorFlow Similarityではnmslibが採用されているのでしょうか？原因はいくつか考えられるのですが、

• NMSLIBはLinux, Mac, Windowsで使える一方、ScaNNはLinuxとMacでないと使えない。
• NMSLIBのver1.0は2014年7月公開であるのに対し、ScaNNは2020年の6月にまだ公開されたばかり。
• NMSLIBはより多くの距離指標に対応していて、コサイン類似度特化型のScaNNより汎用性がある。

あたりが原因ではないかと考えています。

visualizer

visualizerには

• matplotlibによるconfusion_matrixの関数
• 近傍の画像を一列に表示するための関数
• embeddingをUMAPで圧縮し、bokehを用いてインタラクティブな可視化を行うprojector

が実装されています。上記二つについてはコードを読めば何をやっているかはおおよそわかると思うので、最後のprojectorについて説明します。

projectorについて

projectorはembeddingを二次元に圧縮し、ラベル毎に色分けされた散布図を描くためのモジュールです。コードは簡単で、visualizationのノートブックでは以下のようなコードだけで散布図を描くことができるようになっています。

num_examples_to_clusters = 720 #@param {type:"integer"}
thumb_size = 96 #@param {type:"integer"}
plot_size = 800
vx, vy = test_ds.get_slice(0, num_examples_to_clusters)
projector(model.predict(vx), labels=vy, images=vx, class_mapping=breeds, image_size=thumb_size, plot_size=plot_size)

実行すると以下のような散布図が表示され、点の上にカーソルを合わせると画像を見ることができます。

projectorに使われているライブラリ

projectorではUMAPでembeddingを圧縮し、bokehで可視化を行っています。この二つのライブラリについて説明します。

bokeh

　bokehはpythonの可視化のライブラリの一つで、グラフをHTMLで出力することによりインタラクティブな操作を可能にします。bokehで何ができるかは、公式のデモページのmoviesのデモを見るとわかりますが、条件を設定して表示するデータを変えたり、点にカーソルを合わせてデータの情報を取得することができるようになります。

UMAP

　UMAPは次元圧縮の手法の一つで、N次元のベクトルを2、3次元に圧縮し可視化する際によく用いられます。次元圧縮の手法としてはPCAに比べるとかなり遅いものの、umapは教師無しでもデータがラベル毎に分離する傾向がありt-SNEと比べても好まれる傾向にあると思います。公式のドキュメントの可視化パートではbokehを用いてfashion MNISTという画像のデータをそのまま突っ込んで次元圧縮したものが表示されていますが、ラベル毎に分離していることが分かります。

可視化の考察

　embeddingを可視化するツールは他にgoogleのEmbedding Projectorがあります。これはTensor Boardか前述のリンクからブラウザ上で実行できるツールで、非常にパワフルかつインタラクティブにembedding同士の関係性をみることが可能です。一方で起動にはTensorBoardを起動するか、データをブラウザにアップロードしないと使えず、あまり小回りが利かない印象があります。一方でTensorFlow Similarityのprojectorは最小限の労力で一定の可視化ができるようにコードが整備されているので、クイックな検証に向いているように感じます。

UMAP部分の疑問

　TensorFlow Similarityではumapをほぼデフォルトの設定で用いておりprojector関数にはumapの設定についての引数がdensmapしかありません。bokehで散布図を書いてみるくらいのサンプル数であれば計算時間はあまり考えないのかもしれませんが、UMAPは訓練時の距離metricをcosineに指定することもでき、他の設定のデフォルト値がコサイン類似度で指定されている中少しだけ違和感があります。

TensorFlow Similarityはどのように使うべきか

感想

　このライブラリを見た時、最初は感想として細かいチューニングは不向きで、想定していそうなデータセットのサイズも画像がメモリに乗り切るくらいだし実投入できるか？そもそも距離学習に取り組む人はある程度知識あるからオレオレパイプライン使うんじゃないか？と思っていました。
　しかし、この記事を書くにあたって各ライブラリについて調べたりパラメータの検討を行っていくうちに、デフォルトのパラメータのままでもある程度のパフォーマンスが出るように設定されており、実装が面倒な可視化や評価指標の部分も揃っているためベースラインとしては非常に強力なライブラリだと感じるようになってきました。

想定する使い方

　TensorFlow Similarityの使い方は、画像ファイルの置き方、ラベルの置き方をガチガチに固めておいて、新しいラベル付き画像データが来た際にどれくらいの精度が出るかやラベル同士の関係、ベースラインの精度をscript一発で出すパイプラインを構築するのが向いていると思います。
　TensorFlow Similarityのtutorialのvisualizationが良い例で、事前学習済みのEfficientNetをwrapしたEfficientNetSimにより一気に可視化まで行っています。Tutorialには書いていませんが、TensorFlowHubとも相性が良さそうで事前学習済みのEfficientNetV2ならより効率的に動きそうです。

今後の方向性

　公式のIntroducing TensorFlow Similarityという記事では今後の方向性としてBYOLのような半教師あり・自己教師ありの方法をサポートする予定と書いており、ラベル無しのデータでも一定のベースラインが作れるようになれば非常に強力だと思います。
　私としてはその方針に加えて、自然言語の距離学習についても同様に学習から評価、可視化までのパイプラインができるようになると嬉しいです。まだ出たばかりで今後どんどん発展していくライブラリだと思うので、引き続き動向を見守り、できればcontributeもできるといいなと思います。