기술 요약: IMC 친화적 추정을 위한 기능 보존 어텐션 대체 (FAR)

1. 문제 제기

변환기 (Transformer) 아키텍처, 특히 자기 어텐션 (self-attention) 에 의존하는 모델들은 ReRAM 크로스바와 같은 신희 메모리 내 컴퓨팅 (In-Memory Computing, IMC) 하드웨어에 배포될 때 상당한 효율성 병목 현상에 직면합니다. 자기 어텐션은 글로벌 토큰 통합을 가능하게 함으로써 병렬 디지털 프로세서 (GPU/TPU) 에서 탁월한 성능을 발휘하지만, 그 계산 특성은 IMC 실행 모델과 근본적으로 충돌합니다:

• 활성화 간 곱셈 (Activation-to-Activation Multiplications): 자기 어텐션은 $QK^\top$ 및 소프트맥스 정규화와 같은 연산에 크게 의존하는데, 이는 가중치 정지 (weight-stationary) 벡터 - 행렬 곱셈 (VMM) 을 최적화하는 IMC 하드웨어가 아닌, 활성화의 동적이고 비국소적인 혼합을 요구합니다.
• 메모리 오버헤드: 토큰 상호작용의 이차 복잡도 ($O(T^2)$) 와 공간적으로 분산된 활성화를 반복적으로 읽어야 하는 필요성은 조각난 아날로그 연산, 광범위한 아날로그 - 디지털/디지털 - 아날로그 (ADC/DAC) 변환, 그리고 낮은 크로스바 활용도를 초래합니다.
• 현재의 한계: 기존 IMC 가속기는 종종 어텐션을 디지털 유닛으로 오프로딩하여 이차 병목 현상을 해결하지 못하거나, 본질적인 데이터 이동 비용을 해결하지 않은 채 어텐션을 아날로그 어레이에 직접 매핑하려 시도합니다.

동시에 최근 분석들은 추론 중 어텐션 레이어가 상당한 중복성을 보이며, 어텐션 맵이 깊이에 따라 더 매끄럽고 압축 가능해짐을 시사합니다. 이는 어텐션의 기능적 역할이 IMC 국소성과 가중치 정지 특성에 더 잘 부합하는 순차적 모듈로 근사될 수 있다는 가설을 자극합니다.

2. 방법론: FAR 프레임워크

저자들은 사전 학습된 비전 변환기 (특히 DeiT 계열) 의 모든 자기 어텐션 블록을 다중 헤드 양방향 LSTM(BiLSTM) 과 같은 순차적 모듈로 대체하는 FAR(Function-preserving Attention Replacement) 프레임워크를 제안합니다. 이 접근법은 처음부터 학습할 필요 없는 "drop-in" 아키텍처 변환으로 설계되었습니다.

2.1 대체 전략

• 블록 단위 대체: 사전 학습된 변환기의 모든 어텐션 블록 $A_l$은 학습 가능한 대체물 $A'_l$(다중 헤드 BiLSTM) 로 대체됩니다. MLP 블록, 패치 임베딩, 레이어 정규화와 같은 모든 다른 구성 요소는 고정 (frozen) 된 상태로 유지됩니다.
• 증류 (Distillation) 를 통한 기능적 동등성: 대체는 레이어별 증류를 통해 감독됩니다. 각 레이어 $l$에서, 유사성 손실 ($L_{sim}$) 이 원본 어텐션 헤드의 출력과 새로운 BiLSTM 헤드의 출력 간 차이를 최소화합니다:
  $$L_{sim}^{(l)} = ||A_l(x_l) - A'_l(x_l)||_2^2$$
  이는 대체물이 교사의 표현 구조를 복제하도록 보장합니다.
• 학습 프로토콜: 이 과정은 두 단계로 구성됩니다:
  1. 증류 단계: 모델의 나머지 부분은 고정된 채로, 유사성 손실과 작업 손실 (분류) 의 조합을 사용하여 대체 BiLSTM 모듈만 학습됩니다.
  2. 파인튜닝 단계: 잔여 정확도 하락을 회복하기 위해 전체 모델이 공동으로 파인튜닝됩니다.

2.2 다중 헤드 BiLSTM 아키텍처

다중 헤드 어텐션의 귀납적 편향을 보존하기 위해, 대체물은 다중 헤드 BiLSTM으로 구조화됩니다:

• 프로젝션 및 분할: 입력 임베딩은 프로젝션되어 $N$개의 부분 공간으로 분할됩니다 (교사 어텐션 헤드의 수와 일치).
• 양방향 처리: 각 부분 공간은 별도의 BiLSTM 헤드로 처리되어 순방향 및 역방향 토큰 종속성을 모두 포착합니다.
• 연결 및 프로젝션: 출력들은 연결되어 원래 차원으로 다시 프로젝션됩니다.
• IMC 호환성: 어텐션과 달리, 이 구조는 순환적 가중치 재사용과 순차적 데이터 흐름에 의존하여 모든 쌍의 활성화 상호작용을 제거합니다. 이는 가중치가 정지하고 데이터가 순차적으로 흐르는 ReRAM 크로스바에 자연스럽게 매핑됩니다.

2.3 IMC 제약을 위한 구조적 가지치기

리소스 제약이 있는 IMC 어레이에 더 잘 적응하기 위해, 저자들은 BiLSTM 헤드의 숨겨진 차원에 구조적 가지치기를 적용합니다:

• DeepHoyer 의 확장: 그들은 LSTM 숨겨진 유닛을 대상으로 DeepHoyer 프레임워크를 확장합니다. LSTM 게이트 (입력, 망각, 셀, 출력) 가 숨겨진 유닛을 공유하므로, 단일 유닛을 가지치기하려면 모든 관련 가중치 행렬 ($W_{ih}$, $W_{hh}$, 및 프로젝션 가중치) 에서 조정된 제거가 필요합니다.
• 그룹 희소성: 각 숨겨진 유닛에 해당하는 가중치 그룹에 미분 가능한 Hoyer 희소성 측정을 적용하여 네트워크가 어떤 유닛이 중복되는지 학습하도록 장려합니다.
• 레이어 적응형 압축: 가지치기는 헤드와 방향별로 독립적으로 수행되어, 모델이 각 레이어의 특정 중복성에 따라 그 용량을 적응하도록 하며, 결과적으로 IMC 어레이 크기에 맞는 컴팩트한 모듈을 생성합니다.

3. 주요 결과

이 프레임워크는 ImageNet 과 다양한 하류 작업 (CIFAR-10/100, Cars, Flowers, iNaturalist) 에서 DeiT 계열 (Tiny, Small, Base) 에 대해 평가되었습니다.

• 정확도 보존:
  • DeiT-Tiny: FAR-Tiny 는 ImageNet 에서 기준 DeiT-Tiny 보다 +1.2% 높은 성능 (73.3% 대 72.2%) 을 보여, 순차적 대체물이 작은 모델에 유익한 귀납적 편향으로 작용함을 시사합니다.
  • 더 큰 모델: DeiT-Small 과 Base 의 경우, FAR 은 교사 모델과 약 1~2% 의 modest 한 격차로 경쟁력 있는 정확도를 유지하여, 어텐션의 기능적 역할이 순차적 순환을 통해 대폭 재현 가능함을 입증했습니다.
  • 하류 전이: FAR 모델은 하류 분류 작업으로의 강력한 전이 능력을 유지하여, 학습된 시퀀스 - 투 - 시퀀스 매핑이 효과적으로 일반화됨을 나타냅니다.
• 가지치기에 대한 강건성:
  • 구조적 가지치기를 적용하면 평균 숨겨진 차원이 4060% 감소하면서도 정확도 하락은 추가로 12% 만 발생했습니다.
  • 가지치기 패턴은 비균일했으며, 중간 레이어에서는 더 높은 유지율을 보이고 네트워크 끝단에서는 더 공격적인 가지치기가 이루어져, 표현 중복성의 내부 이질성을 반영했습니다.
• 헤드별 전문화: 토큰 상호작용 맵의 시각화는 FAR 헤드가 원본 어텐션 헤드의 고유한 의미론적 역할을 성공적으로 계승하여, 중복된 패턴으로 수렴하는 대신 특정 객체 영역에 집중했음을 보여주었습니다.

하드웨어 효율성 분석

분석적 ReRAM IMC 모델을 사용하여 저자들은 지연 시간과 에너지를 추정했습니다:

• 지연 시간: ReRAM IMC 에서의 FAR 는 이차 어텐션 오버헤드를 피함으로써 ReRAM IMC 에서의 DeiT 보다 약 18 배 더 빠르고, GPU 에서의 DeiT 보다 약 400 배 더 빠릅니다.
• 에너지: ReRAM IMC 에서의 FAR 는 ReRAM IMC 에서의 DeiT 보다 약 3 배 더 에너지 효율적이며, GPU 에서의 DeiT 보다 약 150 배 더 효율적입니다.
• 메커니즘: 이러한 이득은 활성화 간 곱셈을 제거하고, ADC/DAC 변환을 줄이며, 가중치 정지 실행을 가능하게 함에서 비롯됩니다. 저자들은 어텐션의 IMC 에서 중요한 병목 현상인 소프트맥스 비용을 제외했기 때문에 이러한 추정치가 FAR 에 대해 보수적이라고 지적합니다.

4. 중요성 및 주장

이 논문은 FAR 이 모델을 처음부터 학습할 필요 없이 사전 학습된 변환기를 IMC 기반 엣지 가속기에 배포하기 위한 실용적이고 하드웨어 효율적인 대안을 제공한다고 주장합니다.

• 전체 대체의 실현 가능성: 이 연구는 어텐션이 추론 시 끝에서 끝으로 대체되어도 실용적인 정확도를 유지할 수 있음을 입증하여, 고성능 비전 작업에 모든 쌍 토큰 상호작용이 필수적이라는 관념에 도전합니다.
• 하드웨어 정렬: 어텐션을 순차적이고 가중치 정지인 모듈로 대체함으로써, FAR 는 계산 데이터 흐름을 ReRAM 크로스바의 물리적 강점과 정렬시켜 메모리 트래픽과 지연 시간을 크게 줄입니다.
• 사전 학습 지식의 보존: 이 프레임워크는 대규모 사전 학습 동안 학습된 풍부한 의미론적 표현이 블록별 증류를 통해 구조적으로 다른 IMC 친화적 아키텍처로 보존되고 전이될 수 있음을 증명합니다.

요약하자면, FAR 는 어텐션 기반 모델의 지배력과 신희 메모리 중심 하드웨어의 제약 사이의 간극을 메우며, IMC 친화적 순차적 모듈을 중심으로 변환기 추론을 재구성하는 경로를 제공합니다.