How will Bayesian machine learning evolve in the era of LLMs? Discover how GFlowNet may hold the key to unlocking its full potential.
저는 처음에 베이지안들의 연구 방법론과 그들이 추구하는 목표가 다소 허황하게 느껴졌습니다. 베이지안 접근법은 때로 고리타분하고 지나치게 엄격한 가정에 기반하여, 현실 문제에서 효율적인 성능을 보장하지 못하는 경우가 많았습니다. 컴퓨터 과학 분야에서는 단순함이 가장 큰 미덕인데, 베이지안들은 수많은 확률적 도구와 복잡한 수식들을 동원하여 결과적으로 매우 느린 알고리즘을 만들어내는 것처럼 보였습니다.
베이지안 접근법은 모델 추정의 불확실성을 정량화하는 목표를 추구합니다. 근본적으로는 Bayesian posterior inference를 수행하는 것이 그 목적입니다. 이를 수식으로 표현하면 다음과 같습니다:
\[p(z \mid x) = \frac{p(z) \, p(x \mid z)}{\sum_z p(z)p(x \mid z)} = \frac{R(z;x)}{\sum_z R(z;x)}.\]여기서 Bayesian posterior inference는 관측 데이터 $x$에 대해 잠재 변수 $z$의 확률 분포를 업데이트하는 방법을 의미합니다. 만약 $z$가 모델의 파라미터라면, 모델(즉, 함수)에 대한 확률 밀도를 얻을 수 있고, 이를 통해 불확실성을 자연스럽게 측정할 수 있습니다.
그러나 실제 문제에서는 분모에 해당하는 $\sum_z R(z;x)$를 계산하는 것이 현실적으로 불가능한 경우가 대부분입니다. $z$의 경우의 수가 너무 많거나 무한하기 때문에, 모든 경우에 대해 unnormalized density (or reward) $R(z;x)$를 일일이 계산하는 것은 불가능하기 때문입니다.
물론, 이러한 상황을 해결하기 위해 여러 방법들이 개발되었습니다. 그 중 가장 대표적인 방법은 MCMC 방법입니다. 이 방법은 $z$ 값을 조금씩 변경하면서, 정규화되지 않은 분포인 $R(z;x)$의 값을 비교합니다. 확률이 높은 영역에는 $z$가 더 많이 분포하기 때문에, $R(z;x)$ 값이 높아지는 방향으로 마치 산을 오르는 것처럼 $z$ 값을 업데이트합니다. 만약 $R(z;x)$ 값이 낮아진다면, $z$ 값을 그대로 유지합니다. 물론 확률이 낮은 경우에도 때때로 $z$ 값을 업데이트할 수 있는데, 이는 탐색 과정에서 우연히 낮은 값의 상태로 전환되는 경우가 있기 때문입니다.
문제는 대부분의 경우 $R(z;x)$ 분포가 여러 개의 산봉우리를 갖는다는 점입니다. 이러한 분포를 multi-modal distribution이라고 하는데, 여러 산봉우리를 촘촘히 탐색하는 것은 사실상 불가능에 가깝습니다. 산봉우리가 하나뿐이라면 계속 오르막길만 잘 올라가도 언젠가는 정상에 도달하겠지만, 산봉우리가 여러 개라면 때때로 내려가 다른 봉우리를 탐색해야 하므로 탐색이 매우 어려워집니다.
딥러닝 시대에 들어서면서, 뉴럴 네트워크의 강력한 성능을 활용해 이러한 다중 모드(산봉우리) 탐색 문제를 효과적으로 해결할 수 있는 방법들이 등장했습니다. 기존에는 지역적인 정보만을 바탕으로 산등성이를 따라 탐색했다면, 이제는 마치 헬리콥터를 띄워 산 아래 전경을 살펴보듯, 더 넓은 시야를 확보할 수 있습니다.
비유하자면, 헬리콥터에서 물줄기를 쏘아 해당 지점의 깊이를 측정하는 방식과 같습니다. 이처럼 넓은 관점에서 탐색하면, 특정 지역에 국한되지 않고 빠르게 다른 봉우리를 찾아낼 수 있습니다.
이러한 접근법은 에이전트의 순차적 의사결정 문제로 해석할 수 있습니다. 에이전트를 물방울에 비유하면, 이 물방울은 하늘에서 시작해 여러 단계를 거쳐 결국 지면에 도달합니다. 이는 보상이 오직 마지막 단계에서만 주어지는 에피소드 기반 강화학습과 유사합니다.
GFlowNet은 심층 강화학습 기법을 활용하여 Bayesian posterior inference 문제를 해결합니다. 기존 강화학습과 다른 점은, 단순히 가장 높은 산봉우리를 찾는 것이 아니라 여러 모드(봉우리)를 포착하여 그 확률 분포 전체를 학습한다는 점입니다 (GFlowNet에 대한 자세한 방법은 여기서 다루지 않겠습니다).
최근 디퓨전 모델에서도 이러한 접근법의 성공을 엿볼 수 있습니다. 디퓨전 모델은 초기에는 가우시안 노이즈 $N(0,I)$로 시작해, 여러 단계를 거쳐 점진적으로 노이즈를 제거하면서 최종적으로 이미지를 생성합니다. 차이점은 디퓨전 모델이 주로 지도학습 방식으로 학습된다는 점입니다. 이미 모아진 데이터 (일부 모드, 산등성이)에만 탐색이 되도록 친절히 가이드가 되므로 학습이 쉽습니다. 반면, Bayesian 추론에서는 모든 모드를 탐색해야 합니다. 따라서 스스로 산등성이를 탐험을 해야하는 것으로 지도학습으로만 학습할 수가 없습니다. 이런 경우 디퓨전 모델을 backbone으로 하는 GFlowNet 추론방식을 사용하는것이 매우 유망한 방법으로 여겨지고 있고, 이런 분야를 Diffusion Sampler 라고 부릅니다.
요약하자면, 사람이 직접 땅을 밟아 보면서 하는 산 등성이 탐색이 기존 베이지안 방법론이였다면, 딥러닝 생성모델과 강화학습의 발전을 끌어와서 헬리콥터에서 물줄기를 쏘는 탐색 방법이 GFlowNet 추론이라고 할 수 있습니다.
왜 산등성이에 존재하는 여러 모드를 모두 캡처해야 할까요? 그 이유는 모델이 생성하는 샘플의 다양성을 확보하기 위함입니다. 만약 모델의 결정이나 생성 결과가 한 가지 모드에 치우친다면, 해당 모드에 치명적인 문제가 있을 때 전체 시스템에 큰 위험을 초래할 수 있습니다. 극단적인 확신을 가진 단일 모드의 모델보다, 여러 선택지를 제공하는 모델이 훨씬 안전하고 신뢰할 수 있습니다.
생성 모델에게 다양성은 그 자체로 매우 중요한 요소이며, 안전성 측면에서도 큰 의미를 가집니다. AI는 인간을 돕는 도구여야 하며, 하나의 정답만을 내놓는 것보다 여러 가능성을 제시하여 최종 결정을 인간이 내릴 수 있도록 하는 것이 바람직합니다.
GFlowNet 추론은 기존의 베이지안 추론 방법과 달리, 딥러닝의 강력한 도구들(예: LLM, 디퓨전 모델)과 강화학습 기법들을 그대로 활용하여 Bayesian posterior inference 문제를 해결할 수 있는 장점을 지닙니다. 만약 이러한 심층 강화학습의 성공 공식을 GFlowNet이라는 연결고리를 통해 베이지안 추론에 접목시킬 수 있다면, 베이지안 방법론은 단순히 이상적인 접근을 넘어서 신뢰할 수 있는 AGI 구축을 위한 중요한 디딤돌이 될 것입니다.
하지만, 여전히 GFlowNet과 관련기법들은 큰 기술적 어려움들이 존재합니다. 저는 다음과 같은 두가지 난제를 해결해야 한다고 생각합니다 (이것은 강화학습 분야에서도 마찬가지로 제기 되었던 오래된 문제입니다): (1) exploration (2) credit assignment
먼저, Exploration 문제는 계속 언급했던 산을 탐색하는 문제로 해석해서 보면 이해하기 쉽습니다. GFlowNet은 Global 적인 탐색을 하도록 학습되지만 여전히 탐색해야할 영역이 지나치게 넓다면 Multimodal posterior 학습이 크게 어렵게 될 수 있습니다. 일부 영역만 탐색한 채로 그 곳을 전부인거처럼 학습하여 큰 편향이 생길수가 있습니다. 저는 이 문제를 해결하기 위해서는 반드시 분산/병렬 학습이 접목되어야 한다고 생각합니다. 한 노드에서만 탐색하는 것이 아니고 다수의 노드에서 분산적인 탐색과 학습 (즉 Multi agent system)을 하고 이 정보를 공유해서 탐색단계에서의 분산적 scaling을 이루어내야 합니다. 산을 탐색할 때 여러 사람이 동시다발적으로 구역을 나눠서 탐색하는 것이 유리하다는 것을 직관적으로도 너무 당연한 일이죠. 이를 통해 발생하는 연구적인 문제 (분산 탐색과 학습의 asynchronization, communication bottleneck 문제)와 엔지니어링 문제 (다중 gpu 혹은 node 사이의 bandwidth)가 둘다 파생될 수 있고 이를 해결하는데 기여하는 것은 의미가 있다고 생각합니다. 저는 최근에 버클리 리버모어 연구소 연구진들 (Brian R. Bartoldson 등)과 밀라 연구진들과 함께 multi-node 분산학습 방법 을 연구했으나, 여전히 이 방법이 Multimodal posterior inference 관점에서 장점을 가지도록 응용하지는 못한 상태입니다.
그 다음으로, Credit Assignment 문제를 설명하겠습니다. 어떤 결과를 내기 위해 (예를 들어 새로운 분자구조) 수많은 중간 의사결정 (예를 들어 큰 분자구조를 만들기 중간에 작은 블록들을 합성)이 동반되는 경우가 대부분인데, 그로 인한 보상은 매우 희소하게 나타납니다. 그 최종 보상을 중간에 의사결정들에게 어떻게 분배할 수 있을지를 해결하는 것이 바로 Credit Assignment 문제입니다. 이러한 문제는 GFlowNet의 학습을 크게 저하시키는 요소입니다. 저는 포항공대의 연구자들과 (장효순, 안성수) 2024년에 에너지 분배 학습 방법을 제안하여 이 문제를 풀고자 했습니다. 이 방법은 보상 함수를 최대한 균일하게 중간 의사결정 단계에 분배하는 에너지 분배 네트워크 를 추가적으로 학습하여 이를 GFlowNet 학습에 사용하는 방법입니다. 이를 통해 중간 단계 의사결정에 대한 Local 한 보상을 에너지 분배 네트워크를 통해 얻을 수 있어서 decision making step을 끝까지 밟지 않아도 학습이 가능한 장점을 가지게 될 수 있었습니다. 다만, 작은 스케일 문제에서 학습 속도를 크게 높혔지만 여전히 매우 큰 decision making step 수에 대한 scalability는 검증되지 않았으며, 근본적인 credit assignment 문제가 해결 되었다고 보기는 여전히 어렵다고 생각합니다. 이 문제는 여전히 머신러닝의 최대 난제중 하나로 남아있기 때문에 여러가지 다른 시도가 필요할 것이고, 무엇보다도 누군가 이 문제에 대한 Long Term 비전을 제시하는 것이 매우 중요할 것입니다.
저의 스승이자 딥러닝 분야의 대가인 Yoshua Bengio 교수님은 2022년까지 베이지안 접근법에 대해 큰 회의감을 가지고 계셨습니다. 교수님은 90년대 당시 벨 연구소에서 박사후 연구원으로 근무하시던 중, 주위 석학들이 베이지안에 대해 갖고 있던 반감을 직접 목격하셨다고 합니다. 그러나 GFlowNet 추론 방법이 등장한 이후, Bengio 교수님은 베이지안에 대한 견해를 바꾸시고 (2022년), “Why we need a multimodal posterior and why I became Bayesian“이라는 제목의 이메일을 랩 학생들에게 보내셨습니다. 이 메일은 Max Welling(Variational Autoencoder의 저자)과 Yarin Gal(MC Dropout을 통한 베이지안 추론의 저자)에게도 공유되었습니다. 이 메일을 통해 Bengio 교수님은 GFlowNet을 통한 multi-modal posterior inferece의 가능성을 설명하고 Bayesian을 실용적으로 구현하는 Long Term 비전을 공유하였습니다. 이 글은 Bengio 교수님의 당시 메일에 큰 영향을 받아 작성 되었습니다.