꼼꼼하고 이해하기 쉬운 Reformer 리뷰

Reformer가 긴 입력에서 Transformer의 시간·메모리 비용을 줄이는 세 가지 장치를 해설한다. 모든 토큰 쌍을 비교하는 attention 대신 LSH로 비슷한 토큰만 묶어 계산하고, reversible layer로 중간 activation 저장을 줄이며, feed-forward 연산을 chunk 단위로 나눈다. Q와 K 공유, reversible 구조, LSH attention이 성능을 얼마나 유지하는지도 실험으로 검증한다.

핵심 포인트

일반 self-attention은 입력 길이의 제곱에 비례해 토큰 쌍을 계산하므로 512토큰을 넘어 문서·이미지 단위로 확장할 때 비용이 급증한다.
Angular LSH로 방향이 비슷한 벡터를 같은 bucket에 모으고 같은 bucket의 인접 chunk 안에서만 attention을 계산해 복잡도를 선형에 가깝게 낮춘다.
Query와 Key 투영을 같게 두는 가정은 실험에서 분리한 모델과 성능 차이가 작았고 enwiki8에서는 더 빨리 수렴했다.
Reversible residual block은 출력에서 입력을 복원할 수 있어 역전파를 위해 모든 층의 activation을 저장할 필요가 없다.
ImageNet64에서 hash를 8개 이상 쓰면 full attention과 거의 같은 성능을 보였고, enwiki8에서는 입력이 길어져도 단계별 시간이 안정적이었다.

상세 정리

비용의 출발점: scaled dot-product attention은 Q와 K의 모든 조합을 곱한다. 입력 길이가 10배가 되면 attention 행렬은 100배로 커져 긴 문서 처리가 어렵다.
FFN 메모리: feed-forward layer는 각 토큰에 적용되고 내부 차원이 크다. 원 Transformer는 최대 512토큰에 2,048차원 FFN을 사용해 입력이 길수록 이 영역도 큰 메모리를 차지한다.
Activation 저장: N개의 residual block을 역전파하려면 각 층의 입력과 출력을 보관해야 한다. 층 수에 비례하는 activation 메모리가 추가된다.
LSH 직관: 영향력이 큰 토큰 쌍은 임베딩 공간에서 가깝다고 보고, 직접 모든 벡터를 비교하는 대신 locality-sensitive hash가 같은 후보만 찾는다.
Angular LSH: 벡터를 단위 구면에 사상하고 회전된 공간의 사분면 번호를 여러 번 기록한다. 가까운 방향의 벡터일수록 hash 배열을 공유할 확률이 높다.
Q와 K 공유: Reformer는 Query와 Key 투영을 동일하게 둬 하나의 표현을 hash한다. 중요한 토큰은 주고받는 영향이 모두 크다는 가정으로 검색 구조를 단순화한다.
Bucket과 chunk: hash 값으로 정렬한 뒤 불균형한 bucket을 고정 크기 chunk로 다시 나눈다. 같은 bucket이면서 현재 chunk나 바로 앞 chunk에 있는 토큰 쌍만 계산한다.
자기 attention 예외: Q와 K가 같으면 자기 자신과의 내적이 항상 커진다. 후보가 자기 하나뿐인 경우를 제외하면 self-attend를 막아 다른 토큰과의 관계를 학습하게 한다.
복잡도 변화: 토큰 길이를 l, chunk 수를 c라 하면 각 토큰은 최대 약 `2l/c` 후보만 본다. c를 충분히 키우면 전체 계산이 l의 제곱보다 선형에 가까워진다.
Reversible block: 입력을 둘로 나누고 `y1=x1+F(x2)`, `y2=x2+G(y1)`로 계산하면 출력에서 입력을 역산할 수 있다. Transformer에서는 F를 attention, G를 FFN으로 둔다.
FFN chunking: FFN은 토큰 위치 사이 상호작용이 없으므로 전체 시퀀스를 한 번에 처리할 필요가 없다. 작은 chunk를 순차 처리해 한 chunk의 activation만 메모리에 둔다.
실험 구성: ImageNet64와 64K토큰 enwiki8을 인코딩·디코딩하고 bit-per-dim으로 압축 성능을 비교했다.
가설 검증: Q=K와 reversible layer는 원 구조 대비 성능 저하가 미미했다. LSH는 hash 수가 늘수록 full attention에 가까워졌고 12층 이후에는 추가 성능 향상이 작았다.
적용 경계: 짧은 문장에서는 일반 Transformer를 대체할 이득이 작을 수 있다. 문서·책 단위 질의응답처럼 긴 문맥에서 메모리와 속도 이점이 커진다.