Torch를 이용한 matrix inverse(+속도를 빠르게 만들 수 있는 경우)

matrix inverse를 하고자 할 때 numpy든 torch든 간단하게 linalg.inv 함수를 사용하면 끝나니 크게 생각해본 적 없겠지만 matrix inverse는 연산량이 굉장히 큰 operation이다.

 

간단히만 봐도 O(n^3)이므로 matrix 사이즈가 커질수록 폭발적으로 연산량이 증가한다.

 

대략 2000 x 2000 보다 큰 크기의 matrix를 다루기 시작하면 영겁의 시간이 걸린다. 

 

이 글에서는 matrix inverse 방식 별 소요 시간을 대충 비교해보고, matrix 크기가 클 때 조금이라도 inverse operation에 들어가는 시간을 줄일 수 있는 노하우를 공유한다. (설명 전에 미리 말하면, 모든 matrix가 가능한 것은 아니다.)

 

1. Gauss-Jordan elimination based

torch를 이용해서 직접 구현한 matrix inverse다. 선형대수 책을 보면 가장 먼저 나오는 내용으로 Gauss-Jordan method을 이용해서 line by line 처리를 반복하여 inverse를 구한다. 

 

가장 느린 방법이다.

def matinv(A):
    b, N = A.size(0), A.size(1)

    A_float = A.clone()
    A_inv = torch.eye(N).cuda().unsqueeze(dim=0).repeat([b, 1, 1]) # b N N

    # To upper triangular matrix
    for i in range(N-1):
        pivot = A_float[:, i:i+1, i:i+1]
        scalers = -A_float[:, i+1:, i:i+1] / pivot
        A_float[:, i+1:, :] += scalers * A_float[:, i:i+1, :]
        A_inv[:, i+1:, :] += scalers * A_inv[:, i:i+1, :]

    # Diagonalize
    for i in range(N-1, -1, -1):
        pivot = A_float[:, i:i+1, i:i+1]
        scalers = -A_float[:, :i, i:i+1] / pivot
        A_float[:, :i, :] += scalers * A_float[:, i:i+1, :]
        A_inv[:, :i, :] += scalers * A_inv[:, i:i+1, :]

    # Last scaling
    for i in range(N):
        pivot = A_float[:, i:i+1, i:i+1]
        A_float[:, i:i+1, i:i+1] /= pivot
        A_inv[:, i:i + 1, :] /=  pivot

    return A_inv

 

2. LU decomposition based

np.linalg.inv 나 torch.linalg.inv 가 사용하는 방식으로 주어진 matrix를 LU decomposition한 뒤 이를 이용해서 inverse를 구하는 방식이다. 

 

이 방식은 2/3 n^3 + O(n^2)의 연산량을 가지며, Gauss-Jordan 보다는 빠르다. 최적화도 어련히 더 잘 해두었을 것이라 빠르다.

 

torch.linalg.inv(A)

 

3. Cholesky decomposition based

일반적으로 LU decomposition이 가장 간단한 방식으로 알려져있지만 특수한 경우에는 이보다 더 간단하게 구할 수 있다.

 

특수한 경우는, matrix가 symmetric positive semi definite 성질을 만족할 때이다.

 

PSD(positive semi definite)가 무엇인지 기억이 안난다면 대충 matrix가 positive element로만 이루어져있고 diagonal element가 0이 아닐 경우, 그냥 이 방법을 시도해보면 된다. 안되면 말고~

 

symmetric PSD matrix는 일반 matrix보다 극도로 깔끔한 matrix이기 때문에 decomposition이 더 가볍다. 그 방식이 Cholesky decomposition인데 1/3 n^3 + O(n^2)의 연산량일 가진다.

 

2. 방식 대비 1/3 n^3 수준의 연산량 차이이므로 n이 충분히 크다면 엄청나게 연산량을 줄일 수 있다. 

 

u = torch.cholesky(A)
A_inv = torch.cholesky_inverse(u)

 

속도 비교

256x256 크기 matrix 대상

• Gauss Jordan elimination based : 38.69 ms
• LU decomposition based : 1.60 ms
• Cholesky decomposition based: 1.21 ms

512 512 크기 matrix 대상

• Gauss Jordan elimination based : 178.60 ms
• LU decomposition based : 11.20 ms
• Cholesky decomposition based: 5.38 ms

크기가 클수록 효과가 압도적으로 차이나는 것을 볼 수 있다. 

 

웬만하면 torch.linalg.inv로 끝날텐데 혹시 다루는 matrix가 symmetric PSD 라면 무조건 cholesky inverse를 사용하길 추천한다!