Linear層から出した分散共分散行列をMultivariateNormalに渡すときに ~ is zero, singular U のエラーが出る


やったこと

pytorchはv1.5.1を使用.
Linearの出力を(x, y, vxx, vyy, vxy)として,分散共分散行列を作りMultivariateNormalに渡す.

fc = nn.Linear(n, 5)
output = fc(x)
mean = output[:2]
vxx, vyy = nn.Softplus()(output[2:4])
vxy = output[-1]

covariance_matrix = torch.zeros(2, 2)
covariance_matrix[0, 0] += vxx
covariance_matrix[0, 1] += vxy
covariance_matrix[1, 0] += vxy
covariance_matrix[1, 1] += vyy

dist = MultivariateNormal(mean, covariance_matrix)

RuntimeError: cholesky_cuda: For batch 0: U(6,6) is zero, singular U.

MultivariateNormalでは分散共分散行列をコレスキー分解しているため,正定値行列を与える必要がある.このままではcovariance_matrixが正定値行列であることを保証していないので上記エラーが起こる.

解決法: コレスキー分解後の下三角行列を渡す.

fc = nn.Linear(n, 5)
output = fc(x)
mean = output[:2]
a, c = nn.Softplus()(output[2:4])
b = output[-1]

L = torch.zeros(2, 2)
L[0, 0] += a
L[1, 0] += b
L[1, 1] += c

dist = MultivariateNormal(mean,  scale_tril=L)

scale_tril (Tensor) – lower-triangular factor of covariance, with positive-valued diagonal

とあるので対角成分の$a, c$はsoftplusで正の値にしている.

分散共分散行列にするときにはコレスキー分解

\Sigma = LL^{T}

に従い,

covariance_matrix = np.dot(L, L.T)

とすればよい.