기술 요약: 재귀적 추론 모델을 에지(Edge)를 위해 압축할 때 무엇이 살아남는가?

문제 정의

Tiny Recursive Model (TRM)과 같은 재귀적 추론 모델은 여러 재귀 사이클에 걸쳐 잠재적인 "캐리(carry)" 상태를 반복적으로 업데이트함으로써 복잡한 구조적 과제를 해결한다. 이러한 모델은 상대적으로 적은 파라미터(예: 6.83M)로도 경쟁력 있는 성능을 달지만, 자원이 제한된 에지 하드웨어(예: 4–8 MB SRAM)에 배포하려면 공격적인 압축이 필요하다.

단일 패스(single-pass) 자기회귀 모델에서 유도된 표준 압축 직관은 여기에는 적용되지 않는다. 전통적인 트랜스포머에서는 양자화 노이즈가 출력 토큰을 따라 축적된다. 그러나 재귀 모델의 경우, 양자화 오류는 매 캐리 상태 업데이트 시마다 $H$번의 재귀적 추론 사이클을 통해 누적된다. 본 논문은 이러한 모델을 엄격한 SRAM 제약 조건 하에서 압축할 때 무엇이 살아남고 무엇이 파괴되는지를 조사하며, 특히 표준 기술(양자화, 프루닝, 증류)이 전역적 추론 능력을 보존하는지 아니면 단순히 국소적 토큰 예측만을 보존하는지를 묻는다.

방법론

저자들은 세 가지 과제(ARC-2024, Maze-Hard, Sudoku-Extreme)와 두 가지 재귀 아키텍처(TRM-Attention 및 TRM-MLP-Mixing)에 걸쳐 체계적인 압축 연구를 수행했다.

1. 압축 스윕(Compression Sweep): 풀 프리시전 스윕(FP32, FP16, INT8, INT4)과 다양한 재귀 깊이($H$ 및 $n_{sup}$)를 평가했다.
2. 메모리 병목 현상 해결: 퍼즐 임베딩 테이블(99.4 MB)이 메모리 점유율(126 MB 베이스라인의 79%)을 지배한다는 점을 인식하고, 추론 시점에 관련 행(row)인 2 KB만을 스트리밍하는 **싱글 퍼즐 플래시 로딩(single-puzzle flash loading)**을 사용하여 이를 백본 압축과 분리했다.
3. 압축 기술:
   • 사후 학습 양자화 (Post-Training Quantization, PTQ): INT8(per-channel calibration 적용) 및 INT4(naïve per-tensor 및 per-channel calibration 적용)를 적용했다.
   • 구조적 변화: 구조적 프루닝(structured pruning), 지식 증류(knowledge distillation), 선형 어텐션 근사(linear attention approximations)를 평가했다.
   • 복구 시도: 추론 붕괴를 수리하기 위해 토큰 수준의 목적 함수를 사용하는 양자화 인식 학습(QAT)을 테스트했다.
4. 진단: 동일한 입력에 대해 압축된 모델의 최종 캐리 상태와 풀 프리시전 참조 모델 간의 코사인 유사도를 측정하는 레이블 프리(label-free) 지표인 **캐리 궤적 충실도(Carry-Trajectory Fidelity)**를 도입했다. 이는 태스크 평가 없이도 추론 붕괴를 예측하는 역할을 한다.
5. 하드웨어 검증: Qualcomm AI Hub를 통해 세 가지 Snapdragon 플랫폼(Galaxy S24, S22, Dragonwing RB3)에서 온디바이스 레이턴시와 메모리 사용량을 벤치마킹했다.

주요 결과

1. 깊이-정밀도 교차 (The Depth–Precision Crossover)

낮은 정밀도가 일반적으로 성능을 저하시키는 단일 패스 모델과 달리, 저자들은 단일 재귀 사이클($H=1$)에서의 INT8이 풀 뎁스(full-depth) FP32 정확도(ARC-2024에서 35.25% vs. 36.00%)와 일치하면서 FLOPs를 6배 줄인다는 것을 발견했다. 이 상호작作用은 재귀 모델의 경우, 정밀도가 충분하다면(INT8) 전체 뎁스를 유지하는 것보다 뎁스를 줄이는 것이 더 효율적일 수 있음을 시사한다.

2. 구성적 붕괴: 국소적 vs. 전역적 (Compositional Collapse: Local vs. Global)

공격적인 압축(naïve INT4, 프루닝, 증류)은 셀 정확도(국소적 토큰 예측)와 퍼즐-정확도(전역적 제약 조건 만족) 사이의 디커플링을 생성한다.

• naïve INT4, 프루닝, 증류 하에서는 셀 정확도가 높게 유지되더라도(예: >84%), 퍼즐-정확도는 0으로 붕괴된다.
• 이는 압축이 전역적 제약 조건을 충족하는 능력은 파괴하면서 국소적 예측 능력은 보존한다는 것을 나타낸다.
• 교차 엔트로피 손실을 사용하는 QAT를 포함한 토큰 수준의 목적 함수는 개별 셀의 정답 여부에 보상을 줄 뿐 구성적 일관성을 보상하지 않기 때문에, 전역적 추론을 복구하는 데 실패한다.

3. 아키텍처적 취약성 (Architectural Fragility)

이 붕괴는 태스크 자체의 문제가 아니라 아키텍처의 문제이다.

• TRM-Attention(ARC 및 Maze에 사용)은 naïve INT4를 견뎌낸다.
• TRM-MLP-Mixing(Sudoku에 사용)은 naïve INT4 하에서 처참하게 붕괴된다.
• 동일한 태스크(Sudoku)에 대한 절제 실험(ablation study) 결과, MLP-mixing에서 Attention으로 전환하면 INT4에서도 정확도가 유지됨을 보여주었다. 취약성은 토큰 믹싱 메커니즘에 있으며, Attention은 거친 노이즈 속에서도 생존할 수 있는 희소한 관계적 업데이트를 인코딩하는 반면, MLP-mixing은 per-channel 스케일이 필요하다.

4. 캘리브레이션을 통한 복구

Per-channel calibrated INT4(사후 학습, 재학습 없음)는 naïve INT4에서 나타난 처참한 손실을 복구한다.

• Sudoku-Extreme에서 naïve INT4는 정확도가 10.2%로 떨어지지만, calibrated INT4는 이를 71.9%(FP32 베이스라인인 73.8%에 근접)로 회복시킨다.
• 이는 calibrated INT4가 4 MB MCU 타겟 내에 들어오면서도 제약 조건 만족 추론을 위한 실행 가능한 최저 한계선임을 입증한다.

5. 진단 도구로서의 캐리 궤적 충실도

제안된 캐리 궤적 충실도(Carry-Trajectory Fidelity) 지표는 태스크 평가 전에 추론 붕괴를 성공적으로 예측한다.

• 높은 충실도($\phi \ge 0.99$)는 보존된 추론과 상관관계가 있다.
• 낮은 충실도($\phi < 0.8$)는 치명적인 실패를 예측한다.
• 이 지표는 정확도 손실과 단조적인 관계를 가지며, 태스크를 실행하지 않고도 사용 가능한(calibrated) 양자화와 사용 불가능한(naïve) 양자화를 구분할 수 있다.

배포 레시피 (Deployment Recipe)

논문은 에지 장치를 위한 구체적인 배포 전략을 도출한다:

• 임베딩 압축: 싱글 퍼즐 플래시 로딩을 사용하여 99.4 MB의 임베딩 테이블을 2 KB의 활성 메모리로 줄인다.
• 모바일 SoC (8 MB SRAM): INT8 백본 + 플래시 로딩 + 감소된 재귀($H=1, n_{sup}=8$)를 사용하면 6.51 MB 내에 들어가며 6배 적은 FLOPs로 거의 전체 정확도를 달달성한다.
• 마이크로컨트롤러 (4 MB SRAM): Per-channel calibrated INT4 백본 + 플래시 로딩을 사용하면 3.26 MB 내에 들어가며 재학습 없이 FP32에 근접한 정확도를 회복한다.

의의 및 주장

본 논문은 하드한 SRAM 제약 조건 하에서 재귀적 추론 모델 압축에 대한 첫 번째 체계적인 연구를 제공한다고 주장한다. 주요 기여는 다음과 같다:

1. 압축 한계의 재정의: 재귀 모델이 단일 패스 양자화에는 없는, 감소된 재귀가 정밀도 손실을 보상할 수 있는 독특한 '깊이-정밀도 트레이드오프'를 가진다는 것을 입증했다.
2. 실패 모드 식별: 표준 압축 기술이 종종 국소적 토큰 예측은 보존하면서 전역적 추론은 파괴한다는 것을 보여주었으며, 이는 토큰 수준의 학습 목적 함수로도 해결할 수 없음을 밝혔다.
3. 아키텍처적 통찰: 토큰 믹서(Attention vs. MLP)의 선택이 중요한 배포 결정임을 확립했다. MLP-mixing은 재귀적 설정에서 양자화 노이즈에 훨씬 더 민감하다.
4. 레이블 프리 진단: 태스크별 평가 없이 배포 성공 또는 실패를 예측할 수 있는 실용적인 도구로 캐리 궤적 충실도를 도입했다.
5. 실질적 실현 가능성: 고성능 재귀적 추론이 calibrated INT4와 플래시 스트리밍 임베딩을 통해 4 MB 마이크로컨트롤러에 배포될 수 있음을 검증하여, 클라우드 의존성을 제거했다.

저자들은 본 연구가 TRM 제품군과 특정 태스크에 집중되어 있으며, 충실도 진단에는 배포 전 사용 가능한 풀 프리시전 참조 모델이 필요하다는 점을 언급하며 일반화에 대해 신중한 태도를 유지한다. 그들은 모든 재귀적 아키텍처의 압축 문제를 해결했다고 주장하는 것이 아니라, 연구된 특정 클래스의 모델과 제약 조건에 대해 검증된 레시피를 제공하는 것이다.