計算時間の速いt-SNE実装とUMAPを、あくまで速度だけで比べてみた

要約

• $ \mathcal{O}(N^2) $かかるt-SNEの計算を高速な近似に置き換えるアルゴリズムを実装したopenTSNEライブラリを見つけた
• $ \mathcal{O}(NlogN) $でt-SNEに類したことができるUMAPについても、Pythonの著名な実装であるumap-learnライブラリも見つけた
• 扱うデータのサンプルサイズに応じて最速の実装は変わる
  • 件数が少なめだとsklearnのt-SNEが有利
  • 件数が増えるとumap-learnのUMAPが有利
  • 件数に関わらず、openTSNEは無難な性能

t-SNEやUMAPの計算時間の比較結果は、本ページのデータ量に応じて計算時間が変わるのでは？だけ読めばだいたいOKです。

はじめに

$ \mathcal{O}(N^2) $かかるt-SNEの計算を、近似計算によって$ \mathcal{O}(N) $オーダまで削減できるアルゴリズムとして、FIt-SNE(Fast Fourier Transform-accelerated Interpolation-based t-SNE)が知られています。

作者によるFIt-SNEの実装はLinuxでしか動きません。Windowsでも動くFIt-SNEの実装が必要な場合、OpenTSNEを使うのがよいです。

OpenTSNEのサンプルコード作らせてみた

せっかくなのでAnythingLLMでOpenTSNEのサンプルコードを作成しました。

AnythingLLMは、RAG(Retrieval Augmented Generation)と呼ばれる機能を持つチャットボットを簡単に作れるアプリのひとつです。

RAGとは、ユーザから質問された内容に関連する情報を知識データベースから検索した結果をもとに、然るべき回答内容をChatGPTなどの大規模言語モデルに作成させるアーキテクチャのことです。

AnythingLLMには、様々なデータソースに含まれるテキストを知識データベースに格納するためのプラグインが同梱されています。その一つに、GitHubリポジトリ内のテキストファイルをまとめてベクトルストアに投入できるものがあります。GitHubにあるOSSの内容を理解するときにはこの機能がかなり役に立ちます。

ほかにもAnythingLLMでは、Web検索などの外部ツールと連携できるエージェントも構築できます。こうした機能をうまく利用すれば、最新情報に基づいた技術動向の案内をAnythingLLMで実現できるポテンシャルもあります。

そのポテンシャルの検証としてAnythingLLMにOpenTSNEのWeb記事を検索させ、t-SNEを実行する処理をPythonで作成してみました。できたコードは以下です。

from sklearn import datasets
from openTSNE import TSNE

# Load the Iris dataset
iris = datasets.load_iris()
X, y = iris["data"], iris["target"]

# Run openTSNE (FIt-SNE)
tsne = TSNE()
embedding = tsne.fit(X)

print(embedding)

次にOpenDevinでコードの動く環境構築・実行や改良を実施しました。OpenDevinは、LLMによるソフトウェア自動生成に特化したツールで、コード実行用のサンドボックス環境をDockerで構築する機能があります。

ただし、こちらがコード仕様を説明してからOpenDevinが仕様通りのコードをワンパス動かすまでの間には、概ね数十万トークンの入出力が消費されます。GPT-4oを使っていても2ドル位かかります。
そうして作ったコードが以下です。

requirement.txt

numpy==1.26.4
openTSNE==1.0.1
scikit-learn==1.5.0

main.py

import time

from openTSNE import TSNE as OpenTSNE
from sklearn.datasets import make_classification
from sklearn.manifold import TSNE as SKLearnTSNE

# Generate dummy data with 30 informative features out of 500
X, _ = make_classification(n_samples=1000, n_features=500, n_informative=30, random_state=42)

# Run sklearn t-SNE
start_time = time.time()
print('Running sklearn t-SNE...')
sklearn_tsne = SKLearnTSNE(n_jobs=4)
sklearn_embedding = sklearn_tsne.fit_transform(X)
sklearn_time = time.time() - start_time
print('sklearn t-SNE completed.')

# Run openTSNE (FIt-SNE) with adjusted parameters
start_time = time.time()
print('Running openTSNE...')
# NOTE: 第2引数はFIt-SNEを明示的に実行するためのオプション
opentsne = OpenTSNE(n_jobs=4, negative_gradient_method="fft")
opentsne_embedding = opentsne.fit(X)
opentsne_time = time.time() - start_time
print('openTSNE completed.')

print('sklearn t-SNE embedding:')
print(sklearn_embedding)
print('openTSNE embedding:')
print(opentsne_embedding)

print(f'Time taken by sklearn t-SNE: {sklearn_time:.2f} seconds')
print(f'Time taken by openTSNE: {opentsne_time:.2f} seconds')

なお、このコードをさらに改修する作業をOpenDevinにやらせましたが、それだけで10ドル近くかかりました。

動かしてみた

ではOpenTSNEを使うとどれだけ処理が速くなるのでしょう。

…という話をしようと思いましたが、実は上記のコードの設定ではたいして速くはありません。むしろsklearnのものに負けています。

Time taken by sklearn t-SNE: 5.56 seconds
Time taken by openTSNE: 7.77 seconds

sklearnのt-SNE実装

sklearnのt-SNE実装がopenTSNEより速い場合があるのは、おそらくsklearnでもt-SNEを速くする仕掛けが実装されていることに理由がありそうです。

sklearn.manifold.TSNEのソースコードを見ると、以下の引用文献が書かれています。

    [4] Belkina, A. C., Ciccolella, C. O., Anno, R., Halpert, R., Spidlen, J.,
        & Snyder-Cappione, J. E. (2019). Automated optimized parameters for
        T-distributed stochastic neighbor embedding improve visualization
        and analysis of large datasets. Nature Communications, 10(1), 1-12.

これはt-SNEの内部計算処理のハイパラを最適化する、"opt-SNE"という手法を提案した論文です。

この論文が出る前にsklearnのt-SNEを触ったきりで、「t-SNEは遅い」というイメージを持って今に至っていましたが、どっこいsklearnも頑張ってt-SNEをアップデートしていたのです。

UMAPとも比較してみよう

ChatGPTにお願いして、t-SNEよりも計算量が少なくて良いと噂のUMAPの実行時間も計測できるようにしてもらいました。

UMAP自体の解説はこちらに詳しいです。

import time

import umap
from openTSNE import TSNE as OpenTSNE
from sklearn.datasets import make_classification
from sklearn.manifold import TSNE as SKLearnTSNE

# Generate dummy data with 30 informative features out of 500
X, _ = make_classification(n_samples=1000, n_features=500, n_informative=30, random_state=42)

# Run sklearn t-SNE
start_time = time.time()
print('Running sklearn t-SNE...')
sklearn_tsne = SKLearnTSNE(n_jobs=4)
sklearn_embedding = sklearn_tsne.fit_transform(X)
sklearn_time = time.time() - start_time
print('sklearn t-SNE completed.')

# Run openTSNE (FIt-SNE) with adjusted parameters
start_time = time.time()
print('Running openTSNE...')
opentsne = OpenTSNE(n_jobs=4, negative_gradient_method="fft")
opentsne_embedding = opentsne.fit(X)
opentsne_time = time.time() - start_time
print('openTSNE completed.')

# Run UMAP
start_time = time.time()
print('Running UMAP...')
umap_reducer = umap.UMAP(n_jobs=4, n_components=2)
umap_embedding = umap_reducer.fit_transform(X)
umap_time = time.time() - start_time
print('UMAP completed.')

print('sklearn t-SNE embedding:')
print(sklearn_embedding)
print('openTSNE embedding:')
print(opentsne_embedding)
print('UMAP embedding:')
print(umap_embedding)

print(f'Time taken by sklearn t-SNE: {sklearn_time:.2f} seconds')
print(f'Time taken by openTSNE: {opentsne_time:.2f} seconds')
print(f'Time taken by UMAP: {umap_time:.2f} seconds')

結果はというと、ダミーデータ1000行では依然としてsklearnのt-SNE実装が最速でした。

Time taken by sklearn t-SNE: 5.29 seconds
Time taken by openTSNE: 7.44 seconds
Time taken by UMAP: 11.35 seconds

データ量に応じて計算時間が変わるのでは？

よくよく思えば1000件のサンプルサイズだけだと何とも性能を論じ得ません。

…というわけでまたもChatGPTにお願いして、サンプルサイズを変えてsklearnのt-SNE、openTSNEおよびUMAPの計算時間を比較・可視化するコードを作ってもらいました。

import time
import umap
from openTSNE import TSNE as OpenTSNE
from sklearn.datasets import make_classification
from sklearn.manifold import TSNE as SKLearnTSNE
import matplotlib.pyplot as plt

# Sample sizes to test
sample_sizes = [100, 500, 1000, 5000, 10000, 50000]
sklearn_times = []
opentsne_times = []
umap_times = []

for n_samples in sample_sizes:
    print(f'Running for sample size: {n_samples}')

    # Generate dummy data with 30 informative features out of 500
    X, _ = make_classification(n_samples=n_samples, n_features=500, n_informative=30, random_state=42)
    
    # Run sklearn t-SNE
    start_time = time.time()
    print('Running sklearn t-SNE...')
    sklearn_tsne = SKLearnTSNE(n_jobs=4)
    sklearn_tsne.fit_transform(X)
    sklearn_time = time.time() - start_time
    print('sklearn t-SNE completed.')
    sklearn_times.append(sklearn_time)
    
    # Run openTSNE (FIt-SNE) with adjusted parameters
    start_time = time.time()
    print('Running openTSNE...')
    opentsne = OpenTSNE(n_jobs=4, negative_gradient_method="fft")
    opentsne.fit(X)
    opentsne_time = time.time() - start_time
    print('openTSNE completed.')
    opentsne_times.append(opentsne_time)
    
    # Run UMAP
    start_time = time.time()
    print('Running UMAP...')
    umap_reducer = umap.UMAP(n_jobs=4, n_components=2)
    umap_reducer.fit_transform(X)
    umap_time = time.time() - start_time
    print('UMAP completed.')
    umap_times.append(umap_time)

# Plot the computation times with logarithmic scales
plt.figure(figsize=(10, 6))
plt.plot(sample_sizes, sklearn_times, label='sklearn t-SNE', marker='o')
plt.plot(sample_sizes, opentsne_times, label='openTSNE', marker='o')
plt.plot(sample_sizes, umap_times, label='UMAP', marker='o')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Number of samples (log scale)')
plt.ylabel('Computation time (seconds, log scale)')
plt.title('Computation Time vs Number of Samples for Different Dimensionality Reduction Algorithms')
plt.legend()
plt.grid(True, which="both", ls="--")
plt.show()

実行してできたグラフは以下です。両対数グラフにしているので、$ \mathcal{O}(N^2) $ のアルゴリズムはグラフ上に直線として描かれます¹。
sklearnのt-SNE実装の計算時間をみるとグラフ上で直線になっています。ハイパラ最適化はなされているものの基本的には、この実装も$ \mathcal{O}(N^2) $ のアルゴリズムになっていそうです。

まとめ

t-SNEの実装はこれまでsklearnしか知らなかったですが、今回調べてみると高速な実装が他にあったり、umap-learnのUMAPも計算自体はt-SNEより速かったりと、新たな発見がありました。

やはり、勉強は大事です。

なお今回検証したのは計算の速さですが、これは計算の結果の正確性とは基本的に独立なものです。したがってまずは一番軽いUMAPを実行してみて、もともと近くにあったデータ点どうしがUMAPでの射影後にも近くにあるか確認してみるのが良いと思います。

おまけ parametric UMAPも試した && ChatGPTにリファクタさせてみた

umap-learnには、Kerasの作者のFrançois Cholletさんが作ったparametric UMAPというアルゴリズムもあります。

初期学習時の性能においては、普通のUMAPに対するメリットはないようですが、より高付加価値・高精度になっているようです。試しにどれくらいの時間で動くか確認してみました。

加えてparametric UMAPを使うように既存コードを拡張し、さらにChatGPTにリファクタさせてみました。コードは以下です。リファクタした手順は後日紹介します。

import time

import matplotlib.pyplot as plt
import umap
from openTSNE import TSNE as OpenTSNE
from sklearn.datasets import make_classification
from sklearn.manifold import TSNE as SKLearnTSNE
from umap.parametric_umap import ParametricUMAP


def run_algorithm(algorithm_class, X, **kwargs):
    start_time = time.time()
    algorithm = algorithm_class(**kwargs)
    if hasattr(algorithm, 'fit_transform'):
        algorithm.fit_transform(X)
    else:
        algorithm.fit(X)
    return time.time() - start_time

def generate_data(n_samples, n_features=500, n_informative=30, random_state=42):
    return make_classification(n_samples=n_samples, n_features=n_features, n_informative=n_informative, random_state=random_state)

def plot_computation_times(sample_sizes, times, labels):
    plt.figure(figsize=(10, 6))
    for time_list, label in zip(times, labels):
        plt.plot(sample_sizes, time_list, label=label, marker='o')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Number of samples (log scale)')
    plt.ylabel('Computation time (seconds, log scale)')
    plt.title('Computation Time vs Number of Samples for Different Dimensionality Reduction Algorithms')
    plt.legend()
    plt.grid(True, which="both", ls="--")
    plt.savefig("comparison_tsnes_and_umap.svg")
    plt.show()

# Define the algorithms and their specific parameters
algorithms = {
    'sklearn t-SNE': (SKLearnTSNE, {'n_jobs': 4}),
    'openTSNE': (OpenTSNE, {'n_jobs': 4, 'negative_gradient_method': 'fft'}),
    'UMAP': (umap.UMAP, {'n_jobs': 4}),
    'parametric UMAP': (ParametricUMAP, {'n_jobs': 4}),
}

# Sample sizes to test
sample_sizes = [100, 200, 500, 1000, 2000, 5000, 10000, 20000, 50000]

# Initialize lists to store computation times
times = {label: [] for label in algorithms}

for n_samples in sample_sizes:
    print(f'Running for sample size: {n_samples}')
    
    # Generate dummy data
    X, _ = generate_data(n_samples)
    
    # Run and time each algorithm
    for label, (algorithm_class, kwargs) in algorithms.items():
        print(f'Running {label}...')
        times[label].append(run_algorithm(algorithm_class, X, **kwargs))
        print(f'{label} completed.')

# Prepare data for plotting
times_lists = [times[label] for label in algorithms]
labels = list(algorithms.keys())

# Plot the computation times
plot_computation_times(sample_sizes, times_lists, labels)

parametric UMAPはTensorflowを使うので、requirements.txtは以下のようになります。

requirements.txt

matplotlib==3.9.0
numpy==1.26.4
openTSNE==1.0.1
scikit-learn==1.5.0
tensorflow
umap-learn==0.5.6

結果は以下です。tensorflowに依存しているものなのにGPUで動かさなかったため、parametric UMAPの性能はよろしくありません。使えるものならGPU環境で動かしたほうが良いと思います。

1. 理由はこちらを参照 https://mathlandscape.com/log-log-graph/ ↩