Transformer-XL

Transformer-XL은 이전 세그먼트의 은닉 상태를 캐싱하고 다음 세그먼트가 이에 어텐션하도록 하여 고정 길이 텍스트 세그먼트 전반에 걸쳐 문맥을 유지하는 트랜스포머 언어 모델입니다. Carnegie Mellon University와 Google Brain의 Zihang Dai, Zhilin Yang 및 동료들은 ACL 2019에서 Transformer-XL: 고정 길이 문맥을 넘어서는 어텐션 기반 언어 모델을 발표했습니다. 각 세그먼트가 별도의 컨텍스트 윈도우인 고정 세그먼트로 학습한 기본 트랜스포머는 현재 세그먼트의 시작 이전 내용을 볼 수 없습니다. 저자들은 이를 문맥 단편화라고 불렀습니다. Transformer-XL은 순환 신경망(RNN)보다 80%, 기본 트랜스포머보다 450% 더 긴 종속성을 학습했습니다.

세그먼트 수준 순환#

모델이 세그먼트를 처리할 때 각 레이어는 이전 세그먼트의 캐시된 은닉 상태를 현재 세그먼트의 상태와 연결합니다. 쿼리는 현재 세그먼트에서만 가져오지만 키와 값은 두 세그먼트의 상태를 모두 포함합니다. 그래디언트는 캐시에서 멈추므로 세그먼트당 학습 비용은 일정하게 유지됩니다. PyTorch의 단일 어텐션 레이어로 이 패턴을 확인할 수 있습니다.

import torch
from torch import nn

torch.manual_seed(0)
seg_len, dim = 4, 16
attn = nn.MultiheadAttention(dim, num_heads=2, batch_first=True)
hidden = torch.randn(1, 3 * seg_len, dim)  # hidden states for three consecutive segments

memory = torch.empty(1, 0, dim)  # nothing is cached before the first segment
for i, segment in enumerate(hidden.split(seg_len, dim=1), start=1):
    # Queries come from the current segment; keys and values also cover the cached one
    context = torch.cat([memory, segment], dim=1)
    out, weights = attn(segment, context, context)
    print(f"Segment {i}: {weights.shape[1]} queries attend over {weights.shape[2]} positions")
    # Cache this segment for the next one; detach() stops gradients as in Transformer-XL
    memory = segment.detach()

첫 번째 세그먼트의 쿼리 4개는 해당 세그먼트의 토큰 4개만 참조하고, 이후 세그먼트의 쿼리는 각자 자신의 토큰 4개와 캐시된 토큰 4개를 합쳐 8개를 참조합니다. 세그먼트의 레이어 n은 이전 세그먼트의 캐시된 레이어 n − 1 상태를 읽으므로 레이어마다 참조 범위가 한 세그먼트씩 늘어나고, 세그먼트 길이가 N이고 레이어가 L개일 때 약 N × L개의 토큰에 도달합니다.

캐시는 평가 속도도 높입니다. 다음 토큰을 하나씩 예측하는 기본 트랜스포머는 새 토큰마다 전체 윈도우를 다시 계산해야 하는 반면, Transformer-XL은 캐시된 상태를 재사용하므로 평가 중 최대 1,800배 이상 빨라졌습니다. Ultralytics YOLO26 추적도 스크립트가 동영상을 한 번에 한 프레임씩 model.track()로 전달할 때 같은 단편화 문제를 겪습니다. persist=True는 호출 사이에 트래커 상태를 유지하므로 각 객체는 프레임이 바뀌어도 동일한 ID를 유지합니다.

상대 위치 인코딩#

캐시된 상태를 재사용하면 절대 위치가 깨집니다. 캐시된 세그먼트의 첫 토큰과 현재 세그먼트의 첫 토큰이 같은 위치 임베딩을 갖게 되기 때문입니다. Transformer-XL은 대신 쿼리와 키의 내용, 두 항목 사이 거리의 사인파 인코딩, 위치 간에 공유되는 두 개의 학습된 바이어스 벡터(내용용과 위치용)를 바탕으로 각 어텐션 점수를 계산합니다. 점수가 거리에만 의존하므로 평가 시 학습 때보다 더 긴 캐시에 어텐션을 적용할 수 있습니다. 현재 대부분의 대규모 언어 모델(LLM)에서 사용하는 회전 위치 임베딩(RoPE)도 어텐션 점수가 상대 위치에 따라 달라지도록 합니다.

결과와 영향#

Transformer-XL은 다섯 가지 벤치마크에서 새로운 최첨단 결과를 기록했습니다. enwik8에서 문자당 0.99비트, text8에서 1.08비트, WikiText-103에서 18.3, One Billion Word에서 21.8, Penn Treebank에서 54.5의 퍼플렉시티를 파인튜닝 없이 달성했습니다. WikiText-103만으로 학습했지만 수천 토큰 길이의 새로운 기사를 상당히 일관성 있게 생성했습니다. 여러 공동 저자가 개발한 XLNet은 Transformer-XL 아키텍처를 기반으로 자연어 이해 사전 학습을 구축했고, DeepMind의 Compressive Transformer는 오래된 메모리를 버리는 대신 압축해 캐시를 확장했습니다.

Transformer-XL과 관련 모델 비교#

• 장단기 메모리(LSTM): 토큰 단위로 은닉 상태를 업데이트하므로 먼 과거의 정보가 중간의 모든 단계를 거쳐야 합니다.
• Longformer: 순환 없이 윈도우 어텐션과 전역 어텐션을 사용해 4,096개 토큰으로 구성된 문서 전체를 한 번에 읽습니다.
• KV 캐시: 생성 중 이전 토큰의 키와 값을 저장해 다시 계산하지 않도록 합니다. Transformer-XL의 메모리는 학습에도 사용되며, 세그먼트 경계에서 그래디언트가 중단됩니다.
• 상태 공간 모델(SSM): 전체 이력을 고정 크기 상태로 압축합니다.