[논문 리뷰] In-Context Learning for Attention Scheme: from Single Softmax Regression to Multiple Softmax Regression via a Tensor Trick

Yeqi Gao, Zhao Song외 1명|arXiv (Cornell University)|2023. 07. 05.

Tensor decomposition and applicationsMathematics인용 수 3

한 줄 요약

이 논문은 어텐션 메커니즘에서의 인-컨텍스트 학습을 위한 텐서-트릭 기반 프레임워크를 제안하며, 행렬 기반 소프트맥스 회귀를 벡터화된 변환을 통해 재구성하여 안정성과 일반화 능력을 향상시킨다. 손실 함수의 리프시츠 연속성을 확립하고, 모델 가중치나 파라미터의 변동이 출력 분포에 유한한 이동을 유도함을 증명하며, 정규화 및 재스케일링된 소프트맥스 공식화 모두에서 모델 용량과 입력 차원에 대한 지수적 의존성을 고려하여 강건한 인-컨텍스트 적응이 가능함을 보여준다.

ABSTRACT

Large language models (LLMs) have brought significant and transformative changes in human society. These models have demonstrated remarkable capabilities in natural language understanding and generation, leading to various advancements and impacts across several domains. We consider the in-context learning under two formulation for attention related regression in this work. Given matrices $A_1 \in \mathbb{R}^{n imes d}$, and $A_2 \in \mathbb{R}^{n imes d}$ and $B \in \mathbb{R}^{n imes n}$, the purpose is to solve some certain optimization problems: Normalized version $\min_{X} \| D(X)^{-1} \exp(A_1 X A_2^ op) - B \|_F^2$ and Rescaled version $\| \exp(A_1 X A_2^ op) - D(X) \cdot B \|_F^2$. Here $D(X) := \mathrm{diag}( \exp(A_1 X A_2^ op) {\bf 1}_n )$. Our regression problem shares similarities with previous studies on softmax-related regression. Prior research has extensively investigated regression techniques related to softmax regression: Normalized version $\| \langle \exp(Ax) , {\bf 1}_n angle^{-1} \exp(Ax) - b \|_2^2$ and Resscaled version $\| \exp(Ax) - \langle \exp(Ax), {\bf 1}_n angle b \|_2^2 $ In contrast to previous approaches, we adopt a vectorization technique to address the regression problem in matrix formulation. This approach expands the dimension from $d$ to $d^2$, resembling the formulation of the regression problem mentioned earlier. Upon completing the lipschitz analysis of our regression function, we have derived our main result concerning in-context learning.

연구 동기 및 목표

• 소프트맥스 함수의 행렬 공식화를 사용하여 어텐션 기반 회귀에서의 인-컨텍스트 학습을 체계화하기.
• 가중치 공간과 입력 공간 양쪽에서 파라미터 편향에 대한 소프트맥스 손실 함수의 리프시츠 연속성 분석하기.
• 정규화 및 재스케일링된 소프트맥스 회귀에서 모델 가중치나 입력의 변화가 출력 분포에 미치는 영향에 대한 이론적 범위 설정하기.
• 작은 파라미터 업데이트가 모델 출력에 제어 가능하고 유한한 변화를 유도함을 보여주어 안정적인 인-컨텍스트 적응 가능하게 하기.
• 벡터화 및 텐서 기법을 통해 기존 소프트맥스 회귀 연구를 확장하여 매트릭스 형태의 어텐션 메커니즘으로 통합하기.

제안 방법

• 논문은 매트릭스 회귀 문제를 고차원 벡터 공간으로 변환하기 위해 벡터화 기법을 사용하며, 파라미터 차원을 $ d $ 에서 $ d^2 $ 으로 확장한다.
• 두 가지 공식화를 도입한다: 정규화된 형태 $ \min_X \|D(X)^{-1}\exp(A_1 X A_2^\top) - B\|_F^2 $ 와 재스케일링된 형태 $ \|\exp(A_1 X A_2^\top) - D(X)\cdot B\|_F^2 $ 로, 여기서 $ D(X) $ 는 지수 활성화의 행 합을 대각선으로 가지는 행렬이다.
• 분석은 $ u(x), \alpha(x), f(x), h(x), c(x), q(x) $ 및 그 역함수들의 핵심 함수를 유한하게 제한하는 데 기반하며, 리프시츠 연속성을 사용해 민감도를 제어한다.
• 손실 함수의 기울기의 리프시츠 범위를 유도하며, 출력의 변화가 파라미터 변화에 비례하며 요소 $ M = \exp(O(R^2 + \log n)) $ 를 통해 결정됨을 보여준다.
• 프레임워크는 $ x $ 나 $ \mathsf{A} $ 에 대한 편향이 출력 분포에 유한한 이동을 유도함을 증명하며, $ \|\delta_c\|_2 \leq M \cdot \|x_{t+1} - x_t\|_2 $ 를 만족한다.
• 이론적 결과는 매트릭스 지수 함수와 대각선 스케일링의 성질을 활용하여 중간 함수의 노름을 재귀적으로 제한함으로써 도출된다.

실험 결과

연구 질문

• RQ1 매트릭스 파라미터를 $ \mathbb{R}^{d^2} $ 로 벡터화하는 것은 소프트맥스 기반 어텐션 회귀의 안정성과 일반화에 어떤 영향을 미치는가?
• RQ2 파라미터 편향에 대한 정규화 및 재스케일링된 소프트맥스 손실 함수의 기울기의 리프시츠 상수는 무엇인가?
• RQ3 모델 가중치나 입력의 작은 변화가 어텐션 메커니즘의 출력 분포에 얼마나 제한된 변화를 유도하는가?
• RQ4 대규모 언어 모델의 인-컨텍스트 학습 행동은 매트릭스 어텐션 회귀의 텐서-트릭 재구성으로 이론적으로 정당화될 수 있는가?
• RQ5 범위 $ M $ 에서 $ R^2 $ 과 $ \log n $ 에 대한 지수적 의존성이 인-컨텍스트 적응의 강건성에 어떤 영향을 미치는가?

주요 결과

• 정규화된 소프트맥스 손실 함수는 파라미터 벡터 $ x $ 와 매트릭스 $ \mathsf{A} $ 에 대해 리프시츠 연속성을 보이며, 기울기 변화는 $ M \cdot \|x_{t+1} - x_t\|_2 $ 로 유한하게 제한되며, 여기서 $ M = \exp(O(R^2 + \log n)) $ 이다.
• x 기반 업데이트의 경우 출력 이동 $ \|\delta_c\|_2 $ 는 $ M \cdot \|x_{t+1} - x_t\|_2 $ 로 유한하게 제한되어 작은 파라미터 변화에 대한 안정적인 인-컨텍스트 적응을 보장한다.
• A 기반 업데이트의 경우 출력 이동 역시 $ M \cdot \|A_{t+1} - A_t\|_2 $ 로 유한하게 제한되어 가중치 편향에 대해 강건함을 확인한다.
• 이론적 분석은 $ x_t $ 에서 $ x_{t+1} $ 로 또는 $ A_t $ 에서 $ A_{t+1} $ 로의 전이가, 편향된 목표 $ \widetilde{b} $ 나 $ \widehat{b} $ 를 가진 새로운 안정적인 회귀 문제로 이어짐을 확인하며, $ \|\widetilde{b} - b\|_2 \leq M \cdot \|x_{t+1} - x_t\|_2 $ 를 만족한다.
• 범위는 $ c(x) = \alpha(x)^{-1} u(x) $ 와 같은 중간 함수에 대해 리프시츠 성질을 재귀적으로 적용함으로써 도출되며, 지수적 및 다항식 항을 사용해 노름 성장에 엄격한 제어를 한다.
• 이 프레임워크는 모델 업데이트가 출력에 제어 가능하고 예측 가능한 변화만을 유도함을 보여주며, 어텐션 메커니즘에서의 인-컨텍스트 학습에 대한 이론적 기반을 제공함으로써 신뢰할 수 있는 소수의 샘플 일반화를 가능하게 한다.