AHC003の2.926T解法+経緯

https://atcoder.jp/contests/ahc003

![順位表](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1570357/7dfca79e-aa69-41b7-707f-c2533e04d53f.png)


人生初のAHC1位(2,926,767,743,437点: 97.56%)だったので記念に記事を書いてみます。

暫定テストの点数は97,506,167,171点でした。

初日に問題を見たとき、これは自分の得意分野に近いな、と思ったので真面目に取り組むことにしました。

## 問題

詳しくは [公式の問題文](https://atcoder.jp/contests/ahc003/tasks/ahc003_a) を参照してください。要点だけ整理すると、

* 辺の長さの分かっていない30×30のグリッドグラフ上の2点が1組ずつ与えられ、逐次経路を出すことが求められる
* 経路を出したらその経路の合計の長さに±10%の誤差が乗ったものが返ってくる
* 出した経路の合計の長さに対する最短経路長の比の重み付き平均が得点になる
    * 重みは後半ほど高くなる
* 真の辺の長さは、ランダムに決められるが、大雑把には
    * $M=1$: 60本の行と列(以下直線と言うことにします)それぞれの上の辺は基本大体同じ長さで、最大$\pm D$の一様ランダムな誤差がついている
    * $M=2$: それぞれの直線はランダムに2分割され、それぞれの側は基本大体同じ長さで最大$\pm D$の一様ランダムな誤差がついている
   
    の2パターンがそれぞれ1/2の確率で出現する。ただし、$D$も100~2000の範囲で一様ランダムに決められる。




## 最終解法

先に最終的な解法を書きます。この解法に至ったまでの経緯は後に書くので、そっちの方に興味がある人はそっちから読むと良いかも。

大雑把にやったこととしては

* メインは$2\times30\times29=1740$変数の(多変数)正規分布モデルによる辺の長さのベイズ推定
* 8通りの$(M, D)$の組で上の正規分布モデルを作り、対数尤度の値によってモデル間の事後分布を評価
* $M=2$で$D$が小さいと推定された時のみ、正規分布を仮定しないモデルでマルコフ連鎖モンテカルロ法(MCMC)も回してその結果を採用する
* 各辺の長さの事後分布の期待値と分散がもとまったら、最終的な辺のコストを$\text{(期待値)}-\gamma_t\sqrt{\text{(分散)}}$ のモデル間の事後分布による線型結合にして、Dijkstra法で最短路を求めて出力した
* ...というのを全部最後までやってたら時間が足りないので、事後確率が小さいモデルから枝刈りした

この節の以降の項では、それぞれのパートについて詳しく説明していきます。

### 正規分布モデル
$M, D$が与えられていることを仮定する。全ての$h_{i,j}, v_{i, j}$をまとめてベクトル$X$と書くことにする(1740次元)

$X$の期待値$x_0:=\mathbb E[X]$は各要素が5000のベクトル。
$X$の分散共分散行列$\Sigma_0:=\mathbb E[(X-x_0)(X-x_0)^T]$も厳密に求められる:

* 一直線上にない辺は独立なので、それらの間の共分散は0
* 一直線上にある2つの辺の端から数えたindexを$i, j (= 0...28)$とすると、それらの間の共分散は
  $$\begin{cases}\frac13D^2\delta\_{ij} + \frac13(4000-D)^2 & (M=1)\\\\
    \frac13D^2\delta\_{ij} + \frac13(4000-D)^2(1-\frac{1}{28}|i-j|) & (M=2)\end{cases}$$
  ただし、$\delta\_{ij}$は$i=j$の時に1、そうでない時に0をとる。

$X$は本来正規分布ではないはずだが、大胆に$X\sim \mathcal N(x_0, \Sigma_0)$と正規分布で生成されると仮定する。

$t$回目の経路を渡すたびに、長さ$y_t := e_t b_t$が返ってくる。$y_t$も本来は正規分布じゃないが、大胆に$y_t\sim\mathcal N(b_t, R_t)$の正規分布とみなす。(ただし、$R\_t = \frac13 (0.1 y_t)^2$)
$t$回目の経路が通る辺のところが1,そうでないところが0なベクトルを$c_t$とすると、$b_t = c_t^T X$である。

このモデルの下で、事後分布$p(X=x\mid y_1,...,y_t) = p(y_t \mid x)p(X =x\mid y_1,...,y_{t-1})$の計算は、次のように逐次的に行える(Woodburyの恒等式を用いた詳しい導出は省略して[^1]、答えだけ書く):

* 事前分布$p(x)$と観測分布$p(y_t \mid x)$を両方とも正規分布と仮定したので、全ての$t$について$p(x\mid y_1,...,y_t)$は正規分布$p(x \mid y_1,...,y_t) = \mathcal N(x; x_t, \Sigma_t)$で表せる。
* $$\begin{align}(S\_{yx})\_t &:= c\_t^T \Sigma\_{t-1} \\\\
(S_{xy})\_t &:= (S_{yx})\_t^T \\\\
(S_{yy})\_t &:= c_t^T \Sigma_{t-1} c_t + R_t
\end{align}
$$とおくと、
    $$\begin{align}x\_t &= x\_{t-1} + (S_{xy})\_t (S_{yy})\_t^{-1} (y_t - c_t x_{t-1}) \\\\
\Sigma\_t &= \Sigma_{t-1} - (S_{xy})\_t(S_{yy})\_t^{-1}(S_{yx})\_t
\end{align}$$
    と漸化式を用いて計算できる
* さらに、対数尤度$l_t:=\log p(y_t \mid y_1, ..., y_{t-1})$もあとで必要になるので計算しておくと、
      $$l_t = -\frac{1}{2}\left(\log(2\pi)^m\det(S_{yy})\_{t-1}+(y_t-c_tx_{t-1})^T(S_{yy})\_{t-1}^{-1}(y_t-c_tx_{t-1})\right)$$
      ここで、$m$は$y_t$の次元数で、今回は1。

上の計算を普通にやると、行列$\Sigma_t$が$n\times n$のサイズで、行列の引き算が$O(n^2)$かかる。これは嫌なので、$\Sigma_t$を直接持つのではなくて$(S_{xy})\_t$と$(S_{yy})\_t^{-1}$の形のまま持って、

```math
c^T_t\Sigma_{t-1} = c^T_t\Sigma_0 - \sum_{s=1}^{t-1}c_s^T(S_{xy})_s(S_{yy})_s^{-1}(S_{yx})_s
```
  を計算するようにした。今回$\Sigma_0$がブロック対角行列で$t$は1,...,1000なので計算時間が1/4ほどになるが、それよりも「$t$が若いうちは計算量が小さい」という性質が後の枝刈りと相性が良かった、ということの方が重要かも。

ここの掛け算が計算量のボトルネックなので、AVX512を使って速くしたら4倍程度になった。doubleではなくfloatを使うことで8倍になったが、予測精度が落ちたのでやめた。


[^1]: 様々な導出をしている文献があって驚かされますが、[このブログ](https://cookie-box.hatenablog.com/entry/2021/01/31/223006)の導出は比較的素直な導出だと思います(上のブログでは観測の平均が0と仮定してしまっているので$x_t$の更新則は出てこないのですが、同じように式変形することで観測の平均が0でない時も導出できると思います)。ここに書いた$S_{xy}, S_{yx}, S_{yy}$はまさにこのブログの導出で出てくる大きな行列のブロック成分となっています。制御工学や信号処理の分野だと同じ処理が[カルマンフィルタ](https://ja.wikipedia.org/wiki/%E3%82%AB%E3%83%AB%E3%83%9E%E3%83%B3%E3%83%95%E3%82%A3%E3%83%AB%E3%82%BF%E3%83%BC)と呼ばれていたりして、非常に広く使われています。
        
### モデル間の事後分布
$(M, D)$の組が$\\{1,2\\}\times \\{250, 750, 1250, 1750\\}$の計8通りの正規分布モデルを作り、対数尤度を用いてモデル間の事後分布を求めた。

$m=0...7$をモデルの番号とし、$L^{(m)}\_{t}:=\log p(y_1, ..., y_m \mid m) = \sum\_{s=1}^{t}l^{(m)}\_{s}$とおく。ただし、$l^{(m)}\_{s}$はモデル$m$での$y_s$の対数尤度$l\_{s}$である。

モデルの事前分布$P(m)$を一様分布としたので、各モデルの事後確率$P^{(m)}\_{t}:=P(m\mid y_1,..., y_t)$は、

```math
P_t^{(m)} = \frac{\exp L^{(m)}_t}{\sum_{m=0}^{7} \exp L^{(m)}_t}
```
で求められる。

![分類結果](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1570357/959e24cb-7594-c396-ee52-2f1bc941dc50.png)

$(M, D)$を指定して生成した入力に対して、テスターを走らせた後に最も事後確率が高いモデルをプロットするとこんな感じになった。そこそこいい感じに分離できていることがわかる。

### MCMC
$M=2$で$D$が小さいとき、本当は素早く辺の長さを学習できて欲しいのに正規分布モデルだとうまく学習できるまでに余分にデータが必要になっていたことに気づいた。正規分布を考える限り「辺の長さがジャンプするのが一箇所しかない」ということをうまく使えないことが原因だった。

![gaussian_only 2.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1570357/f9256762-b732-681f-2f07-85c808535097.jpeg)

![image.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1570357/a8265a22-167e-1b17-a8e1-53cdc617ac35.jpeg)
↑正規分布の限界: 慣れないとみづらいが、太い線が推定値、細い線が真値を表す。太い線の太さが推定分散(太いほど分散が大きい)を表す。これは100ステップ後の推定値を表す。
4行目に注目: 16列目で急激に辺の長さが変化しているのが正しいが、推定値はなだらかになってしまっている。


そこで、正規分布を仮定せずに辺の長さを推定できる方法として[MCMC(マルコフ連鎖モンテカルロ法)](https://ja.wikipedia.org/wiki/%E3%83%9E%E3%83%AB%E3%82%B3%E3%83%95%E9%80%A3%E9%8E%96%E3%83%A2%E3%83%B3%E3%83%86%E3%82%AB%E3%83%AB%E3%83%AD%E6%B3%95)、その中でも[ギブスサンプリング](https://ja.wikipedia.org/wiki/%E3%82%AE%E3%83%96%E3%82%B9%E3%82%B5%E3%83%B3%E3%83%97%E3%83%AA%E3%83%B3%E3%82%B0)を用いた方法を実装した。計算量が正規分布モデルよりも多くなるはずなので、さすがに1740変数は無理だと思い、180変数で我慢した。

使ったモデルは次の通り:

* $z_i \sim \mathrm{Uniform}(1, 28)$: 各直線$i$の区切れの位置
* $X_{1, i}\sim \mathcal N(5000, \frac{4000^2}{3})$: 各直線$i$の左(上)側の辺の長さ
* $X_{2, i}\sim \mathcal N(5000, \frac{4000^2}{3})$: 各直線$i$の右(下)側の辺の長さ
* 観測される経路の長さ$y_t$は平均が予測値、分散が$\frac13 (0.1 y_t)^2 + \frac13 1000^2(\text{経路の通った辺の数})$の正規分布に従う
    * 経路の通った辺の数の項は、$D$による分散の分を適当に近似したもの。

ガッツリ正規分布を使っているじゃないか、と思われるかもしれないが、$X_{1, i}$と$X_{2, i}$が正規分布でも$z_i$で区切れると指定しているので辺の長さ全体の同時分布が正規分布にはならない。

このモデルで、$X_{1, i}, X_{2, i}, z_i$の事後分布のサンプルを得るために次のようにギブスサンプリングを行った:

* まず何も情報がない時の初期値は$X_{1, i} = X_{2, i} = 5000, z_i = 15$に固定
* $y_1,..., y_t$が得られたとき、各$i=0...59$について以下を順に繰り返す
    * 直線$i$以外の辺の長さと区切り位置を現在の値で固定し、さらに$z_i$と$(X_{1, i}, X_{2, i})$を仮定した上で$y_1, ..., y_t$が得られる確率密度($(X_{1, i}, X_{2, i})$の尤度)は$(X_{1, i}, X_{2, i})$に関してはガウス関数になる
    * $(X_{1, i}, X_{2, i})$の事前分布も2変数正規分布なので、$z_i$を固定した上での$(X_{1, i}, X_{2, i})$の事後分布も2変数正規分布となり、解析的に$(X_{1, i}, X_{2, i})$を周辺化した確率($z_i$に関する尤度)が求まる。
    * $z_i$の事前分布は一様分布だったので、この$z_i$に関する尤度はそのまま$z_i$の事後分布になる
    * 各区切り位置$z_i=1...28$についてこの事後分布を求め、そこからサンプリングする
    * サンプリングされた$z_i$についての$(X_{1, i}, X_{2, i})$の事後分布から$(X_{1, i}, X_{2, i})$をサンプリングする
    * $(X_{1, i}, X_{2, i})$が1000...9000の範囲から超えていたら値をクリッピングする(本当は分布を歪ませる良くない方法。でも制約付きの2変数正規分布からサンプリングするのが面倒だったのでこうした)

これを$N$周回すと、$z_i, X_{1, i},X_{2, i}$の分布が求める事後分布に収束していくので、期待値と分散をとった。
何周回すかが問題になるが、これは$t\le 200$まで$N=10$にし、それ以降は$N=2$にしたら良かった。(こんなに少なくていいのかと自分も思ったが、このくらいじゃないと計算時間が収まらなかったしこれで点数が伸びたのでよしとした。)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1570357/9e52b69a-235b-4f04-893c-14352a72dca8.png)
MCMCだけを使った場合はこのように$M=2$かつ$D$が小さいときに顕著に効果が出る。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1570357/272c67c2-a583-91ed-85b5-eae92e728865.png)
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1570357/ba456d71-e0ee-b653-acb1-3aca3bccc7a4.png)

上が正規分布モデル(上に貼った図の一部)で下がMCMC。MCMCの方が正規分布モデルより鋭い推定をしていることがわかる。


ただし、$M=1$の時や$M=2$で$D$が大きいときに使うと逆効果だった(モデルが真のモデルから遠くなるのでそれはそう)ので、上のモデル間の事後分布で$(M, D)=(2, 250), (2, 750)$の確率の分だけMCMCを信用することにした。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1570357/259c1b51-07c9-8419-801f-7c8a5c46ba4b.png)
モデルを組み合わせた結果、このように点数が上がって欲しいところだけ上がってくれた。

本当はMCMCで尤度$l_t$も推定してモデル間ベイズ推定に組み込みたかったが、尤度を安定して推定できる自信がなかったのでモデル間ベイズ推定は正規分布モデルの尤度を使うことにした。

### 経路の出力
最終的な各辺のコストは、

```math
\sum_{m=0}^{7} P^{(m)}_t\cdot\left(x_t^{(m)} - \gamma_t \sqrt{\mathrm{diag} \Sigma_t^{(m)}}\right)
```
とし、この上でDijkstra法で最短経路を求めた。ただし、$(M, D)=(2, 250), (2, 750)$のモデルの分については期待値と分散はMCMCの出力で置き換えた。

$\gamma_t$は大きくするとまだコストがはっきりとわかっていない辺を優先的に辿るようになる。今回は$\gamma_t = 1.5(1-\frac{t}{1000})^2$を使った。

### 時間管理
8個の正規分布モデル+MCMCを全部最後まで回すのは時間が足りないので、次のようにモデルの枝刈りをした:

* モデルの対数尤度 $L^{(m)}\_t$ が最も良いモデルの対数尤度 $L^{(m^*)}\_t$ と比べて15以上小さくなったら、モデル$m$は捨てる
* それぞれのモデルの計算時間を大雑把に予測して、一つのモデルで最後まで計算する程度の時間しか残ってなかったらその時点で最も尤度が高いモデルに絞った
* MCMCは途中で計算をやめても後から再開できるので、事後分布が0.1%を下回ったら計算を保留しておいてまた0.1%を上回ったときに保留した計算を処理するようにした。

本当はこれで完璧に動いて欲しかったが、実際はMCMCの計算時間は動的に変わる(経路が通る直線の数が変わる)ので、これでもTLEしてしまうケースがあった。そのため実行時間が1.9秒を超えたらその時点で辺のコストを固定しDijkstra法だけを回してクエリに答えるTLE緊急回避モードを実装した。

## 試したがうまくいかなかったこと

### 「情報量」によるパスの選択
辺のコストを

```math
\sum_{m=0}^{7} P^{(m)}_t\cdot\left(x_t^{(m)} - \gamma_t \sqrt{\mathrm{diag} \Sigma_t^{(m)}}\right)
```
にするのは、一見ごもっともに見えるが、標準偏差の和は和の標準偏差ではないので、実はよくわからないことをしている(別に経路の合計の長さの (期待値$-\gamma_t$√分散) がこの値の和になるわけではない)。しかもこのコストは辺同士の共分散を考えていないので、「情報量はあるが分散が小さいわけではない」辺を誤って選択してしまう可能性がある。

分散以外の情報量の指標は、いくつか考えられる。真っ先に思いつくのはエントロピー(正規分布の場合は$\frac12\log \det \Sigma_t + \mathrm{const.}$)だが、これも適切とは言えない: 例えば、一つの変数の分散が0に極端に近づくことで、他の変数の情報が事前分布から全く更新されていなくてもエントロピーは$-\infty$に発散してしまう。「なんでも良いから情報をくれ」という時にはエントロピーを減らすことは役に立つかもしれないが、今は「なんでも良いから情報をくれ」という状況ではない。

より適切な指標として考えられるのは、ある予め決めておいた関数$f(x)$について、「$f(X)$の各成分の推定分散の和」という指標である。具体的には、例えば$f(x)=x$とすれば、$f(X)$の推定分散は$\Sigma_t$そのものなので、各成分の推定分散の和というのは$\mathrm{tr}(\Sigma_t)$となる。$\Sigma$の更新によってこれが最も減れば良いので、正規分布モデルの更新式から、それは

```math
\mathrm{tr}((S_{xy})_t(S_{yy})_t^{-1}(S_{yx})_t) = \frac{\|c_t^T \Sigma_{t-1}\|^2}{c_t^T \Sigma_{t-1} c_t + R_t}
```
が最も大きい$c_t$を考えることになる。(実際にはこの「情報量ゲイン」に$\gamma_t$をかけて経路の長さを足したものを経路全体のコストとする。)

残念ながら、この$f(x)=x$の場合の「情報量ゲイン」を1740次元の変数について計算するのは計算量が高く、dijkstra法の各更新で計算するのは大変だった。そのため、関数$f(x)$の次元を減らすことでなんとかできないか考えてみた。

結局やったのは以下の二つ:
1. $f(x)$を$x$の各直線での辺の長さの平均とする(60次元)
2. $f(x)$を$x$の各直線での辺の長さの「片側に偏った重み付き平均」と「もう片側に偏った重みつき平均」とする(120次元)

1の気持ちとしては、「各直線の辺の長さの平均が精度良く分かれば、マップを大体分かっていることになろう」という感じ。それれに対して2の気持ちは、「いやいや、$M=2$の場合とかは各直線の辺の長さの平均だけわかっててもダメで、左側の平均と右側の平均が分かってる必要があろう」という感じ。

このようにしても、上の式と似たような式が出てきて、それを最大化すれば良いとなる。ただし、次元数が減るので、計算量が減って現実的になる。

これらを経路のコストに加えて、Dijkstraを回すように試してみた。ただし、これらのコストは経路依存になってしまうので、Dijkstra法で最適解が求まらない。そのためここではDijkstra法で最初に見つけた(最適とは限らない)解で満足することにした。

結果として、確かに情報量が大きそうなパスの取り方をするようになり、推定分散が大きく残っている辺が少なくなり、スコアの分散が減った。しかし、同時にそれはコストが高そうな辺も分散が減るまで探索することを意味していて、スコアの平均値はむしろ下がってしまった。


### 初めに辺の長さを低めに見積もっておくこと

情報を得るために探索させる方法でもっと単純に考えられるものとして、初期推定値を低めに設定しておくということがある。そうすれば通った辺の推定値が上がってくれて、結果的に通っていない道を好むようになるだろう。

これは正規分布モデルが一つの時にはうまくいった。しかし、複数のモデルを使うようになると、このように嘘の事前分布をモデルに伝えると間違ったモデルを選択してしまうようだった。そのため、モデルに嘘はつかない方が良いと思って、正しい平均値を初期推定値にするようにした。

## 最終解法に至るまでの経緯

思い出話気分でできるだけ思考過程全てを書き下したので、長いです。そういうのが読みたい人向けです。

### Day 1

問題文を読んだ。今回はいつものマラソンみたいに焼きなまし大会じゃない、これは良問な匂いがする、と思って、興味を引かれた。あと問題のストーリーが好きだった。

今回の問題は、自分がロボコンサークルでやってきたオンラインの状態推定に問題設定が似ているな、と感じた。でも特に$M=2$の場合に正規分布から外れている度合いが強く、一筋縄ではいかなさそう。何より、真面目に解くと状態が1740次元になってしまってつらい。(注: 私が今までロボコンで扱ってきた問題は、観測の方が多くても基本的に状態は10次元とかせいぜいそのくらいで、$O(n^3)$の解法とかも余裕で使える世界にいた。しかも基本的に大体正規分布でよく近似できる場合しか扱ってこなかった。)

それでも、$M=1$の場合に限って考えれば問題は結構簡単になりそう、と思って、初submitは$M=1$に特化した解法を考えた。

まず、各直線に対して一つの平均的辺の長さを推定することにし、$D$による各辺の長さの誤差は観測誤差の方に吸収させる。そして平均的長さの事前分布と観測誤差の分布を正規分布と仮定すれば(これは$M=1$の時は多分そこまで悪くない仮定)、60変数の正規分布モデルで楽勝になる。

これを初日の夜に実装して、94,352,328,036点を得た(一応この時は$D$を固定した。また事前分布を低めに与えることで、見ていない辺を優先的に探索させた)。一瞬だけ暫定2位になれたが、すぐに他の人々に抜かされた。

提出した後、手元の評価scriptとoptunaを使ったパラチュンscriptの整備を始めた。その夜に寝ている間optunaを回し続けることに成功したが、手元の500ケースで評価したスコアは微小にしか上がらなかった。

(注: 終わった後の人の感想を見ると、500ケースで評価というのはずいぶん少なかったっぽい。でも手元で評価すると500ケース8並列でも数十分かかるし辛いんですよ... うん千ケースに関して評価をしたという人の計算環境が羨ましい)

### Day 2

$M=2$の場合にどうにかして対処したい。そうだ、辺の長さを各直線で定数と仮定するのではなく、両端の辺の長さを推定してその間を線形補間すれば良いのではないか。まぁ一箇所でジャンプしている形からお世辞にも近いとはいえないけど、定数よりはましかな。そう思って、120変数の正規分布モデルを実装した。

結果は96,297,636,439点。暫定1位になれた。(その後tomerunさんに抜かれることになる)



ここまでは$D$による誤差は観測誤差に含めて、完全に忘れてしまっていたが、$D$が大きいような問題では「この辺は長いが次の辺は短い」みたいな情報を覚えておくことが有利になるだろう。それを実装するには、1740変数のモデルが必要になる。

しかし、この時点で実装していたアルゴリズムの計算量は$O(n^2Q)$。これだと$n=1740, Q=1000$で2秒に間に合うか微妙。てか、普通のコンテストの感覚だと間に合わない。でも今回のボトルネックは行列演算であり、行列演算に関しては世の中の人々がめっちゃ頑張って高速化しているはずだ。だから一か八か、AVX512を使って行列演算を実装してみて間に合うか試してみようじゃないか。

ということで、AVX512のintrinsicsを学びながら、必要な行列演算を実装した。手元のPCが古くAVX512が動かなかったので、AtCoderのコードテスト環境をありがたく使わせていただいて、ベンチマークを行った。すると、floatの行列演算がなんと8倍早くなるじゃないですか！！！

これだと時間が間に合いそう、と興奮して1740変数の正規分布モデルを実装した。初期分散$\Sigma_0$は、とりあえず$M=2$の場合決め打ちで出した。結果は97,196,447,062点で、また暫定１位に復活することができた。(その後ハイパラチューニングをすることで97,243,242,935点が得られた。)

### Day 3

しかし、初期分布の分散$\Sigma_0$が$M=2$の場合決め打ちだと、$M=1$の場合により速く推定できるはずなのにそれをうまく使うことができない。

実行時間が0.85秒程度で時間の余裕が多少あったので、正規分布モデルを$M=1$に対応する事前分布と$M=2$に対応する事前分布で2つ回して、尤度を使ってモデルの間でもベイズ推定することにした。

しかし、なんか常に片方のモデルが選択されてしまう。初めの方でモデルの事後分布が大きく傾いてしまい、そこから修正できていないような症状だった。

その原因はどうやら最初に盤面を探索させるために初期推定値を低くしていたことっぽい。嘘のモデルを使うのは数学的に美しい正規分布モデルの中で大きな汚点になるだけでなく、モデルの尤度を大きく歪ませてしまうのだった。

より筋が良い方法で盤面を探索させることはできないか？

せっかく分散も推定しているんだから、分散が大きいところを優先的に探索させれば良いじゃないか、と思いついた。ということで、辺の長さから分散の平方根(単位を合わせるため)の定数倍を引いたものを2つのモデルの確率分布を使って線形結合し、それを辺のコストとした。($\gamma_t$に相当するものであるが、この時点では定数にした。)

これをして嘘のモデルを使うことをやめたら、うまくモデルが推定できていそうで感動した。97,266,801,933点。


そこからさらに前に進むためには、二つのモデルが推定している辺の長さと分散を表示するビジュアライザーが必要だな、と思ったので、その日の残りの時間はビジュアライザーを作ることにした。私はRustは数回しか書いたことがなかったが、それでもWeb版のビジュアライザのコードに30行くらい追加するだけでできるならそれが早そうと思ったので、そうした。結局いろいろ(wasmの作り方を1から勉強したり、Arithmetic overflowの罠に引っ掛かったりとか)苦労して6時間くらいかかってしまった。


![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1570357/86e5b447-5c6f-40d0-7763-2d735966d93b.png)
↑成果物(最終解法の図ではありません)

その後、なんとかして$M$だけでなく$D$の推定もできないかと夜明けまで考察してしまった。

### Day 4

まず$D$の推定をするのにあたって、決定的に足りないものはやはり時間だったので、定数倍高速化を頑張ることにした。$\Sigma_t$を陽に持たなくすれば今回の場合計算量が大体1/4になることに気づいたので、それを実装した。

これは等価なコードの書き換えをしただけなはずで、実際に手元で数ケース実行してみても点数は変わらなかったので、良さそうと思って提出した。しかし、点数が下がってしまった。

手元でテストする時はAVX512を使わずに等価なより遅い計算で頑張るので、AVX512のコードにバグがあるのではないか、ということを真っ先に疑った。しかし、AtCoderのコードテスト環境でテストするかぎり、そんなことはなさそうだった。

手元で評価スクリプトを回して点数が減ったかどうか確かめる。すると、なんと手元でも点数が下がっていた。確かに、大体のケースにおいては厳密に出力が一致する。しかし、20ケースに一つくらいの割合で、2行程度だけ出力が違うものがある。もしやと思って、ソースコードでfloatと書かれているところをdoubleに置換して、もう一度提出してみた。すると、なんと元のコードの点数を超える点数(97,277,276,885点)が出た。笑顔になった。元のコードでも、floatの計算誤差によって精度が落ちていたのだ。

floatをdoubleにして2倍程度の時間がかかるようになってしまったが、それでも4つくらいのモデルを並列に走らせることができるようになった。よし、と思いながら、それぞれの$M$について2通りの$D$でモデルを作成し、試してみた。しかし、一向に手元のテストで結果が良くならない。なぜだろうと思っていろいろパラメタを調整したが、結局良かったのは$M=1$と$M=2$で一つずつモデルを作ることだった。

いろいろ悩んだ後、もしかして$\gamma_t$が定数なのが悪いのかもしれない、と思い至った。適当に$\gamma_t = \frac{1.5}{1 + t/300}$みたいなスケジュールを指定してみたら、97,399,878,369点にまで上がった。


複数のモデルでベイズ推定をすると、途中からモデルの事後分布が1:1e-10などと大きく偏ることが多いことがわかった。このような場合、1e-10まで事後分布が落ちてしまったモデルは計算量を使って推定を続ける意味がほぼない。ある程度まで事後確率が下がってしまったものを完全に捨てることによって、初めは4つだけでなくより多くのモデルがあったとしても時間内に計算を終わらせることができるはず、と気付いた。

もちろん事後確率が下がらない場合も考慮しないといけない。そういうときはギリギリのところで強制的に最も良いモデルを選ぶ必要がある。それをするために、AtCoderのコードテスト環境で毎回の計算時間のベンチマークをし、一回の計算でどれくらいの時間がかかるかを大雑把に推定できるようにした。

そして、ついに8つのモデルを持っても時間内に走らせることができた。しかし、手元でテストしてみると、なぜか点数がまた下がってしまった。いろいろ試した結果$(M, D) = (2, 1750)$のモデルを除いた時が手元の500ケースでは前のものより良かったので、これを提出したが、暫定テストの点数は下がってしまった。それでも暫定テストは100ケースしかないから手元の方を信用することにした。

### Day 5

モデルとしては明らかに$(M, D) = (2, 1750)$があった方が良いはずなのに、なぜ手元のテストだとない方が良いことになるのか、ずっと気になっていた。$\gamma_t$の形が悪いのかといろいろ試したが、$(M, D) = (2, 1750)$を抜いたモデルになかなか勝てなかった。

特に$D$が大きい場合、テストケースごとの分散が大きかったので、それが問題では?、全てのモデルで$\gamma_t$が共通なのは無理があるのでは？と考えて、各モデルごとに$\gamma_t$の値を調整してみた。しかし、全然良くならなかった。

ビジュアライザーを見ると、どうやら最後までほとんど通らず、推定分散が大きかった辺が実は短い辺だったりすると点数が悪くなるらしい。もっと強い探索手法が必要なのかもしれない。もしや嘘のモデルを使うことはモデル選択には悪いが探索には良いんだとしたら、モデル選択用の正しいモデルと探索用の嘘のモデルの両方を持っておけばうまくいくのでは?と思ったが、それも微妙だった。

他に良い探索手法はないか？ということで、上に書いた「情報量によるパス選択」のアルゴリズムを考えた。

### Day 6

「情報量によるパス選択」のアルゴリズムを実装した。greedyなdijkstra法を実装する必要があったので、実装は結構重かった。しかし、あまりうまくいかなかった。

ビジュアライザーを見ると、確かに情報量は増えていて明らかに最後まで分散が大きく残る辺が減ったが、同時に「長そうな辺だからやめとこ」と敬遠していたところも探索して潰していたので、そこでの犠牲が大きかったのかもしれない。しかし、各モデルごとに別のパラメタを調整することでどうにかなる気がしてしまったので、optunaを使って頑張った。1日くらいかけたが、結局あんまり良くならなかった。

### Day 7

なぜうまくいくモデルとうまくいかないモデルがあるのかを分析するために、デフォルトの入力ジェネレータではなく **既知の$(M, D)$から生成された入力** を使うべきだ、ということに気付いたのはこの日だった。もっとずっと前に気づくべきだった。

この日はモデルの分析の環境を整えるのに使った。入力ジェネレータを改造して各$(M, D)$の組について4つずつ入力データを作成した。モデルの評価を並列で行えるようにscriptを修正し(それまでは1つずつ走らせていて遅かった)、さらに評価中に*リアルタイムで*$D$と得点の関係がプロットされるようにした。リアルタイムでプロットされると、500個全ての評価が終わるはるか前にモデルの傾向がわかって捗った。

自分で$(M, D)$を指定した入力を食わせてみると、$D$が大きい場合の得点の分散が酷いのは仕方なくて、どっちかというと$M=2$で$D$が小さい時の得点が思ったより伸びないことが問題であることがわかった。

これは辺の長さが正規分布を仮定すると逃げられない制約だと思った(最終解法のMCMCの節の図を参照)。正規分布だけでなく他の分布でも扱える推定方法としてMCMCが浮かんだが、自分は一回も実装したことがなかったので自信はなかった。あと2日しかないときに完全に新しいことをやることに躊躇いも感じたが、ここまではっきりと問題が見えてきたらやらなきゃ後悔すると思って次の日にMCMCを実装することを決意した。

### Day 8

MCMCを実装した。初めは上述のMCMCのアルゴリズムのみで答えを出すプログラムを評価してみた。案の定、$M=2$で$D$が小さい場合の性能がもとよりかなり上がり、$M=1$の時の性能はかなり下がった。

これはDay 4に実装したモデルに組み合わせると有意に性能が上がりそうだな、と思ってテンションが上がった。しかしDay 4のモデルはまだ$(M, D)=(2, 1750)$を抜いた汚いモデルだったので、もしやと思って$(2, 1750)$を戻して一度提出してみたら、なんと暫定テストではこっちの方が97,448,033,028点と元のモデルを大幅に上回る点数を出してしまった。手元の500ケースが運が悪かったか、暫定テストの100ケースが運が良かったかもしれない。

せっかくなので、$(2, 1750)$を入れた綺麗な方のモデルにMCMCをマージすることにした。

すると、見事に手元のテストではこれまでのモデルを大きく上回って97.65Gの点数がでた。興奮した。

最後に残った問題は計算量だった。MCMC単体で1秒かかっていて、愚直に正規分布モデルと組み合わせたのでAtCoderのコードテストで試してみたら2.8秒の時間がかかってしまっていた。これを2秒以内に抑えるために、MCMCの計算時間をベンチマークしてうまく正規分布モデルの時間管理の機構に乗っける必要があった。

それっぽく書いて、提出したが、TLEを食らった。どうやらMCMCは計算量があまり一定ではなく、入力によって変わるみたい。使って良い合計「計算量」のパラメタを減らして余裕を持たせようとしたが、またTLEを食らった。

仕方なかったので、使った時間が1.9秒を超えたときにTLE緊急回避モードが発動するようにした。するとACをもらえて、最初に書いた97,506,167,171点が得られた。97.5Gを超えることができて、満足した。

https://twitter.com/contramundum2/status/1398715655436656640

### Day 9

人々の最後の追い上げをビクビクしながら見守っていた。結局何もしなかった。

### 終わった後、他の人の解説を見て

2位のyosssさんのコードとtwitter上の解説を見て、私の使った正規分布モデルとほぼ同じことをしていて、仲間だ〜と思った。

5位のprime_numberさんを含め、(正則化付き)最小二乗法っぽい方法を使っている人を多数観測した。正則化付き最小二乗法も上に書いた正規分布モデルで言う$\Sigma_0$に単位行列などを入れたものとして解釈できるので、仲間。私がAVX512とか使って高速化を頑張っているところを[CG法](https://ja.wikipedia.org/wiki/%E5%85%B1%E5%BD%B9%E5%8B%BE%E9%85%8D%E6%B3%95)を使って素早く解いているの、感心した。

4位のtoameさん、8位のgasinさんを含め、通った辺に差分を更新する方針の人も結構いたっぽい。これも最小二乗法を勾配法などで解いていると解釈することもできるのかも(そこまでちゃんと考えたわけではないのでそうではないかもしれない)。 周辺の辺の長さと平均化するのも、$\Sigma_0$に単位行列以外の分散共分散行列を指定していると解釈できるかも。

3位のKomakiさんがNeural networkでこの点数を出しているのを見て、とても驚いた。6位のsaharanさんを含めSigmoidなどの非線形関数を使ってモデルを組んでいた人も多そう。非線形なモデルを組むとどうやって最適化するかが重要になってくるわけだが、確率的勾配法やAdamのような亜種を使った人がいたり焼きなましている人がいたり普通に勾配法を使う人がいたり、という感じでしょうか？ 非線形なモデルを最適化するのは私はコンテスト中に思いつかなかったが、確かに試してみるべきだったかもしれない。

## 感想

今回の問題は賛否両論ありそうな気もしますが、私にとっては**神回**でした。もちろん自分の得意分野に近かったということに多少影響されているかもしれないですが、それを抜きにしても非常に良い問題だったと思います。

というのも、この問題、一応用いるアルゴリズムの分野としては統計・機械学習ということになるはずですが、全然機械学習の問題で典型的という感じではなく、むしろとても"AtCoder的な"機械学習の問題だと思うからです。

普通の機械学習の問題では、データが一部与えられ、ほぼ無制限の計算資源を使ってそのデータを学習し、残された検証データに対して予測した答えを提出し、その精度で評価される、というような問題が多いと思います。そのような問題においては、「いかにデータにあったモデルを作るか」ということが本質となるはずです。しかし、そのような問題は現実世界に存在するデータこそが正義であり、より"AtCoder的な見方"からすれば若干ill-definedな問題という感じがしてしまうかもしれません。

一方、今回の問題は最初から完全に正しいモデルが与えられています。もし仮に十分な時間が与えられていたとすれば、推定パートは(1740+α)変数のMCMCなどでほぼ自明に解けてしまう(事後分布が完全に求まる)と言っても良いと思います。それでも問題として成立するのは、他ならぬ2秒の制約があるからで、だからこそ例えば私の場合は正規分布を使って近似したりMCMCの変数の数を減らして近似したりして「サボる」必要がありました。統計や機械学習の概念を使いながら、結局のところ高速な近似アルゴリズムを構成しろというマラソンの原点に帰着されるのは、まさに「機械学習の分野上でのAtCoderの問題」のような感じがしてとても好きでした。

もちろん個人的にもMCMCを初めて実装できたとか、AVX512のintrinsicを初めて勉強できたとかoptunaを初めて使ってみれたとか様々なことを学ぶことができました。Day 1からDay 8までずっと参加しましたが、全く「人生の8日間を無駄にした」という気持ちにならないのは非常に良かったと思います。

長文を最後まで読んでいただきありがとうございました。間違い等があればぜひご指摘お願いします。