KSTU AI Lab

Mixture-of-Depths: Тереңдетилген сабак

Трансформерлердин эсептөө ресурстарын кантип акылдуулук менен бөлүштүрүп, ылдамдыгын жана натыйжалуулугун арттырарын билип алыңыз.

1. Көйгөй: "Бир Калыптагы" Эсептөө

Стандарттык Трансформер моделдеринде (мисалы, LLaMA, GPT) ар бир сөз (токен) моделдин бардык катмарларынан өтүүгө милдеттүү. Тексттеги "жана", "менен" сыяктуу жөнөкөй токендер да, "архитектура", "философия" сыяктуу татаал токендер да бирдей көлөмдөгү эсептөөнү талап кылат.

Бул өтө натыйжасыз. Муну ар бир суроо үчүн бардык окуучуларды — башталгыч класстагыларды да, мыктыларды да — мектептин 12 классынын баарынан өткөргөнгө салыштырса болот. Эсептөө баасы туруктуу жана жогору:

Эсептөө баасы ∝ (Токендердин саны) × (Катмарлардын саны)

Бул ыкма жөнөкөй токендер үчүн ашыкча ресурс сарптап, моделдин иштешин жайлатат.

2. Чечим: Mixture-of-Depths (MoD) менен Динамикалык Тереңдик

Mixture-of-Depths (MoD) — бул ар бир токенге канчалык деңгээлде эсептөө керектигин динамикалык түрдө чечкен акылдуу механизм. Башкача айтканда, маанилүү же татаал токендер моделдин терең катмарларына жөнөтүлсө, жөнөкөй токендер кээ бир катмарларды "секирип" өтүп, ресурсту үнөмдөйт.

A. Багыттоочу (Router) жана Top-k тандоо

MoD архитектурасында моделдин ар бир N-катмарынан кийин атайын багыттоочу (router) — кичинекей нейрон тармагы — жайгашат. Бул багыттоочу ар бир токенди анализдеп, анын "татаалдыгын" баалайт.

Андан соң, ал эң жогорку балл алган токендердин белгилүү бир бөлүгүн (мисалы, top-k, k=12.5%) тандап алат. Тандалган токендер гана кийинки катмарлар блогунда толук иштетилет.

B. Калдыктуу байланыш (Residual Connection)

Тандалбаган, "жөнөкөй" деп табылган токендер эмне болот? Алар жок кылынбайт. Алар кийинки катмарлар блогунан калдыктуу байланыш (residual connection) аркылуу түз өткөрүлүп, иштетилген токендердин жыйынтыгына кошулат. Бул маалыматтын жоголбоосун камсыздайт.

Бул процесс модель боюнча бир нече жолу кайталанат. Натыйжада, ар бир токен өзүнүн татаалдыгына жараша уникалдуу "жол" менен жүрөт.

3. Жыйынтык: Ылдамдык жана Натыйжалуулук

Mixture-of-Depths ыкмасы эсептөө формуласын өзгөртөт:

Эсептөө баасы ∝ (Иштетилген Токендердин саны) × (Катмарлардын саны)

Белгилүү бир катмарларда токендердин аз гана бөлүгү иштетилгендиктен, жалпы эсептөө (FLOPs) көлөмү кескин азаят. Бул бир нече артыкчылыктарды берет:

  • Ылдам үйрөтүү: Моделдер ошол эле жабдыкта бир нече эсе тезирээк үйрөтүлөт.
  • Натыйжалуу инференция: Жооп берүү убактысы кыскарат.
  • Күчтүүрөөк моделдер: Ошол эле эсептөө бюджети менен тереңирээк же кененирээк моделдерди курууга болот.

Mixture-of-Experts (MoE) менен айырмасы: MoE токенди кайсы экспертке (нейрон тармакка) жөнөтүүнү чечсе, MoD токенди кийинки катмарга жөнөтүү же жөнөтпөөнү чечет.

1. Проблема: "Один Размер для Всех" в Вычислениях

В стандартных моделях-трансформерах (например, LLaMA, GPT) каждый токен (слово) обязан пройти через абсолютно все слои сети. И простые токены, такие как "и" или "с", и сложные, как "архитектура" или "философия", требуют одинакового объема вычислений.

Это крайне неэффективно. Это можно сравнить с тем, чтобы заставлять каждого ученика — и новичка, и гения — проходить все 12 классов школы для ответа на любой вопрос. Вычислительная стоимость фиксирована и высока:

Стоимость вычислений ∝ (Количество токенов) × (Количество слоев)

Такой подход тратит ресурсы впустую на простые токены и замедляет работу модели.

2. Решение: Динамическая Глубина с Mixture-of-Depths (MoD)

Mixture-of-Depths (MoD) — это умный механизм, который динамически решает, сколько вычислений требуется каждому отдельному токену. Другими словами, важные или сложные токены отправляются на обработку в глубокие слои модели, в то время как простые токены "пропускают" некоторые слои, экономя ресурсы.

A. Маршрутизатор (Router) и отбор Top-k

В архитектуре MoD после каждых N слоев располагается специальный маршрутизатор (router) — небольшая нейронная сеть. Этот маршрутизатор анализирует каждый токен и оценивает его "сложность" или "важность".

Затем он выбирает определенную долю токенов с наивысшими оценками (например, top-k, где k=12.5%). Только эти выбранные токены отправляются на полную обработку в следующий блок слоев.

B. Остаточное соединение (Residual Connection)

Что происходит с токенами, которые не были выбраны и сочтены "простыми"? Они не отбрасываются. Они передаются напрямую через следующий блок слоев с помощью остаточного соединения (residual connection) и добавляются к результату обработанных токенов. Это гарантирует, что информация не будет потеряна.

Этот процесс повторяется несколько раз по всей модели. В результате каждый токен следует по уникальному "маршруту", зависящему от его сложности.

3. Итог: Скорость и Эффективность

Метод Mixture-of-Depths изменяет формулу вычислений:

Стоимость вычислений ∝ (Количество обработанных токенов) × (Количество слоев)

Поскольку на определенных слоях обрабатывается лишь малая часть токенов, общее количество вычислений (FLOPs) значительно сокращается. Это дает несколько преимуществ:

  • Быстрое обучение: Модели обучаются в несколько раз быстрее на том же оборудовании.
  • Эффективный инференс: Время генерации ответа сокращается.
  • Более мощные модели: Можно создавать более глубокие или широкие модели в рамках того же вычислительного бюджета.

Отличие от Mixture-of-Experts (MoE): MoE решает, какому эксперту (нейросети) отправить токен, тогда как MoD решает, отправлять ли токен на следующий слой вообще.

1. The Problem: "One-Size-Fits-All" Computation

In standard Transformer models (e.g., LLaMA, GPT), every single token is forced to pass through all layers of the network. Simple tokens like "and" or "the" require the same amount of computation as complex tokens like "architecture" or "philosophy".

This is highly inefficient. It's like making every student—from beginner to genius—go through all 12 grades of school for every single question. The computational cost is fixed and high:

Computational Cost ∝ (Number of Tokens) × (Number of Layers)

This approach wastes resources on simple tokens and slows down the model.

2. The Solution: Dynamic Depth with Mixture-of-Depths (MoD)

Mixture-of-Depths (MoD) is a smart mechanism that dynamically decides how much computation each token needs. In other words, important or complex tokens are sent for processing in deeper layers, while simple tokens can "skip" some layers, saving computational resources.

A. The Router and Top-k Selection

In the MoD architecture, a special router—a small neural network—is placed after every N layers. This router analyzes each token and assigns it a "complexity" or "importance" score.

It then selects a certain fraction of tokens with the highest scores (e.g., the top-k, where k=12.5%). Only these selected tokens are fully processed by the next block of layers.

B. The Residual Connection

What happens to the tokens that are not selected and deemed "simple"? They are not discarded. They are passed directly across the next block of layers via a residual connection and are added to the output of the processed tokens. This ensures no information is lost.

This process is repeated multiple times throughout the model. As a result, each token follows a unique "path" based on its complexity.

3. The Result: Speed and Efficiency

The Mixture-of-Depths approach changes the computation formula to:

Computational Cost ∝ (Number of Processed Tokens) × (Number of Layers)

Since only a fraction of tokens are processed at certain layers, the total computational cost (FLOPs) is significantly reduced. This provides several benefits:

  • Faster Training: Models can be trained several times faster on the same hardware.
  • Efficient Inference: The time it takes to generate a response is reduced.
  • More Capable Models: It's possible to build deeper or wider models within the same computational budget, leading to better performance.

Difference from Mixture-of-Experts (MoE): MoE decides which expert (neural network) to send a token to, whereas MoD decides whether to send a token to the next layer at all.

Библиография