Kaggle NeurIPS - Ariel Data Challenge 2024の2位ベイズ解法の解説

はじめに

皆さん、こんにちは、良いKaggleライフをお過ごしでしょうか。
本日は、NeurIPS - Ariel Data Challenge 2024の2位解法をご紹介しようと思います！
コンペのタスクは知っている事を前提として説明します。
コンペのタスクについては以下のような資料を参考にしてください。
素晴らしいコンペ振り返り記事を書いてくださっている方々、ありがとうございます！

• Kaggleコンペページ
• NeurIPS - Ariel Data Challenge 2024振り返り(moto)
• NeurIPS - Ariel Data Challenge 2024の振り返り(takaitoさん)
• Arielコンペ負け犬の遠吠え(junkodaさん)
  

ベイズ推定のことは知っていた方が良いですが、あまり細かい数式の話はしないと思うので、そこまで詳しくなくても本記事を読むうえでは大丈夫だとは思います。

私は最初にこのsolutionを読んでもあまりよく理解できなかったので、自分でもまとめるというモチベーションで本記事を書いています。
この解法を作った方は詳細についてはあまり説明する気はなく、私もできるだけ最善を尽くしましたがコードで一部理解できなかった箇所があることはご了承ください。細かい数式計算を理解できなかったというだけで大まかな解法については正しく説明できると思っています。
データの前処理などは省くことにします。コードを読んで頂いてもすぐ理解できると思います。

多変量ガウス分布について

今更多変量ガウス分布について説明する意味はあまりないかもしれませんが、良い機会なので私の多変量ガウス分布に対する理解をまとめておきます。
平均$\mu \in \mathbb{R}^d$、共分散行列$\Sigma \in \mathbb{R}^d$の多変量ガウス分布の確率密度関数は

f(x) = \frac{1}{(2\pi)^{d/2} |\Sigma|^{1/2}} \exp\left(-\frac{1}{2} (x-\mu)^T \Sigma^{-1} (x-\mu) \right)

になります。($|\cdot|$は行列式を表しています。$\Sigma$は半正定値対称行列)
平均$\mu \in \mathbb{R}^d$、共分散行列$\Sigma \in \mathbb{R}^d$と言われる所以は実際に計算してみると$\mu$が多変量ガウス分布の平均値、$\Sigma$が多変量ガウス分布の共分散行列になることです。

一見、$1$次元のガウス分布によりもだいぶ複雑な対象に見えますが、対称行列は直交行列による対角化ができることから、$\Sigma$は$\Sigma=P\Lambda P^{T}$と対角化をすることができ、$y=P^{T}x$と座標変換をすれば多変量ガウス分布の確率密度関数は独立な$d$個の$1$次元ガウス分布の確率密度関数の積で表せます。
つまり、多変量ガウス分布は単に独立な$d$個のガウス分布を直交座標変換したものと考えることができ、とても単純な対象であることが分かります。

直交行列を使うと座標変換としては分かりやすいですが、直交行列でなくてもいいのであれば正定値対称行列に対してはコレスキー分解と呼ばれる分解$\Sigma=LL^{T}$($L$は下三角行列)が存在し、$y=L^{T} x$と変数変換したら$y$は共分散行列が単位行列の多変量ガウス分布に従います。
実際の数値計算ではコレスキー分解の方がよく用いられます。
正定値行列でなければうまくいかないのでうまくいかない場合は対角成分に微小値を足してコレスキー分解を行います。
どのような場面でコレスキー分解がよく用いられるかというと、$Ax=b$となる$x$を求める場合や$\det A$を求めるような場合に$A$に対してコレスキー分解が用いられます。

余談でコレスキー分解の導出の話をざっくりします。
LU分解をした後にL、Uの対角成分が全て1になるようにしてLDU分解を導出すると、対称行列を分解していることから$U=L^{T}$となり、$LD^{1/2}$($D^{1/2}$は$D$の対角成分に対して全て√を取った行列)を$L$と置きなおすとコレスキー分解が導出できます。
対称行列の直交行列による対角化も個人的に好きな証明があるのですが、ちょっと話が長くなるので省略します。

ガウス過程回帰について

一応ざっくりとだけガウス過程回帰の復習をします。
ガウス過程回帰ではカーネル関数$k(\cdot, \cdot)$というのを定義し、データ点$x_1, x_2, \cdots, x_n$とラベル$y_1, y_2, \cdots, y_n$を与えた時に次の仮定を置きます。
$K=(k(x_i, x_j)) \in \mathbb{R}^{n \times n}$で行列を定義した時、

\begin{pmatrix}
y_1\\
\vdots\\
y_n
\end{pmatrix}\sim N(0, K)

と仮定する。($N$は多変量正規分布、$y_1, \cdots, y_n$は平均値が$0$になるように定数を足して調整する)
このような分布を事前分布と言います。
さらにデータ点$x'$が与えられたときに予測値$y'$を求めたい場合は、
$k'=(k(x', x_{i}))\in \mathbb{R}^{n}$, $k''=k(x', x') \in \mathbb{R}$とし、

K'=\begin{pmatrix}
K&k'\\
k'^T&k''
\end{pmatrix}

とします。
このとき、

\begin{pmatrix}
y_1\\
\vdots\\
y_n\\
y'
\end{pmatrix}\sim N(0, K')

が成り立ちます。(ガウス過程回帰の仮定)
ここから、$p(y'|x_1,\cdots, x_n, x', y_1,
\cdots, y_n)$を求めることによってガウス過程回帰の推論ができます。
なお、計算過程は省きますが、

y=\begin{pmatrix}
y_1\\
\vdots\\
y_n
\end{pmatrix}

とすると、
$p(y'|x_1,\cdots, x_n, x', y_1, \cdots, y_n)$が$N(k'^{T}K^{-1}y, k''-k'^{T}K^{-1}k')$に従うことが知られており、ガウス分布の平均値と標準偏差によって予測値と予測値の自信度を求めることができます。(ちなみに、ベイズの定理と平方完成に似た計算で導出できます)
こちらの分布を事後分布と言います。
事前分布と事後分布は$2$次元の場合はこんな感じで図示できます。
$y_1$と$y_2$に正の相関があれば、$y_1>0$が分かっているの時の$y_2$の確率分布は正の方へと平行移動し、分散が小さくなります。

ガウス過程回帰では、$x_i$と$x_j$の類似度のようなものを表すカーネル関数で$y_i$と$y_j$の共分散を定めており、しばしば$x$が似ていれば$y$が似ているというように表現されます。

ガウス過程回帰について詳しく勉強したい場合はガウス過程と機械学習などを参照してください。
標準的なガウス過程回帰はこのようにして求められますが、残念ながら2位の方の推論方法と標準的なガウス過程回帰の推論方法との対応はきちんととれませんでした…
何が分からなかったのかについては、分からなかったことの章を見てください…

解法について

この解法は与えられた信号を複数のガウス過程回帰を用いてモデリングする解法になっています。
具体的には、与えられた信号を再現するようにnoise、star spectrum、drift、transitのモデルを学習します。
以下の図は2位解法ページからの引用です。

noiseは信号に混ざったノイズを、star spectrumは波長ごとの信号への依存性、driftは時間や波長による信号のズレ、transitはtransit現象による信号の変化を表しています。
star spectrumはtargetの光の波長の値から求められます。

driftとtransitに関しては、さらに次のように分解しています。


spectral driftというのはAIRSのデータ(時間方向には平均を取ってある$2$次元データ。波長$\times$空間が残っている)から推論されるスペクトルごとのノイズで、non-spectral noiseは時間方向のズレです。non-spectral noiseは時間の値、spectral driftは時間と波長の値から求められます。
transitはtransitが生じている間だけ値が$1$、それ以外の時には$0$となるtransit_windowとスペクトルごとのtransitの深さを表すtransit_depthから構成されます。いずれにしてもtargetの光の波長の長さを入力として受け取ります。
(transitについてはPCAを使ってplanetごとにベースとする深さを取る項が実はありますが、細かいことは元のdiscussionやコードを参照してください。)

この解法を作った方は、観測される値はこのような構成要素に分解できると考えてモデリングをしているわけです。
もし、このようなモデリングに成功したら、transit_depthをモデリングするガウス過程回帰を用いることで、transit depthの情報を予測することができ、一件落着です。ちなみに、sigmaはガウス過程回帰の事後分布からサンプリングして標準偏差を求めることで予測します。
(ここでいうtransit depthとは、NeurIPS - Ariel Data Challenge 2024振り返り(moto)の資料における$\frac{d}{1+d}$の値の事です。)

ただ、ガウス過程回帰を勉強したことがある方は、ガウス過程回帰の和や積ってどうやって学習するねんとは思うでしょう…

ガウス過程回帰の和や積を学習する方法

まずは、$y$がガウス分布の事後分布としたら、定数行列$J$、定数ベクトル$z$に関して、$Jy+z$がどのようになるかを考えましょう。
$y$が$N(\mu, \Sigma)$に従うと推論したら、行列$Jy+z$は$N(J\mu+z, J\Sigma J^{T})$に従います。

次に、$y_1, y_2$がガウス分布の事後分布としたら、$y_1+y_2$がどのような値になるかを考えましょう。実は、先程の線形倍のケースに帰着させることができます。
$y_1$が$N(\mu_1, \Sigma_1)$で、$y_2$が$N(\mu_2, \Sigma_2)$と推論したとします。

\mu'=\begin{pmatrix}
\mu_1\\
\mu_2
\end{pmatrix},
\Sigma'=\begin{pmatrix}
\Sigma_1&O\\
O&\Sigma_2
\end{pmatrix}

としたら$(y_1, y_2)$は$N(\mu', \Sigma')$に従い、$I$を単位行列としたら$y_1+y_2$は$(y_1, y_2)$の$\begin{pmatrix}I&I\end{pmatrix}$倍が従う確率分布に従います。

次に、$y_1, y_2$がガウス分布の事後分布としたら、$y_1y_2$(要素積)がどのような値になるかを考えましょう。

どうやら局所的な線形化による近似をすることで対処をしたようです。
$y_1$は$N(\mu_1, \Sigma_1)$に、$y_2$は$N(\mu_2, \Sigma_2)$に従うとしたら、$y_1y_2$を$\mu_2y_1+\mu_1y_2-\mu_1\mu_2$と近似しました。
$\mu_2y_1+\mu_1y_2$の項は先に決めて、平均値が$\mu_1\mu_2$となるように定数項を調整しています。

このようにして、ガウス過程回帰の線形倍に話を帰着させることができ、学習を行えます。
ただし、積の部分は近似値なので、一発で事前分布やカーネル関数などのパラメーターの最適値を求めることはできず、繰り返し最適化をする必要が生じます。

最適化の目的関数

パラメーターごとに最適化の方法が微妙に異なりますが、事後分布をサンプリングして対数尤度を最適化することは多くのパラメーターで共通です。
(transit曲線の形状のハイパーパラメーターについては2位の方に聞いてみても本質的な話は無いと断られてしまいました。transitが生じたタイミングは導関数を見ることで最適化しています。)
基本的にはハイパーパラメーターはtrain dataで学習されていますが、transit depthの大きさのフィッティングはtest dataで行っています。
事前分布のノイズ以外の部分が$N(\mu, P)$、ノイズの部分が$N(0, N)$($N$は対角行列)とし、ノイズ無し信号の予測値を$x'$としたときに、
(事後分布の平均値を$N(\mu, P)$の対数密度関数に代入した値)
+(観測された信号を$N(x', N)$の対数密度関数に代入した値)
を対数尤度としています。

コードについて

クラスの継承関係はこうなっています

• Superclass
  • Observable
  • PriorMatrices
  • Model
    • FixedShape
      • FixedValue
    • Passthrough
      • ParameterScaler
      • Sparse2D
    • Compound
      • CompoundNamed
    • ValueHolder
      • FeatureSelector
        • FixedBasis
        • Uncorrelated
        • UncorrelatedVaryingSigma
        • SquaredExponentialKernelMulti

FeatureSelectorを継承しているクラスたちは共分散行列(このコードでは共分散行列の逆行列の精度行列を主に使っていますが…コードでの変数名はprior_precision_matrix)などを計算するクラスです。
ガウス過程の線形行列倍の行列名はdesign_matrix、定数項はobservable_offsetになっていて、PriorMatricesを継承する全クラスが持っています。
CompoundNamedクラスのmodelというプロパティには名前をkey、FeatureSelectorを継承したクラスをvalueとしたdictを入れます。
モデル毎に入力が異なるので、.featuresで指定しています。

時間計算量と空間計算量をともに減らすためにスパース行列を用いています。

先程から線形行列倍をするという話に帰着させ続けてきましたが、2D画像の入力に対して入力を間引いてガウス過程回帰に渡して出力では逆に線形補間をするという話の線形補間の部分を線形行列倍のところで実装できます。(どこで使っているかなどの詳細はNeurIPS - Ariel Data Challenge 2024の2位解法をご覧ください。)
また、出力で同じ値になる箇所に関しては必要になるまで同じ値を出力しないで必要になったら成分を増やすような実装を線形行列倍の部分でしています。(恐らく計算量を落とすため)

波長ごとにノイズの分散を変えるUncorrelatedVaryingSigmaクラスや、複数のハイパーパラメーターによるsquared exponential kernel(exponentiated quadratic kernel、Gaussian kernel、radial basis function kernelなどとも知られているらしいです。(引用元))の和をカーネル関数として使うSquaredExponentialKernelMultiクラスなどが実装されています(どこで使っているかなどの詳細はNeurIPS - Ariel Data Challenge 2024の2位解法をご覧ください。)

コレスキー分解は失敗することがあるので、失敗した場合は左右から対角行列でスケール倍をした後、対角成分に微小量を足してからもう一度コレスキー分解を試みる実装をしています。(robust_cholesky関数)
robust_cholesky関数はコレスキー分解後の行列を扱うクラスとスケール倍の行列を返すので、robust_cholesky関数の戻り値をcholQとしたらcholQ[1]@(cholQ[0].solve_A(cholQ[1]@rhs)のような実装が多数見受けられます。
(左右から掛けたスケール倍を戻すためにcholQ[1]が必要。そのクラスが扱っている行列をAとしたときにsolve_AはAx=bとなるxを返す関数)

残りはdiscussionを参照

元のdiscussionには他にも色々な説明がされています。
noise、star spectrum、drift、transitのモデルがどのようなカーネル関数を用いているかなどです。
実際にコードを読んでみると、元のdiscussionのsolutionで説明されている箇所も多く感じました。
この記事で書いたことが理解できていれば、元のdiscussionのsolutionを読むのもそこまで難しくないと思います！
NeurIPS - Ariel Data Challenge 2024の2位解法へのリンク

分からなかったこと

↓このツイートの返信部分に分からなかった箇所と何が関係しそうかの予想があるので、分かった方がいれば教えて頂けたら助かります。