Background

대규모 언어 모델들은 그럴듯한 텍스트들을 만들어내는데 탁월하지만, 실제로 텍스트 생성하는 상황에서는 어느정도 통제되어야 할 필요가 있습니다. 생성할 문장들이 전체적으로 특정 감정이나 주제에 맞게 생성되어야 하거나, 더 나아가 특정 syntactic structure 와 같은 복잡한 구조에 맞는 문장이 생성할 수 있어야 하기 때문입니다.

이를 해결하기 위한 방법으로 제일 먼저 떠오르는 생각은 이러한 언어 모델들을 지도학습 세팅으로 fine-tuning하는 것입니다. 하지만, 대규모 언어 모델에 맞는 지도학습 corpus를 찾는 것도 어렵고 비용적으로 큰 모델을 재학습하는 것은 어렵습니다. 이를 위해 제안된 방식들 중 하나가 plug-and-play controllable generation 방식입니다. (논문에서는 prompt-based learning을 언급하거나 비교하고 있지는 않으나, 해당 방식도 충분히 가능한 방법이라고 생각됩니다.)

plug-and-play controllable generation 방식은 먼저 언어모델을 freeze시키고, external classifier를 사용하여 해당 언어모델의 생성 방식을 조종하는 방식으로 볼 수 있습니다. 기존 연구들은 단순한 attribute-level의 control은 가능했으나, syntax parse tree나 syntax length와 같은 복잡한 syntactic structure-wise control은 다루고 있지 않습니다. 이 문제를 해결하기 위해, 해당 연구는 diffusion model을 기반으로 한 non-autoregressive language model인 Diffusion-LM을 제안합니다. 해당 모델을 plug-and-play controllable generation 방식으로 기존의 단순한 control뿐만 아니라 fine-grained control들을 수행할 수 있다고 합니다.

Approach 1: Constructing a Diffusion-based Language Model

먼저, language model을 만들기 위해 기존 diffusion 모델을 그대로 사용하게 되면, 크게 2가지 문제를 고려해야 합니다.

첫번째는, forward diffusion process에서 x0 벡터는 연속적이지만 최종적으로 mapping해줘야 하는 텍스트는 discrete하다는 점입니다. (기존의 vision domain에서는 pixel 값들을 continuous domain으로 간주 할 수 있기 때문에 직접적으로 적용할 수 있었습니다.) 이는 Embedding Function을 통해 해결할 수 있고, learnable하게 하는 것이 가장 효과적이라고 본 연구의 Section 4.1에서는 말하고 있습니다.

마찬가지로 두번째 문제는, backward diffusion process에서 denoising에 의해 생성된 x0가 discrete한 단어들에 mapping될 수 있어야 한다는 점입니다. 이는 rounding function을 통해 해결할 수 있고, argmax-rounding과 같은 방법으로 통해 discrete한 단어들로 mapping시킬 수 있습니다.

예를 들어, 기존의 텍스트가 N개 단어의 나열이라면, x0는 Embedding function에 의해 N*d 차원의 벡터가 되고, 학습 파라미터가 없는 forward diffusion process을 통해 Gaussian Noise 벡터로 향하게 됩니다. 반대로, backward diffusion process를 통해 Gaussian Noise 벡터에서 출발하여 x0를 샘플링할 수 있고, 1부터 N의 각 위치 i에 해당하는 x0[i]의 d 차원 벡터마다 가장 가까운 단어 벡터를 찾아 mapping 시켜주는 방식으로 볼 수 있습니다.

위 2가지 고려사항을 해결해도 발생하는 큰 문제가 있습니다. 위 방법을 통해 설계된 diffusion-based model을 학습하더라도, 샘플링된 x0가 애초에 하나의 텍스트로 수렴하지 않는다는 점입니다. 이는 기존의 diffusion training objective가 각 diffusion step t에 해당하는 학습 목표로 구성되어 있다보니, 단어 임베딩에 가까운 x0로 향할 수 있는 기회가 적다는 것입니다.

이를 해결하기 위해, 1) 모든 diffusion step t에 해당하는 학습 목표를 local하게 학습하는 것이 아니라, x0를 직접적으로 예측하게끔 학습 목표를 수정하게 됩니다. 2) 또한 inference 과정에서, 매 diffusion step마다 clamping이라는 방법을 적용하여 decoding 하게 됩니다. 이 두가지 해결방법을 통해, 샘플링된 데이터 분포가 원하는 단어 임베딩 벡터로 잘 수렴할 수 있게 합니다.

Approach 2: Decoding and Controllable Generation with Diffusion-LM

구축한 Diffusion-LM 모델이 controllable generation task를 plug-and-play 방식으로 수행하려면 어떻게 하면 될까요? 바로 매 (t-1)-step에 생성된 latent variable x_{t-1}에 gradient gudiance로 update하고 x_{t}를 sampling하면 됩니다. 자세한 수식은 다음과 같습니다. 우측의 첫번째 term은 fluency와 관련지어 볼 수 있고, 두번째 term은 control으로 볼 수 있습니다.

추가적인 테크닉으로, 자연스러운 텍스트 생성을 위해 fluency regularization을 적용했고, 빠른 생성을 위해 mutiple gradient steps를 적용했습니다. 또한, 최종적으로 좋은 품질의 텍스트 생성을 위해, Minimum Bayes Risk (MBR) 디코딩 방식을 따르고 있습니다. 자세한 내용은 페이퍼를 참조하시면 될 것 같습니다.

Questionnaires

Q1: 왜 하필 diffusion-based model이어야 하는가?

A1: 제 1저자의 설명에 따르면, 이 연구의 초기 목표는 plug-and-play 방식으로 complex/fine-grained control까지도 할 수 있는 방법론을 제안하는 것이었다고 합니다. 따라서, 기존의 autoregressive LM들을 continuous relaxation을 통해 많이 시도해봤다고 합니다. 하지만, complex/fine-grained control들을 수행할 수 없었고, 이에 대해 rounding error가 가장 큰 bottleneck이었다고 합니다. 이를 위해, diffusion based model로 바꿨고, 바꾼 motivation은 diffusion step들마다 rounding을 시켜서 학습하고, denoise를 잘할 수록 결국 rounding을 잘하는 것으로 관계 지어 설명할 수 있다고 합니다.

References

[1] Li, Xiang Lisa, et al. “Diffusion-LM Improves Controllable Text Generation.” arXiv preprint arXiv:2205.14217 (2022).

[2] https://github.com/XiangLi1999/Diffusion-LM/issues/17 (Asking about necessity of using diffusion model)