WNet: 효율적인 토큰 혼합을 위한 이산 웨이블릿 변환

트랜스포머(Transformer)에서 토큰 혼합(token mixing)은 각 토큰이 다른 토큰으로부터 정보를 가져올 수 있게 하는 단계이며, 긴 시퀀스를 인코딩하는 비용의 대부분을 차지합니다. 셀프 어텐션(Self-attention)은 이 혼합을 매우 잘 수행하는데, 모든 토큰이 내용에 따라 다른 모든 토큰에 가중치를 부여하여 강력한 문맥적 모델링을 제공합니다. 이러한 전쌍 비교(all-pairs comparison)가 바로 비용이 시퀀스 길이에 대해 이차적으로 증가하는 이유이기도 합니다. 우리는 이산 웨이블릿 변환(DWT, discrete wavelet transform) 기반의 토큰 혼합으로 셀프 어텐션을 대체하는 트랜스포머 인코더인 WNet을 소개합니다. 세 가지 어텐션-프리 믹서(attention-free mixers)가 스케일을 재조합합니다: 선형 융합(linear fusion), 학습된 게이팅(learned gating), 또는 각 토큰이 자신의 스케일을 선택하도록 함으로써입니다. 하이브리드 모델은 마지막 레이어에서만 셀프 어텐션을 추가합니다. 수용장 분석(receptive-field analysis)에 따르면, Haar와 같은 두 탭 필터(two-tap filters)로 구축된 웨이블릿 믹서는 네트워크가 아무리 깊더라도 고정된 블록 외부의 토큰 간 관계를 갖지 못하며, 이는 필터가 학습되더라도 마찬가지입니다. 더 긴 필터는 두 개의 레이어 내에서 전체 시퀀스에 도달합니다. 우리는 C4의 고정된 토큰 부분 집합을 사용하여 마스크드 언어 모델링(masked language modeling)으로 모든 모델을 사전 훈련하고, 크기가 일치하는 BERT 및 FNet 베이스라인과 토큰 혼합이 불가능한 제어 그룹을 사용하는 하나의 통제된 설정으로 GLUE에서 미세 조정합니다. 토큰 게이티드 믹서(token-gated mixer)는 256개의 토큰에서 어텐션만큼 빠르게 훈련되며, 4,096개에서는 2.7배 더 빠릅니다.