어텐션 잔차 연결 설명 — 트랜스포머 깊이를 낭비하지 않는 새 구조와 대규모 학습 인프라 설계

핵심 메시지

• 잔차 연결은 기울기가 지나갈 고속도로를 만들어 깊은 신경망 학습을 가능하게 했지만, 현대 트랜스포머에서는 깊이를 제대로 쓰지 못하게 하는 부작용이 있다.
• 층마다 출력이 계속 더해지면서 잔차 스트림의 크기가 커지고, 새 층의 기여가 상대적으로 희석되는 현상을 프리놈 희석(pre-norm dilution)이라고 부른다.
• 각 층이 바로 앞 층의 벡터 하나만 본다는 점에서 잔차 연결은 순환 신경망의 은닉 상태 병목과 구조가 같다. 어텐션이 RNN을 개선했듯 잔차 연결도 개선할 수 있다는 것이 이 연구의 출발점이다.
• 고정값 1 대신 학습 가능한 어텐션 가중치로 이전 층 출력들을 섞으면 검증 손실이 낮아지고 층별 기울기 크기가 고르게 맞춰진다.
• 논문은 같은 손실에 도달하는 데 필요한 연산이 20% 줄어드는 1.25배의 연산 이점을 보고했고, 어텐션 잔차를 쓰면 넓은 모델보다 깊고 좁은 모델이 유리해진다.

쉽게 이해하기

표준 트랜스포머는 층의 출력을 입력에 그대로 더하는 잔차 연결을 쓴다. 이 식을 재귀적으로 펼쳐 보면 깊은 층의 표현은 얕은 층의 표현에 그 사이 모든 변환의 합을 더한 형태가 된다. 역전파 관점에서는 곱셈으로 줄어들지 않는 직통 경로가 생기므로 기울기 소실이 완화되고 아주 깊은 모델도 학습할 수 있게 된다.

문제는 현대 트랜스포머가 학습 안정화를 위해 프리 레이어 정규화를 쓴다는 데서 시작된다. 층마다 보정이 계속 얹히면서 표현의 크기가 깊이에 따라 자연스럽게 커지고, 각 층의 기여는 불어난 잔차 스트림에 비해 점점 작아진다. 그 결과 깊은 층은 작은 기울기를, 초기 층은 상대적으로 큰 기울기를 받으며 신경망의 깊이가 제대로 활용되지 못한다.

또 다른 문제는 정보 병목이다. 각 층은 직전 표현 하나만 보고, 그 이전에 일어난 모든 일은 벡터 하나로 압축돼 있다. 이는 순환 신경망이 지금까지 본 것을 은닉 상태 하나로 요약하느라 겪는 병목과 같은 구조다. 어텐션은 과거를 하나로 압축하는 대신 필요한 입력을 직접 참조해 이 병목을 없앴다. 잔차 연결을 90도 돌려 놓고 보면 같은 처방을 적용할 수 있다는 착상이 나온다.

구체적인 방법은 단순하다. 이전 층 출력들을 더할 때 쓰던 고정 가중치 1을 데이터에 따라 달라지는 값으로 바꾼다. 층마다 학습 가능한 질의 벡터를 두고, 이전 층 출력들을 키와 값으로 삼아 내적을 계산한 뒤 소프트맥스로 정규화한다. 층 출력마다 크기가 크게 다를 수 있어 키 벡터는 정규화한 뒤 내적한다. 질의 벡터는 0으로 초기화해 학습 초반에는 단순 평균, 즉 기존 잔차 연결과 거의 같은 동작에서 출발하게 만든다.

이 방식은 모든 층의 출력을 순전파와 역전파 내내 살려 둬야 해서 대규모 학습에서는 메모리와 통신 비용이 문제가 된다. 그래서 층들을 묶어 블록 단위 요약 벡터로 압축하는 방식이 함께 제안됐다. 12층 모델이라면 4층씩 세 블록으로 묶고, 진행 중인 블록에 대해서는 순전파 도중 누적되는 부분합을 함께 참조한다. 부분합은 그때그때 더해 나가므로 추가 저장 공간이 필요 없다. 블록 8개면 전체 어텐션 잔차의 성능을 거의 회복한다.

논문의 상당 부분은 인프라 설계에 할애돼 있다. 깊이 방향으로 모델을 쪼개는 파이프라인 병렬화에서는 GPU들이 서로를 기다리는 빈 구간(파이프라인 버블)이 생기는데, 배치를 잘게 나누고 각 GPU가 떨어져 있는 여러 층 묶음을 맡는 인터리브 방식으로 이를 줄인다. 블록 어텐션 잔차는 은닉 상태 외에 블록 요약까지 전송해야 해 통신량이 늘지만, 이미 그 GPU에 있는 요약을 다시 보내지 않고 캐시해 재사용하면 예시 상황에서 전송 횟수가 19회에서 11회로 42% 줄었다.

주요 인사이트

• 이 연구가 흥미로운 지점은 새 아이디어라서가 아니다. 영상도 이를 분명히 한다. 이전 층 출력을 이어 붙이는 덴스넷, 층별 결합 가중치를 학습하는 덴스포머, 입력에 따라 달라지는 가중 벡터를 쓰는 딥 크로스 어텐션 등 유사한 시도가 이미 있었다.
• 차별점은 구현과 인프라다. 대규모 학습에서 실제로 감당 가능한 비용으로 돌아가게 만들었기 때문에 이 구조의 장점을 규모에서 입증할 수 있었다는 것이 영상의 결론이다.
• 구조 변경이 최적 모델 형태까지 바꾼다는 결과도 눈에 띈다. 같은 연산 예산이라면 어텐션 잔차 모델은 폭을 넓히기보다 층을 더 쌓는 쪽이 유리했다. 기존 잔차 연결에서는 층을 쌓아도 기여가 희석돼 수확이 줄었기 때문이다.
• 학습된 어텐션 분포를 보면 각 층이 바로 앞 층에 가장 강하게 주목하는 대각선 패턴이 뚜렷하면서도, 대각선을 벗어난 선택적 참조가 함께 나타난다. 모델이 실제로 건너뛰기 경로를 활용하고 있다는 뜻이다.
• 질의 벡터를 0으로 초기화해 기존 구조에서 출발하도록 만든 설계는 새 구조를 도입할 때 흔히 쓰는 안전장치다. 학습 초반의 불안정을 피하면서 필요할 때만 새 자유도를 쓰게 한다.

자주 묻는 질문

관련 AI 소식