Reformer

Reformer는 UC Berkeley와 Google Research의 Nikita Kitaev, Łukasz Kaiser, Anselm Levskaya가 ICLR 2020에서 발표한 Reformer: 효율적인 트랜스포머의 메모리 효율적인 트랜스포머입니다. Reformer는 시퀀스 길이를 제한하는 두 가지 비용을 줄입니다. 지역성 민감 해싱(LSH) 어텐션은 전체 n × n 어텐션 행렬을 대체하며, 길이가 L인 시퀀스에 대해 약 L log L의 비용이 듭니다. 가역 잔차 레이어는 역전파에 필요한 활성값을 레이어마다 저장하는 대신 전체 네트워크에 대해 한 번만 저장합니다. Google은 Reformer가 메모리 16 GB를 탑재한 단일 가속기에서 최대 100만 단어의 컨텍스트 윈도우를 처리할 수 있다고 보고했으며, 장편 소설 죄와 벌 전체를 단일 학습 예제로 처리했습니다.

LSH 어텐션#

소프트맥스 적용 후에는 각 쿼리의 어텐션이 쿼리와 가장 유사한 몇 개의 키에 집중되므로 Reformer는 해당 키만 찾습니다. 쿼리와 키에 하나의 공유 투영을 사용하고, 각 벡터에 각도 기반 지역성 민감 해싱을 적용합니다. 무작위 행렬 R을 곱한 뒤 [xR, −xR]에서 가장 큰 값의 위치를 버킷으로 지정하므로, 방향이 비슷한 벡터는 대체로 같은 버킷에 들어갑니다. 그런 다음 토큰을 버킷별로 정렬하고, 정렬된 시퀀스를 동일한 크기의 청크로 나누며, 각 토큰은 자신이 속한 청크와 바로 앞 청크 내에서만 어텐션을 수행합니다. 여러 해시 라운드를 병렬로 실행하면 한 라운드에서 서로 다른 버킷으로 나뉜 유사 쌍도 포착할 수 있습니다.

LSH는 근사 최근접 이웃 벡터 검색의 고전적인 방법입니다. 정확한 검색보다 훨씬 적은 비용으로 쿼리의 k-최근접 이웃 대부분을 찾습니다. 버킷화 단계는 단독으로 실행할 만큼 간단합니다:

import torch

torch.manual_seed(0)
dim, n_buckets, rounds = 64, 16, 1000
query = torch.randn(dim)
vectors = {"query": query, "near": query + 0.3 * torch.randn(dim), "far": torch.randn(dim)}

# Angular LSH as in Reformer: random projection R, then bucket = argmax over [xR, -xR]
R = torch.randn(rounds, dim, n_buckets // 2)
buckets = {}
for name, x in vectors.items():
    proj = torch.einsum("d,rdb->rb", x, R)
    buckets[name] = torch.cat([proj, -proj], dim=-1).argmax(dim=-1)

for name in ("near", "far"):
    cosine = torch.cosine_similarity(query, vectors[name], dim=0)
    shared = (buckets[name] == buckets["query"]).float().mean()
    print(f"{name}: cosine similarity {cosine:.2f}, same bucket as the query in {shared:.0%} of rounds")

쿼리와 코사인 유사도가 0.96인 근접 벡터는 1,000회 라운드 중 72%에서 쿼리와 같은 버킷에 들어가며, 유사도가 0.04인 관련 없는 벡터는 9%에서 같은 버킷에 들어갑니다. 따라서 버킷 내로 제한된 어텐션은 중요한 쌍 대부분을 유지합니다.

가역 레이어와 청킹#

표준 신경망은 역방향 패스를 위해 각 레이어의 활성값을 모두 유지하므로, 메모리 사용량은 깊이와 길이에 따라 증가합니다. Reformer는 가역 잔차 네트워크(RevNets)를 응용합니다. 활성값을 두 부분으로 나누고, 각 부분을 다른 부분을 입력으로 사용하는 함수로 업데이트합니다.

역방향 패스 중 각 레이어는 출력에서 입력을 복원합니다. x2 = y2 − FeedForward(y1) 다음에 x1 = y1 − Attention(x2)를 계산하므로 최종 활성값만 저장하고 각 레이어는 한 번씩 재계산합니다. 그래디언트 체크포인팅도 활성값을 몇 개의 체크포인트에 저장하고 그 사이의 레이어를 재계산하는 방식으로 비슷한 절충을 합니다. Reformer는 모델보다 몇 배 더 넓은 은닉 레이어를 사용하기 때문에 피드포워드 레이어도 위치별 청크로 나누어 실행합니다. 활성값 메모리는 Ultralytics YOLO26 학습 시 배치 크기를 제한하는 요소입니다. batch=-1을 사용하면 AutoBatch가 여러 배치 크기에서 GPU 메모리를 측정하고 직선을 적합한 뒤, 사용 가능한 메모리의 60%를 채울 것으로 예측되는 배치 크기를 선택합니다.

결과 및 구현#

Reformer는 enwik8의 문자 단위 텍스트와 ImageNet-64의 픽셀 단위 이미지 생성에서 표준 트랜스포머와 비슷한 성능을 냈으며, 메모리를 훨씬 적게 사용하고 긴 시퀀스에서 더 빠르게 실행됐습니다. Hugging Face Transformers는 Reformer를 ReformerModel로 구현하며, enwik8과 죄와 벌로 학습한 체크포인트를 공개했습니다.

Reformer와 관련 모델 비교#

• Longformer: 어텐션 패턴을 미리 고정합니다. 각 토큰 주변의 슬라이딩 윈도우와 일부 전역 토큰을 사용합니다.
• 희소 어텐션: 일부 토큰 쌍에 대해서만 점수를 계산하는 방법의 총칭입니다. LSH 어텐션은 그중 콘텐츠 기반 방식입니다.
• 선형 어텐션: 소프트맥스를 커널 특징 맵으로 대체하여 어텐션을 선형 시간에 계산할 수 있도록 합니다.
• Flash attention: 점수 행렬을 저장하지 않고 GPU 온칩 메모리 타일에서 정확한 어텐션을 계산하므로, 점수 개수는 제곱에 비례하지만 메모리 사용량은 선형적으로 증가합니다.