eubinecto/the-clean-transformer

multi-head attention, split head 구현

jl749 opened this issue · 3 comments

jl749 commented

안녕하세요

보통 multi-head attention 구현시 N, L, H --> N, L, heads, H//heads
처럼 머리를 나누는것으로 알고 있습니다

image

위 이미지 처럼 루프와 concat 을 통한 구현할시 문제가 있을까요?
값은 정상적으로 출력됩니다...

++ 추가 (개인적인 생각입니다..)
헤드를 나눈다면 grouped convolution 처럼 속도 개선이 가능하다고 생각합니다
하지만 저희가 원하는건 attention 이 다르게 적용된 다양한 feature map 을 확보하고 concat 후 의미값들을 linear 레이어를 통해 다시 함축하는것인데

헤드를 나눌시 같은 H 에 다양한 질문을 던지는 방식보단 각각 다른 H sub parts 에 대해 각기 다른 질문들을 던지는 것이기에 (hidden vector size 를 질문의 갯수에 비유 한다면)

이렇게 구현한다면 single head attention 이랑 다른 점이 있을까 라는 생각이 들었습니다

eubinecto commented

안녕하세요

보통 multi-head attention 구현시 N, L, H --> N, L, heads, H//heads 처럼 머리를 나누는것으로 알고 있습니다

image

위 이미지 처럼 루프와 concat 을 통한 구현할시 문제가 있을까요? 값은 정상적으로 출력됩니다...

++ 추가 (개인적인 생각입니다..) 헤드를 나눈다면 grouped convolution 처럼 속도 개선이 가능하다고 생각합니다 하지만 저희가 원하는건 attention 이 다르게 적용된 다양한 feature map 을 확보하고 concat 후 의미값들을 linear 레이어를 통해 다시 함축하는것인데

헤드를 나눌시 같은 H 에 다양한 질문을 던지는 방식보단 각각 다른 H sub parts 에 대해 각기 다른 질문들을 던지는 것이기에 (hidden vector size 를 질문의 갯수에 비유 한다면)

이렇게 구현한다면 single head attention 이랑 다른 점이 있을까 라는 생각이 들었습니다

안녕하세요! 이런 조용한 리포에 갑자기 질문 이라니! 감사합니다 ㅎㅎ

위 이미지처럼 루프와 concat을 통한 구현시 문제가 있을까요?

image

루프와 concat을 통한 구현시 논리적인 문제는 전혀없습니다. 다만 병렬화가 불가능하다는 기술적인 문제가 발생합니다. For loop을 명시하게 되면 동일한 연산을 GPU로 동시처리하고 싶어도 할수가 없습니다. 하지만 동일한 루프를 행렬연산으로 정의하면 (이를 "vectorization" 이라고 부릅니다) pytorch/tensorflow 등의 힘을 빌려 GPU에서 병렬처리할 수 있습니다. 때문에 For loop을 행렬연산으로 정의하는 방법이 있다면 항상 그렇게 정의를 해주는 편이 학습을 훨씬 빠르게 할 수 있도록 도와줍니다. 그게 Best Practice에요 :)

        result = AttentionLayer(self.hidden_size, self.encoding_size, self.masked)(q, k, v)  # (N, L * heads, E)
        for _ in range(self.heads - 1):
            head = AttentionLayer(self.hidden_size, self.encoding_size, self.masked)(q, k, v)  # (N, L, E)
            result = torch.cat((result, head), dim=2)

아, 그리고 구현하신 코드를 한번 살펴보았는데 (위), 구현 방법에 오류가 있어요. 이렇게 forward가 호출될때마다 AttentionLayer 객체를 생성하게되면 사실상 AttentionLayer 속 가중치는 계속 랜덤한 값으로 재생성돼요. 최적화가 전혀 이루어지지 않습니다. 레이어를 생성하는 것은 MultiHeadAttentionLayer 객체를 생성할 때 멤버변수로 등록할 때 한번하는 것이고, forward에서는 등록된 가중치로 연산해야 올바르게 최적화가 됩니다. 예를 들면 다음과 같이 구현하셔야합니다:

def __init__(self, ...):
  ...
  self.layers = torch.nn.ModuleList([
            AttentionLayer(...)
            for _ in range(heads)
        ])

def forward(self, ....):
  ... 
  results = list()
  for layer in self.layers:
     head = layer(q, k, v)  # (N, L, E)
     results.append(head)

물론, 저라면 이렇게 루프를 쓰기보단, 행렬연산으로 병렬화를 할 것이지만요! ㅎㅎ

각각 다른 H sub parts 에 대해 각기 다른 질문들을 던지는 것보단 같은 H에 다양한 질문을 던지는 것이 낫지 않을까요?

Inductive Bias의 당위성을 파고드는 아주 좋은 질문이네요. 우선, 각기 다른 H sub parts에 각기 다른 질문을 던지는 것은 singlehead attention과는 다릅니다. Single head attention은 각기 다른 H sub parts에 동일한 질문을 던진다는 점에서 달라요.

한편 연산시간에 대한 고민은 차치했을 때, 같은 H에 다양한 질문을 던지는 것이 더 적절할수는 있어요. 하지만 "히든"벡터이기 때문에 정말로 더 적절한지는 논리적으로 판단할 수 없습니다. H개의 각 차원이 무엇을 뜻하는지는 그 누구도 몰라요. 어차피 각 차원의 의미를 모른다면 새로운 가중치를 추가하는 것보단 이미 존재하는 것을 재활용하도록 유도하면, 더 적은 가중치로 같은 H에 다양한 질문을 던지는 것과 동일한 상태에 이를 수 있는 가능성은 있습니다.

  • 🚀1
eubinecto commented

@jl749 답이 되었다면 이슈를 닫아주시고, 더 궁금한게 있으시면 계속 코멘트 달아주세요~~! ㅎㅎ

jl749 commented

와 정말 감사합니다!
트렌스포머의 장점이 병렬화가 가능하다는건데 GPU 병렬연산에 대해서는 전혀 생각해보지 않았네요
궁금점이 한번에 정리 되었습니다!!

  • 👍1