t-SNEで遊ぶ

はじめに

最近よく見かける「t-SNE」という非線形次元圧縮手法を試してみた。

t-SNEとは

  • t-Distributed Stochastic Neighbor Embedding
  • SNEと呼ばれる次元圧縮手法の問題点を改善した手法
    • SNEは、「各点間の"ユークリッド距離"を、類似度に相当する"条件付き確率"に変換して低次元にマッピングする手法」のこと
    • 各点について、高次元での確率分布と低次元での確率分布が一致するように低次元での点を探している
    • 確率分布の違いは「カルバックライブラー情報量」で、これを損失関数にして勾配法で低次元の点を求める
  • 低次元での分布に「自由度1のt-分布」を利用する
  • さらに、高速化を行った「Barnes-Hut t-SNE」という手法ではO(n log n)で計算できる

詳しい説明は、元論文や紹介記事がたくさんある。


また、実用で使う場合は、以下に各種実装へのリンクがまとめられているので、そちらを使うのがよい。
https://lvdmaaten.github.io/tsne/

コード

元論文にある「Algorithm 1: Simple version of t-Distributed Stochastic Neighbor Embedding」をそのまま実装したつもり。
論文で紹介されている「early compression」「early exaggeration」はいれていない。

#include <iostream>
#include <vector>
#include <cmath>

// 注意: longではなくint(32bit)にすべき
unsigned long xor128(){
  static unsigned long x=123456789, y=362436069, z=521288629, w=88675123;
  unsigned long t;
  t=(x^(x<<11));
  x=y; y=z; z=w;
  return w=(w^(w>>19))^(t^(t>>8));
}
// 注意: int_maxぐらいで割るべき
double frand(){
  return xor128()%1000000/static_cast<double>(1000000); 
}
static const double PI = 3.14159265358979323846264338;
double normal_rand(double mu, double sigma2){
  double sigma = sqrt(sigma2);
  double u1 = frand(), u2 = frand();
  double z1 = sqrt(-2*log(u1)) * cos(2*PI*u2);
  //double z2 = sqrt(-2*log(u1)) * sin(2*PI*u2);
  return mu + sigma*z1;
}

typedef std::vector<double> P;
static const double EPS = 1e-5;

class tSNE {
  int fromDim, toDim;
  double perp;
  int T;
  double eta;

  double dist2(const P& xi, const P& xj){
    double ret = 0.0;
    for(int i=0; i<xi.size(); i++){
      ret += (xi[i]-xj[i]) * (xi[i]-xj[i]);
    }
    return ret;
  }
  
public:
  tSNE(int fromDim, int toDim = 2, double perp = 5.0, int T = 100000, double eta = 100):
  fromDim(fromDim), toDim(toDim), perp(perp), T(T), eta(eta) { }
  
  std::vector<P> reduce(const std::vector<P>& X){
    int N = X.size();
    std::vector<P> Y(N, P(toDim,0));
    
    //compute pairwise affinities p(j|i) with perplexity perp
    std::vector< std::vector<double> > condP(N, std::vector<double>(N,0));
    for(int i=0; i<N; i++){
      //binary search
      double lb = 0.0, ub = 1e10;
      for(int bi=0; bi<200; bi++){
        double sigma2 = (lb+ub)/2.0;
        
        for(int j=0; j<N; j++){
          if(i!=j){
            condP[j][i] = exp( -dist2(X[i],X[j]) / (2.0 * sigma2) );
          }else{
            condP[j][i] = 0.0;
          }
        }
        double sum = 0.0;      
        for(int k=0; k<N; k++){
          if(i != k) sum += condP[k][i];
        }
        for(int j=0; j<N; j++){
          condP[j][i] /= sum;
        }
        //compute perplexity
        double HPi = 0.0;
        for(int j=0; j<N; j++){
          if(i!=j) HPi -= condP[j][i] * log(condP[j][i]);
        }
        double perpPi = pow(2.0, HPi);
        if(perpPi < perp) lb = sigma2;
        else ub = sigma2;
      }
      std::cerr << "sigma^2 of " << i << " = " << lb << std::endl;
    }

    //set p(i,j)
    std::vector< std::vector<double> > jointP(N, std::vector<double>(N,0));
    for(int i=0; i<N; i++){
      for(int j=0; j<N; j++){
        if(i!=j) jointP[i][j] = (condP[j][i] + condP[i][j]) / (2.0 * N);
      }
    }
    
    //sample initial solution Y^(0)
    for(int i=0; i<N; i++){
      for(int j=0; j<toDim; j++){
        Y[i][j] = normal_rand(0.0, 0.001);
      }
    }
    
    //iteration
    std::vector<P> diff(N, P(toDim, 0));
    double alpha = 0.5;
    
    for(int t=1; t<=T; t++){
      std::vector<P> newY = Y;
      //compute low-dimensional affinities q(i,j)
      std::vector< std::vector<double> > jointQ(N, std::vector<double>(N,0));
      double sum = 0.0;
      for(int k=0; k<N; k++){
        for(int l=0; l<N; l++){
          if(k!=l) sum += 1.0 / (1.0 + dist2(Y[k], Y[l]));
        }
      }
      for(int i=0; i<N; i++){
        for(int j=0; j<N; j++){
          if(i!=j){
            jointQ[i][j] = 1.0 / (1.0 + dist2(Y[i], Y[j]));
            jointQ[i][j] /= sum;
          }
        }
      }

      //(compute cost C = KL(P||Q))
      double cost = 0.0;
      for(int i=0; i<N; i++){
        for(int j=0; j<N; j++){
          if(i!=j){
            cost += jointP[i][j] * log( jointP[i][j] / jointQ[i][j] );
          }
        }
      }
      
      //compute gradient
      std::vector<P> delta(N, P(toDim, 0));
      for(int i=0; i<N; i++){
        for(int j=0; j<N; j++){
          if(i!=j){
            double a = (jointP[i][j] - jointQ[i][j]) * (1.0 / (1.0 + dist2(Y[i], Y[j])));
            for(int k=0; k<toDim; k++){
              delta[i][k] += 4.0 * a * (Y[i][k] - Y[j][k]);
            }
          }
        }
      }

      //set Y^(t)
      for(int i=0; i<N; i++){
        for(int j=0; j<toDim; j++){
          newY[i][j] = Y[i][j] - eta * delta[i][j] + alpha * diff[i][j];
          diff[i][j] = newY[i][j] - Y[i][j];
        }
      }

      Y = newY;
      
      if(t%100==0){
        std::cerr << "t = " << t << " : cost = " << cost << std::endl;
      }
      
      if(t >= 250) alpha = 0.8;
    }
    return Y;
  }
};


int main(){
  int N, fromDim, toDim;
  std::cin >> N;
  std::cin >> fromDim >> toDim;

  //input
  std::vector<P> X;
  for(int i=0; i<N; i++){
    P p(fromDim, 0);
    for(int j=0; j<fromDim; j++){
      std::cin >> p[j];
    }
    X.push_back(p);
  }

  //compute
  tSNE tsne(fromDim, toDim, 5.0, 100000, 100.0);
  std::vector<P> Y = tsne.reduce(X);

  //output
  for(int i=0; i<N; i++){
    for(int j=0; j<toDim; j++){
      if(j!=0) std::cout << "\t";
      std::cout << Y[i][j];
    }
    std::cout << std::endl;
  }
  return 0;
}

実験

MNISTの「test」のデータから500件分をプロット。各要素は0〜255だったようなので、256で割った値を入力として利用。
入力の形式は、

データ点数
高次元での次元数 低次元での次元数
高次元でのデータ点1(タブかスペース区切り、高次元での次元数分)
高次元でのデータ点2(タブかスペース区切り、高次元での次元数分)
...

にしている。

$ cat in.txt
500
784 2
0       0       0       0       0       0       0       0 ...
0       0       0       0       0       0       0       0 ...
...

最終状態は以下。一応収束してくれている。


ステップ5000ずつの低次元マッピングの状態をプロットしてみる。mnistの数字(0〜9)ごとに色や記号の形をわけている。

最初のバラバラな状態からすぐに塊ができて、それらが微調整されていっている。
効率化テクニックを入れていないせいか、どんどん大きくなっているし、10万回回してもまだ収束しきっていないようにも見える。。。