기술 요약: 텐서 메모리: 장기 범위 트랜스포머를 위한 고정 크기 순환 상태

1. 문제 제기

트랜스포머는 공간 및 시간 차원을 긴 토큰 시퀀스로 평탄화하여 이미지와 비디오를 처리합니다. 어텐션 메커니즘과 키-값 (KV) 캐싱이 과거 특징을 보존하지만, 시퀀스 범위가 길어짐에 따라 두 가지 주요 한계를 겪습니다:

1. 메모리 확장: 실제 추론은 과거 키와 값을 저장하는 데 의존하며, 이는 토큰 수에 비례하여 선형적으로 증가합니다. 이는 긴 비디오와 문서에 대해 계산적, 메모리적으로 불가능하게 만듭니다.
2. 지속적 상태의 부재: 기존 트랜스포머는 명시적이고 지속적인 상태를 갖지 않습니다. 과거 특징에 주의를 기울이지만, 입력이 변경됨에 따라 업데이트되고 조회될 수 있는 구조화된 표현을 유지하지는 않습니다. 이로 인해 가려짐 (occlusion) 과 누락된 관찰 하에서 취약해지며, 컴팩트하고 지속적인 상태를 유지하는 대신 끊임없이 확장되는 과거 기록에서 글로벌 컨텍스트를 반복적으로 재구성해야 하는 문제가 발생합니다.

스파스 어텐션, 히스토리 압축, 또는 검색 증강 컨텍스트와 같은 기존 접근법들은 종종 컨텍스트 길이를 확장하지만, 정보를 비구조화된 토큰 벡터 또는 슬롯의 집합으로 저장합니다. 이들은 공간 구조를 명시적으로 보존하거나 비디오 이해와 같은 상태 기반 작업에 적합한 컴팩트한 "세계 상태"를 유지하지는 않습니다.

2. 방법론: 텐서 메모리

저자들은 표준 트랜스포머 블록을 고정 크기의 순환 3D 메모리 텐서로 보강하는 경량 모듈인 텐서 메모리를 제안합니다. 이 모듈은 상태 용량을 입력 길이와 분리하면서도 공간적 유추 편향을 보존합니다.

핵심 아키텍처

메모리는 숨겨진 상태 $h_t$와 셀 상태 $c_t$로 구성된 볼륨 쌍으로 정의되며, 차원은 $C \times D \times H \times W$입니다. 여기서 공간 차원은 입력 길이와 무관하게 일정합니다. 이 모듈은 토큰 청크당 사이클로 작동합니다:

1. 청킹: 입력 토큰은 청크로 그룹화됩니다. "읽기 쿼리"는 청크의 첫 번째 토큰에서 유도되며, "쓰기 요약"은 전체 청크의 학습된 투영입니다.
2. 연속 주소 지정: 공유된 좌표 예측기 (연결된 가중치 사용) 가 볼륨 격자 내에서 읽기 ($\mu_{read}$) 와 쓰기 ($\mu_{write}$) 를 위한 연속적인 3D 좌표를 생성합니다.
3. 읽기 (연속 샘플링): 모델은 삼선형 보간을 사용하여 $\mu_{read}$에서 숨겨진 상태 $h_t$를 읽습니다. 검색된 특징은 토큰 차원으로 투영되어 게이트된 잔차 연결을 통해 토큰 스트림에 융합됩니다. 학습된 스칼라 게이트 $\gamma$가 이 융합의 강도를 제어하여, 특정 작업에 메모리 경로가 유익하지 않을 경우 이를 억제할 수 있게 합니다.
4. 쓰기 (미분 가능한 소프트 쓰기): 모델은 콘텐츠 벡터와 확산 매개변수 $\sigma$를 예측하여 정보를 씁니다. 이 콘텐츠는 예측된 3D 위치 $\mu_{write}$ 주변의 가우스 가중 볼륨으로 배치됩니다. 이러한 "소프트 쓰기"는 하드 볼륨 할당을 피하여 기울기가 정보가 배치되는 위치를 형성하도록 합니다.
5. 순환 업데이트: 쓰기 볼륨은 현재 숨겨진 상태와 연결됩니다. 경량 분해된 3D 연산자 (깊이별 3D 컨볼루션 후 포인트별 혼합) 가 ConvLSTM 스타일의 게이트 ($i_t, f_t, o_t, g_t$) 를 계산합니다. 이러한 게이트는 셀 상태와 숨겨진 상태를 업데이트하여 정보가 지속되고 국소적으로 전파되며 관찰이 누락된 경우에도 사용 가능하도록 합니다.

주요 설계 선택

• 고정 크기 상태: 메모리 텐서 크기가 일정하여 KV 캐싱과 관련된 선형 메모리 증가를 방지합니다.
• 공간 구조: 비구조화된 토큰 슬롯과 달리 3D 볼륨 격자는 공간적 관계를 명시적으로 보존합니다.
• 미분 가능한 소프트 쓰기: 가우스 가중 볼륨을 사용하면 이산적이고 하드한 할당이 아닌 부드럽고 연속적인 업데이트가 가능합니다.
• 게이트된 융합: 학습된 게이트 $\gamma$는 모듈이 훈련을 불안정하게 하지 않고 사전 훈련된 모델에 연결될 수 있도록 보장하며, 메모리가 불필요한 경우 게이트가 해당 경로를 억제합니다.

3. 주요 기여

• 새로운 메모리 메커니즘: 토큰이 연속 좌표를 통해 읽고 쓸 수 있는 고정 크기의 순환 3D 볼륨 상태 도입으로 명시적인 공간 상태를 구현합니다.
• 통합: 이 모듈은 표준 트랜스포머 훈련 파이프라인에 깔끔하게 통합되며, 아키텍처 변경 없이 기존 블록에 연결하거나 제거할 수 있습니다.
• 진단 도구 세트: 지속적 상태가 유익한 시나리오 (예: 좌표 바인딩, 가려짐, 지도 구축) 와 그렇지 않은 시나리오를 분리하도록 설계된 제어된 "토이" 진단 도구 세트 개발.

4. 실험 결과

저자들은 언어, 이미지, 비디오 도메인의 표준 벤치마크와 진단 도구 세트에서 텐서 메모리를 평가했습니다.

• 언어 모델링 (WikiText-2): 텐서 메모리는 가장 큰 절대적 개선을 달성하여 테스트 퍼플렉시티를 130.30(기본 트랜스포머) 에서 100.07로 감소시켰습니다. 이는 Transformer-XL, 로컬 어텐션, 선형 어텐션을 포함한 베이스라인들을 능가했습니다.
• 이미지 패치 재구성 (CUB-200-2011): 마스킹된 패치 재구성 작업에서 텐서 메모리 변형 (Spatial ViT) 은 Base ViT, Local ViT, Registers ViT 대비 마스크 비율 (25%, 50%, 75%) 전반에 걸쳐 PSNR 및 SSIM 에서 일관된 개선을 보였습니다.
• 비디오 동작 인식 (UCF-101): 이 모델은 Top-1 (89.21% 대 82.12%) 및 Top-5 정확도에서 결합된 시공간 트랜스포머 베이스라인과 강력한 "레지스터 토큰" 베이스라인 모두를 능가했습니다. 저자들은 이는 단계별 볼륨 순환으로 인해 훈련 벽시계 시간이 상당히 증가하는 대가였다고 지적했습니다.
• 토이 진단:
  • 좌표 바인딩: 어텐션 기반 베이스라인이 거의 무작위 성능 (7–14%) 으로 붕괴되는 반면, 텐서 메모리는 쓰기 횟수가 증가함에 따라 높은 정확도 (75–91%) 를 유지했습니다.
  • 무해성 제어: 메모리가 필요 없는 작업에서 학습된 게이트 $\sigma(\gamma)$는 초기화 값 근처에 머무르며 (메모리 경로 억제), 모든 모델이 100% 정확도를 달성했습니다.
  • 지도 구축: 텐서 메모리는 시간이 지남에 따라 정보를 유지해야 하는 장기 범위 그리드 탐색 작업에서 베이스라인을 크게 능가했습니다.

5. 중요성 및 주장

이 논문은 텐서 메모리가 명시적이고 컴팩트하며 구조화된 상태를 트랜스포머에 제공함으로써 장기 범위 추론을 처리하는 실용적인 방법을 제공한다고 주장합니다.

• 길이와 분리된 용량: 고정 크기 텐서를 사용하여 KV 캐시의 선형 메모리 확장을 피함으로써, 이론적으로 긴 시퀀스에 대해 더 확장 가능하게 만듭니다.
• 가려짐에 대한 강건성: 지속적 상태는 모델이 장면에 대한 "신념 상태"를 유지하도록 하여, 토큰 스트림에서 컨텍스트를 재구성하는 데만 의존하는 모델보다 누락된 관찰이나 가려짐에 대해 더 강건하게 만듭니다.
• 다용도성: 장기 범위 비디오에 동기를 부여했지만, 이 메커니즘은 언어 (글로벌 컨텍스트 축적) 와 이미지 (구조화된 공간 집계) 에도 유익함이 입증되었습니다.
• 미래 작업의 기반: 저자들은 현재 구현을 기반으로 보고 있습니다. 향후 방향으로는 고차원 메모리 (예: 시간 인식 공간 메모리를 위한 4D), 풍부한 동역학 (예: 광학 흐름 왜곡), 그리고 효율성 개선을 위한 스파스 업데이트 등을 제안합니다.

저자들은 현재 한계에 대해 겸손하게 인정하며, 밀집된 볼륨 업데이트가 처리량 병목 현상이며, 고정 크기 상태가 방대한 수의 고유한 세부 정보를 저장해야 하는 작업의 경우 포화될 수 있음을 인정합니다. 그러나 소규모 구조화된 상태를 추가하는 것이 더 능력 있는 장기 범위 지각 및 추론 시스템으로 가는 실현 가능한 길이라고 주장합니다.