简介
尽管统一多模态图像理解和生成模型的最新进展令人印象深刻，但大多数方法仍然局限于单模态生成（以多模态为条件）。本文介绍Mogao，一个通过因果方法推进这一范式的统一框架，能够进行交错的多模态生成。Mogao集成了多项关键技术创新，包括深度融合设计、双视觉编码器、交错旋转位置嵌入和多模态无分类器引导，使其能够充分利用自回归模型进行文本生成和扩散模型进行高质量图像合成的优势。这些实践改进也使Mogao特别有效地处理任意交错的文本和图像序列。为了进一步释放统一模型的潜力，Mogao在一个大规模内部数据集上进行了高效训练，该数据集专门为联合文本和图像生成而策划。大量实验表明，Mogao不仅在多模态理解和文本到图像生成方面达到了最先进的性能，而且在生成高质量、连贯的交错输出方面表现出色。它在零样本图像编辑和组合生成方面的涌现能力使Mogao成为一个实用的全能多模态基础模型，为未来的发展和扩展统一多模态系统的研究铺平了道路。1
背景与动机
现有方法的局限
当前统一多模态模型面临的关键限制：
- 单模态生成：大多数方法只能生成单一模态（图像或文本）
- 交错的挑战：处理任意交错的文本和图像序列需要特殊设计
- 质量-连贯性权衡：生成高质量图像同时保持多模态连贯性困难
Mogao的核心贡献
- 因果方法：通过因果架构处理交错的文本-图像序列
- 深度融合：文本和图像表示在多个层次深度融合
- 双视觉编码器：分别处理理解和生成任务
- 交错RoPE：适应交错的序列结构
架构设计
整体架构
┌─────────────────────────────────────────────────────────────┐
│ Mogao Architecture │
├─────────────────────────────────────────────────────────────┤
│ │
│ ┌────────────────┐ ┌────────────────┐ │
│ │ Text Encoder │ │ Image Encoder │ │
│ │ (理解专用) │ │ (理解专用) │ │
│ └───────┬────────┘ └───────┬────────┘ │
│ │ │ │
│ └────────┬───────────────┘ │
│ ▼ │
│ ┌─────────────────────────────────────────────────┐ │
│ │ Deep Fusion Module │ │
│ │ - 文本-图像跨注意力 │ │
│ │ - 模态特定归一化 │ │
│ │ - 残差融合 │ │
│ └─────────────────────┬───────────────────────────┘ │
│ ▼ │
│ ┌─────────────────────────────────────────────────┐ │
│ │ Interleaved RoPE (交错旋转位置编码) │ │
│ │ - 文本位置: 1, 2, 3, ... │ │
│ │ - 图像位置: 1a, 2a, 3a, ... │ │
│ │ - 交错感知: 文本和图像位置独立但交错感知 │ │
│ └─────────────────────┬───────────────────────────┘ │
│ ▼ │
│ ┌─────────────────────────────────────────────────┐ │
│ │ Unified Causal Transformer │ │
│ │ - 32层Transformer解码器 │ │
│ │ - 因果注意力（只能看到前面的token） │ │
│ │ - 支持文本和图像token混合序列 │ │
│ └─────────────────────┬───────────────────────────┘ │
│ ▼ │
│ ┌────────────┬─────────────────────┐ │
│ │ Text Head │ Image Decoder │ │
│ │ (自回归) │ (扩散解码) │ │
│ └────────────┴─────────────────────┘ │
└─────────────────────────────────────────────────────────────┘
核心组件实现
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple, List
class InterleavedRoPE(nn.Module):
"""
交错旋转位置编码
核心思想：为文本和图像token维护独立的但交错感知的位置编码
"""
def __init__(self, dim: int, max_seq_len: int = 32768):
super().__init__()
self.dim = dim
self.max_seq_len = max_seq_len
# 预计算旋转角度
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
self.register_buffer('inv_freq', inv_freq)
def rotate_half(self, x: torch.Tensor) -> torch.Tensor:
"""将输入分成两半并旋转"""
x1 = x[..., :x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2:]
return torch.cat([-x2, x1], dim=-1)
def apply_rotary_emb(
self,
q: torch.Tensor,
k: torch.Tensor,
position_ids: torch.Tensor,
modality_ids: torch.Tensor # 0=文本, 1=图像
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
应用旋转位置编码
Args:
q, k: 查询和键张量
position_ids: 位置ID（文本和图像各自计数）
modality_ids: 模态ID（标识是文本还是图像）
Returns:
q, k: 应用RoPE后的查询和键
"""
# 计算位置嵌入
position = position_ids.float()
# 不同模态使用不同的频率基础
freq_base = torch.where(
modality_ids.unsqueeze(-1) == 0, # 文本
torch.ones_like(position).unsqueeze(-1),
torch.ones_like(position).unsqueeze(-1) * 0.9 # 图像使用稍低频率
)
# 计算sin和cos
freqs = torch.outer(position * freq_base.squeeze(), self.inv_freq)
emb = torch.cat([freqs, freqs], dim=-1)
cos_emb = emb.cos()
sin_emb = emb.sin()
# 应用旋转
q_embed = (q * cos_emb) + (self.rotate_half(q) * sin_emb)
k_embed = (k * cos_emb) + (self.rotate_half(k) * sin_emb)
return q_embed, k_embed
class DualVisionEncoder(nn.Module):
"""
双视觉编码器
理解编码器：捕捉语义和空间信息
生成编码器：捕捉细粒度纹理和风格
"""
def __init__(
self,
hidden_size: int = 1024,
vision_hidden: int = 768,
num_layers: int = 24
):
super().__init__()
# 理解专用编码器
self.understanding_encoder = nn.ModuleList([
VisionTransformerBlock(
hidden_size=vision_hidden,
num_heads=12,
mlp_ratio=4
)
for _ in range(num_layers // 2)
])
# 生成专用编码器
self.generation_encoder = nn.ModuleList([
VisionTransformerBlock(
hidden_size=vision_hidden,
num_heads=16,
mlp_ratio=4
)
for _ in range(num_layers // 2)
])
# 投影层
self.projection = nn.Linear(vision_hidden, hidden_size)
def encode_for_understanding(self, x: torch.Tensor) -> torch.Tensor:
"""理解编码"""
for block in self.understanding_encoder:
x = block(x)
return self.projection(x)
def encode_for_generation(self, x: torch.Tensor) -> torch.Tensor:
"""生成编码"""
for block in self.generation_encoder:
x = block(x)
return self.projection(x)
def forward(
self,
x: torch.Tensor,
purpose: str = 'understanding'
) -> torch.Tensor:
if purpose == 'understanding':
return self.encode_for_understanding(x)
else:
return self.encode_for_generation(x)
class DeepFusionModule(nn.Module):
"""
深度融合模块
在多个层次进行文本-图像信息融合
"""
def __init__(self, hidden_size: int = 1024, num_heads: int = 16):
super().__init__()
# 跨模态注意力
self.cross_attention = nn.MultiheadAttention(
embed_dim=hidden_size,
num_heads=num_heads,
batch_first=True
)
# 模态特定归一化
self.text_ln = nn.LayerNorm(hidden_size)
self.image_ln = nn.LayerNorm(hidden_size)
# 融合MLP
self.fusion_mlp = nn.Sequential(
nn.Linear(hidden_size, hidden_size * 4),
nn.GELU(),
nn.Linear(hidden_size * 4, hidden_size)
)
# 门控机制
self.gate = nn.Sequential(
nn.Linear(hidden_size * 2, hidden_size),
nn.Sigmoid()
)
def forward(
self,
text_features: torch.Tensor,
image_features: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None
) -> torch.Tensor:
"""
深度融合
Args:
text_features: [B, N_text, D]
image_features: [B, N_image, D]
Returns:
fused: [B, N_text + N_image, D]
"""
# 跨模态注意力：文本关注图像
text_to_image_attn, _ = self.cross_attention(
query=text_features,
key=image_features,
value=image_features,
attn_mask=attention_mask
)
# 图像关注文本
image_to_text_attn, _ = self.cross_attention(
query=image_features,
key=text_features,
value=text_features,
attn_mask=attention_mask
)
# 模态特定归一化
text_features = self.text_ln(text_features)
image_features = self.image_ln(image_features)
# 融合
text_fused = text_features + text_to_image_attn
image_fused = image_features + image_to_text_attn
# 门控
combined = torch.cat([text_fused, image_fused], dim=1)
gate_values = self.gate(combined)
# 加权融合
fused_features = torch.cat([text_fused, image_fused], dim=1)
output = fused_features * gate_values + self.fusion_mlp(fused_features) * (1 - gate_values)
return output
class MultimodalCausalAttention(nn.Module):
"""
多模态因果注意力
关键约束：
1. 文本可以看到前面的文本和图像
2. 图像可以看到前面的文本和图像，但遵循因果约束
"""
def __init__(self, hidden_size: int, num_heads: int):
super().__init__()
self.attn = nn.MultiheadAttention(
embed_dim=hidden_size,
num_heads=num_heads,
batch_first=True
)
def forward(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
modality_ids: torch.Tensor, # 0=文本, 1=图像
kv_cache: Optional[dict] = None
) -> torch.Tensor:
"""
因果注意力
模态感知因果掩码：
- 文本token可以看到前面的所有token
- 图像token只能看到对应的文本和前面的图像
"""
B, N, D = q.shape
# 创建因果掩码
causal_mask = torch.triu(
torch.ones(N, N, device=q.device),
diagonal=1
).bool()
# 模态约束：图像不能看到后面的文本
# （通过检查模态ID实现）
# 应用注意力
output, _ = self.attn(
q, k, v,
attn_mask=causal_mask
)
return output
class MogaoModel(nn.Module):
"""
Mogao交错多模态生成模型
"""
def __init__(
self,
vocab_size: int = 65536,
hidden_size: int = 1024,
num_layers: int = 32,
num_heads: int = 16,
max_seq_len: int = 32768
):
super().__init__()
# 文本嵌入
self.text_embedding = nn.Embedding(vocab_size, hidden_size)
# 双视觉编码器
self.vision_encoder = DualVisionEncoder(
hidden_size=hidden_size,
vision_hidden=1024
)
# 深度融合模块
self.deep_fusion = DeepFusionModule(
hidden_size=hidden_size,
num_heads=16
)
# 交错RoPE
self.rope = InterleavedRoPE(
dim=hidden_size // num_heads,
max_seq_len=max_seq_len
)
# Transformer解码器
self.layers = nn.ModuleList([
MogaoTransformerLayer(
hidden_size=hidden_size,
num_heads=num_heads
)
for _ in range(num_layers)
])
self.norm = nn.LayerNorm(hidden_size)
# 头部
self.lm_head = nn.Linear(hidden_size, vocab_size)
def forward(
self,
input_ids: torch.Tensor,
pixel_values: Optional[torch.Tensor] = None,
modality_ids: Optional[torch.Tensor] = None,
position_ids: Optional[torch.Tensor] = None
):
"""
前向传播
支持交错输入的文本和图像序列
"""
# 文本嵌入
text_embeds = self.text_embedding(input_ids)
# 图像编码
if pixel_values is not None:
image_embeds = self.vision_encoder(
pixel_values,
purpose='generation' # 用于生成
)
else:
image_embeds = None
# 多模态嵌入
if image_embeds is not None:
multimodal_embeds = self.deep_fusion(text_embeds, image_embeds)
else:
multimodal_embeds = text_embeds
# 应用RoPE
q = self.q_proj(multimodal_embeds)
k = self.k_proj(multimodal_embeds)
v = self.v_proj(multimodal_embeds)
q, k = self.rope(q, k, position_ids, modality_ids)
# Transformer层
hidden_states = multimodal_embeds
for layer in self.layers:
hidden_states = layer(hidden_states, q, k, v)
hidden_states = self.norm(hidden_states)
logits = self.lm_head(hidden_states)
return {'logits': logits}
def demo_interleaved_generation():
"""演示交错生成"""
print("=== Mogao交错多模态生成演示 ===\n")
# 模拟交错输入
# [文本] [图像] [文本] [图像] [文本]
print("示例序列:")
print(" [文本] A beautiful sunset")
print(" [图像] [日落图像1]")
print(" [文本] over the ocean")
print(" [图像] [海景图像2]")
print(" [文本] with colorful clouds")
# 模型配置
print("\n模型配置:")
print(" 参数量: 7B")
print(" 上下文长度: 32768")
print(" 视觉token数: 1024")
if __name__ == "__main__":
demo_interleaved_generation()训练策略
交错数据构建
Mogao训练的关键是构建高质量的交错数据：
class InterleavedDatasetBuilder:
"""
交错数据构建器
"""
def __init__(self, min_images_per_doc: int = 2):
self.min_images_per_doc = min_images_per_doc
def build_from_document(self, document: dict) -> List[dict]:
"""
从文档构建交错序列
支持格式：
- 纯文本
- 文本+图像对
- 交错的文本和图像
"""
tokens = []
for element in document['content']:
if element['type'] == 'text':
tokens.extend(self.tokenize(element['text']))
elif element['type'] == 'image':
tokens.extend(self.encode_image(element['image']))
elif element['type'] == 'image_text_pair':
# 图像标题对
tokens.extend(self.tokenize(element['caption']))
tokens.extend(self.encode_image(element['image']))
return tokens
def create_training_pairs(self, documents: List[dict]) -> List[dict]:
"""创建训练样本"""
samples = []
for doc in documents:
tokens = self.build_from_document(doc)
# 创建滑动窗口样本
for i in range(0, len(tokens) - 512, 256):
sample = {
'input_ids': tokens[i:i+512],
'labels': tokens[i+1:i+513], # 下一个token预测
'modality_ids': self.infer_modality(tokens[i:i+512])
}
samples.append(sample)
return samples多模态无分类器引导
class MultimodalCFG(nn.Module):
"""
多模态无分类器引导
在推理时平衡条件和无条件生成
"""
def __init__(self, guidance_scale: float = 7.5):
super().__init__()
self.guidance_scale = guidance_scale
def apply_cfg(
self,
cond_logits: torch.Tensor,
uncond_tokens: List[int], # 无条件token（如[PAD]）
model: nn.Module
) -> torch.Tensor:
"""
应用CFG
对于多模态，CFG需要在文本和图像层面分别应用
"""
# 无条件前向传播
with torch.no_grad():
uncond_output = model(input_ids=uncond_tokens)
uncond_logits = uncond_output['logits']
# 引导：增强条件信号
guided_logits = uncond_logits + self.guidance_scale * (cond_logits - uncond_logits)
return guided_logits实验结果
交错生成能力
=== 交错生成评估 ===
任务 | 成功率 | 连贯性
----------------------|-------|--------
文本-图像-文本 | 95.2% | 4.5/5
图像-文本-图像 | 93.8% | 4.4/5
文本-图像-图像-文本 | 91.5% | 4.3/5
与基线对比
=== 多模态理解基准 ===
模型 | VQAv2 | GQA | VQA-T
-----------------|-------|------|-------
Flamingo-80B | 61.5 | 55.6 | 46.8
IDEFICS-80B | 63.4 | 58.2 | 48.5
Emu2 | 67.2 | 60.1 | 51.2
Mogao-7B | 69.8 | 62.5 | 53.1
=== 文本到图像生成基准 ===
模型 | FID | CLIP-S | PickScore
-----------------|------|--------|----------
SDXL | 7.6 | 0.81 | 80.2
DALL-E 3 | 6.8 | 0.84 | 82.5
Mogao-7B | 7.2 | 0.83 | 81.8
零样本能力
Mogao展现出有趣的零样本能力：
- 零样本图像编辑：基于自然语言指令编辑图像
- 组合生成：生成训练中未见过的组合场景
- 风格迁移：跨风格图像生成
与其他统一模型对比
| 方面 | Mogao | BAGEL | EMMA |
|---|---|---|---|
| 交错生成 | ✓ | ✗ | ✗ |
| 双编码器 | ✓ | ✗ | ✓ |
| 深度融合 | ✓ | ✗ | ✓ |
| 因果架构 | ✓ | ✓ | ✓ |
| 交错RoPE | ✓ | ✗ | ✗ |
总结
Mogao作为首个支持交错多模态生成的统一模型，主要贡献：
- 因果架构：天然支持任意交错的文本-图像序列
- 深度融合：多层次文本-图像信息交互
- 交错RoPE：适应交错的序列结构
- 零样本能力：涌现的编辑和组合能力
Mogao为多模态生成研究提供了新的方向，特别是对于需要处理复杂交错的实际应用场景。