Weighted Reservoir SamplingとStreaming RIS

このページではサンプルの未知のサイズのストリーム入力から、シングルパスで各サンプルのウェイトに比例した確率でサンプルを選択することができるWeighted Reservoir Samplingというアルゴリズムについて紹介します。また、Resampled Importance Samplingと組み合わせることによって生まれるStreaming RISという新たなサンプリングアルゴリズムについても紹介します。

離散分布に沿ったサンプリング

$ M $ 個のサンプル $ x_1, \ldots, x_M $ と対応するウェイト $ w_1, \ldots, w_M $ があるときに、ウェイトの大きさに比例した確率で一つのサンプルを選択する場合、アルゴリズム1に示すような手法が使われます。サンプルサイズが予め決まっていてサンプルの選択も入力がすべて終わってからで良いという場合はこの処理で問題が無いのですが、入力が追加されるたび逐次ウェイトに比例した選択を行いたい場合には毎回入力列を巡回して選択する処理 $ O(N) $ が必要になるため非効率です。CDFを作っておいて逐次更新すれば一回あたりの選択の計算量は $ O(\log N) $ まで落ちますが、それでも毎回その計算量が必要となるのは非効率です。また各サンプルとウェイトを保持しておく必要もあるため(入力列の長さが未知の場合は特に)メモリの使用量も問題になりえます。

function $ \mathrm{sampleDiscreteDistribution} $($ M $ 個のサンプル $ x_1, \ldots, x_M $ とウェイト $ w_1, \ldots, w_M $, $ [0, 1) $ の範囲の一様乱数 $ u $) {
	ウェイトの合計を初期化 $ \wsum \LAR 0 $
	for $ i $ in $ [1, M] $
		ウェイトを加算 $ \wsum \LAR \wsum + w_i $

	$ [0, \wsum) $ の範囲の乱数を得る $ u \LAR \wsum \cdot u $
	ウェイトの累積値を初期化 $ w_{\mathrm{acc}} \LAR 0 $
	for $ i $ in $ [1, M] $ {
		ウェイトを累積 $ w_{\mathrm{acc}} \LAR \wsum + w_i $
		if $ w_{\mathrm{acc}} > u $
			return $ x_i $
	}
}

未知の要素数の入力からのサンプリング

未知の長さのサンプルのストリーム入力から、ウェイトに応じた確率でサンプルを選択する手法にWeighted Reservoir Sampling (WRS)と呼ばれるものがあります。アルゴリズム2にその概要を示します。この手法ではReservoir (発音はレゼヴァーやレゼボアといった感じ)と呼ぶデータ構造に今までの入力のウェイトの総和とひとつのサンプルだけを保持します。新たなサンプル(とウェイト)が入力されるたびに、確率的にサンプルの置き換えを行います。新たなサンプルに置き換わる確率は、ウェイトの総和に対する新たなサンプルウェイトの比によって決定されます。

このように非常に簡単なアルゴリズムで未知の長さのストリーム入力に対してサンプルのウェイトに比例した確率でサンプルを残す、選択することができます。メモリ使用量も非常に少なく済みます。欠点としてはサンプルが入力されるたびに乱数生成が必要なため、決まった長さの入力からサンプルを1回選択するだけで良いのであればアルゴリズム1よりも計算量が多くなること、また過去に入力されたサンプルやウェイトは逐次棄てているため、2度目のサンプリングはできないことが挙げられます。

class $ \mathrm{Reservoir} $ {
	ウェイトの合計 $ \wsum \LAR 0 $
	現在のサンプル $ x_z \LAR \emptyset $

	function $ \mathrm{update} $(新しいサンプル $ x_i $, ウェイト $ w_i $, $ [0, 1) $ の範囲の一様乱数 $ u_i $) {
		ウェイトを累積 $ \wsum \LAR \wsum + w_i $
		if $ u_i < w_i \;/\; \wsum $
			サンプルを更新 $ x_z \LAR x_i $
	}
}

Reservoirの結合

Reservoirの面白く有用な性質として複数のReservoirを適切に結合することで、それぞれの入力列を復元することなく、すべてを結合した入力列からウェイトに応じた確率でサンプルを選択したことに等しくなることがあります。アルゴリズム3にReservoirの結合を示します。これまた非常に簡単で、それぞれのReservoirが保持しているサンプルとウェイトの合計値を、単純に新たな一つのサンプルとウェイトが入力されたかのごとく処理するだけです。

function $ \mathrm{combineReservoirs} $ ($ N $ 個のReservoir $ R_1, \ldots, R_N $ と一様乱数 $ u_1, \ldots, u_N $) {
	結合結果のReservoir $ R_{\mathrm{c}} $
	for $ i $ in $ [1, N] $
		$ R_{\mathrm{c}} \,. \mathrm{update} $($ R_i \,. x_z $, $ R_i \,. \wsum $, $ u_i $)

	return $ R_{\mathrm{c}} $
}

Streaming RIS

Resampled Importance Sampling (RIS)は、複数のサンプリングを実行、それぞれの候補サンプルのウェイトの評価を行い最終的に1つのサンプルを選択(、そしてImportance Samplingに使用)するというアルゴリズムでした。単純な実装では候補サンプル数が多いときに各候補のサンプルとウェイトを保持しておくためのメモリが問題でした。WRSはまさにこの特性・問題にぴったりとあてはまります。RISにWRSを組み合わせた手法をStreaming RISと呼びます[Bitterli2020]。Streaming RISに使用するReservoirをアルゴリズム4に示します。基本的にはアルゴリズム2と同じですが、無駄な再計算をしないためと、後ほど紹介する発展的な応用のために、一部の値のキャッシュなど少し拡張してあります。

class $ \mathrm{RISReservoir} $ {
	ウェイトの合計 $ \wsum \LAR 0 $
	現在のサンプル $ x_z \LAR \emptyset $
	目標分布の値 $ \hat{p}_z \LAR 0 $
	ストリーム長 $ M \LAR 0 $

	function $ \mathrm{update} $(新しいサンプル $ x_i $, ウェイト $ w_i $, 目標分布の値 $ \hat{p}_i $, $ [0, 1) $ の範囲の一様乱数 $ u_i $) {
		ウェイトを累積 $ \wsum \LAR \wsum + w_i $
		ストリーム長を更新 $ M \LAR M + 1 $
		サンプルを採用するか否か $ b_{\mathrm{accept}} = u_i < w_i \;/\; \wsum $
		if $ b_{\mathrm{accept}} $
			サンプルと目標分布の値を更新 $ x_z \LAR x_i $, $ \hat{p}_z \LAR \hat{p}_i $
		return $ b_{\mathrm{accept}} $
	}

	function $ \mathrm{calcReservoirWeight} $() {
		return $ \displaystyle \frac{1}{\hat{p}_z} \Prt{\frac{1}{M} \wsum} $
	}
}

アルゴリズム5にこのReservoirを用いたStreaming RISを示します。RISでは目標分布と提案分布の値の比をウェイトとして候補サンプルを生成、最終的にひとつのサンプルを選択します。ここではそれらのウェイトと候補サンプルをReservoirに逐次入力、入力が終わればReservoirが保持しているサンプルはウェイトに比例した確率で選ばれたものになっています。最後に選ばれたサンプルを使って被積分関数の値を評価、Reservoir全体のウェイト(重点的サンプリングにおけるPDFの逆数に相当)を計算すれば積分結果の推定値が得られます(サンプルサイズ1)。

function $ \mathrm{performStreamingRIS} $() {
	Reservoir $ R $
	for $ i $ in $ [1, M] $ {
		提案分布 $ q $ からサンプル $ x_i \sim q $
		目標分布の値を計算 $ \hat{p}_i = \hat{p}(x_i) $
		ウェイトを計算 $ w_i = \hat{p}_i \;/\; q(x_i) $
		$ R . \mathrm{update} $($ x_i $, $ w_i $, $ \hat{p}_i $, $ \mathrm{rand} $())
	}

	被積分関数の値を計算 $ f = f(R . x_z) $
	PDFの逆数の推定値を計算 $ W = R . \mathrm{calcReservoirWeight} $()
	return $ f \cdot W $
}

異なる目標分布に対して作られたReservoirの結合

function $ \mathrm{combineReservoirs} $ ($ N $ 個のReservoir $ R_1, \ldots, R_N $ と一様乱数 $ u_1, \ldots, u_N $, 新しい目標分布 $ \hat{p}_{\mathrm{new}}(\cdot) $) {
	結合結果のReservoir $ R_{\mathrm{c}} $
	合計のストリーム長を初期化 $ M_{\mathrm{c}} \LAR 0 $
	選択されたサンプルを生成したReservoirのインデックス $ s $
	for $ i $ in $ [1, N] $ {
		Reservoirのウェイト $ W_i = R_i \,. \mathrm{calcReservoirWeight} $()
		目標分布の値 $ \hat{p}_i = \hat{p}_{\mathrm{new}}(R_i \,. x_z) $
		目標分布の違いを考慮したウェイトを計算 $ \displaystyle w_i = \hat{p}_i \cdot W_i \cdot R_i \,. M $ ($ \displaystyle = \frac{\hat{p}_{\mathrm{new}}(R_i \,. x_z)}{R_i \,. \hat{p}_z} \cdot R_i \,. \wsum $)
		if $ R_{\mathrm{c}} \,. \mathrm{update} $($ R_i \,. x_z $, $ w_i $, $ \hat{p}_i $, $ u_i $)
			$ s \LAR i $
		$ M_{\mathrm{c}} \LAR M_{\mathrm{c}} + R_i \,. M $
	}
	ストリーム長を修正 $ R_{\mathrm{c}} \,. M \LAR M_{\mathrm{c}} $

	return $ R_{\mathrm{c}} $, $ s $
}

function $ \mathrm{performMultiReservoirStreamingRIS} $ ($ N $ 個のReservoir $ R_1, \ldots, R_N $ と目標分布 $ \hat{p}_{\mathrm{org}, 1}, \ldots, \hat{p}_{\mathrm{org}, N} $) {
	新たな目標分布 $ \hat{p}_{\mathrm{new}}(\cdot) $
	$ R_{\mathrm{c}} $, $ s $ = $ \mathrm{combineReservoirs} $($ R_1, \ldots, R_N $, $ \Brc{\mathrm{rand}(), \ldots, \mathrm{rand}()} $, $ \hat{p}_{\mathrm{new}}(\cdot) $)

	MISウェイトの分母を初期化 $ \mathrm{denomMisWeight} \LAR 0 $
	for $ i $ in $ [1, N] $
		$ \mathrm{denomMisWeight} \LAR \mathrm{denomMisWeight} + \hat{p}_{\mathrm{org},\, i}(R_c \,. x_z) \cdot R_i \,. M $
	// 現在のサンプルを生成したReservoirの目標分布からMISウェイトの分子を計算
	$ \mathrm{numMisWeight} = \hat{p}_{\mathrm{org},\, s}(R_c \,. x_z) $
	MISウェイトを計算 $ \displaystyle m = \frac{\mathrm{numMisWeight}}{\mathrm{denomMisWeight}} $

	被積分関数の値を計算 $ f = f(R_c \,. x_z) $
	PDFの逆数の推定値を計算 $ \displaystyle W = \frac{1}{\hat{p}_{\mathrm{new}}(R_c \,. x_z)} \Prt{m \cdot R_c \,. \wsum} $
	return $ f \cdot W $
}

Streaming RISはReSTIR (Reservoir-based Spatio-Temporal Importance Resampling) [Bitterli2020, Wyman2021]における根幹のアルゴリズムとして使われます。

参考文献

• [Bitterli2020] Benedikt Bitterli, Chris Wyman, Matt Pharr, Peter Shirley, Aaron Lefohn, Wojciech Jarosz - "Spatiotemporal reservoir resampling for real-time ray tracing with dynamic direct lighting", 2020
• [Marrs2021] Adam Marrs, Peter Shirley, Ingo Wald - "RAY TRACING GEMS II: NEXT GENERATION REAL-TIME RENDERING WITH DXR, VULKAN, AND OPTIX", 2021, Apress
• [Wyman2021] Chris Wyman, Alexey Panteleev - "Rearchitecting Spatiotemporal Resampling for Production", 2021