Retrieval-Centric Deep Learning은 가중치 행렬을 벡터 데이터베이스로 대체한다

요하네스 오스발트(Johannes Oswald)와 공동 연구진은 가중치 행렬 대신 학습 입력과 역전파 기울기를 저장한 데이터베이스를 추론에 활용하는 ‘검색 중심 딥러닝(Retrieval-Centric Deep Learning·RCDL)’을 제안했다. 연구진은 CIFAR 이미지 분류와 합성 언어 모델링 실험에서 가능성을 관찰했지만, 언어 과제에서 성능 차이를 보인 층은 폭이 수백만에 이르러 계산량과 메모리 요구가 매우 컸다고 밝혔다. 언어 모델링 결과는 아직 논문으로 발표되지 않았다.

가중치 행렬을 검색으로 해석하기

연구진은 저렴한 연산 자원과 메모리가 충분해지는 미래에, 경사 하강법으로 가중치를 학습하는 방식보다 나은 AI 모델 학습법이 가능할지 질문하며 프로젝트를 시작했다. 계기는 파라다임스 오브 인텔리전스 팀 세미나에서 발표된 가즈키 이리에(Kazuki Irie)의 논문이었다. 이 논문은 선형 층 y=Wx에서 가중치 행렬에 대한 경사 하강법을 쌍대 관점으로 해석한다. 논문

이 해석에 따르면 학습된 선형 층의 추론 계산은 학습 중 저장한 출력에 대한 역전파 기울기를 입력으로부터 검색해 조합하는 과정으로 볼 수 있다. 입력 x를 키로 삼고, 내적 기반 어텐션으로 각 항목에 가중치를 주는 방식이다. 이는 계산 결과를 바꾸지 않지만, 가중치 행렬 W 대신 해당 층에서 본 모든 입력과 기울기를 데이터베이스 D에 저장해야 한다. 모델의 각 가중치 행렬마다 벡터가 수십억 개에 이를 수도 있다. 연구진은 이 프로젝트에서는 우선 연산량과 메모리 부담을 고려하지 않았다고 설명했다.

이 관점은 학습률 스케줄링, 가중치 감쇠, 모멘텀 같은 기법이 데이터베이스 항목에 어떤 영향을 주는지 분석하는 데도 활용할 수 있다. 다만 연구진은 이런 분석에서 무엇을 알아낼 수 있을지는 아직 확실하지 않으며, 탐구되지 않은 가능성이 있을 수 있다고 밝혔다.

비선형 커널과 성장형 층

연구진은 내적 기반 선형 어텐션 대신 더 강력한 함수로 데이터베이스를 검색하면 어떨지 살폈다. 먼저 방사 기저 함수(RBF) 커널을 시험했고, 이후 양의 정부호 커널(PDK)이라면 어떤 것이든 사용할 수 있다는 점을 확인했다. 연구진은 이 학습 알고리즘이 원리 있는 방법임을 보였으며, 관련 결과가 그럽(Grubb)과 배그넬(Bagnell)의 2010년 논문에도 나와 있다고 소개했다.

그 결과는 학습된 가중치 행렬처럼 역전파 기울기를 검색해 조합하되, 비선형 어텐션 함수로 가중치를 정하는 ‘성장형’ 신경망 층이다. 연구진은 논문의 알고리즘 5로 이를 제시했다.

하지만 소프트맥스 자기 어텐션을 학습 규칙에 그대로 넣으면 문제가 생겼다. 새 데이터베이스 항목을 추가하는 과정이 연속적이지 않기 때문이다. 학습률을 0에 가깝게 낮춰도 추가 항목에 따른 항이 사라지지 않아 함수에 변화가 생긴다. 연구진은 이를 해결하기 위해 가중 소프트맥스 층을 도입했다. 이 방식에서는 데이터베이스 항목을 부드럽게 추가할 수 있다.

연구진은 가중치 w로 표현되는 무한 데이터베이스 공간에서 함수적 경사 하강법을 수행하고, 이를 이산화하면 새 데이터베이스 항목을 추가하는 방식으로 대응할 수 있다고 설명했다. 자세한 수학적 전개는 논문에 제시했다. 이렇게 만든 ‘성장형 가중 소프트맥스 어텐션 층’은 논문의 알고리즘 3이며, 가중치 행렬을 입력과 기울기에 대한 가중 소프트맥스 커널로 대체하고 원리 있는 방식으로 최적화하도록 설계됐다. 연구진은 이 접근을 RCDL이라고 부른다.

연구진은 델타넷(DeltaNet)과 메사넷(MesaNet) 같은 시퀀스 층도 가중치 학습 규칙으로 옮길 수 있다고 설명했다. 이 관점에서는 경사 하강법으로 가중치를 최적화하는 데 초점을 두기보다, 키에서 역전파 기울기인 값으로 잘 대응하는 행렬을 찾는다. 연구진은 이를 새로운 가중치 최적화 기법을 설계하는 관점으로 제시했다. DeltaNet 논문 · MesaNet 논문

회귀·이미지 분류 실험

연구진은 먼저 학생-교사 회귀 실험에서 가중 소프트맥스 층 하나가 잘 최적화되는 결과를 보였다고 밝혔다. 과매개변수화된 ReLU 학생 모델은 교사 모델을 잘 맞췄지만, 학습이 끝나도 교사의 비선형성을 재현하지는 못했다. 연구진의 층은 폭을 완전히 성장한 가중 소프트맥스 층과 맞춘 매개변수화 소프트맥스 기준선보다 성능이 높았다. ReLU 다층 퍼셉트론(MLP) 학생 모델에서는 안장점과 비슷한 거동이 나타났지만, 연구진의 층에서는 그러지 않았다고 덧붙였다.

CIFAR 실험에서는 잔차 MLP의 피드포워드 네트워크(FFN)를 RCDL 층으로 바꿨다. 연구진의 메사 옵티마이저는 뮤온(Muon)과 경쟁력 있는 결과를 보였지만, 이름과 달리 메사 구조와는 관련이 없으며 프레데리크 벤징(Frederik Benzing)의 잘 알려지지 않은 FOOF 옵티마이저를 다시 만든 것이라고 설명했다. RCDL 층은 매개변수 기반 기준선보다 연산 비용이 훨씬 컸고, 기준선과 비슷한 수준에 도달했지만 이를 앞서지는 못했다.

연구진은 여러 하이퍼파라미터와 모델 구조를 집중적으로 탐색했고, 각 방법에서 테스트 세트에 과적합한 최상의 설정을 제시했다고 밝혔다. 이 구조와 데이터셋에서 뮤온은 Adam이나 RMSprop보다 크게 나은 성능을 보였다. 또 RCDL 층은 당시 일반 경사 하강법만 사용했는데, 이는 Adam이나 뮤온으로 학습한 가중치와 비교할 때 불리할 수 있다고 설명했다. 성장형 RCDL에 Adam과 유사한 최적화를 적용하려면 함수 공간에서 작동하는 방법이 필요하지만, 아직 완전히 개발하지 못했다. 논문에는 성장형 RCDL과 매개변수 기반 학습 규칙을 결합한 초기 결과도 담겼다.

논문은 소프트맥스 커널의 온도가 미치는 영향과 매개변수 기반 기준선과 다른 일반화 양상도 분석한다. 연구진은 실험 결과가 엇갈리는 가운데 논문이 NeurIPS에 채택돼 12월 시드니와 파리에서 발표할 예정이라고 밝혔다.

합성 언어 과제와 계산 비용

언어 모델링 실험은 실제 언어가 아닌 합성 언어를 대상으로 했다. 연구진은 니콜라스 주케트(Nicolas Zucchet)가 전기 데이터셋의 학습 역학을 분석한 결과를 출발점으로 삼았다. 이 데이터셋은 제위안 앨런주(Zeyuan Allen-Zhu)와 위안즈 리(Yuanzhi Li)의 ‘언어 모델링의 물리학’ 연구에서 다뤄졌다. 주케트의 논문 · 언어 모델링의 물리학 논문

주케트의 관찰에 따르면 전기 데이터셋으로 학습한 트랜스포머는 세 단계를 거쳐 지식을 습득한다. 먼저 일반적인 속성-값 통계를 배우고, 개인별 지식이 없는 이상적인 모델도 도달할 수 있는 성능 수준에서 정체기를 맞는다. 이후 대상과 속성의 연관성을 학습한다. 이 정체기가 지속되는 기간은 개인 수에 거의 비례했다. 연구진은 Gemma 3 백본 등 최신 최적화와 구조를 사용해 이 결과를 재현했다고 설명했다.

주케트는 새로운 개인 데이터로 미세 조정할 때 강한 ‘망각’도 관찰했다. 새 개인에 관한 지식을 습득하는 속도는 더딘 반면, 사전 학습 중 배운 개인에 대한 성능은 빠르게 떨어졌다. 사전 학습 데이터를 재생하는 고전적 방식은 최종 성능 하락을 일부 줄였지만, 초기 하락은 막지 못했다.

연구진은 언어 실험에 쓰는 현실적인 학습 기간과 모델 구조에서 가중치 행렬을 RCDL 층으로 바꾸면 데이터센터 전체를 사용할 만큼의 계산·메모리 비용이 발생한다고 밝혔다. 이에 부담을 줄일 방법으로 유한 특성 벡터를 사용해 양의 정부호 커널을 근사하는 접근을 검토했다. 이 유한 차원 공간에서 경사 하강법을 수행하면 성장형 RCDL의 학습 규칙을 근사할 수 있다.

성장형 RBF 층을 근사한 결과는 특별한 매개변수화를 적용한, 코사인 비선형성을 사용하는 FFN이 된다. 특성 수가 커질수록 성장형 RBF 층에 가까워진다. 연구진은 무작위 특성과 비모수 커널 방법의 관계는 이미 알려져 있지만, 이를 성장형 RBF 층과 연결한 점은 특히 뒤이은 실험 결과와 함께 볼 때 새로운 기여일 수 있다고 주장했다.

연구진은 Gemma 3 트랜스포머 백본의 FFN을 GeGlu 층, 코사인 보흐너(Cos Bochner) 근사 층, 완전한 성장형 RBF 층으로 교체해 비교했다. 성장형 RBF 층과 공정하게 비교하기 위해 매개변수 기반 GeGlu 및 코사인 보흐너 층의 폭을 조정해, 추론 시 부동소수점 연산량(FLOPs)이 완전히 성장한 RCDL RBF 층과 대체로 비슷하도록 했다.

인구 규모 1만6000명으로 학습한 2층 모델에서는 RBF와 폭이 넓은 코사인 보흐너 층이 성능 정체기를 크게 줄였다. 코사인 보흐너 층은 폭을 키울수록 성장형 RBF 층과 비슷한 수준으로 개선됐고, 연구진은 이것이 폭이 커지면서 RBF 층을 더 잘 근사한다는 점을 시사한다고 봤다. 표준 GeGLU FFN은 빠르게 성능 한계에 도달했다. 다만 연구진은 이런 결과가 나타나는 이유를 아직 알지 못하며, 현상을 신비롭게 여긴다고 밝혔다. 성장형 RBF 층과 매개변수화된 코사인 보흐너 층은 폭이 매우 넓어, 규모가 큰 모델에 바로 넣을 단계는 아니라고 덧붙였다.

이 결과는 다른 인구 규모와 혼합 구조에서도 나타났다. 혼합 구조에서는 모델 중간의 표준 GeGLU FFN을 하나 또는 두 개의 성장형 RBF·코사인 보흐너 층으로 바꿨다. 미세 조정 실험에서도 RBF와 코사인 보흐너 층이 GeGLU보다 나은 듯한 결과를 보였다. 연구진은 RCDL 층이 재생 없이도 망각된 지식을 ‘되찾는’ 현상을 계속 살펴보고 있다고 밝혔다.

남은 과제

연구진은 RCDL을 입력과 역전파 기울기로 구성된 성장형 데이터베이스에서 정보를 검색하고 추론 때 보간하는 방식으로 요약했다. 이 틀은 비선형 커널을 가중치 학습에 도입하지만, 현재는 데이터베이스가 매우 커질 수 있다. 회귀와 이미지 분류에서 흥미로운 양상을 보였고, 성장형 RCDL 층과 근사 층은 합성 언어 과제에서 지식 습득 속도를 높였다. 그러나 뚜렷한 이점을 관찰하려면 폭을 수백만까지 늘려야 해 층의 비용이 매우 크다.

연구진은 언어 모델링 결과를 논문으로 정리할 계획이며, RBF 층을 하위 선형적으로 성장시키고 키-값 캐시(KV-cache) 압축 기법을 사용하는 방안을 찾고 있다. RBF 층에 더 강력한 최적화 기법을 개발하고, 어떤 조건에서 일반 경사 하강법으로 최적화한 가중치 행렬보다 이점이 나타나는지도 이해하는 것이 남은 과제다.

코드와 논문은 각각 GitHub 저장소와 arXiv 논문에서 확인할 수 있다.