Test-Time Supervised Adaptation를 위한 Forward-Only Associative Learning via Closed-Form Fast Weights

사전학습된 표현을 고정한 상태에서 입력-출력 임베딩 간의 연상 매핑을 학습하는 FAAST는 gradient-based 업데이트를 제거하고 단일 forward 패스에서 task-specific 적응을 수행한다. 메모리 조회나 컨텍스트 의존성을 요구하지 않으며, 이미지 분류와 언어 모델링 벤치마크에서 backprop 기반 적응과 경쟁하거나 우수한 성능을 보이고 학습 시간과 메모리 사용을 크게 감소시킨다.

출발점은 고정된 표현(K, V)을 이용해 단순 선형 매핑 W를 학습하는 문제이다. FAAST는 데이터 집합 D = {(xi, yi)}에서 키-값 쌍 K = [k1, ..., kN]⊤, V = [v1, ..., vN]⊤를 구성한 뒤, L(W) = ||KW − V||F^2를 최소화하는 해를 분석적으로 구한다. 최적 해는 W⋆ = K†V이며, 이는 Moore-Penrose 역함수를 이용한 해이다. K의 특이값 분해(K = URΣV⊤)에서 Σ†를 정의하고 W⋆ = RΣ†U⊤V로 표현할 수 있다. 이때 Σ†에 σi가 너무 작으면 불안정하므로 Σ†ε를 도입해 작은 σ를 차단하는 스펙트럴 필터링을 적용한다. FAAST의 이 해는 attention 기반 검색의 한 형태로 해석될 수 있는데, a⋆ = K†q를 통해 질의 q에 대한 가중치를 계산하고 h = a⋆⊤V로 출력을 얻는 형태가 된다. 이로써 메모리 조회가 필요 없이 fast weights를 통해 task-specific 매핑을 단일 forward 패스에서 구성할 수 있다. 이 접근은 pretrained 표현과의 아키텍처적 분리를 유지하면서도, memory 기반 Retrieval이나 ICL 대비 inference-time 비용을 크게 줄이고, memory를 고정된 크기로 압축해 실용성을 높인다. FAAST는 이미지 분류 및 언어 모델링 벤치마크에서 backprop 기반 적응과 유사하거나 더 나은 성능을 보였고, 학습 시간은 대폭 감소한다. 또한 작은 데이터에서도 일반화가 잘되며, 임의의 정의된 라벨에서도 견고한 성능을 보인다.

1. 전체 접근 방식과 아이디어: representation 학습과 associative learning을 계층적으로 분리하고, downstream task 적응을 gradient 없이 analytically 수행한다. 2) 핵심 메커니즘/알고리즘: 데이터에서 K, V를 수집하고 min_W ||KW − V||^2를 최소화하여 W⋆를 구한다. 이를 SVD를 통해 계산하며, W⋆ = K†V로 표현된다. 3) Incremental 업데이트: Nt를 누적하는 업데이트 규칙과 S, T를 이용한 빠른 갱신(W⋆ = S−1T)을 제시한다. 4) 스펙트럴 필터링: σmaxϵ 조건으로 Σ†를 정의하고 ϵ = 1/Nα로 설정해 불안정 방향을 차단한다. 5) pretrained 네트워크와의 통합: classification의 경우 W0와 W⋆의 보간으로 적응하며, language model의 경우 중간 레이어의 kℓ,t, vℓ,t를 바탕으로 per-layer W⋆를 구성하고 readout projection Pℓ를 도입해 memory 출력과의 결합을 수행한다. 6) 이론적 근거: W⋆는 일반적 gradient 기반 학습의 최소-노름 해와 일치하며, pseudoinverse attention으로 해석 가능하다. Softmax attention은 엔트로피 규제를 통한 근사로 해석된다.

메인 벤치마크 결과: 이미지 분류에서 5-shot CIFAR10은 73.8±0.3, full은 86.7±0.2, miniImageNet은 5-shot 88.6±0.2, full 93.0±0.1이다. 텍스트 분류에서 SST-2는 1-shot 78.5±1.1, 5-shot 80.8±1.1, full 87.5±0.9이며, IMDB는 1-shot 86.7±0.9, 2-shot 87.4±0.9, full 90.4±0.8이다. 시퀀스 모델링에서 WikiText-103의 GPT2-XL(1.5B) 기반 실험에서 FAAST의 PPL은 15.35로, readout를 추가 trained하면 13.23으로 감소해 backprop 기반 적응과 동등하거나 우수한 성능을 보인다. BLEU 기반 기계 번역(IWSLT2017)에서 1-shot/full 데이터 설정에서 De-En, En-Fr, Fr-En 등에서 2–3 BLEU 포인트 이상 개선을 달성했다. 학습 비용 측면에서 이미지 분류에서 FAAST의 학습 비용은 38 GPU seconds로, CIFAR10의 풀 파인튜닝(1,512 GPU seconds) 대비 약 97% 감소했다(마이크로 배치 대비). mini-ImageNet의 경우 3 GPU seconds로 감소했다(풀 파인튜닝 212 GPU seconds 대비). 언어 모델링의 경우 대형 모델에서 학습 비용이 0.2 GPU hours로 감소했고(backprop 대비 93.3% 절감), 추론 비용은 메모리 기반 대비 최대 96.5% 절감했다. 또한 memory-sparse 설정에서 memory 크기와 업데이트 discount를 조정함으로써 memory의 효율성과 업데이트 빈도 간의 트레이드오프를 분석했다.

FAAST 모듈은 고정된 표현으로 구성된 키 K, 값 V를 모아 행렬로 만들고, 최소-제곱 문제 L(W) = ||KW − V||^2_F를 최소화하는 해 W⋆를 구한다. 해는 W⋆ = K†V이고, K의 SVD(K = UΣR⊤)를 이용해 Σ†를 정의하면 W⋆ = RΣ†U⊤V가 된다. Σ†ε는 σi가 σmaxϵ 미만인 경우 0으로 대체하는 방식으로, Σ†ε = diag(σ−1_i)에서 I[σi ≥ σmaxϵ]를 곱해 얻는다. 이를 통해 작은 특이값에 의한 과적합을 억제한다. incremental update는 Wt+1 = (Nt/(Nt+1))Wt + (N/(Nt+1))W⋆로, 새 배치의 Kb, Vb를 추가해 S = K⊤K, T = K⊤V를 갱신하고 W⋆ = S−1T를 재계산하는 방식이다. FAAST는 분류기의 W0와 W⋆의 보간으로 기존 모델과의 호환을 유지하며, Lℓ(읽기 읽기용 Pℓ)과 메모리 스코어러를 도입해 informative 쌍의 기여를 가중한다. 시퀀스 모델링의 경우 kℓ,t, vℓ,t를 수집해 Kl, Vl로 확장하고, 각 레이어에 W⋆ℓ를 구성해 잔차 연결(readout projection Pℓ)으로 다음 레이어에 전달한다. 이때 a⋆ = K†q를 이용한 가중치 획득과 h = a⋆⊤V의 형태로 출력을 얻으며, softmax attention은 엔트로피-정규화된 근사로 해석된다.