[논문 리뷰] Simplifying Transformer Blocks

Bobby He, Thomas Hofmann|arXiv (Cornell University)|2023. 11. 03.

Neural Networks and Reservoir Computing인용 수 10

한 줄 요약

이 논문은 skip connections, value/projection 매트릭스, 순차적 서브-블록, 심지어 정규화 계층까지 제거한 단순화된 트랜스포머 블록(SAS 및 SAS-P)을 제시한다. 이는 표준 Pre-LN 트랜스포머와 대등하거나 더 나은 훈련 속도 및 처리량을 달성하면서 매개변수를 최대 약 16%까지 감소시키고 처리량을 대략 15-16% 증가시킨다.

ABSTRACT

A simple design recipe for deep Transformers is to compose identical building blocks. But standard transformer blocks are far from simple, interweaving attention and MLP sub-blocks with skip connections & normalisation layers in precise arrangements. This complexity leads to brittle architectures, where seemingly minor changes can significantly reduce training speed, or render models untrainable. In this work, we ask to what extent the standard transformer block can be simplified? Combining signal propagation theory and empirical observations, we motivate modifications that allow many block components to be removed with no loss of training speed, including skip connections, projection or value parameters, sequential sub-blocks and normalisation layers. In experiments on both autoregressive decoder-only and BERT encoder-only models, our simplified transformers emulate the per-update training speed and performance of standard transformers, while enjoying 15% faster training throughput, and using 15% fewer parameters.

연구 동기 및 목표

• 표준 트랜스포머 블록을 훈련 불안정성과 아키텍처의 취약성 때문에 단순화할 필요성을 제시한다.
• 훈련 속도와 성능에 필수적인 구성 요소(skip connections, value/projection 매트릭스, 순차적 서브블록, 정규화)가 무엇인지 조사한다.
• 신호 전파 이론과 경험적 증거에 guided된 점진적으로 단순화된 블록 변형을 개발한다.
• 단순화된 블록이 자동회귀 모델 및 인코더 전용 모델 전반에서 표준 트랜스포머의 훈련 속도와 성능에 도달하거나 이를 능가함을 입증한다.
• 다양한 깊이에서의 효율성 및 확장성 이점과 다운스트림 작업(GLE)을 포함한 평가를 통해 확인한다。

제안 방법

• Pre-LN 트랜스포머 블록에서 시작하여 신호 전파와 실험에 의해 구성 요소를 하나씩 제거한다.
• self-attention에서 값 및 프로젝션 매트릭스를 고정하거나 제거하여 아이덴티티로 만들어 매 업데이트당 속도를 회복한다.
• 잔차 구성 요소를 재매개변수화하거나 축소하여 실제 스킵 없이도 스킵 연결의 이점을 모방한다.
• 병렬 MHA 및 MLP 서브-블록을 채택하여 병렬화 및 순차 의존성 감소(SAS-P)를 가능하게 한다.
• 정규화 계층을 선택적으로 제거하고 영향 평가를 수행하되, 주된 기준은 SAS와 SAS-P에서 정상화가 있는 상태를 유지한다.
• GLUE 파인튜닝을 포함한 디코더 전용 GPT 유사 모델 및 인코더 전용 BERT 설정에서 평가하고 Pre-LN 기준선과 비교한다.

Figure 1: Comparison between different Transformer blocks. (Left) The standard Pre-LN block. (Top Right) Our most simplified block. (Bottom Right) The parallel block (Wang & Komatsuzaki, 2021 ) . Like the parallel block, our block eschews the need for sequential sub-blocks, but we additionally remov

실험 결과

연구 질문

• RQ1 표준 트랜스포머 블록의 핵심 구성 요소(skip connections, value/projection 매트릭스, 순차적 서브블록, 정규화)를 제거해도 업데이트당 훈련 속도가 손상되지 않는가?
• RQ2 아이덴티티 초기화와 특정 매트릭스의 제한된 업데이트가 skipless 주의 블록의 속도 손실을 완화하는가?
• RQ3 병렬화된 MHA와 MLP 블록(병렬 블록)이 성능을 유지하거나 향상시키면서 훈련 처리량을 지속시키는가?
• RQ4 단순화된 블록이 더 깊은 아키텍처로 확장되고 인코더 전용 모델 및 다운스트림 작업(GLUE)으로의 전달이 가능한가?

주요 결과

• 주의 서브블록의 스킵 연결은 적절한 초기화(Shaped Attention)와 MLP 경로의 축소를 통해 훈련 속도를 보존하며 제거될 수 있다.
• 값 및 프로젝션 매트릭스를 아이덴티티로 고정하거나 제거하면 업데이트당 훈련 속도가 유지되거나 향상되어 매개변수 및 FLOP를 크게 줄일 수 있다.
• MLP 스킵 연결을 제거하는 것은 병렬 서브블록과 결합될 때 가능하며, SAS-P를 가능하게 하여 Pre-LN 속도와 유사한 훈련 속도와 더 적은 매개변수를 달성한다.
• SAS 및 SAS-P 블록은 런타임 측면에서 Pre-LN의 훈련 속도에 맞추거나 능가할 수 있으며 매개변수는 약 13% 감소하고 처리량은 최대 ~16% 증가한다.
• 깊이가 깊어질수록 단순화된 블록의 이점이 커지며 인코더 전용 설정(BERT) 및 GLUE 파인튜닝으로 확장되어 성능을 유지한다.
• GLUE 결과는 SAS 및 SAS-P가 베이스라인 성능과 일치하면서 매개변수 수가 감소하고 처리량이 증가하는 것을 시사한다; V-SkipInit은 업데이트당 속도와 확장성에서 뒤처진다.

Figure 2: Loss of training speed in transformers without attention sub-block skip (He et al., 2023 ) , even with Shaped Attention, Eq. 5 , and MLP skips ( $\alpha_{\text{FF}}=1$ ).