[24’ ICLR] PASTA: Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs
카테고리: Language
태그: Attention
🔍 Abstract
- Problem
- 저자들은 LLM도 우리가 강조한 텍스트, 즉 Bold나 Italic처럼 강조하고자 하는 특정 텍스트에 집중할 수 있어야 한다고 주장한다. 그러나 기존의 방법들, 즉
*를 사용하거나"를 사용하는 방법은 효과적이지 않으며, Few-shot learning은 Prompt의 품질에 매우 민감하다.
- 저자들은 LLM도 우리가 강조한 텍스트, 즉 Bold나 Italic처럼 강조하고자 하는 특정 텍스트에 집중할 수 있어야 한다고 주장한다. 그러나 기존의 방법들, 즉
- Insight & Method
- 저자들은 이를 Internal Mechanism을 조작하여 해결할 수 있다고 생각한다. 즉, 강조하고자 하는 Text 부분의 Attention을 높여 성능을 향상시키고자 한다.
- 그러나, 모든 Head의 Attention을 조작하는 것은 각 Head의 역할이 다르기 때문에 적절하지 않다. 대신, 저자들은 실제 성능 개선에 도움이 되는 Head를 Small Training Set을 통해 찾고, 해당 Head의 Attention을 조작하는 방법으로 접근한다.
- Result
- Instruction Following, Contextual Information Interpretation, Counterfactual Conflict Resolving 등 다양한 Task에서 성능 향상을 보였다. 이 성능 향상은 다른 Baseline에 비해 굉장히 크다.
- 이 방법은 Robust하며, 성능 향상에 도움이 되는 Head를 결정하는 것이 성능에 큰 영향을 미친다.
1. PASTA: Post-hoc Attention Steering Approach
1.1. Multi-Task Model Profiling
PASTA(Past-hoc Attention Steering Approach)는 Attention을 조작하여 성능을 향상시키는 방법이다. 이를 위해 Small Training Set을 이용하여 성능 향상에 도움이 되는 Head를 찾아내고, 해당 Head의 Attention을 조작한다. 이 두 과정을 각각 (1) Multi-Task Model Profiling과 (2) Post-hoc Attention Steering이라고 부른다.
먼저 Multi-Task Model Profiling은 Small Training Set을 이용하여 성능 향상에 도움이 되는 Head를 찾아내는 과정이다. 이 과정은 다음과 같다.
- 각 Task $i$에 대해 small training subset $\mathcal{D}^ {(i)}$를 구한다.
- 각 Layer $l$, Head $h$에 대해 Attention Steering을 진행한다. Attention Steering의 자세한 과정은 다음과 같다.
- 강조하고자 하는 Context $\mathcal{G}$와 그렇지 않은 Context $\mathcal{G}^ -$가 있다고 하자.
- 이때 강조하지 않고자 하는 $j \in \mathcal{G}^ -$에 대한 Attention을 $\alpha$만큼 scaling하여 줄인다. 그리고 나머지는 합이 1이 되도록 정규화한다.
- 이를 통해 강조하고자 하는 $j \in \mathcal{G}$에 대한 Attention을 높일 수 있다.
- 값을 높이는 것이 아니라, 낮추는 방식으로 구현한 것은 구현의 안정성 때문이다. 수식은 1.2절에서 자세히 다룬다.
- 이렇게 구한 Attention Steering의 결과를 통해 성능 향상에 도움이 되는 Head를 찾아낸다. 즉, 성능이 향상된 순서대로 Ranking $R^ {(i)} = [(l_1, h_1), (l_2, h_2), \cdots]$을 구한다.
- 모든 Task에 대하여 top-k 수준으로 성능이 향상되는 Head Set $\mathcal{H} = \cup_ i R^ {(i)} _ {1:k}$를 구한다.
1.2. Post-hoc Attention Steering
Post-hoc Attention Steering은 Attention Steering을 통해 구한 성능 향상에 도움이 되는 Head Set $\mathcal{H}$에 대해 Attention을 조작하는 과정이다. 이 과정은 다음과 같은 수식에 의해 이루어진다.
\[\mathbf{H}^ {(l, h)} = \mathcal{T} (\mathbf{A}^ {(l, h)}) \mathbf{V}, \quad \text{where } [\mathcal{T} (\mathbf{A})]_ {ij} = \begin{cases} \alpha \mathbf{A}_ {ij} / \mathbf{C}_ i & \text{if } j \in \mathcal{G}^ - \\ \mathbf{A}_ {ij} / \mathbf{C}_ i & \text{otherwise} \end{cases}\] \[\mathbf{C}_ i = \sum_ {j \in \mathcal{G}} \mathbf{A}_ {ij} + \sum_ {j \in \mathcal{G}^ -} \alpha \mathbf{A}_ {ij}\]합을 1으로 하는 것은 분포를 유지하기 위하여 중요하다. 일부는 값을 곱하는 것이 아니라 더하는 것을 선호할 수도 있다. 그러나, 일정한 값을 더하는 것은 Attention Score를 Uniformly distributed하게 만들어 불필요한 정보 손실을 일으킬 수 있다. 따라서 곱하는 방식을 사용하는 것이 더 효과적이다.
1.3. Algorithm
전체 알고리즘은 위와 같다.
2. Experiments
2.1. Main Results
전체 결과는 위와 같고, 알 수 있듯이 PASTA는 다른 Baseline에 비해 성능 향상이 굉장히 크다. 저자들은 대략 $k=400$ 에서 $k=500$ 정도로 설정하였고, 따라서 $50 \leq \vert \mathcal{H} \vert \leq 100$ 정도로 설정하였을 때 성능 향상이 가장 컸다고 보고하였다. 예시 결과는 아래와 같이 확인할 수 있다.
2.2. Ablation Study
Model Profiling. 모든 Layer의 모든 Head를 Steering하거나, 하나의 Layer 전체를 Steering하는 것, 혹은 하나의 Head만 Steering하는 것은 모두 성능 향상에 뚜렷한 영향을 미치지 않았다. 따라서 성능 향상에 도움이 되는 Head를 찾는 것이 성능 향상에 중요하다.
Varying the Number of Heads. Head의 수가 늘어날수록 Instruction Following 즉 강조한 것에 대한 결과는 좋아지지만 전반적인 Performance, Fluency가 떨어지는 것을 확인할 수 있다. 따라서 적절한 Trade-off를 만족하는 Head의 수를 찾는 것이 중요하다. 한편, Scaling Factor $\alpha$는 값이 굉장히 많이 바뀌어도 성능에 큰 영향을 미치지 않았다. 즉, $\alpha$는 Robust하다.
💡 Summary
지금까지 24’ ICLR에 발표된 PASTA에 대해 알아보았다. 해당 논문의 내용을 요약하면 다음과 같다.
- User-specified Text를 강조하고자 해당 부분의 Attention을 높이는 Attention Steering을 도입하였으며, 성능 향상에 도움이 되는 Head를 찾아내어 선택적으로 Steering을 적용하는 것이 중요함을 확인
- Steering의 정도를 결정하는 Scaling Factor $\alpha$는 Robust하여 큰 영향을 미치지 않음
댓글 남기기