일본 AI 기업 사카나 AI가 역전파를 쓰지 않고도 1,000층 규모 신경망을 학습시키는 방법 ‘PC-ALM(증강 라그랑주 예측 부호화)’을 공개했다. 회사 연구 페이지와 아카이브(arXiv) 논문으로 함께 발표했으며, 9월 18일 공개한 사내 연구 조직 프런티어 인텔리전스 그룹(FIG)의 성과로 소개됐다. 사카나 AI의 최고기술책임자(CTO) 리온 존스는 트랜스포머 구조를 만든 공동 저자 중 한 명이다.

역전파는 현재 거의 모든 딥러닝 모델이 쓰는 학습 방식이다. 신경망이 산출한 답과 정답의 차이를 마지막 층에서 첫 층까지 거꾸로 전달해 각 층의 가중치를 고친다. 정확하지만 순전파가 끝나기를 기다리고 역전파가 끝나기를 다시 기다린 뒤에야 가중치를 고칠 수 있어, 신경망 전체가 엄밀하게 맞물려 돌아가야 한다. 사카나 AI는 뇌에서는 이런 방식이 성립하지 않는다고 설명한다.

예측 부호화는 이를 대신하는 방법으로 오래 연구돼 왔다. 각 층이 자기 주변 정보만 보고 오차를 줄이는 방식이라 층끼리 동시에 계산할 수 있지만, 신경망이 깊고 좁아질수록 오차 신호가 앞쪽 층까지 제대로 전달되지 않아 역전파보다 성능이 떨어졌다.

사카나 AI는 각 층에 피드백 제어기 역할을 하는 ‘쌍대 뉴런’을 더해 이 문제에 대응했다. 쌍대 뉴런은 그 층에 쌓인 예측 오차를 기록해 두는 변수로, 오차 신호가 열처럼 서서히 퍼지지 않고 앞쪽 층까지 고르게 전달되게 만든다. 선형 신경망에서는 이웃 층과 주고받는 계산만으로도 역전파와 같은 기울기에 도달한다는 것이 논문의 결론이다.

회사는 이 방식으로 최대 1,000층 신경망을 이웃 층과의 정보 교환만으로 학습시켰고, 정확도가 역전파에 거의 맞먹었다고 밝혔다. 뇌가 국소 정보만으로 어느 뉴런에 오차의 책임이 있는지 판단하는 방식은 신경과학의 미해결 문제로 남아 있다.

성능은 신경망의 폭과 깊이를 바꿔 가며 측정한 모든 조건에서 역전파와 같은 수준을 기록했다. 특히 기존 예측 부호화가 약했던 깊고 좁은 신경망에서 차이가 컸고, 1,000층짜리 잔차 다층 퍼셉트론 학습에도 성공했다.

이 방식이 주목받는 이유는 계산을 나눠 처리할 수 있기 때문이다. 역전파는 신경망 전체가 하나로 묶여 있어 여러 장치에 쪼개 학습시키기 어렵지만, 층 단위로 계산이 끝나는 방식은 장치를 나눠 쓰기 쉽다. 사카나 AI는 분산 프로세서와 뉴로모픽 하드웨어에서 새로운 학습 방식으로 이어질 수 있다고 설명했다. 뉴로모픽 하드웨어는 뇌의 구조를 본떠 만든 반도체다.

사카나 AI는 JAX 기반 구현 코드를 깃허브에 공개했다. 다만 이번 실험은 다층 퍼셉트론을 대상으로 했고, 대형언어모델 규모에서 같은 결과가 나오는지는 아직 확인되지 않았다.

자세한 내용은 사카나 AI(Sakana AI)에서 확인할 수 있다.

이미지 출처: 사카나 AI