[Paper Review] TOKENFORMER: RETHINKING TRANSFORMER SCALINGWITH TOKENIZED MODEL PARAMETERS

안녕하세요. 오늘은 Transformer의 Layer 구조를 변형하여 Model Parameter를 Token으로 활용하여 모델의 확장성을 높이고, 학습 비용을 줄인 TOKENFORMER를 리뷰하겠습니다. 이름에서 부터 직관적으로 Transformer의 구조를 TOKEN화 한것 같은 느낌을 줍니다. Transformer 논문을 보고 오시면 이해가 쉽습니다.

- 논문 관련 문서 -

arXiv 30 Oct 2024. [Paper] [Page] [Github]

- 목차 -

1. Background

2. Token-Parameter Attention(Pattention)

3. Progressive Model Scaling

4. Result

---

1. Background

• Transformer의 한계점: Linear projection
• 본문 내용: Conversely, token-parameter computations rely heavily on linear projections (Dunford & Schwartz, 1988), where input tokens are multiplied by a fixed set of parameters. This prescribed design limits scalability because increasing the model size requires altering core architectural components, often necessitating retraining the entire model from scratch.
• 분석: Transformer의 매개변수 계산은 Linear Projection에 의존하는데, 여기서 입력 토큰은 고정된 매개변수 집합으로 곱해진다고 합니다. 아래 예시와 같이 고정된 매개변수인 W를 사용하는데, 이때 모델의 크기를 확장하기 위해서는 W의 크기를 키워야합니다. 그러면, 차원의 크기가 맞지 않아 연산이 불가능하여, 기존 모델에서 사용할 수 없죠. 그래서 모델을 확장하게 되면 처음부터 다시 학습해야합니다.

Linear Projection 행렬 연산

• 이런 모델 확장에 대한 유연성을 가지기 위해 고안된 것이 Tokenformer입니다.
• 본문 내용: Specifically, to project features with input and output dimensions D1 and D2, we employ two sets of parameters, each comprising N learnable tokens with channel dimensions of D1 and D2,
  respectively. In this formulation, input tokens serve as queries, and model parameters as keys and
  values. This flexibility renders our model’s parameters inherently scalable with variable N, allowing
  for efficient expansion by continuously adding new key-value parameter pairs.
• 분석: 입/출력 차원이 D1/D2 일때, Linear projection을 하면 W 크기가 고정된다. 하지만, Token 형식으로 Key-Value를 구성하면 N개만큼 블록으로 쌓아 버리면, 모델을 그냥 확장해버릴 수 있다. 아래 내용을 보면, 초기 모델에서 계산된 Key-Value Tokens을 그대로 두고, N만큼 새롭게 추가된 Key-Value Tokens을 쌓아서 출력을 내면 W가 고정되지 않고 N만큼 더 행렬 연산을 수행할 수 있게 된다.

초기 모델 파라미터

초기 모델 계산 과정

모델 확장할때, N=1으로 가정 했을때, 파라미터 정보

모델 확장 - N이 늘어나서 계산되는 과정

2. Token-Parameter Attention(Pattention)

• TokenFormer: Transformer 구조에서 Q, K, V Projection 부분을 Pattention layer로 대체한 구조이다.
• Pattention Layer: Key, Value 파라미터를 Tokens으로 사용하여 Input Tokens과 Key Tokens을 dot product하여 Attention Scores를 구하고 거기에 Value Tokens를 Weighted Sum하여 최종 Output Tokens을 산정한다.
• Parameter Reusing in Pattention: 사전 훈련된 Old Key-Value Param Tokens 세트가 장착된 기존 Tokenformer 모델을 고려해 보면, New Key-Value Param Tokens를 추가하여 이 모델을 확장합니다. 이러한 스케일링 방식은 입력이나 출력 차원을 변경하지 않고도 임의의 수의 매개변수를 적분하는 것을 허용합니다.
• Git repo [2]: TokenFormer/blob/main/megatron/model/tokenformer.py 이 코드를 보면, Pattention 클래스가 있는데, 여기서 초기화 init 함수에서 파라미터로 param_token_num을 활용하여 모델을 확장할 수 있는 것으로 추정된다.

3. Progressive Model Scaling

 

• 본문 내용: While these approaches can be effective, they often disturb the pre-established distribution of the smaller model, increasing the risk of losing pre-trained knowledge and slowing convergence. In contrast, our model allows for parameter scaling in a natural and seamless manner and preserves the integrity of the existing model.
• 분석: 모델이 무결성을 가질 수 있다고 하는데, 사전 학습된 Key-Value param Tokens 값은 고정해버리고 추가되는 Key-Value param Tokens만 업데이트하면 무결성을 보장할 수 있는 것이다.
• Progressive Model Scaling: Tokenformer를 사용하는 우리의 점진적 스케일링 방법론은 처음부터 학습한 Transformer 모델과 비슷한 성능을 달성하는 동시에 학습 예산을 크게 줄입니다.
• 왼쪽 그래프: 누적 계산 예산을 통해 모델 확장 비용 평가. Transformer 기준선은 처음부터 독립적으로 수행한 각 개별 확장 단계에 대한 비용이 발생하는 반면 Tokenformer는 처음에 124M 모델을 학습하고 점진적으로 354M, 757M, 1.4B 매개변수로 확장하는 것을 포함하여 모든 확장 단계에 걸쳐 비용을 집계합니다. 오른쪽 그래프: 각 확장 단계에 필요한 예산을 측정하여 모델 확장 비용 평가. 사용된 Transformer 기준선은 30B 및 300B 토큰으로 학습한 그림 3에 나와 있는 것과 일치합니다. 마찬가지로 Tokenformer의 경우 비용은 더 작은 것에서 각 증분 확장 단계에 필요한 예산입니다. 모든 실험은 TPU v4 하드웨어에서 수행되었습니다.

4. 결론 [1]

TokenFormer

Rethinking Transformer Scaling with Tokenized Model Parameters

 

TokenFormer는 토큰-토큰 및 토큰-매개변수 상호작용의 계산을 전적으로 어텐션 메커니즘을 사용하여 통합하고 신경망의 유연성을 극대화하는 fully attention-based neural network입니다.

데이터뿐만 아니라 모델 매개변수를 Key, Value로 토큰화하여 모델 개념을 Input Tokens와 매개변수 토큰(Key, Value Tokens)간의 상호작용으로 대체합니다. 이를 통해 가변적인 수의 매개변수를 처리할 수 있고 기존 Transformer보다 더 큰 유연성을 제공합니다.

Incremental Model Scaling: 이 모델은 New Key-Value Tokens 매개변수 쌍을 추가하여 처음부터 다시 학습하지 않고도 점진적 스케일링을 허용하여 학습 비용을 크게 줄입니다. 이 이점이 이 논문의 중요성! 입니다.
Device-Cloud Collaboration: 기기-클라우드 협업에서 클라우드 지식 기반 역할을 할 수 있으며 각 키-값 매개변수 토큰 쌍은 학습 가능한 패턴을 나타냅니다.
Sparse Inference (MoE): Tokenformer를 각 Key-Value 쌍이 전문가 역할을 하는 extreme Mixture of experts으로 해석하여 추론 비용을 크게 줄입니다.
Parameter Efficient Tuning: 새로운 작업이나 데이터 세트에 직면했을 때 모델은 이러한 새로운 매개변수 토큰을 통합하여 사전 훈련된 매개변수를 증강할 수 있으므로 특정 작업 요구 사항에 빠르게 적응할 수 있습니다.
Integrating Vision and Language Models: 사전 훈련된 시각 및 언어 모델의 Key-Value Tokens을 병합하여 시각 언어의 원활한 통합을 가능하게 합니다.

 

[Reference]

[1] https://haiyang-w.github.io/tokenformer.github.io/
[2] https://github.com/Haiyang-W/TokenFormer