Mixture of Experts: когда нейросеть учится делегировать

Привет, чемпионы!

Представьте, что у вас есть большой и сложный проект, и вы наняли двух управленцев: Кабан-Кабаныча и Руководителева. Вы даете им одинаковую задачу: набрать штат сотрудников и выполнить ваш проект. Вся прибыль вместе с начальным бюджетом останется у них.

Кабан-Кабаныч решил, что нет смысла платить отдельным специалистам по DevOps, backend, ML и другим направлениям, и нанял всего одного сотрудника за 80 монеток. Этот бедняга работал в стиле «один за всех» и, естественно, быстро выгорел и «умер». Кабан-Кабаныч, не долго думая, нанял еще одного такого же сотрудника. В итоге вы вернулись и увидели печальную картину: задачу никто не решил, остался лишь Кабан-Кабаныч и кладбище несчастных сотрудников.

А вот Руководителев поступил иначе: он распределил бюджет на несколько похожих сотрудников, но сначала не понимал, кто из них в чём лучше. Тогда он стал давать им небольшие задачи и внимательно наблюдать за результатами. Через некоторое время он понял, что сотрудник №1 на 70% лучше справляется с задачами по ML, сотрудник №2 на 80% эффективнее в backend-разработке и так далее. Так Руководителев постепенно сформировал команду экспертов, сам став управляющим (или «gating»-узлом), который распределяет задачи на основе знаний о возможностях каждого сотрудника. Сотрудники углубляли экспертизу в своих направлениях, а Руководителев становился всё эффективнее в распределении задач.

Внезапно мы пришли к интересному решению:

• Руководителев — это gating network, который распределяет задачи, исходя из предыдущих успехов сотрудников.

• Сотрудники — это local experts, каждый из которых специализируется на своей части задач.

Таким образом, мы экономим ресурсы, получаем сильных специалистов и достигаем отличных результатов за короткое время.

Именно так в 1991 году и появилось решение Adaptive Mixtures of local Experts

Этот подход доказал эффективность, сокращая время обучения моделей почти вдвое.

Как работает MoE?

Представьте модель, у которой есть входные и выходные данные, а между ними набор экспертов. Этих экспертов организует управляющая сеть (gating network), определяющая, какие эксперты могут лучше справиться с конкретной задачей. Gating-сеть, которая присваивает веса результату каждого эксперта, объединяя их в итоговый ответ.

Звучит красиво, но не всё так просто… Во время обучения возникают интересные и даже «ломающие мозг» ситуации, особенно когда осознаёшь, что созданная тобой модель может «вынести» тебя самого.

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

Ключевая особенность — разреженность. С помощью MoE можно масштабировать модель без пропорционального увеличения вычислительной нагрузки. Это очень важно, ведь позволяет обучать огромное количество экспертов, используя при этом только нужных. В этом нам помогает важный гиперпараметр — top_k, определяющий, сколько лучших экспертов будет выбрано для каждого входа.

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

Возникает закономерный вопрос: зачем тогда вообще нужны остальные эксперты?

Как с этим бороться? В своём коде я добавил трекер распределения данных по экспертам, чтобы контролировать, не «залип» ли гейт. Также я внедрил несколько хитрых решений, подсмотренных на профессиональных форумах.

Давайте кратко резюмируем:

Технология MoE выгодна за счёт разреженности и гибкости использования экспертов. Однако это «сделка с дьяволом», поскольку возникают сложности:

• Сложная балансировка работы экспертов.

• Функция потерь должна учитывать как производительность экспертов, так и маршрутизатора.

• Количество гиперпараметров (количество экспертов, архитектура gating-сети) усложняет настройку модели.

Где сейчас используют MoE?

Почти все современные LLM используют MoE. Например, недавно вышедшая модель Llama4 Scout с 16x17B параметрами — это 16 экспертов по 17 миллиардов параметров каждый. То есть на инференсе вы используете не все 272 млрд параметров, а только top_k выбранных. Впечатляющее снижение вычислительных затрат, правда?

Также технология активно применяется в компьютерном зрении, и сейчас мы её протестируем на простом примере V-MoEs.

Тест драйв технологии

Итак, для обучения возьмем простенький датасет CIFAR100 и обучим на нем нашу кастомную V-MoEs для классификации изображений.

Сама по себе архитектура будет состоять из следующего:

Классический VIT, но ее часть классификатора мы обернем в decoder блок, где у нас будет применена MOE

Начнем с маршрутизатора, в нашем случае он был реализован следующим образом

import torch import torch.nn as nn import torch.nn.functional as F  class GatingNetwork(nn.Module):     def __init__(self,                  input_dim = 151296,                  num_experts=4,                  top_k=2,                  use_noise=True,                  noise_std=1e-2,                  temperature=1.0):         super().__init__()          self.num_experts = num_experts         self.top_k = top_k         self.use_noise = use_noise         self.noise_std = noise_std         self.temperature = temperature          self.gate = nn.Sequential(             nn.Linear(input_dim, 512),             nn.ReLU(),             nn.Dropout(0.2),             nn.Linear(512, 256),             nn.ReLU(),             nn.Dropout(0.2),             nn.Linear(256, num_experts)         )      def forward(self, x):         logits = self.gate(x)  # (B, num_experts)          if self.use_noise and self.training:             scale = logits.std(dim=1, keepdim=True).clamp(min=1e-3)               noise = torch.randn_like(logits) * self.noise_std * scale             logits = logits + noise          topk_vals, topk_indices = torch.topk(logits, self.top_k, dim=1)          gates = F.softmax(topk_vals / self.temperature, dim=1)  # (B, top_k)          return topk_indices, gates 

Он берёт входной вектор x, оценивает, какие эксперты из num_experts лучше подойдут для каждого примера в батче, и возвращает top_k лучших экспертов с их весами.

То есть это — классическая Gating Network, которая решает, каким экспертам дать поработать с входом.

Обратите внимание, что тут есть noisy gating — это один из способов избежать «залипания гейта» на одном и том же эксперте. Во время тренировки шум масштабируется и в зависимости от поставленной нами пропорции влияет на решение о том какого эксперта повыбирать. Иными словами мы влияем на «результатова», чтобы он давал шансы большему числу экспертов, а не выбирал любимчиков.

Создадим экспертов

import torch.nn as nn   class FFNExpert(nn.Module):     def __init__(self, input_dim, hidden_dims, output_dim, dropout_prob=0.5):         super(FFNExpert, self).__init__()          layers = []         self.linears = nn.ModuleList()          prev_dim = input_dim         for hidden_dim in hidden_dims:             linear = nn.Linear(prev_dim, hidden_dim)             self.linears.append(linear)             layers.append(linear)             layers.append(nn.LayerNorm(hidden_dim))             layers.append(nn.ReLU())             layers.append(nn.Dropout(dropout_prob))             prev_dim = hidden_dim          final_linear = nn.Linear(prev_dim, output_dim)         self.linears.append(final_linear)         layers.append(final_linear)          self.network = nn.Sequential(*layers)         self._initialize_weights()      def _initialize_weights(self):         for linear in self.linears:             nn.init.xavier_uniform_(linear.weight)             if linear.bias is not None:                 nn.init.zeros_(linear.bias)      def forward(self, x):         return self.network(x)   class FFNExpertSmall(FFNExpert):     def __init__(self, input_dim, output_dim):         super(FFNExpertSmall, self).__init__(input_dim, hidden_dims=[256, 128], output_dim=output_dim, dropout_prob=0.3)   class FFNExpertMedium(FFNExpert):     def __init__(self, input_dim, output_dim):         super(FFNExpertMedium, self).__init__(input_dim, hidden_dims=[512, 256, 128], output_dim=output_dim,                                               dropout_prob=0.4)   class FFNExpertLarge(FFNExpert):     def __init__(self, input_dim, output_dim):         super(FFNExpertLarge, self).__init__(input_dim, hidden_dims=[1024, 512, 256, 128], output_dim=output_dim,                                              dropout_prob=0.5)   class FFNExpertVeryLarge(FFNExpert):     def __init__(self, input_dim, output_dim):         super(FFNExpertVeryLarge, self).__init__(input_dim, hidden_dims=[2048, 1024, 512, 256, 128],                                                  output_dim=output_dim, dropout_prob=0.6) 

Тут в целом все просто, мы набросали 4 эксперта с разными параметрами и посмотрим на то как они будут обучаться.

Начнем собирать модель

import torch.nn as nn import timm  class ViT_backbone(nn.Module):     def __init__(self):         super().__init__()         self.backbone = timm.create_model('vit_base_patch16_224',                                           pretrained=True)          for param in self.backbone.parameters():             param.requires_grad = False          self.embed_dim = self.backbone.head.in_features          self.backbone.reset_classifier(0)          self.ln = nn.LayerNorm(self.embed_dim)         self.ln2 = nn.LayerNorm(self.embed_dim)         self.attn = nn.MultiheadAttention(embed_dim=self.embed_dim,                                           num_heads=8,                                           batch_first=True)      def forward(self, x):         skip = self.backbone.forward_features(x)  # [B, N, D]         x_ln = self.ln(skip)         attn_out, _ = self.attn(x_ln, x_ln, x_ln)         x_attn = attn_out + skip         x_final = self.ln2(x_attn).flatten(1)  # [B, N*D]         return x_final

Тут все просто возьмем классическую модель VIT и добавим к ней слои нормализации после Multihead Attention и skip connection.

После сделаем наше объединение и наконец-то MOE

import torch import torch.nn as nn import torch.nn.functional as F from model.gating_network import GatingNetwork from model.Vit_model import ViT_backbone  class MoECNN(nn.Module):     def __init__(self,                  experts,                  input_for_gating = 151296,                  top_k=2,                  output_dim=100,                  use_aux_loss=True,                  aux_loss_weight=0.01,                  warmup_iters=500,                  noise_std = 0.5):         super().__init__()         self.num_experts = len(experts)         self.top_k = top_k         self.output_dim = output_dim         self.use_aux_loss = use_aux_loss         self.aux_loss_weight = aux_loss_weight         self.warmup_iters = warmup_iters         self.iter = 0          self.backbone = ViT_backbone()         self.experts = nn.ModuleList(experts)         self.gating = GatingNetwork( input_dim = input_for_gating,                                      num_experts=self.num_experts,                                      top_k = top_k,                                      noise_std=noise_std)          self.register_buffer("expert_usage",                              torch.zeros(self.num_experts))      def forward(self, x):         batch_size = x.size(0)         device = x.device         x = self.backbone(x)          if self.training and self.iter < self.warmup_iters:             random_indices = torch.randint(0,                                            self.num_experts,                                            (batch_size, self.top_k),                                            device=device)             gates = torch.full((batch_size, self.top_k),                                1.0 / self.top_k,                                device=device)             topk_indices = random_indices             self.iter += 1         else:             topk_indices, gates = self.gating(x)          output = torch.zeros(batch_size, self.output_dim, device=device)         self.expert_usage.zero_()          for i in range(self.top_k):             idx = topk_indices[:, i]             for expert_idx in torch.unique(idx):                 expert_mask = (idx == expert_idx)                 if expert_mask.sum() == 0:                     continue                 x_sel = x[expert_mask]                 y_sel = self.experts[expert_idx](x_sel)                 gate_weight = gates[expert_mask, i].unsqueeze(1)                 output[expert_mask] += gate_weight * y_sel                  self.expert_usage[expert_idx] += expert_mask.sum()          aux_loss = None         if self.use_aux_loss and self.training:             usage = self.expert_usage / batch_size             aux_loss = ((usage - usage.mean()) ** 2).mean() * self.aux_loss_weight          return output, aux_loss 

Первое на что обратим внимание это это warmup_iters. Тут у нас это число итераций где мы как-бы отключаем gating-сеть , чтобы избежать коллапса распределения (один эксперт выбирается чаще остальных до того, как сеть обучится разумно маршрутизировать входы). Это дает нам «разогреть» экспертов передавая им равномерно данные и далее мы уже начинаем более тонко избирать экспертов за счет gating network.

Второй момент это добавление use_aux_loss. Данный параметр позволяет нам учитывать в общем лоссе неравномерное распределение по экспертам в общий loss.

Как итог модель выбирает tok_k экспертов и на основе их предсказаний делает взвешенную сумму, после чего выдает результат и loss по распределению.

Что в итоге?

При простом «на коленке» мы смогли получить f1 на тесте 89.%. Более явно поиграв с гипперпараметрами, типами экспертов и некоторыми изощренностями думаю, что можно получить результат лучше. Самое главное, что

Давайте теперь проведем модель через один батч и посмотрим, что там на одном батче, что произошло по графику и посмотрим на первые 10 сэмплов батча.

Как можем увидеть, у нас 2 эксперт оказался в данной итерации не востребован, а использовали мы с 0,1,3 эксперта в разной пропорции.

Можно сказать, что вот: «второй эксперт переобучился или обучился плохо». Однако давайте глянем глубже! Мы ведь отслеживаем все через clearml 🙂

На тестовом датасете в среднем можно заметить, что все эксперты примерно вышли на какую-то свою зону ответственности. Хотя конечно от первого эксперта хотелось ожидать побольше!

Теперь давайте посмотрим на визуализацию результатов:

Несмотря на шакальность(мы работаем с CIFAR100 напоминаю) мы получили весьма неплохие результаты. И теперь вишенка на торте — это отслеживание по экспертам. Их собственно говоря мы итак логируем и сейчас можем провести на маленьком сэмпле аналитику. Если у нас есть очень большой эксперт и он не пригодился в использовании в вычислениях, то мы можем сэкономить очень много памяти.

Подводя итоги основным концептом было показать какие проблемы бывают и сложности при работе с технологией, а также ее возможности и потенциал, который уже сейчас очень успешно реализуется!

🔥 Ставьте лайк и напишите какие темы было бы интересно разобрать дальше! Самое главное — пробуйте и экспериментируйте!

✔️ Присоединяйтесь к нашему Telegram-сообществу @datafeeling, чтобы первыми применять на практике передовые технологии!