リザーバーサンプリング - コンピュータフリーク

リザーバーサンプリング（貯水池サンプリング）とは、巨大なデータ集合から一様ランダムなサンプルをオンラインで抽出する方法です。データがメモリに乗らない場合や、データサイズが予め分かっていない場合にも適用することができます。

問題設定

データの列 $\{x_1, x_2, \cdots, x_t, \cdots\}$ を前から順番に観測していきます。$x_t$ は一行のテキストや、一件のログデータなどに対応します。ある時刻 $T \ge r$ において、サンプルを下さいという問い合わせが与えられるので、$\{x_1, \cdots, x_T\}$ の中から $r$ 件の一様ランダムサンプルを抽出して出力してください。件数 $r$ は予め決まっていますが、いつ問い合わせが来るかの $T$ は知らされません。また、$T$ は非常に大きいため、$\{x_1, \cdots, x_T\}$ を全て保存することはできません。

一様ランダムサンプルについては、非復元抽出、つまり重複のないサンプリングを主に考え、最後に復元抽出の場合も考えます。

応用例: ログデータの解析

ウェブページへのアクセスログを監視しているとします。非常にアクセス数の多いページなので、全データを保存することはできません。適当なタイミングで過去の訪問者を分析したくなり $r = 100$ 件程度のランダムなログが要求される場合が考えられます。

応用例: 機械学習用のデータ保持

システムのログを元に機械学習モデルを構築しようとしています。しかし、ログの量があまりに多く、全データを使ってモデルを訓練することができません。$r = 10000$ 件程度のランダムなログをサンプリングしてからそれらを用いてモデルを訓練する場合が考えられます。

応用例: オンラインアルゴリズムの構成要素として

オンラインアルゴリズムとはデータが順次与えられ逐次的に処理していくアルゴリズム一般を指します。まさにリザーバーサンプリングの問題設定もこの範疇に入ります。より複雑高度な問題設定において、まずリザーバーサンプリングで一様サンプリングを行い、得られたサンプルに対してバッチ処理をかけることで、バッチアルゴリズムをオンラインアルゴリズムに変換する方法がしばしば行われます。リザーバーサンプリングは一様サンプリングであることが保証されるため、この方法による理論性能が解析しやすいことが特徴です。

解法: 時刻 $T$ が予め分かっている場合

時刻が予め分かっている場合は非常に簡単です。

まず最初に $\{1, 2, \cdots, T\}$ 上の一様分布から $r$ 件 i.i.d. サンプリングを行い、対応するインデックスのデータを保存して最後に出力すればよいです。

解法: ナイーブな解法（アルゴリズム R）

アルゴリズム R は最も単純なオンラインアルゴリズムです。計算速度がネックでない場合にはこのアルゴリズムで十分な場合も多いです。

アルゴリズム R では、リザーバー（貯水池）とよばれるサンプル集合 $\mathcal{R}_t$ を保持し、各時刻 $t$ において、$\mathcal{R}_t$ は $\{x_1, x_2, \cdots, x_t\}$ からの一様サンプルであるようにします。データの流れが雨に相当して、その一部分を貯水池に貯めているようなイメージです。問い合わせが来た時点で $\mathcal{R}_T$ を出力すれば完了です。問題は $\mathcal{R}_t$ をどのように構築するか、というところにあります。

時刻 $r$ までは全てのデータを使うことになるので単純に $\mathcal{R}_t$ に要素を追加していけばよく、時刻 $r$ の時点ではちょうど $\mathcal{R}_r = \{x_1, x_2, \cdots, x_r\}$ となります。これが帰納法のベースステップに対応します。

時刻 $t > r$ の開始時点で $\mathcal{R}_{t-1}$ が $\{x_1, x_2, \cdots, x_{t-1}\}$ からの一様サンプルから成り立っているとしましょう。データ $x_t$ が到着したとき、このデータが $\mathcal{R}_t$ に含まれるべき確率は $\frac{r}{t}$ です。よって

確率 $\frac{r}{t}$ で表が出るコインを振る
裏が出れば何もせず $\mathcal{R}_t = \mathcal{R}_{t-1}$ とする
表が出れば、$\mathcal{R}_{t-1}$ の要素を一つランダムに抽出し、$x_t$ で置き換えて $\mathcal{R}_t$ とする。

とします。このとき、明らかに $x_t$ が $\mathcal{R}_t$ に含まれる確率は $\frac{r}{t}$ です。また、$x_i ~(i < t)$ については、$$\begin{align*}\text{Pr}[x_i \in \mathcal{R}_t] &= \text{Pr}[x_i \in \mathcal{R}_{t-1} \land \text{置き換えられない}] \\ &\stackrel{\text{(a)}}{=} \frac{r}{t-1} \text{Pr}[\text{置き換えられない} \mid x_i \in \mathcal{R}_{t-1}] \\ &\stackrel{\text{(b)}}{=} \frac{r}{t-1} (1 – \frac{r}{t} \cdot \frac{1}{r}) \\ &= \frac{r}{t}\end{align*}$$ となります。ここで、(a) は $\mathcal{R}_{t-1}$ が $\{x_1, x_2, \cdots, x_{t-1}\}$ からの一様サンプルであることから、(b) は、コインが表かつ $r$ 要素の中から選ばれたときに置き換えられることから従います。

よって、$\{x_1, x_2, \cdots, x_t\}$ のそれぞれが $\mathcal{R}_t$ に含まれる確率は一様となり、帰納法より常に $\mathcal{R}_t$ には一様にサンプルが含まれることになります。

解法: ナイーブな解法（優先度付きキュー）

優先度付きキューを用いた解法では、データ $x_t$ を観測するたびに、$[0, 1]$ の一様乱数をサンプリングし、これを $x_t$ の優先度としてキューに追加します。もしキュー内の要素が $r$ 件を超えれば、最も優先度の小さい要素を破棄します。優先度は完全に独立に定まるので、最終的に残る要素が一様ランダムであることは明らかです。実装が単純であり、乱数の分布をデータごとに変えることで重み付きサンプリングに拡張ができることも嬉しい点です。計算量は優先度付きキューの保持コスト分だけわずかに劣ることとなります。

解法: 高速な解法（アルゴリズム L）

アルゴリズム R や優先度付きキューを用いた解法では、$t$ が非常に大きいとき、新しい要素がリザーバーに追加される確率は非常に小さくなります。たとえばアルゴリズム R では、コインが表になる回数の期待値は $$\begin{align*}\mathbb{E}[\text{コインが表になる回数}] &= \sum_{t = r+1}^T \frac{r}{t} \\ &\le \sum_{t = 2}^T \frac{r}{t} \\ &\le r \log T \end{align*}$$ であるので、ほとんどの結果が裏であることが分かります。それにも関わらず、毎回コインを振るのは無駄な労力です。

アルゴリズム L では、次に新しい要素がリザーバーに追加される時刻を表す分布からサンプリングを行い、その時刻までのデータを全て読み飛ばします。

ベースとなるアルゴリズムは優先度付きキューを用いたアルゴリズムです。時刻 $r$ の時点で $\mathcal{R}_r = \{x_1, x_2, \cdots, x_r\}$ とし、これを帰納法のベースステップとします。

時刻 $t$ でデータが追加されたとし、優先度付きキューに含まれる優先度の最小値が $q \in [0, 1]$ であるとしましょう。このとき、新しいデータの優先度が $q$ を上回りリザーバーに追加される確率は $1 – q$ です。つまり、

次にリザーバーに要素が追加される時刻が $t + 1$ である確率は $1 – q$
次にリザーバーに要素が追加される時刻が $t + 2$ である確率は $q (1 – q)$
次にリザーバーに要素が追加される時刻が $t + 3$ である確率は $q^2 (1 – q)$
…

となります。これは幾何分布に他なりません。よって、時刻 $t$ において幾何分布からサンプリングを行い、得られた値を $k$ とすると、時刻 $t+1, t+2, \cdots, t+k-1$ のデータは読み飛ばし、時刻 $t+k$ のデータを処理すれば良いことになります。これにより効率が大幅に向上します。

また、上記のプロセスにおいては、各データの優先度の値自体は保持する必要がありません。必要なのは、優先度付きキュー内の最小の優先度 $q$ のみです。よって、アルゴリズム L では、各データの優先度を保持することも、優先度付きキューを保持することも無く、リザーバー $\mathcal{R}$ と（仮想的な）最小の優先度 $q$ のみを保持し、

幾何分布 $\text{Geom}(1 – q)$ から $k$ をサンプリング
$\mathcal{R}$ の要素の要素を一つランダムに抽出し $x_{t + k}$ で置き換える
優先度 $q$ を更新
$t \leftarrow t + k$

という等価で、より単純な処理で置き換えます。優先度 $q$ の更新については、$q$ 以上であるという条件の下での $r$ 個の要素の最小値の分布であるので、一様分布の極値分布からのサンプリングにより行うことができます。

復元抽出の場合

復元抽出、つまり重複ありで一様サンプリングを行う場合を考えます。

最も単純な方法は、リザーバーの大きさを $r = 1$ とし、上述のアルゴリズムを並列に複数実行するやり方です。明らかに一様復元抽出になっており、非復元抽出の場合のプログラムが存在する場合には流用できるので実装も容易です。計算量についてもアルゴリズム L を使えば非常に効率良くなります。

必要なサンプル数 $r$ が比較的大きい場合には、アルゴリズム R を改良した以下のような方法も考えられます。まずはナイーブにリザーバーの大きさを $1$ としたアルゴリズム R を並列に動作させることを考えます。データ $x_t$ が到着したとき、確率 $\frac{1}{t}$ のコインを $r$ 回振ることになります。しかし、アルゴリズム L の冒頭で述べたように、これはほとんどの場合裏となります。そこで、