이 논문은 변환기 기반 확률적 모델 (Transformer Probabilistic Models) 에서 발생하는 효율적인 자기회귀 (Autoregressive, AR) 추론 문제를 해결하기 위해 제안된 새로운 아키텍처, 인과적 자기회귀 버퍼 (Causal Autoregressive Buffer) 에 관한 연구입니다. ICLR 2026 에 발표된 이 논문은 신경 프로세스 (Neural Processes), 사전 적합 네트워크 (Prior-Fitted Networks), 테이블 기반 파운데이션 모델 등 다양한 분야에서 적용 가능한 방법을 제시합니다.

다음은 논문의 상세한 기술 요약입니다.

1. 문제 제기 (Problem Statement)

• 배경: 신경 프로세스 (NPs) 나 테이블 파운데이션 모델과 같은 집합 기반 (Set-based) 모델은 가변 크기의 컨텍스트 (관측 데이터) 를 한 번에 처리하여 새로운 타겟에 대한 주변 분포 (Marginal Distribution) 를 효율적으로 예측하는 데 탁월합니다.
• 한계: 많은 응용 분야 (신호 보간, 행동 데이터 모델링, 다중 열 예측 등) 는 여러 타겟 간의 연관성을 고려한 결합 분포 (Joint Distribution) 가 필요합니다.
• 기존 방식의 비효율성:
  • 기존 모델은 결합 분포를 얻기 위해 자기회귀 (AR) 방식으로 타겟을 순차적으로 생성합니다.
  • 그러나 기존 변환기 아키텍처는 모든 입력 (컨텍스트 + 이전 예측값) 에 대해 양방향 자기주의 (Bidirectional Self-Attention) 를 사용하므로, 새로운 예측값이 추가될 때마다 전체 컨텍스트를 다시 인코딩 (Re-encoding) 해야 합니다.
  • 이로 인해 시간 복잡도가 $O(K(N+K)^2)$ (여기서 $N$은 컨텍스트 크기, $K$는 타겟 개수) 로 급격히 증가하여, 대규모 컨텍스트나 긴 시퀀스에서는 실용적이지 않습니다.

2. 제안 방법: 인과적 자기회귀 버퍼 (Methodology)

저자들은 한 번만 인코딩된 컨텍스트와 가벼운 순차적 예측을 분리하는 인과적 자기회귀 버퍼 (Causal Autoregressive Buffer) 를 제안합니다.

핵심 메커니즘

1. 컨텍스트 캐싱 (Context Caching):
   • 초기 컨텍스트 집합 $C$를 한 번만 인코딩하여 키 (Key) 와 값 (Value) 을 메모리에 저장 (캐싱) 합니다. 이 캐시는 읽기 전용 (Read-only) 으로 고정됩니다.
2. 인과적 버퍼 (Causal Buffer):
   • 생성된 타겟들을 별도의 버퍼에 저장합니다.
   • 버퍼 내의 토큰들은 인과적 마스크 (Causal Mask) 를 적용받아, 오직 자신보다 이전에 생성된 버퍼 토큰과 고정된 컨텍스트만 참조할 수 있습니다.
   • 이는 언어 모델의 KV 캐싱 (Key-Value Caching) 아이디어를 확률적 모델에 적용한 것입니다.
3. 예측 과정:
   • 새로운 타겟을 예측할 때, 모델은 캐시된 컨텍스트와 현재까지의 버퍼 상태를 모두 참조합니다.
   • 컨텍스트를 다시 인코딩할 필요가 없으므로, 계산 비용이 크게 절감됩니다.

복잡도 분석

• 기존 AR 방식: 매 단계마다 전체 집합을 재인코딩하므로 $O(K(N+K)^2)$.
• 제안 방식:
  • 컨텍스트 인코딩: $O(N^2)$ (1 회만 수행)
  • 교차 주의 (Cross-attention): $O(NK)$
  • 버퍼 자기주의 (Self-attention): $O(K^2)$
  • 총 복잡도: $O(N^2 + NK + K^2)$
  • 이는 $N$과 $K$가 클 때 기존 방식보다 훨씬 효율적입니다.

훈련 전략

• 마스킹 어텐션 (Masked Attention): 훈련 시 컨텍스트, 버퍼, 타겟 토큰을 하나의 시퀀스로 합치고, 구조화된 어텐션 마스크를 사용하여 각 역할 (Context, Buffer, Target) 간의 정보 흐름을 제어합니다.
• 버퍼 크기 커리큘럼: 일부 타겟은 컨텍스트만 참조하고, 일부는 컨텍스트와 버퍼의 일부 prefix 를 참조하도록 훈련하여, 테스트 시 버퍼 상태 (비어있음 또는 가득 참) 에 관계없이 견고한 성능을 발휘하도록 합니다.

3. 주요 기여 (Key Contributions)

1. 아키텍처 혁신: 변환기 기반 확률적 모델의 자기회귀 결합 추론 비용을 $O(K(N+K)^2)$에서 $O(N^2 + NK + K^2)$로 줄이는 인과적 자기회귀 버퍼를 도입했습니다.
2. 통합 훈련 전략: 마스크 어텐션과 버퍼 크기 커리큘럼을 결합하여, 단일 모델이 효율적인 주변 예측과 가속된 자기회귀 샘플링/밀도 평가를 모두 수행할 수 있도록 했습니다.
3. 광범위한 적용성: TNP (Transformer Neural Processes), PFN (Prior-Fitted Networks), TabICL (Tabular Foundation Models) 등 다양한 모델에 적용 가능함을 입증했습니다.

4. 실험 결과 (Results)

저자들은 합성 함수, EEG 시계열 데이터, 베이지안 모델 비교 작업, 테이블 회귀 등 4 가지 도메인에서 실험을 수행했습니다.

• 성능 (Predictive Performance):
  • 제안된 방법은 전체 컨텍스트를 재인코딩하는 기존 AR 방식 (TNP-D-AR, TNP-A) 과 동등한 예측 정확도를 달성했습니다.
  • 버퍼가 비어있는 경우 ($K=1$) 는 기존 모델과 정확히 일치하며, 버퍼가 사용될 때 ($K=16$) 도 성능 저하가 거의 없었습니다.
• 속도 (Speed):
  • 결합 샘플링 (Joint Sampling): 기존 방식 대비 최대 20 배 빠른 속도를 기록했습니다.
  • 밀도 평가 (Density Evaluation): 병렬 처리가 가능한 TNP-A 와 유사한 속도를 내면서도, 순차적 AR 방식 대비 K 배 더 빠릅니다.
• 메모리 효율성 (Memory Efficiency):
  • 큰 컨텍스트 크기 ($N=1024$) 에서 최대 7 배 적은 VRAM을 사용했습니다. 이는 컨텍스트 캐시를 재사용하고 버퍼만 동적으로 관리하기 때문입니다.
• 구체적 작업별 성과:
  • EEG 데이터: 보간 및 예측 작업에서 기존 AR 방식과 유사한 성능을 내면서도 계산 비용이 획기적으로 감소했습니다.
  • 베이지안 모델 비교: 다중 감각 인과 추론 모델에서 로그 주변 가능도 (LML) 추정 시, ground truth 와 매우 높은 상관관계 ($R^2 \approx 1.0$) 를 보이며 모델 선택 능력을 입증했습니다.
  • 테이블 파운데이션 모델 (TabICL): 6 개의 UCI/Kaggle 데이터셋에서 독립 예측 및 표준 AR 방식과 비교하여, 버퍼를 사용한 AR 방식이 낮은 컨텍스트 환경과 예측 (Forecasting) 작업에서 우수한 성능을 보였습니다.

5. 의의 및 결론 (Significance)

• 효율성과 표현력의 균형: 기존에 상충관계로 여겨졌던 '집합 기반 모델의 효율성'과 '자기회귀 모델의 표현력 (결합 분포 모델링)'을 동시에 달성했습니다.
• 실용성: 대규모 컨텍스트와 긴 시퀀스가 필요한 실제 응용 분야 (의료 신호 처리, 금융 예측, 과학 시뮬레이션 등) 에서 확률적 모델의 실용성을 크게 높였습니다.
• 확장성: 이 버퍼 메커니즘은 플러그 앤 플레이 (Plug-and-play) 방식으로 기존 모델에 쉽게 통합될 수 있으며, 잠재적 병목 (Latent Bottleneck) 이 있는 아키텍처 (예: Perceiver) 로도 확장 가능함을 보였습니다.

결론적으로, 이 논문은 변환기 기반 확률적 모델의 추론 병목 현상을 해결하고, 대규모 데이터와 복잡한 의존성 모델링이 필요한 현대적 AI 응용 분야에서 실시간에 가까운 효율적인 베이지안 추론을 가능하게 하는 중요한 기술적 진전을 제시합니다.