Paper name
SegMoE: Segmind Mixture of Diffusion Experts
Paper Reading Note
Blog URL: https://blog.segmind.com/introducing-segmoe-segmind-mixture-of-diffusion-experts/
Project URL: https://huggingface.co/blog/segmoe
Code URL: https://github.com/segmind/segmoe
Intro Video URL: https://www.youtube.com/watch?v=6Q4BJOcvwGE
TL;DR
2024 年 Segmind 研发的全球首个用于 Stable Diffusion 的开源专家混合（Mixture of Experts，MoEs）框架。这是一种能够将多个稳定扩散模型动态组合在一起的框架，无需训练即可在短时间内创建更大的 MoE 模型。
Introduction
背景
- 专家混合是一类稀疏深度学习模型，包含由线性层和门控层组成的稀疏线性层。这些门控层会将输入路由到少数几个线性层，从而节省计算资源并减少推理时间。这种架构使得模型在保持推理时间较低的同时，能够拥有大量的参数。
- 受 Mixtral 模型的启发，创建一个类似于 mergekit 的框架，以结合 Stable Diffusion 模型，这将显著增加模型的知识量，同时保持接近单个 Stable Diffusion 模型的推理时间。
本文方案
- 我们的框架 SegMoE 将 Stable Diffusion 1.5 和 Stable Diffusion XL 模型结合成一个专家混合风格的模型，提升了模型的提示依从性和多样性。未来，我们计划支持更多模型，并支持 SegMoE 模型的训练，这可能进一步提高模型的质量，并为文本生成图像提供一个新的 SOTA（最先进）模型。
什么是 SegMoE？
SegMoE 模型遵循与 Stable Diffusion 相同的架构。与 Mixtral 8x7b 类似，SegMoE 模型包含多个模型。这种方式是通过用稀疏 MoE 层替换部分前馈层来实现的。MoE 层包含一个路由器网络，用于选择哪个专家模型最有效地处理哪些 token。你可以使用 segmoe 包来创建你自己的 MoE 模型！整个过程只需几分钟。
关于名称
SegMoE 模型的命名为 SegMoE-AxB，其中 A 代表 MoE 组合在一起的专家模型数量，B 代表参与每个图像生成的专家数量。根据配置设置，模型的某些层（前馈块、注意力层或全部）会被复制，而其余参数则与 Stable Diffusion 模型相同。
推理
我们在 Hub 上发布了三个合并模型：
- SegMoE 2x1 有两个专家模型。
- SegMoE 4x2 有四个专家模型。
- SegMoE SD 4x2 有四个 Stable Diffusion 1.5 专家模型。
效果
- SegMoE 4x2
免责声明及正在进行的工作
- 速度较慢：如果每个 token 的专家数量大于 1，MoE 将在多个专家模型之间执行计算。这使得它比单个 SD 1.5 或 SDXL 模型更慢。
- 高 VRAM 使用：MoEs 推理速度很快，但仍需要大量的 VRAM（因此需要昂贵的 GPU）。这使得它们在本地设置中使用具有挑战性，但它们非常适合多 GPU 部署。作为参考，SegMoE-4x2 在半精度下需要 24GB 的 VRAM。
- SegMoE 已全面集成到 Hugging Face 生态系统中，并得到 diffusers 的支持。
代码分析
- 主要分析代码：https://github.com/segmind/segmoe/blob/main/segmoe/main.py
SparseMoeBlock
- SparseMoeBlock 的实现与 Mixtral 基本一致，稍微做了一定简化
# 受 transformers.models.mixtral.modeling_mixtral.MixtralSparseMoeBlock 的启发
class SparseMoeBlock(nn.Module):
def __init__(self, config, experts):
super().__init__()
self.hidden_dim = config["hidden_size"] # 隐藏层的维度
self.num_experts = config["num_local_experts"] # 可用专家的数量
self.top_k = config["num_experts_per_tok"] # 每个 token 选择的专家数量
self.out_dim = config.get("out_dim", self.hidden_dim) # 输出的维度，如果未指定则默认为 hidden_dim
# 门控机制，通过线性层生成每个 token 的专家选择概率
self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
self.experts = nn.ModuleList([deepcopy(exp) for exp in experts]) # 复制专家层
def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
batch_size, sequence_length, f_map_sz = hidden_states.shape # 获取批量大小、序列长度和特征维度
hidden_states = hidden_states.view(-1, f_map_sz) # 展平批量和序列维度
# 通过门控机制计算路由 logits
router_logits = self.gate(hidden_states)
_, selected_experts = torch.topk(
router_logits.sum(dim=0, keepdim=True), self.top_k, dim=1
) # 选择 top_k 专家
routing_weights = F.softmax(
router_logits[:, selected_experts[0]], dim=1, dtype=torch.float
) # 计算选择的专家的路由权重
# 将路由权重转换回输入的 dtype
routing_weights = routing_weights.to(hidden_states.dtype)
final_hidden_states = torch.zeros(
(batch_size * sequence_length, self.out_dim),
dtype=hidden_states.dtype,
device=hidden_states.device,
) # 初始化最终的隐藏状态
# 遍历所有选中的专家，并在每个专家上执行计算
for i, expert_idx in enumerate(selected_experts[0].tolist()):
expert_layer = self.experts[expert_idx]
# 计算当前专家的隐藏状态
current_hidden_states = routing_weights[:, i].view(
batch_size * sequence_length, -1
) * expert_layer(hidden_states)
# 由于 `index_add_` 仅支持使用 torch 张量进行索引，因此我们使用 `top_x` 张量。
final_hidden_states = final_hidden_states + current_hidden_states
final_hidden_states = final_hidden_states.reshape(
batch_size, sequence_length, self.out_dim
) # 恢复最终隐藏状态的形状
return final_hidden_states # 返回最终的隐藏状态
SegMoEPipeline
- 定义了一个名为 SegMoEPipeline 的类，用于处理和管理混合扩散专家模型。
以下是对 SegMoEPipeline 中涉及的各个函数的定义和功能的解释，算法方面重点函数是 self.get_gate_params 的实现：
-
remove_all_forward_hooks: 这个函数用于删除模型中的所有前向钩子。它会遍历模型的所有子模块，如果子模块有前向钩子，就会将它们删除。
-
SparseMoeBlock: 这个类定义了一个稀疏的Mixture of Experts（MoE）块。它包含一个门控网络和多个专家网络。在前向传播过程中，它会根据输入的隐藏状态选择合适的专家网络进行计算，并将结果组合起来。
-
getActivation: 这个函数用于获取模型中的激活值。它接受一个模型、输入和输出作为参数，并返回一个钩子函数。这个钩子函数会在模型的前向传播过程中被调用，并将激活值保存到一个字典中。
-
SegMoEPipeline: 这个类定义了SegMoEPipeline，它是一个用于动态组合多个Stable Diffusion模型的管道。它接受一个配置文件或路径作为参数，并根据配置文件中的设置加载和组合模型。
-
load_from_scratch: 这个函数用于从头开始加载SegMoEPipeline。它会根据配置文件中的设置加载基础模型和专家模型，并将它们组合起来。
-
call: 这个函数用于调用SegMoEPipeline进行推理。它会将输入传递给基础模型，并返回推理结果。
-
create_empty: 这个函数用于创建一个空的UNet2DConditionModel模型。它会根据配置文件中的设置初始化模型，并将模型中的层替换为稀疏的MoE层。
-
save_pretrained: 这个函数用于将SegMoEPipeline保存到磁盘上。它会将模型的权重和配置保存到指定的路径下。
-
cast_hook: 这个函数用于在模型的前向传播过程中获取隐藏状态。它会为模型中的每个层注册一个前向钩子，并在前向传播过程中将隐藏状态保存到一个字典中。
-
get_hidden_states: 这个函数用于获取模型中的隐藏状态。它接受一个模型、正向提示和负向提示作为参数，并返回一个包含隐藏状态的字典。
-
get_gate_params: 这个函数用于获取门控网络的参数。它接受一个专家模型列表、正向提示列表和负向提示列表作为参数，并返回一个包含门控网络参数的字典。
算法实现思路
定义专家和对应的 positive_prompt/negative_prompt
- 定义 experts 的时候会给定该 expert 的 positive_prompt 和 negative_prompt，类似下面。一个完整示例见：https://github.com/segmind/segmoe/blob/main/segmoe_config_4x2.yaml
- 这些 positive_prompt 和 negative_prompt 会用来计算 MoE router 的权重
base_model: Base Model Path, Model Card or CivitAI Download Link
num_experts: Number of experts to use
moe_layers: Type of Layers to Mix (can be "ff", "attn" or "all"). Defaults to "attn"
num_experts_per_tok: Number of Experts to use
type: Type of the individual models (can be "sd" or "sdxl"). Defaults to "sdxl"
experts:
- source_model: Expert 1 Path, Model Card or CivitAI Download Link
positive_prompt: Positive Prompt for computing gate weights
negative_prompt: Negative Prompt for computing gate weights
- source_model: Expert 2 Path, Model Card or CivitAI Download Link
positive_prompt: Positive Prompt for computing gate weights
negative_prompt: Negative Prompt for computing gate weights
- source_model: Expert 3 Path, Model Card or CivitAI Download Link
positive_prompt: Positive Prompt for computing gate weights
negative_prompt: Negative Prompt for computing gate weights
- source_model: Expert 4 Path, Model Card or CivitAI Download Link
positive_prompt: Positive Prompt for computing gate weights
negative_prompt: Negative Prompt for computing gate weights
将多个 expert 转换为一个 MoE 模型
- 把多个模型的 ffn 和 attn 中的 to_q/to_k/to_v 都会替换为 SparseMoeBlock，也即模型每一层都能自动选择不同的专家模型的参数
MoE router 参数初始化
- 因为 router 是没有经过训练的，需要通过给定的模型 postive prompt 和 negative prompt 进行 MoE router 参数初始化。算法逻辑主要是从多个专家模型中提取隐藏状态，并对这些状态进行归一化处理和汇总。这种思路其实就是根据不同专家模型输出的 hidden_states 中参数整合后作为 MoE router 参数，这样如果是与 hidden_states 更相关的输入可以达到更高的 MoE router 激活值（其实是有很强的先验假设在这里）。
@torch.no_grad
def get_gate_params(
self,
experts, # 专家模型列表
positive, # 正样本列表，与专家模型一一对应
negative, # 负样本列表，与专家模型一一对应
):
gate_vects = {} # 用于存储每一层的门控向量
for i, expert in enumerate(tqdm.tqdm(experts, desc="Expert Prompts")):
expert.to(self.device) # 将专家模型移动到指定设备
expert.unet.to(
device=self.device,
dtype=self.torch_dtype,
memory_format=torch.channels_last,
) # 将专家模型的unet部分移动到指定设备，并设置数据类型和内存格式
hidden_states = self.get_hidden_states(expert, positive[i], negative[i])
# 获取专家模型的隐藏状态
del expert # 删除专家模型以释放显存
gc.collect() # 手动进行垃圾回收
torch.cuda.empty_cache() # 清空CUDA缓存
for h in hidden_states:
if i == 0:
gate_vects[h] = [] # 初始化门控向量列表
hidden_states[h] /= (
hidden_states[h].norm(p=2, dim=-1, keepdim=True).clamp(min=1e-8)
) # 对隐藏状态进行L2归一化
gate_vects[h].append(hidden_states[h]) # 将归一化后的隐藏状态添加到门控向量列表中
for h in hidden_states:
gate_vects[h] = torch.stack(
gate_vects[h], dim=0
) # 将门控向量堆叠成张量，形状为 (num_expert, num_layer, hidden_size)
gate_vects[h].permute(1, 0) # 对张量进行维度置换，形状为 (num_layer, num_expert)
return gate_vects # 返回包含所有层门控向量的字典
总结
- 无需训练的多模型 MoE 方案，代码完全开源。
- MoE router 这里虽然是无需训练的，但是这里的初始化方法或许在从头训练的 MoE router 参数初始化上也能用上。
- 这种无需训练的 MoE 方案目前 follow 的工作不多，可能还是有一定先天劣势，比如模型的生成质量上其实没有明显比单专家模型会更好，同时会占用更多的显存，速度也会变慢。
扫一扫
1173
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
