Prompt Highlighter - Interactive Control for Multi-Modal LLMs
University of Hong Kong에서 CVPR’24에 publish한 Prompt Highlighter: Interactive Control for Multi-Modal LLMs를 읽고 정리한 논문 Review 입니다.
Contribution
해당 논문의 Contribution을 먼저 짚고 넘어간다. 논문의 Contribution은 크게 두 가지로 구분 할 수 있다.
- Inference 과정 중에, LLM 모델에게 Attention weight를 통한 강조 표시된 (highlighted) 토큰을 가이드 하면 사용자가 원하는 더 바람직한 결과를 낼 수 있다.
- LLM에서 제어 가능한 (Controllable) 생성을 위한 Token-level의 User interactions (사용자-상호작용)
- 이 논문에서 참고한 prior work : https://arxiv.org/abs/2306.17806
Preliminaries
논문을 이해하기 위해서 두 가지 기본적인 Preliminaries들을 정리하고 넘어가자. - Classifier-free Guidance in Diffusion Process - How does LLM generate text?
Classifier-free Guidance in Diffusion Process
위 Figure는 논문의 Main Idea를 설명하는 그림이다.
논문에서 LLM은 주어진 prompt에 따라 텍스트를 생성한다. 이 때, 사용자는 특정 토큰(특정 단어)을 강조하고 싶을 수 있다. 이 때, LLM은 강조된 토큰을 참고하여 User의 의도에 맞게끔 더 나은 결과를 생성할 수 있다. 이러한 과정을 저자들은 Classifier-free Guidance in Diffusion Process를 통해 가능하다고 설명한다.
위 Figure에서 빨간색 박스친 부분에서 CFG re-weight 이라는 단어가 나오는데 이를 이해하기 위해서 우선 실제 Diffusion Process에서의 CFG(Classifier-free Guidance)를 이해해야 한다.
P Dhariwal and A Nichol, Diffusion Models Beat GANs on Image Synthesis, NeurIPS 2021
J Ho and T Salimans, Classifier-Free Diffusion Guidance, NeurIPS 2021 Workshop
이 그림은 Classifier Guidance을 통한 Diffusion Process를 설명하는 그림이다. 내용은 단순하다. Classifier Guidance는 Diffusion Process에서 Forward process (Image -> Noise)가 최종적으로 종료된 시점에서 원본 이지미의 특정 Class 정보로 Guidance를 주입한다. 반면 Classifier-free Guidance는 Forward process 과정에서 Guidance를 주입한다. 이 때, Guidance는 이미지의 특징이 포함된 Embedding 형태로 주입된다. “주입된다.” 라는 말은 이해를 돕고자 사용한 표현이다. 실제로는 해당 과정은 수식적으로 다음과 같다.
Classifier Guidance: Let \(\epsilon_\theta(x_t, t)\) be the noise prediction model, \(p_\phi(y|x_t, t)\) be the classifier model, and \(w\) be the guidance scale. The classifier guidance equation is:
\(\epsilon_\theta(x_t, t) \longleftarrow \epsilon_\theta(x_t, t) + w \cdot \nabla_{x_t} \log p_\phi(y|x_t, t)\)
The noise prediction is updated by adding the gradient of the log-probability of the desired class \(y\) with respect to the noisy image \(x_t\), scaled by the guidance scale \(w\).
Classifier-Free Guidance (CFG): Let \(\epsilon_\theta(x_t, t, c)\) be the conditional noise prediction model, where \(c\) is the conditioning information (e.g., text prompt). The CFG equation is:
\(\epsilon_\theta(x_t, t, c) \longleftarrow (1 + w) \cdot \epsilon_\theta(x_t, t, c) - w \cdot \epsilon_\theta(x_t, t, \emptyset)\) The noise prediction is updated by taking a weighted sum of the conditional noise prediction \(\epsilon_\theta(x_t, t, c)\) and the unconditional noise prediction \(\epsilon_\theta(x_t, t, \emptyset)\), where \(\emptyset\) represents an empty conditioning. The weight is determined by the guidance scale \(w\).
In both cases, the guidance scale \(w\) controls the strength of the guidance effect. Higher values of \(w\) lead to generated images that more closely match the desired class (classifier guidance) or the conditioning information (CFG).
These guided noise predictions are then used in the reverse diffusion process to generate the final image. The key difference is that classifier guidance requires a separate classifier model, while CFG only relies on the conditional noise prediction model.
How does LLM generate text?
위 Figure는 Language Models(LMs)이 텍스트를 생성하는 과정을 설명하는 그림이다.
LM은 주어진 단어 시퀀스에서 다음에 올 가능성이 가장 높은 단어를 예측한다.
LLM의 응답의 두 번째 단어는 일반적으로 프롬프트와 예측된 응답의 첫 번째 단어를 분석하여 생성된다. 그런 다음 LLM은 새로운 시퀀스를 분석하여 다음 단어를 예측한다. 이 과정은 complete response가 형성될 때까지 반복된다.
여기서 특징적인 것은 같은 프롬프트를 사용해도 결과는 매우 다양할 수 있다는 것이다. 이는 LLM의 의도적으로 동적인 특성 때문일 수 있고, 아래에 설명된 응답 처리 단계 때문일 수 있다.
중요한 점은 LLM은 텍스트를 생성할 때, 이전 단어들을 참고하여 다음 단어를 예측하는데 이 때 생성된 단어도 예측에 사용된다.
Language Modeling with CFG guidance
다시 논문의 Overview를 살펴보자. CFG는 수식적으로 아래와 같이 간단하게 표현된다.
(1) Sampling process of CFG (Bayse rule) \[ P_\Theta(x|c) \propto P_\Theta(x|c)^\gamma/P_\Theta(x)^{\gamma-1} \]
다음으로 Auto-regressive의 수식은 아래와 같이 표현된다.
(2) Auto-regressive Manner \[ P_\Theta(x) = \prod_{i=1}^N P_\Theta(x_i|x_{j<i}) \]
식 (1)과 식 (2)를 결합하면 아래과 같이 표현된다. 이 식은 G Sanchez et al, Stay on topic with classifier-free guidance. (LLM-CFG) 에서 참고한 내용이라고 한다.
(3) Fusion - G Sanchez et al, Stay on topic with classifier-free guidance. (LLM-CFG) \[ P_\Theta(x|c) \propto \frac{P_\Theta(x|c)^\gamma}{P_\Theta(x)^{\gamma-1}} \propto \prod_{i=1}^N \frac{P_\Theta(x_i|x_{j<i}, c)^\gamma}{P_\Theta(x_i|x_{j<i})^{\gamma-1}} \]
식 (3)을 log로 변환하면 아래와 같이 표현된다. (Logarithm rule with eq.(3)) \[ \log P_\Theta(x|c) = \sum_{i=1}^N \gamma \log P_\Theta(x_i|x_{j<i}, c) - (\gamma-1) \log P_\Theta(x_i|x_{j<i}) \]
여기서 c는 condition이자 강조되는 prompt 토큰을 의미한다. 이로써 사용자는 Controllable Generation in LLMs을 할 수 있게 된다.
User can control input as masked
논문에 따르면 위 그림과 같이 사용자는 특정 토큰을 강조하고 싶을 때, 해당 토큰을 Masking하여 입력할 수 있다. 이 때 모델은 Two-branch로 동작하여 normal context branch와 highlighted context branch로 나뉘어진다.
이 아이디어가 논문의 핵심 아이디어라고 할 수 있는데 CFG를 잠재적으로 부여하기 위해서 이런 방식을 채택한 것이다.
논문에서는 바로 위 Figure에 대해서
“the additional conditional context \(c\) is naturally incorporated as the difference between \(s\) and \(\bar{s}_{j<i}\)”
라고 설명한다.
Attention Activation
논문에서는 Masking 기법을 이용해서 특정 토큰을 강조하는 방법을 제시하고 이를 CFG의 이론을 사용해서 가능하게 하였다. 하지만 논문에서 추가로 제시한 문제점은 Long-range sequence에서 Highlight된 부분이 잘 눈에 띄지 않는다는 것이다. 논문에서는 이러한 이유를 근본적으로 파악한다.
이 그림은 생성된 Token (y-axis) 과 입력 Token (x-axis) 간의 Attention score를 내적하여 어떤 관계성이 있는지 파악한 그림이다.
문제 상황은 다음과 같다.
“… fix the grammar and keep LaTeX format, make it compact …” 라는 prompt를 입력하고, LLM에게 문장을 생성하라고 했을 때 생성된 문장과 입력된 prompt의 관계성에서 나타난 결과는 특정 일부 단어에만 (Grammar, LaTeX)와 Correlation이 강하다는 것을 알 수 있다. (위 그림에서 빨간색 박스)
저자들은 이러한 상황은 LLM이 입력 Token들을 골고루 참고하지 않고, 특정 Token들에만 집중하는 것으로 해석한다.
예를 들어서 사용자가 아무리 Compact 라는 단어를 강조하고 싶어도 LLM은 이를 무시하고 Grammar, LaTeX 에만 집중할 수 있다는 것이다.
이러한 문제를 해결하기 위해서 논문에서는 Attention Activation의 Balance를 맞추는 해결 방안을 제시한다.
\[ h_i = \log{\beta}\cdot m_i + k_i \]
\(\beta\): Scale factor
\(m_i\): Masked-condition
\(k_i\): Original Attention score
위 식에서 \(h_i\)는 re-valued attention score를 의미한다. 이를 통해서 LLM은 Highlight된 Token에 대한 Attention을 더욱 강조할 수 있게 된다.
\(h_i\)를 구하는 식 자체는 그렇게 어렵지 않다. 그냥 Scaling 하고 나서 원본 Attention score에 그 값을 더하는 것 뿐이다.
하지만 이 방법을 통해서 Balanced Attention Activation을 이루어 낼 수 있고 논문의 요점인 Highlight Token에 대한 기능을 더 잘 살릴 수 있게 됐다고 생각한다.
\[ p_i = \frac{\exp^{h_i}}{\sum_{j=1}^{N} exp^{h_j}} = \frac{\beta^m_i \cdot \exp^{k_i}}{\sum_{j=1}^{N} \beta^m_j \cdot \exp^{k_j}} \]
최종 attention score는 위와 같이 SoftMax 확률로 변환하는 방식으로 계산된다.
unconditional branch에서는 \(-\delta\) 만큼 negative-scaled 된 attention score를 사용한다. 따라서 the normal context branch에서만 re-valued attention score \(p_i\)를 사용한다고 생각하면 된다.
추가로, \(\mathbf{\tilde{P}}_\Theta(x_i \vert \bar{s}_{j<i}, -\delta\mathbf{m})\) 는 unconditional branch에서의 inference 과정을 의미하는 수식이다.
Highlighting Visual Tokens
논문은 Text에 대한 Highlighting을 설명하고 나서 Visual Tokens에 대한 Highlighting에 대해서도 설명한다. 방법은 Text와 크게 다르지 않다.
Text에서는 Masking을 통해서 Highlighting을 했고, Visual Tokens에서도 Masking을 통해서 Highlighting을 한다.
논문에서 Visual Materials를 다루는 모델로는 BLIP-2의 Q-Former를 사용했다고 한다. 저자들은 Q-former로 Visual 정보를 다룰 때 다음과 같은 문제가 있다고 설명한다.
This process obscures the direct correlation between image patches and input tokens
이게 무슨 말이냐면, 위 사진은 1st cross-attention layer of the Q-Former의 32개의 query token에 대한 attention score를 나타낸 것이다.
이 attention score에서는 원본 이미지에서 특징이 되는 semantic area에 대한 correlation을 찾기 어렵다.
Q-Former에서의 query는 이미지 embedding을 LLM-friendly 하게 변환하는데 이 과정에서 Image의 의미가 분할되지 않기 때문에 이미지안에 존재하는 각각의 정보를 명확히 전달하기 어렵다는 것이 저자들의 주장이다.
Process of Highlighting Visual Tokens
저자들은 이러한 Q-Former의 문제점을 꼬집고 Query vector를 간단하게 수정하는 방식으로 문제를 해결한다. Visual Tokend을 User controllable 하도록 처리하는 과정은 다음과 같다.
위 그림과 같이 이미지는 Patch-wise 하게 표현할 수 있고 Masking 하고자 하는 각 Patch에 대해서 embedding을 생성할 수 있다.
그 다음으로 각각의 key, value embed function \(f_k\), \(f_v\) 을 통해서 Q-Former의 \(f_k(\mathbf{q}^{\text{im}})\) 와 \(f_v(\mathbf{q}^{\text{im}})\)를 구할 수 있다.
그 다음으로는 Masking된 Embedding을 어떻게 처리하는 지를 살펴봐야하는데 이전에 보았던 Text를 처리하는 방식은 아래 식으로 표현 가능했다.
\[ \bar{s}_i = (\alpha - 1)m_i \cdot f(x_i) + f(x_i) \]
이를 Visual Token에 대해서도 적용하면 아래와 같이 표현 가능하다. 그 과정을 논문의 그림과 같이 살펴보자면 다음과 같다:
\[ \mathbf{\hat{q}} = \text{softmax}((\log{\beta} \cdot \mathbf{w} + \mathbf{Q} \mathbf{K}^\top) / \sqrt{d_k}) V \]
위 수식은 Q-Former의 Query Vector의 변형된 형태인 \(\mathbf{\hat{q}}\)를 구하는 식이다. 여기서 \(\mathbf{w}\)는 다른게 아니고 masked condition \(\mathbf{m}\)의 Reshaped를 의미한다. (이미지이기 때문에 Reshape)
Q-former의 Learnable Q-vector를 highlight mask를 기준으로 수정하여 Inference를 했을때의 결과물이다. Text는 LLM이 생성한 Text를 나타내고 이미지는 그 당시 Inference 과정에서의 Attention Score를 Visualization한 것이다.
Experiments
추가적으로 실험 결과를 살펴보겠다. 신기한 결과가 많아서 분석하는 재미가 있었다.
프롬프트에서 computer vision 이라는 토큰을 강조하자 Feature Detection, Object Detection, Image Segmentation 과 같은 실제 computer vision에서 사용되는 전문적인 용어의 결과물이 생성된 것을 볼 수 있다.
다음은 사진에 Highlighting을 한 결과를 나타낸다.
(a)부분 (남녀 커플)과 (b)부분 (선박)을 highlight 함으로써 LLM이 해당 Object에 대한 설명을 더 잘하는 것을 볼 수 있다.
정량적인 성능도 다른 모델들에 비해서 높은것으로 측정되었다고 한다. 근데 이건 뭐 비교대상이 되는 Prior work가 적고 엄청난 의미가 있다고 보이진 않는다.
나는 이런 평가과정은 처음봐서 신기했는데, 생성된 프롬프트를 다시 DALLE-3에 넣어서 생성된 이미지를 평가하는 방식이다. 평가 Metric은 CLIP Score로 평가한다.
Appendix
Pseudocode of Text generation process with Highlighting.
이 결과물은 Appendix에 나와있던건데 코드를 생성하는데 있어서도 Highlighting을 통해서 더 나은 결과물을 생성할 수 있었다. 한마디로 불필요하게 반복되는 코드는 생성하지 않고 목적에 맞는 코드만을 잘 생성하게 된다.
위 사진은 실제 Github 소스코드에서 Attention Activation이 어떻게 이뤄지는지에 대한 코드이다.
Conclusion
이 논문은 Training 없이 사용자가 원하는 토큰을 강조하고 싶을 때, LLM이 그 토큰을 참고하여 더 나은 결과물을 생성할 수 있도록 하는 방법을 제시한 논문이다. 나로서는 CFG를 이용한 Highlighting 방법이 가장 신기했고, 이론적으로 LLM-CFG가 뒷받침 되어있어서 더욱 설득력이 있었다. 추가로 Two-branch를 통한 Highlighting 방법이 참신하다고 느꼈다.
나머지 Solution들은 첫 번째 Controllable Guidance이 Impact가 컸기 때문에 Technical Contribution이 눈에 띄는 정도는 아니었다.
하지만 LLM이라는 연구주제에서 어떻게 저런 문제상황들(Imbalanced Attention Score)을 발견했는지 Motivation이 궁금하게 하는 Solution이었다.
WoLF 프로젝트 이후로 처음 읽은 논문인데, WoLF 이후의 논문도 LLM 관련 논문일것이기 때문에 LLM에 어떤 문제들이 있는지 더 면밀하게 조사해봐야겠다.