Pytorch Distributions Shape 의 이해(Understanding Shapes in Pytorch Distribution)

 Pytorch에서는 다양한 distribution에 대해서 지원을 하고 있습니다 (https://pytorch.org/docs/stable/distributions.html). 최근 일을 하면서 distribution을 정말 많이 사용하고 있는데 정리가 잘 되어 있으면 좋겠다는 생각에 공부도 할 겸 포스팅으로 남기려고 합니다. 

 최근 저는 transformer 를 활용을 하면서 decoder를 거치지 않고 encoder의 latent vector를 사용해서 distribution을 통해 추론하고자 하는 값을 sampling 하는 방식을 사용하는 경우가 많이 있습니다. 이때, multivariate 인 경우와 univariate 인 경우에 따라서 Pytorch의 distribution을 다양하게 선택하고 활용해야하는데 그때마다 찾아보는 것도 어렵고 docs는 복잡한 경우가 많아서 이 기회에 저같이 어려움을 겪는 분들이 없기를 바라면서 정리를 시작하려고 합니다.

 이 포스팅에서는 docs에서 자주 사용되는 3가지 용어에 대해서 살펴보려고 합니다. 

 

Three types of shapes 

---

 Pytorch Distribution에 대한 document를 보면 제일 많이 등장하는 단어 3개가 sample shape, batch shape, event shape 입니다. 그런데 너무 당연하게 안다는 것을 가정하고 docs에서는 설명을 하고 있어서 이해하기가 매우 어렵습니다. 아래 그림에 잘 설명이 되어있지만 시계열의 ML문제로 가정을 해서 조금 더 쉽게 설명을 하자면 다음과 같습니다.

1. Sample shape : batch size의 수와 같은 개념
2. Batch shape : multivariate time series 문제로 가정하고 보면, 택시의 수요 예측을 위해서 홍대라는 지역의 "시간별 유동인구수", "시간별 버스운행 수", "시간별 택시운행 수"라는 3개의 시계열이 입력으로 들어갈 때, 각각의 시계열 입력은 독립적 (independent)하지만 모두 다른 distribution argument (ex. mean 과 std)로 부터 값이 발생할 것이므로 "not identical" 한 distribution에서 나왔다고 이해할 수 있습니다. 즉, 위 예시의 경우에 batch_shape는 "3"이 됩니다. 
3. Event shape : 하나의 시계열 데이터가 차지하는 dimension의 크기. "시간별 유동인구수"가 5명, 3명 이렇게 1개의 값을 가질 수 도 있지만 (남자 2명, 여자 3명), (남자 1명, 여자 2명) 이렇게 2차원씩으로 구성될 수도 있다고 이해할 수 있습니다. 이런 경우 event shape는 "2"가 되겠죠.

Figure 1: Three groups of shapes. Reproduced from Dillon et al. (2017).

 

 위의 설명도 이해가 안가실 수 있는데, 저도 처음에는 잘 안갔습니다. 그래서 제가 distribution 공부를 하면서 제일 좋았던 그림을 가져와봤습니다. 주의하실 점은 위에서 설명드린 1, 2, 3의 순서가 그림의 (-) 순서와 반대라는 것에 주의하면서 살펴보셔야 합니다.

 Event Shape인 1개의 데이터가 가지는 dimension을 의미한다고 말씀드렸습니다. 아래의 그림에서 Draws가 가로로 붙으면 이 Event Shape가 증가하는 것을 의미합니다.  Tensor로 가정하면 (1, )이 (2, ) 와 같이 증가하게 되는 거라고 생각하시면 됩니다.

 똑같이, 세로로 붙으면 Sample Shape가 1개 증가하는 것을 의미하며 다시말해 1개의 샘플을 더 만들어 낸 것이므로 (S, B, E)가 (1, ?, 1) 이었던 경우가 (2, ?, 1)로 변했다고 생각하시면 됩니다.

뒤로 붙는 경우에는 Batch Shape가 증가하게 됩니다.  

 그러면, Pytorch Distribution에서 제공하는 Normal과 MultivarateNormal 그리고 Independent Class를 사용해서 어떻게 위의 shape들이 변하는지 살펴보도록 하겠습니다. 

 Normal Distribution의 input은 loc과 scale입니다. 아래와 같이 코드를 입력하게 되면 loc은 (2, ) shape를 가진 Tensor, scale은 (4, 1) shape를 가진 Tensor를 할당해 준 것입니다. 그러면, 이렇게 생성된 Normal Distribution의 batch_shape와 event_shape는 어떻게 될까요? 

n = Normal(torch.randn(2), torch.randn((4, 1)))
print(n.batch_shape, n.event_shape)

#torch.Size([4, 2]) torch.Size([])


from torch.distributions import Gamma, Beta
n = Gamma(torch.randn(2), torch.randn((4, 1)))
print(n.batch_shape, n.event_shape)

n = Beta(torch.randn(2), torch.randn((4, 1)))
print(n.batch_shape, n.event_shape)

# torch.Size([4, 2]) torch.Size([])
# torch.Size([4, 2]) torch.Size([])

 정답은 (4, 2)와 () 입니다. 우선 Normal Distribution을 사용하게 되면 pytorch는 Univariate distribution을 가정하기 때문에 event_shape는 무조건 ()가 됩니다. 그러므로, loc과 scale에서의 dimension들은 모두 batch_shape로 할당되게 됩니다. 

 그런데, 왜 loc은 (2, ) 이고 scale이 (4, 1)인데 batch_shape가 (4, 2)가 되는거지? 라고 생각하실 수 있습니다. 이는 code를 보면 이해하실 수 있습니다. pytorch의 Normal은 내부적으로 log_prob이나 sample 계산시에 loc / scale.pow(2) 와 같이 loc과 scale을 곱하거나 나누는 연산을 수행합니다. 이때, pytorch는 row vector와 column vector의 연산을 수행하게 되는데, 이때 아래와 같이 연산의 결과가 생성됩니다. 

 그러므로, Normal Distribution이든 어떤 Univariate 를 가정하는 Distribution을 사용하게 된다면 위와 같은 batch_shape를 가지게 됩니다. 

row_vector = torch.Tensor([1.0, 2.0])
column_vector = torch.Tensor([[3.0],[4.0]])

row_vector * column_vector

#tensor([[3., 6.],
#        [4., 8.]])

row_vector / column_vector

#tensor([[0.3333, 0.6667],
#        [0.2500, 0.5000]])

 그렇다면, MultivariateNormal은 어떻게 다를까요? 기본적으로 MultivariateNormal에서는 Event_shape가 어떠한 값을 가질 수 있다는 차이가 있습니다. Normal의 scale 위치에는 covariance matrix (공분산 행렬)을 넣어주게 됩니다. Normal 의 경우에는 torch.eye(5)를 넣는다고 하더라고 event_shape가 ()인 것을 알 수 있습니다. 하지만, MultivariateNormal은 event_shape가 (5,) 가 됩니다. 

mvn = MultivariateNormal(torch.randn(5), torch.eye(5))
print(mvn.batch_shape, mvn.event_shape)

# torch.Size([]) torch.Size([5])

n = Normal(torch.randn(5), torch.eye(5))
print(n.batch_shape, n.event_shape)

# torch.Size([5, 5]) torch.Size([])

 그렇다면, 1-dimensional matrix를 covariance matrix위치에 넣을 수 있는지 궁금할 수 있습니다. 만약 1-dim의 matrix를 입력으로 주게되면 아래와 같이 오류를 반환합니다. 

mvn = MultivariateNormal(torch.randn(5), torch.randn(5))
print(mvn.batch_shape, mvn.event_shape)

# covariance_matrix must be at least two-dimensional, with optional leading batch dimensions

 만약 batch_shape를 주면서 event_shape도 주고 싶다면 다음과 같이 줄 수 있습니다. loc에 해당하는 matrix의 row 갯수는 batch_shape가 되며, column의 수가 event shape가 되는 것을 알 수 있습니다. 이때, scale matrix 자리에는 항상 square matrix가 들어가야한다는 것을 다시한번 잊지 말아야 합니다. 

mvn = MultivariateNormal(torch.randn((3, 4)), torch.eye(4))
print(mvn.batch_shape, mvn.event_shape)

# torch.Size([3]) torch.Size([4])

mvn = MultivariateNormal(torch.randn((3, 4)), torch.randn(4, 2))
print(mvn.batch_shape, mvn.event_shape)

# A must be batches of square matrices, but they are 2 by 4 matrices

 

Independent Class

---

 그런데, 가끔 구글링으로 pytorch distribution에 관련된 코드를 검색하다보면 Independent class를 활용하는 코드들이 등장합니다. 실제로 mxnet이나 pytorchTS에서는 Univariate Distribution 을 생성한 뒤에 Indepedent Class를 사용해서 multivariate distribution으로 변환시키는 방법을 사용하고 있습니다. 

 Independent Class는 어떻게 동작하는 것인지 간단히 살펴보도록 하겠습니다.

 먼저 Normal 을 다음과 같이 선언해서 batch_shape가 (3, 5)인 분포를 얻습니다. Indepedent는 인자로 base_distribution과 reinterpreted_batch_ndims를 인자로 받습니다. 이때 두번째 인자인 reinterpreted_batch_ndims는 base distribution의 batch_shape를 뒤에서부터 몇개나 event_shape로 가져올 것이지를 의미합니다. 

 다시말해, base_distribution의 batch_shape가 (3, 5) 일때 reinterpreted_batch_ndims가 "1"이라면 뒤에서 부터 1개인 (5)를 가져온다는 것을 의미하며, reinterpreted_batch_ndims가 "2"라면 뒤에서부터 2개인 (3, 5)를 가져오게 됩니다. 

n = Normal(torch.randn((3, 5)), torch.randn(5))
print(n.batch_shape, n.event_shape)

# torch.Size([3, 5]) torch.Size([])

normal_ind_1_ind_1 = Independent(n, 1)
print(normal_ind_1_ind_1.batch_shape, normal_ind_1_ind_1.event_shape)

# torch.Size([3]) torch.Size([5])

normal_ind_1_ind_1 = Independent(n, 2)
print(normal_ind_1_ind_1.batch_shape, normal_ind_1_ind_1.event_shape)

# torch.Size([]) torch.Size([3, 5])

---------------------------------------------------------

n = Normal(torch.randn((1, 3, 5)), torch.randn(5))
print(n.batch_shape, n.event_shape)

# torch.Size([1, 3, 5]) torch.Size([])

normal_ind_1_ind_1 = Independent(n, 3)
print(normal_ind_1_ind_1.batch_shape, normal_ind_1_ind_1.event_shape)

# torch.Size([]) torch.Size([1, 3, 5])

 위와 같은 특성을 통해서 주로 Multivariate Distribution 을 구현하고자 할때 Univariate Distribution을 정의하고 neural network 모델의 output을 우리가 원하는 event_shape가 마지막 차원이 되도록 구현한 뒤에 실제로 distribution을 거친 output을 구하는 시점에 Indepedent(univariate_distr, 1) 을 추가하는 방법으로 하나의 batch 가 multi-dimesion을 가진 것으로 변환시키는 방법으로 사용되게 됩니다.

if is_1d:
  dist = Normal(0., 1.) 
else:
  dist = Independent(Normal(torch.zeros(dim), torch.ones(dim)), 1)

References

---

https://bochang.me/blog/posts/pytorch-distributions/

https://ericmjl.github.io/blog/2019/5/29/reasoning-about-shapes-and-probability-distributions/