DeepMind는 Transformer를 업그레이드하여 순방향 패스 FLOP를 최대 절반까지 줄일 수 있습니다.

하이브리드 깊이를 도입한 DeepMind의 새로운 디자인은 Transformer 효율성을 크게 향상시킬 수 있습니다.

Transformer의 중요성은 말할 필요도 없습니다. 현재 이 혁신적인 기술을 개선하기 위해 노력하는 많은 연구팀이 있습니다. 중요한 개선 방향 중 하나는 Transformer에 적응형 컴퓨팅 기능을 제공하는 등 효율성을 높이는 것입니다. 이렇게 하면 불필요한 계산이 절약됩니다.

Transformer 아키텍처 제안자 중 한 명이자 NEAR 프로토콜의 공동 창립자인 Illiya Polosukhin은 얼마 전 Jen-Hsun Huang과의 대화에서 다음과 같이 말했습니다. "적응형 컴퓨팅은 다음에 나타나야 합니다. 우리는 구체적으로 문제에 얼마나 많은 컴퓨팅 리소스가 사용되는지 주의하세요. "

사실 인간은 적응형 컴퓨팅 능력을 갖고 태어났습니다. 사람들은 다양한 문제를 해결할 때 자연스럽게 다른 시간과 에너지를 할당합니다. .

언어 모델링에서도 마찬가지입니다. 정확한 예측 결과를 얻기 위해 모든 토큰과 시퀀스에 동일한 시간이나 리소스를 투자할 필요는 없습니다. 그러나 Transformer 모델은 순방향 전달에서 각 토큰에 대해 동일한 양의 계산을 소비합니다. 이것은 사람들을 한탄하게 만듭니다. 대부분의 계산이 낭비됩니다!

이상적으로는 불필요한 계산을 수행할 수 없다면 Transformer의 계산 예산을 줄일 수 있습니다.

조건부 계산은 필요할 때만 계산을 수행하여 전체 계산량을 줄이는 기술입니다. 이전에도 많은 연구자들이 언제 계산이 수행되는지, 얼마나 많은 계산이 사용되는지를 평가할 수 있는 다양한 알고리즘을 제안해 왔습니다.

그러나 이 어려운 문제의 경우 일반적으로 사용되는 솔루션 형식은 동적 계산 그래프를 도입하는 경향이 있기 때문에 기존 하드웨어 제한에 잘 대처하지 못할 수 있습니다. 대신, 가장 유망한 조건부 계산 방법은 현재 하드웨어 스택을 일관되게 사용하고 하드웨어의 최대 활용도에 따라 선택된 정적 계산 그래프와 알려진 텐서 크기의 사용을 우선시하는 방법일 수 있습니다.

최근 Google DeepMind는 Transformer에서 사용하는 계산량을 줄이기 위해 더 낮은 컴퓨팅 예산을 사용하여 이 문제를 연구했습니다.

• 논문 제목: Mixture-of-Depths: 동적으로 컴퓨팅 할당 변환기 기반 언어 모델
• 논문 주소: https://arxiv.org/pdf/2404.02258.pdf

그들은 다음을 구상합니다. 각 계층에서 네트워크는 각 토큰에 대해 결정을 내리는 방법을 학습하여 사용 가능한 컴퓨팅 예산을 동적으로 할당해야 합니다. 특정 구현에서 총 계산 노력은 훈련 전에 사용자에 의해 설정되며 네트워크가 작동하는 동안 내려지는 결정의 함수가 아니라 결코 변경되지 않습니다. 이를 통해 하드웨어 효율성 향상(예: 감소된 메모리 공간 또는 전달 패스당 FLOP 감소)을 미리 예상하고 활용할 수 있습니다. 팀의 실험은 전체 네트워크 성능을 저하시키지 않고도 이러한 이점을 얻을 수 있음을 보여줍니다.

DeepMind의 이 팀은 전체 네트워크 깊이에 걸쳐 동적 토큰 수준 라우팅 결정이 수행되는 MoE(Mixed Expert) Transformer와 유사한 접근 방식을 채택합니다.

MoE와는 달리 여기에서 선택 사항은 토큰에 계산을 적용하거나(표준 Transformer와 동일) 잔여 연결을 통해 우회(동일하게 유지, 계산 저장)하는 것입니다. MoE와의 또 다른 차이점은 이 라우팅 메커니즘이 MLP와 다중 헤드 주의 모두에 사용된다는 것입니다. 따라서 이는 네트워크에서 처리하는 키와 쿼리에도 영향을 미치므로 경로는 업데이트되는 토큰뿐만 아니라 주의할 수 있는 토큰도 결정합니다.

DeepMind는 각 토큰이 Transformer 깊이에서 서로 다른 수의 레이어 또는 모듈을 통과한다는 사실을 강조하기 위해 이 전략을 MoD(Mixture-of-Depths)라고 명명했습니다. 여기서는 이를 "혼합 깊이"로 번역합니다(그림 1 참조).

MoD는 사용자가 성능과 속도를 평가할 수 있도록 지원합니다. 한편, 사용자는 일반 Transformer와 동일한 훈련 FLOP를 사용하여 MoD Transformer를 훈련할 수 있으며, 이는 최종 로그 확률 훈련 목표에서 최대 1.5% 향상을 가져올 수 있습니다.반면에 MoD Transformer는 더 적은 계산을 사용하여 일반 Transformer와 동일한 훈련 손실을 달성합니다. 즉, 순방향 패스당 FLOP가 최대 50% 더 적습니다.

이러한 결과는 MoD Transformer가 지능적으로 라우팅하는 방법을 학습할 수 있음을 보여줍니다(예: 불필요한 계산 건너뛰기).

Implementing Mixed Depth (MoD) Transformer

요약하면 전략은 다음과 같습니다.

• 정적 컴퓨팅 예산을 기존의 동등한 수준보다 낮게 설정합니다. Transformer에 필요한 계산 방법은 모듈 계산(즉, self-attention 모듈 및 후속 MLP)에 참여할 수 있는 시퀀스의 토큰 수를 제한하는 것입니다. 예를 들어 일반 변환기는 시퀀스의 모든 토큰이 self-attention 계산에 참여하도록 허용할 수 있지만 MoD 변환기는 시퀀스에서 토큰의 50%만 사용하도록 제한할 수 있습니다.
• 각 토큰에는 스칼라 가중치를 제공하는 라우팅 알고리즘이 있습니다. 이 가중치는 모듈 계산에 참여할지 아니면 우회할지 여부와 같이 각 토큰의 라우팅 선호도를 나타냅니다.
• 각 모듈에서 상위 k개의 가장 큰 스칼라 가중치를 찾으면 해당 토큰이 모듈 계산에 참여합니다. k개의 토큰만 이 모듈의 계산에 참여해야 하므로 계산 그래프와 텐서 크기는 훈련 과정에서 정적입니다. 이러한 토큰은 라우팅 알고리즘에 의해 인식되는 동적 및 컨텍스트 관련 토큰입니다.

라우팅 체계

팀에서는 두 가지 학습된 라우팅 체계(그림 2 참조)인 토큰 선택과 전문가 선택을 고려했습니다.

토큰 선택적 라우팅 방식에서 라우팅 알고리즘은 계산 경로(예: MoE Transformer의 전문가 ID 전체)에서 각 토큰에 대한 확률 분포를 생성합니다. 그런 다음 토큰은 선호하는 경로(즉, 확률이 가장 높은 경로)로 전송되며, 보조 손실은 모든 토큰이 동일한 경로로 수렴되지 않도록 보장합니다. 토큰 선택 라우팅에는 토큰이 가능한 경로 간에 적절하게 분할되는지 보장되지 않기 때문에 로드 밸런싱 문제가 있을 수 있습니다.

전문가 선택적 라우팅은 위의 방식을 반대로 합니다. 즉, 토큰이 선호하는 경로를 선택하도록 하는 대신 각 경로는 토큰 선호도에 따라 상위 k 토큰(top-k)을 선택합니다. 이는 각 경로가 항상 k개의 토큰을 보장하므로 완벽한 로드 밸런싱을 보장합니다. 그러나 이로 인해 일부 토큰이 여러 경로의 상위 k에 있을 수 있고 다른 토큰에는 해당 경로가 없을 수 있으므로 일부 토큰이 과도하게 또는 과소 처리될 수도 있습니다.

DeepMind의 선택은 세 가지 이유로 전문가 선택 라우팅을 사용하는 것입니다.

첫째, 보조 잔액 손실이 필요하지 않습니다.

둘째, 상위 k를 선택하는 작업은 라우팅 가중치의 크기에 따라 달라지므로 이 라우팅 방식을 사용하면 현재 모듈 계산에 가장 필요한 토큰을 결정하는 데 도움이 되는 상대 라우팅 가중치를 사용할 수 있습니다. 라우팅 알고리즘은 가장 중요한 토큰이 상위 k 중에 있도록 가중치를 적절하게 설정할 수 있습니다. 이는 토큰 선택적 라우팅 방식으로는 수행할 수 없는 작업입니다. 특정 사용 사례에는 기본적으로 null 작업인 계산 경로가 있으므로 중요한 토큰을 null로 라우팅하는 것을 피해야 합니다.

셋째, 라우팅은 두 개의 경로만 통과하므로 단일 top-k 작업으로 토큰을 상호 배타적인 두 세트(각 계산 경로당 한 세트)로 효율적으로 나눌 수 있으며, 이는 위에서 언급한 사항에 대처할 수 있습니다. 또는 처리 중인 문제.

이 라우팅 방식의 구체적인 구현에 대해서는 원본 논문을 참조하세요.

Sampling

전문가 선택 라우팅에는 많은 장점이 있지만 명백한 문제도 있습니다. 즉, top-k 작업은 인과적이지 않습니다. 즉, 주어진 토큰의 라우팅 가중치가 상위 k에 있는지 여부는 그 뒤의 라우팅 가중치 값에 따라 달라지지만 자동회귀 샘플링을 수행할 때는 이러한 가중치를 얻을 수 없습니다.

이 문제를 해결하기 위해 팀에서는 두 가지 방법을 테스트했습니다.

첫 번째는 간단한 보조 손실을 도입하는 것입니다. 실습을 통해 언어 모델링의 주요 목표에 미치는 영향은 0.2%−0.3%이지만 모델이 자동 회귀 방식으로 샘플링할 수 있음이 입증되었습니다. 그들은 라우팅 알고리즘의 출력이 로짓을 제공하는 이진 교차 엔트로피 손실을 사용했으며 이러한 로짓의 상위 k를 선택하여 대상을 제공할 수 있습니다(즉, 토큰이 상위 k에 있으면 1, 그렇지 않으면 0).

두 번째 방법은 작은 보조 MLP 예측기(다른 라우팅 알고리즘과 마찬가지로)를 도입하는 것입니다. 이 예측기의 입력은 라우팅 알고리즘의 입력과 동일하지만(스톱 그래디언트 포함) 출력은 예측 결과입니다. 시퀀스의 top-k에 있는지 여부입니다. 이 접근 방식은 언어 모델링 목표에 영향을 미치지 않으며 실험에 따르면 이 단계의 속도에는 큰 영향을 미치지 않습니다.

이러한 새로운 방법을 사용하면 라우팅할 토큰을 선택하여 자동 회귀 샘플링을 수행하거나 향후 토큰에 대한 정보에 의존하지 않고 라우팅 알고리즘의 출력을 기반으로 모듈을 우회할 수 있습니다. 실험 결과에 따르면 이는 99%의 정확도를 빠르게 달성할 수 있는 비교적 간단한 보조 작업입니다.

Results

Training, isoFLOP 비교

먼저, 팀은 최적의 하이퍼파라미터를 결정하기 위해 상대적으로 작은 FLOP 예산(6e18)으로 일부 모델을 트레이닝했습니다(아래 그림 3 참조). .

전반적으로 MoD Transformer가 기준선 isoFLOP 곡선을 오른쪽 아래로 드래그하는 것을 볼 수 있습니다. 즉, 최적의 MoD 변환기는 최적의 기준 모델보다 손실이 낮고 매개변수도 더 많습니다. 이 효과는 운이 좋은 결과를 가져옵니다. 비록 하이퍼파라미터 설정에서 자체적으로 isoFLOP 최적은 아니지만 최적의 기본 모델보다 더 나은 성능을 발휘하는 일부 MoD 모델이 있습니다. 예를 들어, 220M 매개변수(그림 3의 모델 번호 3)가 있는 MoD 변형은 isoFLOP 최적 기준 모델(역시 220M 매개변수, 그림 3의 모델 번호 1)보다 약간 우수하지만 이 MoD 변형은 훈련 중 단계는 다음과 같습니다. 60% 이상 빨라졌습니다.

아래 그림 4는 총 FLOP가 6e18, 2e19 및 1e20일 때 isoFLOP 분석을 보여줍니다. 볼 수 있듯이 FLOP 예산이 더 커질수록 이러한 추세는 계속됩니다.

아래 그림 5는 인터리브 라우팅 모듈을 사용하여 훈련된 MoD Transformer의 라우팅 결정을 보여줍니다. 많은 수의 모듈 바이패스에도 불구하고 이 MoD Transformer는 여전히 일반 Transformer보다 더 나은 성능을 달성합니다.

자동 회귀 평가

MoD 변형의 자동 회귀 샘플링 성능도 평가했으며 결과는 아래 그림 6에 나와 있습니다. 이러한 결과는 MoD 변환기를 통해 달성된 계산 절감 효과가 훈련 설정에만 국한되지 않음을 보여줍니다.

MoDE(혼합 깊이 및 전문성)

MoD 기술은 MoE 모델과 자연스럽게 통합되어 소위 MoDE 모델을 형성할 수 있습니다. 아래 그림 7은 MoDE와 MoDE가 제공하는 개선 사항을 보여줍니다.

MoDE에는 단계적 MoDE와 통합 MoDE의 두 가지 변형이 있습니다.

단계적 MoDE는 자체 주의 단계 전에 라우팅 우회를 수행하거나 토큰 작업에 도달하는 반면, 통합 MoDE는 일반 MLP 전문가 간에 "비작업" 전문가를 통합하여 MoD 라우팅을 구현합니다. 전자의 장점은 토큰이 self-attention 단계를 건너뛸 수 있다는 것이고, 후자의 장점은 라우팅 메커니즘이 간단하다는 것입니다.

팀은 통합 방식으로 MoDE를 구현하는 것이 전문가의 능력을 직접적으로 줄이고 잔여 라우팅 설계를 구현하기 위해 토큰을 폐기하는 데 의존하는 것보다 훨씬 낫다는 점을 알아냈습니다.

위 내용은 DeepMind는 Transformer를 업그레이드하여 순방향 패스 FLOP를 최대 절반까지 줄일 수 있습니다.의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!