대형 언어 모델(LLMs)이 다양한 작업에서 탁월한 성공을 거두는 데는 Transformer 아키텍처의 핵심 혁신인 자기 주의 메커니즘이 큰 역할을 한다. 쿼리와 키 간의 쌍별 유사도 점수를 계산하여, 이 메커니즘은 모델이 동적으로 문맥에 민감한 표현을 구축하게 한다. 기본 원리는 의미적으로 관련된 토큰에 더 큰 가중치를 부여하고 그렇지 않은 토큰에는 작은 가중치를 부여함으로써, 시퀀스 내의 토큰 간의 의미 있는 의존성을 포착하는 것이다.

그러나 경험적 증거는 표준 주의 메커니즘이 이상적인 동작에서 벗어난다는 것을 보여주며, 두 가지 널리 관찰되는 현상인 “표현력 붕괴"와 “주의 침전"을 일으킨다. 표현력 붕괴는 상대적으로 긴 문맥 시나리오에서 주의 메커니즘이 너무 많은 토큰으로부터 정보를 모으려고 하여 과압축이 발생하고 결국 구별할 수 없는 최종 표현을 초래한다. 반면, 주의 침전은 초기 토큰이 높은 가중치를 비례적으로 받는다는 것을 의미하며, 이러한 토큰들은 의미적 중요성이 적음에도 불구하고 주의 분포를 안정화하는 역할을 한다. 이 두 현상은 모델 성능을 저하시키고 효율적인 배포를 방해하며, 현재의 주의 설계에 깊은 구조적 결함이 있음을 시사한다.

기존 연구들은 이러한 문제들을 각각 다룬다. 붕괴 문제가 발생할 때 일반적으로 차원이나 수치 정밀도를 증가시키는 방법을 사용하여 표현 손실을 완화하지만, 이는 훈련과 추론 비용이 크게 증가하는 대신이다. 침전 현상에 대해서는 주로 “침전 토큰"이나 “침전 편향"을 명시적으로 보존하거나 정규화를 제거하여 해결하려고 한다. 그러나 이러한 접근법은 모델의 유연성을 저하시키거나 성능 안정성에 손상을 초래할 수 있다. 이러한 방법들은 개별 문제를 완화하지만, 본질적으로 분산되어 있으며 단일 프레임워크 내에서 두 가지 문제를 모두 해결하는 데 실패한다.

본 논문에서는 위의 문제들에 대한 통합적인 관점을 제공한다: 자기 주의 메커니즘의 핵심 제한 사항은 부적절한 주의 할당이며, 이는 두 가지 극단으로 나타난다(Figure 1). 주의 과부하는 밀도가 높은 문맥에서 토큰 간에 주의를 너무 넓게 분산시켜 의미적 특징을 평균화하고 구별할 수 없는 표현을 초래한다. 반면, 주의 부족은 이전 문맥과 거의 관련이 없음에도 불구하고 소프트맥스 정규화가 요구하는 가중치의 합이 1이 되도록 하는 경우, 침전 현상이 발생한다. 이러한 두 가지 실패 모드는 주의 메커니즘이 중요한 의미에 더 잘 맞춰진 선택적인 집중을 할당하도록 개선할 필요성을 강조한다.

따라서 우리는 Lazy Attention이라는 접근법을 제안하며, 이는 집중 할당을 개선하기 위해 두 가지 보완적인 메커니즘을 통합한다: 1) 위치별 차별화: 이 메커니즘은 RoPE와 학습 가능한 주의 편향을 결합하여 각 주의 헤드와 헤드 차원 간에 위치 특징의 차별화를 강화한다. 이를 통해 토큰 표현의 혼동을 줄이고 주의 과부하를 완화시킨다. 2) 탄성 소프트맥스: 이 변종은 주의 가중치에 오프셋을 도입하여 표준 소프트맥스의 엄격한 정규화 제약을 완화한다. 이를 통해 모델이 무관심한 토큰에게 0 가중치를 할당하고 침전 문제를 해결할 수 있다. 이러한 메커니즘은 더 의미적으로 관련된 표현에 맞춘 집중적이고 희소성 있는 주의를 촉진하여 언어 모델링의 효과성을 개선한다.

우리의 기여는 다음과 같이 요약된다:

• 우리는 주의 과부하와 부족을 두 가지 기본적인 주의 할당 실패 모드로 특징짓는 통합적 관점을 제공한다.
• Lazy Attention을 제안하며, 이는 (i) RoPE와 학습 가능한 주의 편향을 결합하여 헤드와 차원 간의 위치별 차별화를 강화하고 주의 과부하를 완화시키며, (ii) 탄성 소프트맥스를 사용하여 무시할 수 있는 가중치를 필터링하여 침전 문제를 억제한다.
• 다양한 벤치마크에서 실험을 통해 Lazy Attention은 주의 침전을 완화하고 주의 가중치에서 평균 59.58%의 희소성을 달성하며 언어 모델링 성능을 개선함으로써, 주의 설계에 강화된 집중이 중요하다는 것을 확인한다.

/>

제반 사항

Transformer 아키텍처

Transformer는 $`L`$ 개의 층으로 구성되어 있으며, 각 층은 multi-head attention (MHA) 모듈과 feed-forward network (FFN)를 포함하고 있다. 잔차 연결(residual connections)이 두 모듈에 모두 적용된다. 입력 시퀀스가 층 $`l`$에서 $`\boldsymbol{X}^l = [\boldsymbol{x}_1, \boldsymbol{x}_2, ..., \boldsymbol{x}_n] \in \mathbb{R}^{n \times d}`$이고, 여기서 $`n`$은 시퀀스 길이이며 $`d`$는 은닉 차원이라면, 층 업데이트는 다음과 같이 작성된다:

\begin{equation}
\hat{\boldsymbol{X}}^l = \boldsymbol{X}^l + \text{MHA}(\text{LN}(\boldsymbol{X}^l)), 
\quad
\boldsymbol{X}^{l+1} = \hat{\boldsymbol{X}}^l + \text{FFN}(\text{LN}(\hat{\boldsymbol{X}}^l)),
\end{equation}

여기서 $`\text{LN}(\cdot)`$는 층 정규화를 나타낸다. MHA 메커니즘은 모델이 여러 표현 하위 공간에서 동시에 정보를 포착할 수 있게 한다. 각 주의 헤드 $`h \in [H]`$에 대해 쿼리, 키, 값은 다음과 같이 계산된다: $`\boldsymbol{Q}_h = \boldsymbol{X}^l \boldsymbol{W}_h^Q, \boldsymbol{K}_h = \boldsymbol{X}^l \boldsymbol{W}_h^K, \boldsymbol{V}_h = \boldsymbol{X}^l \boldsymbol{W}_h^V,`$ 여기서 $`\boldsymbol{W}_h^Q, \boldsymbol{W}_h^K, \boldsymbol{W}_h^V \in \mathbb{R}^{d \times d_h}`$는 투영 행렬이다. 마지막으로 주의 연산은 다음과 같다:

\begin{equation}
\text{Attention}(\boldsymbol{Q}_h, \boldsymbol{K}_h, \boldsymbol{V}_h) = \text{softmax}\!\left(\boldsymbol{Q}_h \boldsymbol{K}_h^\top/\sqrt{d_h}\right)\boldsymbol{V}_h,
\end{equation}

여기서 $`d_h = d/H`$. 주의 가중치는 각 쿼리 위치에 대해 시퀀스에 대한 확률 분포를 형성하는 $`\text{softmax}(\boldsymbol{Q}_h \boldsymbol{K}_h^\top / \sqrt{d_h}) \in \mathbb{R}^{n \times n}`$이다.

위치 인코딩

자기 주의 메커니즘은 순서에 무관하므로, 위치 정보를 모델에 통합해야 한다. Transformer는 정현 함수를 사용하여 고정된 위치 인코딩을 생성하며 이는 종종 **절대 위치 인코딩(APE)**이라고 불린다: $`\boldsymbol{P}_{(i, 2t)} = \sin(i/10000^{2t/d})`$와 $`\boldsymbol{P}_{(i, 2t+1)} = \cos(i/10000^{2t/d})`$. 여기서 $`i`$는 토큰의 위치, $`t`$는 차원, 그리고 $`d`$는 은닉 차원 크기를 나타낸다. GPT와 같은 후속 연구에서는 학습 가능한 위치 임베딩을 모델 매개변수로 제공한다. 두 가지 변형 모두에서 위치 임베딩 $`\boldsymbol{P} \in \mathbb{R}^{n \times d}`$는 토큰 임베딩에 추가되어 첫 번째 Transformer 층으로 입력된다: $`\text{TokenEmbed}(\boldsymbol{X}) + \boldsymbol{P}`$. 그러나 APE는 학습 시퀀스를 벗어난 길이 일반화에서 부족하다.

APE의 한계점을 해결하기 위해 **상대 위치 인코딩(RPE)**이 도입되었다. RPE는 토큰 간의 쌍별 거리를 모델링한다. 임베딩을 추가하는 대신, RPE는 주로 주의 점수를 수정한다. 중요한 방법론에는 Transformer-XL이 포함되어 있으며, 이는 내용과 상대 위치를 분리시키는 미세한 주의 공식을 사용하고, ALiBi는 단순히 주의 점수에 정적 편향(예: $`-m \cdot |i-j|`$)을 추가한다. 최근에는 **회전 위치 임베딩(RoPE)**이 LLMs에서 채택되었다. RoPE는 회전 변환을 적용하여 상대 위치를 인코딩하며 다음과 같이 간결하게 표현된다: $`\mathbf{q}'_i = \boldsymbol{R}_i \mathbf{q}_i,\; \mathbf{k}'_i = \boldsymbol{R}_i \mathbf{k}_i`$, 여기서 $`\boldsymbol{R}_i`$는 블록 대각 회전 행렬이며 주의 점수가 상대 위치를 반영하도록 한다. 더 관련된 연구들은 부록 8.1에 상세히 설명되어 있다.

요약하자면, 위치 인코딩 방법은 절대적에서 상대적이고 회전 설계로 진화하여 더욱 풍부하고 유연한 위치 의존성을 포착하는 경향을 보여준다.

주의 침전과 위치 인코딩에 대한 재고

주의 침전의 본질

/>

/>

/>

주목 침전은 추론 중 초기 토큰이 대량의 주의를 받는 현상을 말한다. 관련 논의는 부록 8.2에 제공된다.

, 우리는 학습 중 특정 위치에 고정된 [Mask] 토큰을 삽입하여 침전 위치를 조사한다. 이를 “Mask@k“로 표기하며 k는 입력에서의 위치를 나타낸다. 그림 2에 따르면, 고정된 토큰이 시작 부분 근처(예: Mask@2)에 배치될 때 첫 번째 위치에서 주의를 받는다. 그러나 시작부분으로부터 약 16개 이상의 토큰 뒤 (예: Mask@16)에 삽입되면 첫 번째 토큰이 다시 침전 현상을 보인다.

한편, 지적했듯이 포지셔널 임베딩 없이도 트랜스포머의 은닉 상태 변동에는 잠재적인 위치 정보가 포함되어 있다. 이는 그림 3과 4의 비교와 일치한다. 침전 현상이 존재할 때, 침전 토큰의 은닉 상태 변동은 다른 토큰보다 훨씬 크다. 마찬가지로, 주의 계산에서 침전 토큰 값의 변동도 다른 토큰보다 훨씬 작으며, 이는 부록 13에 상세히 설명되어 있다.

위치 인코딩의 역할

/>

/>

/>

/>

/>

/>

/>

/>

/>

/>

위의 발견을 바탕으로 우리는 위치 인코딩이 주의 행동에 미치는 영향을 검토한다. “the the the…“와 같은 반복된 토큰 시퀀스를 제공함으로써, 의미적 차이와 토큰 임베딩 변동성을 제거하고 침전 토큰을 구별하는 데 도움이 된다. 그림 16은 다섯 가지 위치 인코딩을 사용하는 대표적인 LLM들에서 주의 가중치를 시각화한다.

APE (예: GPT-2)는 반복된 입력에서도 침전 현상이 지속된다(그림 11). 대조적으로, RoPE와 ALiBi를 사용하는 Llama2와 Bloom과 같은 RPE를 채택한 LLM들은 침전 현상을 완전히 잃는다. 그들의 주의 분포는 위치 인코딩의 장기 감소에 의해 조절되는 균일 패턴으로 붕괴된다. 이는 RoPE, ALiBi와 같은 RPE 방법이 주의 가중치를 본질적으로 위치에 민감하게 제공하며, 은닉 상태에 위치 정보를 포함시키지 않고도 이를 전 층에 전파할 수 있음을 나타낸다.

두 모델은 특히 독특한 행동을 보인다. Transformer-XL (그림 16a.2)는 첫 번째 토큰이 아니라 구두점 및 기타 특수 토큰에서 침전 현상을 보인다. 이는 블록 재귀 학습 방식으로 인한 것이다. 대조적으로, GPT-OSS은 반복된 토큰 입력에서도 침전을 유지한다. 이것은 GPT-OSS이 소프트맥스 분모에 학습 가능한 편향을 추가하여 가상의 침전 토큰으로 기능하기 때문이다. 이러한 편향은 동일한 토큰 간에도 약간의 주의 가중치 차이를 유발하여 모델이 위치 구별성을 유지할 수 있게 한다.

Limit to 15k chars for stability