Demystify Transformer의 컨텍스트 기반 학습 기능(feat. TransNormer)

GPT와 같은 LLM(대형 언어 모델)의 성공은 Softmax 어텐션 메커니즘과 불가분의 관계가 있지만 이 메커니즘도 높은 비용과 같은 몇 가지 단점이 있습니다. 최근 Shanghai Artificial Intelligence Laboratory와 OpenNLPLab의 연구팀은 Softmax 기반의 어텐션 메커니즘을 완전히 버리고 새로 제안된 선형 어텐션을 대신 사용하는 새로운 대규모 언어 모델 TransNormerLLM을 제안했습니다. 보고서에 따르면 TransNormer LLM은 선형 주의를 기반으로 하는 최초의 대규모 언어 모델(LLM)로, 정확도와 효율성 측면에서 기존의 Softmax 주의 기반 모델을 능가합니다. 연구원들은 사전 훈련된 모델의 오픈 소스 버전도 출시할 예정입니다.

논문: https://arxiv.org/abs/2307.14995
모델: https://github.com/OpenNLPLab/TransnormerLLM
관련된 포스팅
paperwithcode Transformer 소개

대규모 언어 모델은 자연어 처리(NLP) 분야에 혁명을 일으켰습니다. 컴퓨터 프레임워크에서 인간의 언어를 이해하고, 생성하고, 상호 작용하는 능력을 향상시키면서 다양한 유형의 작업에 탁월합니다. 언어 모델링의 이전 개발은 주로 Transformer 아키텍처를 중심으로 진행되었으며, 그 중 기둥 모델에는 기본 Transformer, GPT 시리즈, BERT 및 BART 및 기타 획기적인 모델이 포함됩니다. Transformer 아키텍처의 성공은 데이터 기반 모델 체계에서 입력 토큰 간의 종속성을 식별할 수 있고 글로벌 위치도 인식할 수 있는 softmax 어텐션 메커니즘을 기반으로 하므로 모델이 장기적으로 효과적으로 처리할 수 있습니다. -자연어의 범위 역학.

그렇더라도 기존 트랜스포머에는 여전히 한계가 있습니다. 무엇보다도 시퀀스 길이와 관련하여 2차 시간 복잡성으로 인해 확장성이 제한되고 교육 및 추론 단계에서 계산 리소스와 시간 효율성이 떨어집니다. 이 2차 시간 복잡도를 선형 복잡도로 단순화하기 위해 많은 연구자들이 다양한 시퀀스 모델링 방법을 제안했습니다. 그러나 이러한 방법은 다음 두 가지 이유로 LLM에 사용하기 어렵습니다.

본 논문에서 제안하는 TransNormerLLM은 선형 어텐션을 기반으로 한 최초의 LLM으로 정확도와 효율성 모두에서 전통적인 소프트맥스 어텐션을 능가합니다. TransNormerLLM의 구성은 이전의 선형 어텐션 아키텍처인 TransNormer를 기반으로 하며 성능을 개선하기 위해 일부 수정이 이루어졌습니다. TransNormerLLM의 주요 개선 사항에는 위치 임베딩, 선형 주의 속도 향상, 게이팅 메커니즘, 텐서 정규화 및 추론 속도 향상이 포함됩니다.

특별한 주의를 기울여야 하는 개선 사항 중 하나는 TransNormer의 DiagAttention을 선형 주의로 대체하여 전반적인 상호 작용 성능을 향상시킬 수 있다는 것입니다. 연구원들은 또한 희석 문제를 해결하기 위해 지수적 감쇠가 있는 LRPE를 도입했습니다. 또 연구진은 라이트닝 어텐션(Lightning Attention)이라는 신기술도 도입해 훈련 중 선형 주의 속도를 2배로 높일 수 있고, 지각 IO를 통해 메모리 사용량도 4배까지 줄일 수 있다고 말했다. 뿐만 아니라 GLU 및 정규화 방법을 단순화했으며 후자는 전체 속도를 20% 높였습니다. 그들은 또한 서로 다른 시퀀스 길이에서 수치적 안정성과 일정한 추론 속도를 보장할 수 있는 강력한 추론 알고리즘을 제안하여 훈련 및 추론 단계에서 모델의 효율성을 개선했습니다.

연구원들은 TransNormerLLM의 효과를 검증하기 위해 6TB 이상의 대용량 코퍼스와 2조 토큰 이상을 신중하게 수집했습니다. 데이터의 품질을 보장하기 위해 수집된 코퍼스를 필터링하기 위한 자체 정리 전략도 개발했습니다. 표 1에서 볼 수 있듯이 연구원들은 원래의 TransNormer 모델을 확장하고 매개변수가 3억 8,500만에서 1,750억에 이르는 여러 TransNormerLLM 모델을 얻었습니다. 그런 다음 그들은 새로운 대규모 코퍼스를 기반으로 포괄적인 실험과 제어 변수 연구를 수행했으며 결과는 새로운 방법이 softmax Attention 기반 방법보다 성능이 우수하고 훈련 및 추론 속도가 더 빠르다는 것을 보여줍니다.

Transformer TransNormer — *표 1: TransNormerLLM 모델의 변형*

LLM 분야의 연구를 촉진하기 위해 Shanghai Artificial Intelligence Lab과 OpenNLPLab의 연구원들은 자체 사전 교육 모델을 소스로 공개할 것입니다. 효율적인 변환기 구조”.

Transformer LLM

아키텍처 개선

다음은 TransNormerLLM의 다양한 모듈과 연구원이 제안한 몇 가지 개선 조치를 간략하게 소개합니다.

개선 1: 위치 인코딩

TransNormer의 하위 계층은 희석 문제를 피하기 위해 DiagAttention을 사용합니다. 그러나 이로 인해 토큰 간의 글로벌 상호 작용이 부족합니다. 이 문제를 해결하기 위해 연구원들은 TransNormerLLM에 대한 지수 감쇠가 있는 LRPE(Linearized Relative Position Encoding)를 사용하여 하위 계층에서 완전한 관심을 유지했습니다. 연구원들은 이 접근법을 LRPE-d라고 불렀습니다.

개선 2: 게이팅 메커니즘

게이팅은 모델의 성능을 향상시키고 학습 프로세스를 원활하게 할 수 있습니다. 연구원들은 TransNormerLLM에 대해 “Transformer quality in linear time” 논문의 Flash 방법을 사용했고 토큰 혼합에서 GLA(Gated Linear Attention) 구조를 사용했습니다.

모델의 속도를 더욱 향상시키기 위해 게이트 자체가 비선형성을 도입할 수 있기 때문에 원래 GLU 구조의 활성화 기능을 제거한 SGLU(Simple GLU)도 제안했습니다.

개선 3: 텐서 정규화

연구원들은 TransNormer에 도입된 NormAttention을 사용했습니다. TransNormerLLM에서는 RMSNorm을 새로운 단순 정규화 함수인 SimpleRMSNorm(약칭 SRMSNorm)으로 대체했습니다.

전체 프레임

그림 1은 TransNormerLLM의 전체 구조를 보여줍니다.

이 아키텍처에서 입력 X의 업데이트는 두 가지 순차적 단계로 수행됩니다. 먼저 SRMSNorm 정규화를 사용하여 GLA(Gated Linear Attention) 모듈을 통과합니다. 그런 다음 SGLU(Simple Gated Linear Unit) 모듈을 통해 SRMSNorm 정규화를 사용합니다. 이 전체 아키텍처는 모델의 성능을 향상시키는 데 도움이 됩니다. 이 전체 프로세스의 의사 코드는 다음과 같습니다.

훈련 최적화

light attention

Attention 계산 속도를 높이기 위해 연구원들은 새로 제안된 선형 Attention을 IO(입력 및 출력) 처리에 더 적합하게 만드는 Lightning Attention 알고리즘을 도입했습니다.

알고리즘 1은 라이트닝 어텐션의 포워드 패스 구현 세부 사항을 보여주고 알고리즘 2는 백워드 패스입니다. 연구원들은 기울기를 더 빠르게 계산할 수 있는 구현 버전도 가지고 있으며 이는 향후 출시될 것이라고 말했습니다.

모델 병렬화

컴퓨터 클러스터 전체에 걸쳐 모든 모델 매개변수, 그래디언트 및 옵티마이저 상태 텐서를 배포하기 위해 연구자들은 완전 분할 데이터 병렬(FSDP/완전 분할 데이터 병렬)을 사용했습니다. 이 전략적 파티셔닝 접근 방식은 각 GPU의 메모리 공간을 줄여 메모리 활용을 최적화합니다. 효율성을 더욱 향상시키기 위해 역방향 전달 중에 메모리에 캐시된 활성화 수를 줄이는 활성화 검사점을 사용했습니다. 대신, 이러한 기울기는 계산할 때 제거되고 다시 계산됩니다. 이 기술은 컴퓨팅 효율성을 개선하고 리소스를 절약하는 데 도움이 됩니다. 또한 GPU 메모리 소비를 줄이면서 계산 속도를 높이기 위해 연구원들은 AMP(Automatic Mixed Precision)도 사용했습니다.

위의 성과 외에도 연구원들은 Nvidia의 Megatron-LM 모델의 병렬화에서 크게 영감을 받은 선형 변압기 실행 모델을 병렬화하여 시스템 엔지니어링을 더욱 최적화했습니다. 어텐션 모듈 다음에 2층 MLP(다층 퍼셉트론) 모듈이 있습니다. Megatron-LM 모델 병렬성을 사용할 때 이 두 모듈에서 독립적으로 사용됩니다. 마찬가지로 TransNormerLLM 구조도 SGLU와 GLA의 두 가지 주요 모듈로 구성되며 이 둘의 모델 병렬화는 별도로 수행됩니다.

강력한 추론

이를 통해 TransNormerLLM은 RNN의 형태로 추론을 수행할 수 있습니다. 알고리즘 3은 이 프로세스의 세부 사항을 제공합니다. 그러나 거기에는 수치적 정밀도의 문제가 있습니다.

이러한 문제를 피하기 위해 연구원들은 강력한 추론 알고리즘을 제안했습니다(알고리즘 4 참조).

원래 추론 알고리즘과 강인한 추론 알고리즘으로 얻은 결과는 동일합니다.

구조

연구원들은 총 700TB가 넘는 방대한 양의 공개 텍스트를 인터넷에서 수집했습니다. 수집된 데이터는 그림 2와 같이 데이터 전처리 절차를 거쳐 처리되어 약 2조 개의 토큰이 포함된 6TB의 클린 코퍼스를 남겼습니다. 더 나은 투명성을 제공하고 사용자가 더 잘 이해할 수 있도록 돕기 위해 데이터 소스를 분류했습니다. 표 2는 특정 범주를 제공합니다.

실험

연구원들은 Metaseq 프레임워크에서 PyTorch 및 Trition을 사용하여 TransNormerLLM을 구현했습니다. 모델은 Adam 옵티마이저를 사용하여 학습되며 FSDP는 모델을 NVIDIA A100 80G 클러스터로 효율적으로 확장하는 데에도 사용됩니다. 또한 모델 병렬 처리를 적절하게 사용하여 성능을 최적화합니다.

건축 절제 실험