R言語とStanでガウス過程画像復元モデルを実装してみた

# はじめに

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

https://dippeople.dip-net.jp/19888/

今回の記事では、前回の[記事](https://qiita.com/Gotoubun_taiwan/items/75d82ec15f85ab4ae803)の応用として、ガウス過程の近似を用いた画像の解像度向上（復元）手法を紹介します。

なお、私の専門は政治学方法論と計量経済学であり、情報工学系の画像処理の新しい手法の開発には正直あまり関心がありません。本記事で紹介した内容が「画像復元」と呼べるものなのかも正直わかりません。なんか復元できているぞーと思って「画像復元」というタイトルをつけただけです。本記事の目的は、あくまでガウス過程が2次元情報をどのように学習するかを示すことにあります。そのため、画像処理の専門家の視点から見ると至らない点があるかもしれませんが、あらかじめご了承ください。

また、ガウス過程と本記事で利用するガウス過程の近似法の考え方は省略します。前回の記事を参照していただければと思います。

# なぜやるのか

我々はデータ分析において、しばしば次のようなモデルを推定します：

$$
y \sim Normal(\alpha + f(\text{緯度}, \text{経度}) + X\beta, \sigma)
$$

例えば、政治学においては、$y$ がある政策の支持フラグであり、確率分布がベルヌーイ分布になる場合を考えます。このとき、緯度と経度を回答者の住所、$X$ を回答者のその他の情報（学歴、年収、所属政党など）とし、ガウス過程を用いて$f(\text{緯度}, \text{経度})$を学習させることで、他の条件を一定とした場合の地理的要因が政策支持に与える影響をノンパラメトリックに推定できます。

この推定結果は、既存の行政区画と一致する場合もあれば、そうでない場合もあります。こうした分析を通じて、政策の支持状況の地理的分布を可視化し、新たな仮説を構築する手がかりを得ることができます。

マーケティングにおいても、他の条件を一定とした場合、どの地域のユーザーのCVRが最も高いかを分析する際に、同様のモデル構造を応用できます。

ただし、内生性などの問題はいったん脇に置くとして、直接観測が難しい$f(\text{緯度}, \text{経度})$をガウス過程の近似によって正確に推定できるのかを慎重に考える必要があります。そのため、まずは$f(x_{1}, x_{2})$の学習結果の妥当性を、より簡単に視覚的に確認できる MNIST の手書き数字データを用いて検証したいと思います。

また、$f(x_{1}, x_{2})$が本当に数字の書き方に関する正しいロジックを学習できた場合、ピクセルを細分化して、学習データに含まれていないピクセルの位置に関する状態を正しく予測（解像度を上げる）できるはずです。これを合わせて検証したいと思います。

# 実装コード

前回の[記事](https://qiita.com/Gotoubun_taiwan/items/75d82ec15f85ab4ae803)のコードを少し修正した内容ですので、詳細の説明は省略します：

```cpp:
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_lpmf(
    array[] int pixel_id,
    int start, int end,
    
    array[] real y,
    array[] int pixel_begin, array[] int pixel_end,
    matrix S, array[] vector pixel,
    
    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[pixel_begin[i]:pixel_end[i]] | gaussian_process_approximation(pixel[i], S, 1.3, alpha, l, betas), sigma);
      count += 1;
    }
    return sum(log_likelihood);
  }
}
data {
  int approximation_type;
  int approximation_type_pow_K;
  
  int pixel_type;
  int N;
  int K;
  
  matrix[approximation_type_pow_K, K] S;
  
  array[pixel_type] vector[K] pixel;
  array[pixel_type] int pixel_id;
  array[pixel_type] int pixel_begin;
  array[pixel_type] int pixel_end;
  array[N] real y;
  
  int val_N;
  array[val_N] int val_pixel_id;
  
  int fine_grained_pixel_type;
  array[fine_grained_pixel_type] vector[K] fine_grained_pixel;
}
parameters {
  real<lower=0> alpha;
  real<lower=0> l;
  vector[approximation_type_pow_K] 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_lupmf, pixel_id, grainsize,
    
    y,
    pixel_begin, pixel_end,
    S, pixel,
    
    alpha, l, betas, sigma
  );
}
generated quantities {
  vector[pixel_type] f;
  vector[fine_grained_pixel_type] f_fine_grained;
  array[val_N] real predict;
  for (i in 1:pixel_type){
    f[i] = gaussian_process_approximation(pixel[i], S, 1.3, alpha, l, betas);
  }
  for (i in 1:val_N){
    predict[i] = normal_rng(f[val_pixel_id[i]], sigma);
  }
  for (i in 1:fine_grained_pixel_type){
    f_fine_grained[i] = gaussian_process_approximation(fine_grained_pixel[i], S, 1.3, alpha, l, betas);
  }
}

```

# 前処理

本記事の分析では、MNISTの数字「6」の画像データに対して、次のような生成モデルを想定しています：

$$
色の濃さ_{i} \sim Normal(f(ピクセル_{i, 1}, ピクセル_{i, 2}), \sigma)
$$

ここで、$ピクセル_{i, 1}$ と $ピクセル_{i, 2}$ はそれぞれピクセルの位置を示しており、地図上の緯度経度のようなイメージで使用します。もちろん、色の濃さもピクセルの位置も標準化されます。

また、異なる画像でも、ピクセルの位置が同じであれば、$f(ピクセル_{i, 1}, ピクセル_{i, 2})$の値が変わらないような設定のモデルであるため、Stan 言語のベクトル化機能をうまく活用して、計算の高速化を図るための前処理も行っています。


```R
mnist <- dslabs::read_mnist()

# ピクセル位置のマスター表
base_df <- matrix(mnist$train$images[1,], nrow=28) |> 
  as.data.frame() |> 
  tibble::rowid_to_column() |> 
  tibble::tibble() |> 
  tidyr::pivot_longer(!rowid, names_to = "colid") |> 
  dplyr::mutate(
    colid = colid |> 
      purrr::map_int(
        \(x){
          stringr::str_remove_all(x, "V") |> 
            as.numeric()
        }
      )
  ) |> 
  dplyr::arrange(colid) |>
  dplyr::select(rowid, colid) |>
  dplyr::mutate(
    rowid = rowid/max(rowid) - 0.5,
    colid = colid/max(colid) - 0.5,
    id = dplyr::row_number()
  )

# 解像度向上させたピクセル位置のマスター表
fine_grained_df <- matrix(nrow = 84, ncol = 84) |>
  as.data.frame() |> 
  tibble::rowid_to_column() |> 
  tibble::tibble() |> 
  tidyr::pivot_longer(!rowid, names_to = "colid") |> 
  dplyr::mutate(
    colid = colid |> 
      purrr::map_int(
        \(x){
          stringr::str_remove_all(x, "V") |> 
            as.numeric()
        }
      )
  ) |> 
  dplyr::arrange(colid) |>
  dplyr::select(rowid, colid) |>
  dplyr::mutate(
    rowid = rowid/max(rowid) - 0.5,
    colid = colid/max(colid) - 0.5,
    # idはピクセルの位置を一意に識別するキー
    id = dplyr::row_number()
  )

set.seed(12345)
six_df <- (mnist$train$labels == 6) |>
  which() |>
  purrr::map(
    \(i){
      this_df <- base_df |>
        dplyr::bind_cols(value = mnist$train$images[i,])

      # ゼロが多いため、ダウンサンプリングを実施
      use_id <- c(
        which(this_df$value != 0),
        sample(which(this_df$value == 0), length(which(this_df$value != 0)))
      )
      
      this_df[use_id,]
    }
  ) |>
  dplyr::bind_rows() |>
  dplyr::mutate(
    value = value/255
  ) |>
  dplyr::arrange(id)


pixel_begin_end_df <- six_df |>
  dplyr::mutate(
    rid = dplyr::row_number()
  ) |>
  dplyr::summarise(
    begin = min(rid),
    end = max(rid),
    .by = id
  )

```

次に、Stanに渡すためのデータを整形します。一つの次元の「近似用関数の数」を30に指定します。ここで注意しなければならないのは、2次元データの場合、「近似用関数の数」の総数が30の2乗、すなわち900になることです。したがって、本記事のように注意してStanのコードをチューニングしないと、計算が非常に重くなってしまいます。


```R
set.seed(123456)
val_id <- sample(
  1:nrow(six_df),
  100
)

data_list <- list(
  approximation_type = 30,
  approximation_type_pow_K = 30^2,
  
  pixel_type = nrow(image_begin_end_df),
  N = nrow(six_df),
  K = 2,
  
  S = tibble::tibble(m1 = 1:30, m2 = 1:30) |> tidyr::complete(m1, m2) |> as.matrix(),
  
  pixel = base_df |> dplyr::select(rowid, colid) |> as.matrix(),
  pixel_id = 1:nrow(pixel_begin_end_df),
  pixel_begin = pixel_begin_end_df$begin,
  pixel_end = pixel_begin_end_df$end,
  y = six_df$value,
  
  val_N = nrow(six_df[val_id,]),
  val_pixel_id = six_df$id[val_id],
  
  fine_grained_pixel_type = nrow(fine_grained_df),
  fine_grained_pixel = fine_grained_df |> dplyr::select(rowid, colid) |> as.matrix()
)
```

# モデル推定

まずはモデルのコンパイルを実施します：

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

次に、変分推論でモデルの学習を行います：

```R
> m_hgp_estimate <- m_hgp_init$variational(
     seed = 123456,
     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. 
------------------------------------------------------------ 
Rejecting initial value: 
  Gradient evaluated at the initial value is not finite. 
  Stan can't start sampling from this initial value. 
Rejecting initial value: 
  Gradient evaluated at the initial value is not finite. 
  Stan can't start sampling from this initial value. 
Rejecting initial value: 
  Gradient evaluated at the initial value is not finite. 
  Stan can't start sampling from this initial value. 
Rejecting initial value: 
  Gradient evaluated at the initial value is not finite. 
  Stan can't start sampling from this initial value. 
Gradient evaluation took 0.105025 seconds 
1000 transitions using 10 leapfrog steps per transition would take 1050.25 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) 
Iteration: 250 / 250 [100%]  (Adaptation) 
Success! Found best value [eta = 0.1]. 
Begin stochastic gradient ascent. 
  iter             ELBO   delta_ELBO_mean   delta_ELBO_med   notes  
   100     -2870565.039             1.000            1.000 
   200     -2077834.316             0.691            1.000 
   300     -1513619.835             0.585            0.382 
   400     -1181948.248             0.509            0.382 
   500     -1086060.683             0.425            0.373 
   600      -888654.605             0.391            0.373 
   700      -807666.383             0.349            0.281 
   800      -719985.639             0.321            0.281 
   900      -655143.633             0.296            0.222 
  1000      -631384.282             0.270            0.222 
  1100      -572076.721             0.181            0.122 
  1200      -553309.856             0.146            0.104 
  1300      -537135.717             0.112            0.100 
  1400      -522542.363             0.086            0.099 
  1500      -503129.521             0.082            0.099 
  1600      -491253.955             0.062            0.039 
  1700      -476790.233             0.055            0.038 
  1800      -465017.312             0.045            0.034 
  1900      -454768.950             0.037            0.030 
  2000      -445912.083             0.036            0.030 
  2100      -437288.441             0.027            0.028 
  2200      -430847.334             0.025            0.025 
  2300      -429357.054             0.023            0.024 
  2400      -423744.755             0.021            0.023 
  2500      -420686.019             0.018            0.020 
  2600      -418253.866             0.016            0.020 
  2700      -413752.208             0.014            0.015 
  2800      -412083.759             0.012            0.013 
  2900      -412068.000             0.010            0.011   MEAN ELBO CONVERGED 
Drawing a sample of size 1000 from the approximate posterior...  
COMPLETED. 
Finished in  5011.1 seconds.
```

最後に、結果を保存します：

```R
m_hgp_summary <- m_hgp_estimate$summary()
```

# 結果の可視化

さて、学習データに含まれた二つの「6」に合わせて、モデルが学習した「6」（$f(ピクセル_{i, 1}, ピクセル_{i, 2})$）の形と、ピクセル数を増やした場合の解像度アップされた「6」の形を可視化します。

```R
par(mfrow = c(2, 2), family = "HiraKakuProN-W3")

image(1:28, 1:28, matrix(mnist$train$images[which(mnist$train$labels == 6)[100],], nrow=28)[, 28:1], 
      col = colorRampPalette(c("white", "blue"))(5000), 
      xlab = "", ylab = "", main = "元データ1", axes = FALSE)

image(1:28, 1:28, matrix(mnist$train$images[which(mnist$train$labels == 6)[200],], nrow=28)[, 28:1], 
      col = colorRampPalette(c("white", "blue"))(5000), 
      xlab = "", ylab = "", main = "元データ2", axes = FALSE)

image(1:28, 1:28, matrix(posterior_six, nrow=28)[, 28:1], 
      col = colorRampPalette(c("white", "blue"))(5000), 
      xlab = "", ylab = "", main = "ガウス過程が学習した「6」", axes = FALSE)

image(1:84, 1:84, matrix(posterior_six_fg, nrow=84)[, 84:1], 
      col = colorRampPalette(c("white", "blue"))(5000), 
      xlab = "", ylab = "", main = "ガウス過程が学習した「6」（解像度アップ）", axes = FALSE)
```
![six_gp.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/d593136d-690b-b224-5636-d33ad43c65ee.png)

このように、近似されたガウス過程で推定された $f(ピクセル_{i, 1}, ピクセル_{i, 2})$ は、元の学習データの解像度での「6」をうまく推定しただけでなく、ピクセル数を増やしても形がスムーズに保たれ、解像度向上のタスクにも成功しています。したがって、$f(ピクセル_{i, 1}, ピクセル_{i, 2})$ は「6」の書き方のロジックをある程度正確に学習できたと言えます。

# 結論

いかがでしたか？

本記事で紹介したモデルを利用すれば、他の条件を一定とした場合に、2 次元データの独立変数が従属変数に与える影響を確認することができます。

今後は、実際に地理データにこの手法を応用したいと考えています。

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

https://www.wantedly.com/projects/1955309