[CV] Hidden dimension이 너무 클 때 flatten 하지 말고 똑똑하게 layer 추가하기

요즘 모델 아키텍처적인 설계를 많이 하구있다. 사실 내가 가장 못하는 구현이기도 하다. ㅜㅜ 연구를 점점 해나갈수록 내가 상대적으로 잘하는 것, 못하는 것이 슬슬 구분되기 시작하는데 이건 내가 못하는 것에 속한다 ㅎㅋ.. 

예전에 친한 언니가 딥러닝은 차원맞추기라고 하는 얘기를 들었던 적이 있는데, 요즘 공감되구 있다. 

---

암튼 ViT 아키텍처를 고치다가, weight matrix를 중간에 추가해야할 일이 생겼다. 그런데 dimension이 맞지 않는거다. 예를 들어서 아래와 같은 상황이다. 

• [1, 65536, 64] 텐서에 특정 weight matrix를 곱해서, 결과적으로 [1,8] 크기의 행렬을 만들어야 함 

 

# Flatten? 

나는 아키텍처 설계에 익숙치 않기 때문에, 위와 같은 상황에서 처음에는 아래와 같이 생각했다. 

• [1,65536,64]의 텐서를 flatten 해서 [1, 65536x64] 행렬을 만든다. 
• 그리고 [65536x64, 8]의 weight matrix를 만든다. (nn.Parameters)
• 위 두개 행렬을 행렬곱(@) 한다. 
• [1, 8] 크기의 행렬을 얻는다. 

위 방법으로도 당연히 계산이 가능하다. 하지만 문제점은.. [65536x64, 8] 이 어마어마한 크기의 파라미터를 만들거냐는거다. 특히 transformer의 경우 block이 굉장히 많은데, 나는 block마다 위와 같은 연산을 했어야 하는 상태였다. 

따라서 업데이트 해야하는 파라미터 수가 폭발해 OOM이 발생하는 이슈가 있었다 ㅋㅋ

---

# 대신 nn.Linear 추가로 해결하기 

위와 같은 무식한 방법 말고, 조금 더 효율적이지만 같은 결과를 내는 방법을 발견해 기록한다. nn.Parameter로 파라미터를 직접 지정하지 말고 ,아래와 같은 nn.Linear layer의 가중치를 이용하는 것이다. (누군가에겐 당연한 방법일 수 있지만, 아키텍처 변형 초보인 나에겐 나름 센세이션 했다) 

• [1, 65536, 64]를 [65536, 64]로 flatten 한다. 
• 아래와 같은 nn.Linear layer를 만든다. (Input node = 64, output node = 8) 

self.f_gate = nn.Sequential(nn.Linear(64, 8, bias=False))

• 위 layer에 [65536, 64]를 통과시킨다. 그럼 [65536, 8] 크기의 행렬이 출력된다. (64x8 = 512개만의 파라미터를 이용한다!)
• [64, 65536] x [65536, 8] = [64, 8] 행렬이 나오고, 크기 [1, 8]을 만들기 위해 필요에 따라 argmax 등을 추가한다. 

 

그럼 아주 미라클적이게도 33554432 개의 파라미터가 512개로 줄어든다!! 필요에 따라 nn.Linear layer를 적절히 추가해 이용해보자. 

마지막으로 같은 결과를 얻은 모습 ㅜㅜ