[논문 리뷰] White-Box Transformers via Sparse Rate Reduction

Yaodong Yu, Sam Buchanan외 6명|arXiv (Cornell University)|2023. 06. 01.

Advanced Neural Network Applications인용 수 23

한 줄 요약

논문은 트랜스포머 계열의 층들을 희소 속도 감소 목표를 최적화하는 풀 언롤링된 단계로 통합하여 완전한 해석 가능 화이트박스 아키텍처(CRATE)를 도출하며, 토큰 표현을 압축하고 희소화하고 대규모 시각 데이터에서 엔지니어링된 트랜스포머와의 경쟁력 있는 성능을 보인다.

ABSTRACT

In this paper, we contend that the objective of representation learning is to compress and transform the distribution of the data, say sets of tokens, towards a mixture of low-dimensional Gaussian distributions supported on incoherent subspaces. The quality of the final representation can be measured by a unified objective function called sparse rate reduction. From this perspective, popular deep networks such as transformers can be naturally viewed as realizing iterative schemes to optimize this objective incrementally. Particularly, we show that the standard transformer block can be derived from alternating optimization on complementary parts of this objective: the multi-head self-attention operator can be viewed as a gradient descent step to compress the token sets by minimizing their lossy coding rate, and the subsequent multi-layer perceptron can be viewed as attempting to sparsify the representation of the tokens. This leads to a family of white-box transformer-like deep network architectures which are mathematically fully interpretable. Despite their simplicity, experiments show that these networks indeed learn to optimize the designed objective: they compress and sparsify representations of large-scale real-world vision datasets such as ImageNet, and achieve performance very close to thoroughly engineered transformers such as ViT. Code is at \url{https://github.com/Ma-Lab-Berkeley/CRATE}.

연구 동기 및 목표

• 표현 학습의 동기를 데이터 분포를 저차원 부분공간의 혼합으로 압축하고 표현을 희소화하는 것으로 삼는 것.
• 손실 인코딩 속도와 희소성을 결합하여 컴팩트한 토큰 표현을 학습하는 단일한 희소 속도 감소 목표를 도입한다.
• 최적화 단계의 풀언롤링으로 트랜스포머 유사 층을 도출하여 어텐션과 MLP 블록에 대한 수학적 해석 가능성을 제공한다.
• CRATE (Coding RAte TransformEr)를 데이터에서 학습된 분포와 사전으로 계층별 확률 모델을 갖는 화이트박스 아키텍처로 제안한다.
• CRATE가 대규모 시각 데이터에서 표현을 압축하고 희소화하는 학습이 가능하며 ViT 스타일의 성능에 근접함을 입증한다.

제안 방법

• 통합 목표를 정의한다: Z = f(X)로 하여 rate reduction과 ℓ0 희소성 페널티를 결합하고, 희소한 rate reduction을 최대화한다.
• 레이어별로 학습된 기저 U[K]를 갖는 저차원 부분공간의 혼합으로 토큰 분포를 모델링한다.
• MSSA에 대해 부분공간 혼합에 대한 코딩 속도를 최소화하기 위한 그래디언트 스텝으로 자기주의(셀프 어텐션) 유사 업데이트를 도출한다.
• Z의 희소성을 촉진하기 위해 학습된 사전 D에 대항하는 ISTA 유사 업데이트로 희소화를 표현한다.
• 레이어별로 MSSA 기반 압축을 수행하고 이어 ISTA 기반 희소화를 수행하는 층을 차례로 쌓아 CRATE를 구성하며, 레이어별 U[K]와 D는 엔드-투-엔드로 학습된다.
• 복제를 위한 코드 링크 제공: https://github.com/Ma-Lab-Berkeley/CRATE

실험 결과

연구 질문

• RQ1 희소성을 가진 rate-reduction 목표가 컴팩트하고 해석 가능한 토큰 표현을 생성할 수 있는가?
• RQ2 풀언롤링 최적화에서 파생된 화이트박스 트랜스포머 층이 대규모 시각 작업에서 경쟁력 있는 성능을 달성하는가?
• RQ3 자기 주의와 MLP 블록이 단일 프레임워크 내에서 노이즈 제거/압축 및 희소 코딩 단계로 재해석될 수 있는가?
• RQ4 레이어별로 학습된 부분공간 기저와 사전이 표현 품질과 전이성에 미치는 영향은 무엇인가?
• RQ5 제안된 MSSA 및 ISTA 블록이 학습 중에 의도된 최적화 목표와 얼마나 충실히 일치하는가?

주요 결과

• CRATE 계층은 점진적 최적화를 구현하여 토큰 분포를 부분공간의 혼합으로 압축하고 표현을 희소화한다.
• MSSA 구성요소는 부분공간에 대한 rate-reduction 노즈 제거에서 유도된 그래디언트-스텝과 같은 동작으로 자기 주의와 유사한 업데이트에 대응한다.
• ISTA 기반 희소화 층은 학습된 사전을 이용해 희소성을 촉진하여 rate 기반 다양성의 계산 가능 근사치를 가능하게 한다.
• ImageNet-1K에 대한 실험에서 CRATE가 표현을 압축하고 희소화하는 학습을 수행하며 ViT 등 엔지니어링된 트랜스포머에 근접한 성능을 보인다.
• 레이어별 분석은 압축과 희소화가 레이어를 거치며 모두 향상됨을 보여주며 의도된 목적 주도 설계를 뒷받침한다.