Reformer - 효율적 Transformer
Digest (CISELQ): 표준 Transformer는 시퀀스 길이 L에 대해 어텐션 메모리/시간이 O(L²)로 폭증하고 각 층의 활성값을 모두 저장해야 해서 long-context 학습이 사실상 불가능하다(Context). 저자들은 “softmax 어텐션은 쿼리마다 유사도 상위 소수 키만이 실질적으로 기여한다”는 관찰과, “활성값은 역전파 시 재계산 가능하다”는 두 가지 통찰을 결합한다(Insight). 해법으로 (1) LSH(Locality-Sensitive Hashing) 어텐션: 쿼리/키를 공유한 뒤(shared-QK) 랜덤 회전 해시로 버킷팅하여 같은 버킷 내에서만 어텐션을 계산하고, (2) Reversible Residual Layer: RevNet 구조로 활성값 저장을 제거하며, (3) FFN chunking으로 큰 MLP를 시간축으로 분할 처리한다(Solution). enwik8(64K 컨텍스트)과 imagenet64에서 표준 Transformer와 동등한 bpd/성능을 유지하면서 메모리/속도 측면에서 수십 배 개선을 보이고, L=64K에서 full attention은 OOM이지만 Reformer는 단일 GPU에서 학습된다(Evidence). 다만 LSH 근사 품질은 해시 라운드 수에 민감하며(round ↑ → 정확도↑/속도↓), 짧은 시퀀스에서는 오버헤드가 상대적 이득을 상쇄할 수 있고, 디코딩 단계에서는 병렬 이점이 제한된다(Limitations). 이 논문은 “패턴을 미리 정하지 않고 content 기반으로 sparsity를 유도”하는 어텐션 계열을 열어, 이후 Longformer/BigBird(구조 기반 sparse)와 대비되는 축을 형성한다(Open Questions: 해시 라운드 자동 조절, 학습된 해시, reversible과 다른 효율화 기법의 결합).
섹션별 요약
Introduction
- 배경: Transformer는 번역/언어모델링에서 SOTA지만, 파라미터 수(수억~수십억)와 시퀀스 길이(수만) 확장 시 단일 가속기 학습이 불가능할 정도로 메모리 부담이 큼.
- 세 가지 병목: (1) N-레이어 활성값 저장 비용, (2) FFN 중간 차원 d_ff가 d_model보다 훨씬 큼, (3) 어텐션 O(L²) 메모리/연산.
- 기여: LSH 어텐션 + 가역 잔차 + FFN 청크화를 결합한 Reformer. 표준 Transformer와 성능 동등성을 유지하면서 길이 L에 대해 O(L log L) 메모리/시간.
Methods
- Shared-QK Attention: Q와 K를 동일한 선형 투영으로 공유. 성능 저하 없이 해싱 설계를 단순화.
- LSH Attention: 쿼리 벡터에 랜덤 회전행렬 R을 곱한 뒤
argmax([xR; -xR])로 해시 버킷 할당(angular LSH). 같은 버킷의 토큰끼리만 어텐션. - 정렬과 청킹: 버킷 인덱스로 시퀀스를 정렬한 뒤 고정 크기 청크로 분할. 각 쿼리는 자기 청크 + 이전 청크(경계 처리)만 attend.
- Multi-round Hashing: 한 해시로 놓치는 이웃을 보완하기 위해 n_rounds개의 독립 해시를 병렬로 수행 후 결합.
- Reversible Residual Layer: y1 = x1 + Attention(x2), y2 = x2 + FFN(y1). 역전파 시 x를 y로부터 복원 → 층별 활성값 저장 불필요.
- Chunked FFN: 위치별로 독립인 FFN을 시간축 c개 청크로 처리하여 d_ff 중간 활성 메모리를 1/c로 감소.
Results
- enwik8 / text8: 표준 Transformer와 동등한 bpd (≈ 1.05 bpd, 논문 Figure 5). 64K 컨텍스트 학습 가능.
- imagenet64 generation: 12K 토큰 컨텍스트(=64×64×3)에서 정상 학습. full attention은 OOM.
- LSH rounds ablation: n_rounds=1은 성능 저하, 4~8에서 full attention에 수렴(Table/Figure 해시 라운드 실험).
- 속도: L이 커질수록 full attention 대비 선형 로그 스케일링(Figure 3).
| Model | Dataset | Metric | Score | vs. Baseline |
|---|---|---|---|---|
| Full Transformer | enwik8 | bpd | ~1.05 | — |
| Reformer (shared-QK) | enwik8 | bpd | ~1.05 | 동등 |
| Reformer (LSH, 8 rounds) | enwik8-64K | bpd | ~1.05 | full과 동등 |
| Full Transformer | imagenet64 | — | OOM | — |
| Reformer | imagenet64 (12K ctx) | bits/dim | 정상 학습 | 가능 |
Discussion
- 한계: 해시 라운드 수에 성능 민감, 짧은 시퀀스에선 청킹/정렬 오버헤드가 비효율, 디코딩(autoregressive) 단계에서 실제 wall-clock 이득 제한.
- 향후 방향: 학습 가능한 해시, 다른 sparsity 방식과 결합, reversible 기법의 다른 도메인 확장.
Insights
- 주목할 점: 어텐션 sparsity를 **content 기반(해시 유사도)**으로 부여 → 구조적 패턴(sliding window 등)에 의존하지 않음.
- 연결 고리: RevNet(가역 네트워크)과 LSH(근사 최근접 탐색)를 NLP 어텐션에 이식.
- 시사점: 긴 문서/유전체/음성 등 초장문 모달리티에 Transformer 적용 가능성 확대.
- 비판적 코멘트: multi-round hashing은 결국 상수 배 증가이며, BigBird처럼 이론적 근사 보장은 약함.
Discussion Points
- 논쟁점: content-based vs structural sparsity 중 어떤 것이 언어 모델링에 적합한가.
- 검증 필요 가정: shared-QK가 모든 태스크에서 full-QK와 동등한가?
- 후속 연구: Routing Transformer(cluster-based), Performer(kernel 근사) 등으로 이어지는 흐름.
메타데이터
| 항목 | 내용 |
|---|---|
| 제목 | Reformer: The Efficient Transformer |
| 저자 | Nikita Kitaev, Łukasz Kaiser, Anselm Levskaya |
| 소속 | UC Berkeley, Google Research |
| 연도 | 2020 |
| 발표 | arXiv:2001.04451, ICLR 2020 |
| 링크 | arXiv, GitHub (trax) |
| 키워드 | LSH attention, reversible residual, sparse attention, long context |
왜 이 연구를 하는가?
핵심 질문
“수만 토큰 시퀀스를 단일 가속기에서 학습 가능한 Transformer를 어떻게 설계할 것인가?”
기존 접근법의 한계
| 한계 | 설명 |
|---|---|
| O(L²) 어텐션 | L=64K면 어텐션 행렬 40억 엔트리, 메모리 폭증 |
| N-layer 활성값 저장 | 모든 중간 활성값을 backprop 위해 저장 → N배 메모리 |
| 거대한 d_ff | FFN 중간 차원이 d_model의 4배 이상, 위치별 활성 저장이 비쌈 |
| 기존 sparse 방법 | Sparse Transformer는 fixed 패턴 → content-aware sparsity 부재 |
핵심 통찰
- softmax 어텐션은 상위 소수의 키에만 집중한다 → 가까운 이웃만 계산해도 충분.
- 유사한 벡터를 같은 버킷으로 보내는 LSH가 어텐션 이웃 검색에 자연스럽게 맞는다.
- 잔차 네트워크를 가역으로 만들면 활성값을 재계산으로 대체 → 저장 불필요.
방법 (Method)
프레임워크 개요
flowchart TB X[입력 시퀀스 x1..xL] --> Q[Shared-QK Projection\n q_i = k_i = Wx_i] Q --> H[LSH Hashing\n h_i = argmax of xR;-xR] H --> S[해시 버킷 기준 정렬\n sort by bucket id] S --> C[고정 크기 청크로 분할\n length L/chunk] C --> A[청크 내 + 이전 청크와만\n Attention 계산] A --> M[Multi-round Hashing\n n_rounds 반복 후 결합] M --> R[Reversible Residual\n y1=x1+Attn, y2=x2+FFN] R --> F[Chunked FFN\n 시간축 c-chunk 처리] F --> O[출력 표현]
핵심 구성요소
- Shared-QK: q_i = k_i 공유 투영.
- Angular LSH: 회전행렬 R ∈ R^{d×b/2}, 버킷 인덱스 = argmax of concat(xR, -xR).
- 정렬+청킹: 버킷 id 정렬 후 길이 m 청크. 경계 이웃 처리 위해 이전 청크도 attend.
- Multi-round: 독립 해시 n_rounds회 → 결과 합산(softmax over union).
- Reversible Layer: x=(x1,x2) → y1=x1+Attention(x2), y2=x2+FFN(y1). 역계산: x2=y2-FFN(y1), x1=y1-Attention(x2).
- Chunked FFN: 위치 독립성 활용해 시간축 c개 청크로 분할 처리.
Key Equations
(1) Angular LSH 버킷 할당
(2) Shared-QK 가정
(3) LSH 어텐션 (버킷 내 softmax)
(4) Reversible Residual
(5) 복잡도
(청크 크기 m ≈ L/버킷 수, 버킷 수 ≈ L/log L.)
복잡도 분석 (O(L log L))
| 축 | 표준 Transformer | Reformer |
|---|---|---|
| 어텐션 메모리 | O(L²) | O(L log L) (버킷 크기 고정) |
| 어텐션 시간 | O(L²·d) | O(L log L · d · n_rounds) |
| 층별 활성값 | O(N·L·d) | O(L·d) (reversible) |
| FFN 메모리 | O(L·d_ff) | O(L·d_ff / c) (chunked) |
발견 (Findings)
주요 결과
| 모델 | 데이터셋 | 컨텍스트 | bpd / 결과 |
|---|---|---|---|
| Full Transformer | enwik8 | 12K | ~1.05 bpd |
| Reformer (LSH-8) | enwik8 | 64K | ~1.05 bpd (동등) |
| Full Transformer | imagenet64 | 12K | OOM |
| Reformer | imagenet64 | 12K | 정상 학습 |
핵심 발견
- 성능 동등성 + 메모리 수십 배 절감이 동시에 성립.
- 해시 라운드 4~8이면 full attention과 수렴; 라운드를 늘릴수록 근사 품질 ↑.
- Reversible layer의 성능 페널티 사실상 0.
이론적 의의
Content-based Sparsity 패러다임 제시
고정 패턴 없이 유사도 기반으로 동적으로 sparsity를 형성하는 길을 열었다. 이는 이후 Routing/Clustered Attention, Performer의 kernel 근사 등과 같은 계열로 확장된다.
메모리 효율 네트워크 설계 원칙
“저장 대신 재계산(reversible)” 아이디어를 Transformer에 도입하여, 깊이 N이 메모리에 미치는 선형 영향을 상수로 바꾸었다. 이는 이후 대규모 모델의 gradient checkpointing 설계와 상호 보완적이다.
재현성 및 신뢰도 평가
| 항목 | 등급 | 비고 |
|---|---|---|
| 코드 공개 | ✅ | Google trax 저장소 공개 |
| 데이터 공개 | ✅ | enwik8/imagenet64 공용 데이터 |
| 하이퍼파라미터 | ✅ | 논문 본문/부록에 상세 기재 |
| 실험 환경 | ⚠️ | 단일 GPU/TPU 환경 일부 모호 |
| 통계적 신뢰도 | ⚠️ | 다중 시드 실험은 제한적 |
| 종합 등급 | A | 오픈소스 구현 다수 존재 |
주장별 신뢰도
| # | 주장 | 근거 | 신뢰도 |
|---|---|---|---|
| 1 | Shared-QK는 성능 손실이 없다 | enwik8 ablation, 표준 bpd 동등 | 🟢 |
| 2 | LSH 어텐션(8 rounds)은 full과 동등 | rounds ablation, 64K 실험 | 🟢 |
| 3 | Reversible layer는 성능 페널티가 없다 | 동일 설정 baseline 비교 | 🟡 |
| 4 | 실제 wall-clock이 L log L로 개선 | Figure 3 scaling plot | 🟡 (디코딩은 제한) |
읽기 난이도: ⭐⭐
Transformer 기본 구조, LSH(근사 최근접 이웃), RevNet(가역 네트워크)에 대한 사전 지식 필요.
관련 연구 비교 매트릭스
| 축 | Reformer | Sparse Transformer (Child 2019) | Longformer (Beltagy 2020) | BigBird (Zaheer 2020) |
|---|---|---|---|---|
| Sparsity 방식 | Content-based (LSH 버킷) | Structural (strided/fixed) | Structural (sliding window + global) | Structural (window + global + random) |
| 패턴 결정 | 입력 의존, 동적 | 사전 정의 | 사전 정의 + few global | 사전 정의 + 랜덤 |
| 복잡도 | O(L log L) | O(L√L) | O(L) | O(L) |
| 이론적 근사 보장 | 약함 (확률적) | 없음 | 없음 | 보편 근사 증명 |
| 메모리 최적화 | Reversible + LSH + FFN chunk | 없음 | gradient checkpoint | 없음 |
| 대표 컨텍스트 | 64K | 12K | 4K~16K | 4K~8K |
| 주요 도메인 | LM / 이미지 생성 | LM / 이미지 | 장문 문서 QA | 장문 QA / 유전체 |
| 코드 공개 | ✅ (trax) | ✅ | ✅ | ✅ |
핵심 대비: Reformer는 “content가 누가 이웃인지 결정”하는 반면, Longformer/BigBird/Sparse Transformer는 “구조가 이웃을 미리 정한다”. 전자는 일반성/동적성이 강점이고, 후자는 구현 단순성과 이론 보장이 강점.
관련 연구
- Attention Methods — 본 문서와 동일 카테고리 허브
- Sparse Transformer (Child et al., 2019) — 구조적 sparse 어텐션의 대표
- Longformer (Beltagy et al., 2020) — 슬라이딩 윈도우 + global
- BigBird (Zaheer et al., 2020) — 구조적 sparse 중 이론 보장
- RevNet (Gomez et al., 2017) — 가역 잔차 네트워크 원류
- Routing Transformer (Roy et al., 2021) — content-based cluster 어텐션, Reformer 후속
원자적 인사이트 (Zettelkasten)
💡 Softmax Attention의 희소성은 content-aware 근사의 정당성을 준다
출처: Reformer - The Efficient Transformer (Kitaev et al., 2020)
유형: 이론적
Softmax는 지수함수적으로 top-k 키에 확률 질량을 집중시키므로, 하위 키들을 0으로 근사하는 비용이 매우 작다. 따라서 “유사 벡터를 찾는 문제”로 어텐션을 환원할 수 있으며, LSH 같은 근사 최근접 이웃(ANN) 알고리즘을 직접 이식할 수 있다.
핵심 조건/맥락: 쿼리/키가 학습되어 유사도 구조를 형성할 때 성립. 초기 학습이나 비정상 분포에서는 깨질 수 있음.
연결: Sparse Attention, Routing Transformer
활용 가능성: 모든 efficient attention 연구의 출발점으로 인용 가능.
💡 “저장 대신 재계산” 원리로 깊이(N)의 메모리 비용을 상수화
출처: Reformer - The Efficient Transformer (Kitaev et al., 2020)
유형: 방법론적
Reversible residual은 층 개수 N에 비례하던 활성값 메모리를 O(1)로 만든다. 이 원리는 어텐션 sparsity와 직교하므로 다른 효율화 기법과 자유롭게 결합 가능하다.
핵심 조건/맥락: 계산량은 약 1.3~2배로 늘어남(재계산 오버헤드). 통신 병목이 큰 환경에선 절약이 덜할 수 있음.
연결: RevNet, Gradient Checkpointing
활용 가능성: 초장문/초심층 모델 설계의 기본 블록.
💡 Content-based vs Structural sparsity는 상보적이다
출처: Reformer - The Efficient Transformer (Kitaev et al., 2020)
유형: 연결
Reformer(content) vs Longformer/BigBird(structural)의 대비는 “동적 일반성 vs 정적 단순성·이론 보장”이라는 트레이드오프를 드러낸다. 두 접근은 서로 대체재라기보다 결합 가능한 축으로 이해해야 한다.
핵심 조건/맥락: 태스크가 지역 의존성이 강하면 structural, 장거리 content 매칭이 중요하면 content-based가 유리.
연결: Longformer, BigBird, Performer
활용 가능성: 새 효율적 어텐션을 설계할 때 두 축을 명시적으로 분리 후 하이브리드화.
💡 해시 라운드 수는 정확도-비용 트레이드오프의 명시적 다이얼
출처: Reformer - The Efficient Transformer (Kitaev et al., 2020)
유형: 실험적
n_rounds는 근사 품질과 연산 비용을 직접 제어하는 하이퍼파라미터이다. 라운드=1이면 눈에 띄는 성능 저하, 4~8이면 full attention 수렴. 이는 근사의 “누락 확률”을 확률적 이론으로 관리할 수 있다는 실용적 증거다.
핵심 조건/맥락: 학습 초기엔 낮은 라운드로 시작해 점진적으로 올리는 커리큘럼도 가능.
연결: LSH, Approximate Nearest Neighbor
활용 가능성: 추론 시 지연에 따라 라운드를 동적 조절하는 adaptive inference 설계.
핵심 용어 정리
| 용어 | 정의 |
|---|---|
| LSH (Locality-Sensitive Hashing) | 유사한 벡터가 같은 해시값을 가질 확률이 높은 해시 함수족. 근사 최근접 이웃 검색에 사용. |
| Angular LSH | 코사인 유사도에 맞춘 LSH. 랜덤 회전으로 벡터를 투영 후 부호로 버킷팅. |
| Shared-QK | Query와 Key를 동일 선형 투영으로 공유하는 설계. |
| Reversible Residual | 출력에서 입력을 역산할 수 있는 잔차 연결. 활성값 저장 불필요. |
| FFN Chunking | 위치 독립 FFN을 시간축 청크로 분할 처리하여 중간 활성 메모리 절감. |
| bpd (bits per dim) | 밀도 추정 모델의 성능 척도. 낮을수록 좋음. |
| Sparse Attention | 전체 L² 어텐션 대신 일부 (쿼리, 키) 쌍만 계산하는 기법. |
| Content-based Sparsity | 입력 내용(유사도)에 따라 어텐션 대상을 동적으로 선택. |
| Structural Sparsity | 입력과 무관하게 사전 정의된 패턴(sliding window 등)을 사용. |
태그
paper #2020 attention reformer lsh reversible sparse-attention efficient-transformer long-context