はじめに

こんにちは、事業会社で働いているデータサイエンティストです。

今回の記事では、Gabriel Riutort-Mayol, Paul-Christian Bürkner, Michael R. Andersen, Arno Solin and Aki Vehtari（2023）の論文で紹介された、ヒルベルト空間を用いたガウス過程の高速化手法をStan言語で実装する方法を紹介します。

なお、Riutort-Mayol et al. (2023) は、元となる理論を提唱した Arno Solin and Simo Särkkä (2020) を実証分析者向けに分かりやすく解説した論文です。本記事では、それをさらに転用しやすい形でコード化し、実際の分析に利活用できるよう提供することを目的とします。

なぜやるのか

ガウス過程（Gaussian Process）は、ノンパラメトリックベイズの代表的な手法であり、関数の形状そのものを推定できる強力な手法です。しかし、学習データが増えると計算コストが急激に増大します。

その主な理由は、ガウス過程が予測を行う際に、学習データの特徴量（独立変数）と予測対象の特徴量との相関を利用するためです。言い換えれば、ガウス過程はデータの構造を柔軟に捉えられる一方で、事実上、全ての学習データとの関係を記憶し続ける必要があるため、データ量が増えるほど保持すべき情報も増え、計算負荷が高くなります。

私は業務でガウス過程をベイズモデルの一部として組み込むことを検討したことがあります。しかし、5年分の日次KPIデータ（約2000サンプル）程度でも、計算に数時間かかることがあり、実用上困難でした。そのため、やむを得ず別の手法を利用したり、新たな手法を考案したりしています。

自作のモデルを活用することは、データサイエンティストの腕の見せ所の一つですが、ガウス過程の優れた統計学的性質や、ハイパーパラメータの事前分布設計に関する有用な文献の存在を考えると、それを捨てるのはなかなか惜しいものがあります。

同じ理由でガウス過程の利用を諦めた方もいるかもしれませんが、本記事を参考にすることで、計算時間を大幅に削減しつつ、その優れた性質を活用できるようになります。ぜひ試してみてください！

考え方

理論的な詳細については、正直なところ私自身まだ十分に理解しきれていない部分もあるため、あまり深く立ち入りません。しかし、基本的な考え方としては、「全ての学習データとの関係を記憶」する巨大な行列 $K$を、$m$個の近似用関数の線形結合として表現することで再現します。ここでは学習データの総数よりもはるかに小さい（例えば10程度）値でも十分に良い近似が得られるため、計算時間を大幅に削減できるのです。

実装コード

実装用のコードは以下となります。

モデルは単純で、

$$
Y \sim Normal(f(X), \sigma)
$$

で、$f(x)$を指数2次カーネルのガウス過程で表現するだけなので、詳細の説明は省略します。

gaussian_process_approximation はガウス過程の近似を行う関数で、観測値の$x$（独立変数）やその他のパラメータを渡すと、対応する$f(x)$を返します。

また、lambda、phi、s_inf、gaussian_process_approximation をそのままコピペすれば、別のモデルにも簡単に転用できます。ただし、現在の実装では単一の観測値に対してのみ動作するため、複数の観測値に対応させたい場合は、ご自身でループ処理を追加してください。

データとして渡す必要があるものは、「近似用関数の数」（approximation_type）および「近似用関数の数の独立変数数乗」（approximation_type_pow_K。R言語などのプログラミングでapproximation_type ^ Kを使って計算可能）です。

そして、Sは独立変数が一つしかない場合、R言語言語でいうとmatrix(1:近似用関数の数)を渡せばいいです。独立変数が複数ある場合、tibble::tibble(m1 = 1:近似用関数の数, m2 = 1:近似用関数の数) |> tidyr::complete(m1, m2) |> as.matrix()を渡すことで対応できますが、独立変数が2つ以上の場合についてはまだ検証していないため、参考程度にご覧ください。

コード内では直接記述してしまいましたが（エンジニアの方に怒られそう💦）、Lという、近似の目的を表すハイパーパラメータもあります。コード内では1.3と指定しているため、標準化後のxが-1.3と1.3の間であるyの値の近似を目指すことになります。

hilbert_gp.stan

functions {
  vector lambda(
    matrix S, real L, int j
  ){
    int dimension = dims(S)[2];
    vector[dimension] output;
    for (i in 1:dimension){
      output[i] = ((pi() * S[j, i])/(2 * L))^2;
    }
    return output;
  }
  
  real phi(
    vector x, matrix S, real L, int j
  ){
    int dimension = dims(S)[2];
    vector[dimension] output;
    for (i in 1:dimension){
      output[i] = sqrt(1/L) * sin(sqrt(lambda(S, L, j)[i]) * (x[i] + L));
    }
    return prod(output);
  }
  
  real s_inf(
    real alpha, real l, vector w
  ){
    int D = size(w);
    return alpha * (sqrt(2 * pi())^D) * (l ^ D) * exp(-0.5 * (l ^ 2) * w'*w);
  }
  
  real gaussian_process_approximation(
    vector x, matrix S, real L, real alpha, real l, vector betas
  ){
    int m = dims(S)[1];
    vector[m] output;
    for (i in 1:m){
      output[i] = (s_inf(alpha, l, sqrt(lambda(S, L, i))) ^ 0.5) * phi(x, S, L, i) * betas[i];
    }
    return sum(output);
  }
  
  real partial_sum_lpdf(
    array[] real y,
    int start, int end,
    
    matrix S, array[] vector x,
    
    real alpha, real l, vector betas, real sigma
  ){
    vector[end - start + 1] log_likelihood;
    int count = 1;
    for (i in start:end){
      log_likelihood[count] = normal_lpdf(y[count] | gaussian_process_approximation(x[i], S, 1.3, alpha, l, betas), sigma);
      count += 1;
    }
    return sum(log_likelihood);
  }
}
data {
  int approximation_type;
  int approximation_type_pow_K;
  int N;
  int K;
  
  matrix[approximation_type_pow_K, K] S;
  
  array[N] vector[K] x;
  array[N] real y;
  
  int val_N;
  array[val_N] vector[K] val_x;
}
parameters {
  real<lower=0> alpha;
  real<lower=0> l;
  vector[approximation_type] betas;
  
  real<lower=0> sigma;
}
model {
  alpha ~ gamma(0.1, 0.1);
  l ~ gamma(0.1, 0.1);
  betas ~ normal(0, 1);
  
  sigma ~ gamma(0.1, 0.1);
  
  int grainsize = 1;
  
  target += reduce_sum(
    partial_sum_lupdf, y, grainsize,
    
    S, x,
    
    alpha, l, betas, sigma
  );
}
generated quantities {
  array[val_N] real predict;
  for (i in 1:val_N){
    predict[i] = normal_rng(gaussian_process_approximation(val_x[i], S, 1.3, alpha, l, betas), sigma);
  }
}

データ生成

ここでは、まずランダムウォークに従う3万件のyを含むデータフレームを生成します。次に、xはx/max(x) - 0.5で標準化し、その値が-0.5から0.5の間になるようにします。yは(y - mean(y))/sd(y)で標準化します：

set.seed(123456)
df <- 30000 |>
  rnorm(n = _)|>
  cumsum() |>
  tibble::tibble(
    y = _
  ) |>
  dplyr::mutate(
    x = dplyr::row_number(),
    y_std = (y - mean(y))/sd(y),
    x_std = x/max(x) - 0.5
  )

Nが3万件なので、通常のガウス過程であれば、変分推論やreduce_sum関数などを活用しても、数時間かからないと計算が終わらないことになります。

次は、2000件の検証データを指定し、Stan言語にわたすための整形を実施します：

val_id <- sample(1:nrow(df), 2000)

data_list <- list(
  approximation_type = 10,
  approximation_type_pow_K = 10,
  N = nrow(df[-val_id,]),
  K = 1,
  
  S = matrix(1:10),
  
  x = matrix(df$x_std[-val_id]),
  y = df$y_std[-val_id],
  
  val_N = nrow(df[val_id,]),
  val_x = matrix(df$x_std[val_id])
)

推定結果

ここではモデルコンパイルと推定に入ります：

m_hgp_init <- cmdstanr::cmdstan_model("hilbert_gp.stan",
                                      cpp_options = list(
                                        stan_threads = TRUE
                                      )
)

> m_hgp_estimate <- m_hgp_init$variational(
     threads = 16,
     data = data_list
 )
------------------------------------------------------------ 
EXPERIMENTAL ALGORITHM: 
  This procedure has not been thoroughly tested and may be unstable 
  or buggy. The interface is subject to change. 
------------------------------------------------------------ 
Gradient evaluation took 0.032744 seconds 
1000 transitions using 10 leapfrog steps per transition would take 327.44 seconds. 
Adjust your expectations accordingly! 
Begin eta adaptation. 
Iteration:   1 / 250 [  0%]  (Adaptation) 
Iteration:  50 / 250 [ 20%]  (Adaptation) 
Iteration: 100 / 250 [ 40%]  (Adaptation) 
Iteration: 150 / 250 [ 60%]  (Adaptation) 
Iteration: 200 / 250 [ 80%]  (Adaptation) 
Success! Found best value [eta = 1] earlier than expected. 
Begin stochastic gradient ascent. 
  iter             ELBO   delta_ELBO_mean   delta_ELBO_med   notes  
   100        -6147.231             1.000            1.000 
   200        -5305.277             0.579            1.000 
   300        -4841.198             0.418            0.159 
   400        -4822.964             0.315            0.159 
   500        -4093.398             0.287            0.159 
   600        -3601.774             0.262            0.159 
   700        -3588.400             0.225            0.136 
   800        -3263.764             0.210            0.136 
   900        -3098.738             0.192            0.099 
  1000        -2874.989             0.181            0.099 
  1100        -2914.376             0.082            0.096 
  1200        -2890.984             0.067            0.078 
  1300        -2897.236             0.058            0.053 
  1400        -2873.490             0.058            0.053 
  1500        -2806.786             0.043            0.024 
  1600        -2824.581             0.030            0.014 
  1700        -2773.042             0.031            0.019 
  1800        -2767.543             0.021            0.014 
  1900        -2802.820             0.017            0.013 
  2000        -2750.181             0.011            0.013 
  2100        -2780.155             0.011            0.011 
  2200        -2790.060             0.011            0.011 
  2300        -2762.715             0.011            0.011 
  2400        -2736.907             0.012            0.011 
  2500        -2838.712             0.013            0.011 
  2600        -2815.478             0.013            0.011 
  2700        -2738.936             0.014            0.011 
  2800        -2725.963             0.014            0.011 
  2900        -2760.536             0.014            0.011 
  3000        -2743.057             0.013            0.010   MEDIAN ELBO CONVERGED 
Drawing a sample of size 1000 from the approximate posterior...  
COMPLETED. 
Finished in  252.0 seconds.

なんと！4分で終わりました！！！

結果を整形します：

m_hgp_summary <- m_hgp_estimate$summary()

最後の、検証データでの精度を可視化しましょう：

m_hgp_summary |>
    dplyr::filter(stringr::str_detect(variable, "^predict\\["))|>
    dplyr::bind_cols(
        x = df$x[val_id],
        y = df$y_std[val_id]
    ) |>
    ggplot2::ggplot() + 
    ggplot2::geom_point(ggplot2::aes(x = x, y = y)) + 
    ggplot2::geom_line(ggplot2::aes(x = x, y = mean), color = "blue") +
    ggplot2::geom_ribbon(ggplot2::aes(x = x, ymin = q5, ymax = q95), fill = "blue", alpha = 0.3)

黒い点は検証データを、青い線は予測値の平均を、青い区間は信用区間を示しています。

我々がガウス過程に期待するトレンドのようなものを抽出できたと思います。

結論

いかがでしたか？

本記事で紹介されたコードを利用すれば、ガウス過程をよりスケールさせることができます。個人的には、緯度と経度を特徴量として渡し、地理的な特性が従属変数に与える影響を確認することに一番可能性を感じます。

最後に、私たちと一緒に、データサイエンスの力で社会を改善したい方はこちらをご確認ください：

参考文献

Riutort-Mayol, Gabriel, Paul-Christian Bürkner, Michael R. Andersen, Arno Solin and Aki Vehtari. "Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming." Statistics and Computing 33.1 (2023): 17.

Solin, Arno, and Simo Särkkä. "Hilbert space methods for reduced-rank Gaussian process regression." Statistics and Computing 30.2 (2020): 419-446.