统一多模态架构

统一多模态架构

统一多模态架构（unified multimodal architecture）用单一系统替代了各自分离的专用模型，使其能跨文本、图像、音频和视频进行阅读、推理与生成。本文件涵盖任意到任意模型（CoDi、NExT-GPT）、原生多模态 LLM（Gemini、GPT-4o）、多模态分词策略，以及统一化在架构上的取舍。

为何要统一

• 想象一位会说五种语言的翻译官，能在句子中间切换语言而毫不迟疑。早期多模态系统更像五位各自独立的翻译官，坐在不同的房间里，每人处理一种语言，通过墙上的小窗口递纸条交流。**统一多模态架构（unified multimodal architecture）**就是那位通晓多语者：一个权重共享的模型，在一次前向传播中就能跨文本、图像、音频、视频甚至动作进行阅读、写作和推理。

• 这样做的动机既有实践上的也有理论上的。在实践层面，为每对模态（文生图、图生文、音生文等）维护各自的专用模型，会导致组合爆炸：k 种模态最多需要 k(k-1) 条有向流水线。统一模型把所有这些都坍缩进一个系统。在理论层面，人类认知并不是在相互隔离的模块中分别处理视觉和语言的；跨模态绑定发生得早且深入，统一化试图镜像这一点。

• 共享权重鼓励跨模态迁移（transfer across modality）。一个在文本中学到时间模式（主语在动词之前、原因在结果之前）的 transformer，可以把同样的注意力回路复用于视频中的时间模式（物体先出现再运动）或音频中的时间模式（起音在延音之前）。这正是你在第 7 章看到的多模态版迁移学习——语言模型微调，以及第 8 章的 ImageNet 预训练。

• 形式上，令 \mathcal{M} = \{m_1, m_2, \ldots, m_k\} 为一组模态。统一模型定义一个单一的可参数化函数 f_\theta，把输入模态的任意子集映射到输出模态的任意子集：

• 其中 \mathcal{P}(\mathcal{M}) 是模态集合的幂集（所有子集）。关键约束是 \theta 大部分共享；只有薄薄的、模态专属的适配器层有所不同。

• 统一的承诺伴随着一个根本性的张力：模态在结构上是不同的。文本是离散词元的 1D 序列。图像是连续像素值的 2D 网格。音频是 1D 连续波形，时间尺度与文本截然不同。视频又给图像加了一个时间轴。把这些迥异的结构调和进一个 transformer 能消化的单一序列，是这个领域的核心工程挑战。

任意到任意模型

• 把它想象成一个万能遥控器，能通过同一个界面操作你的电视、空调和音响。**任意到任意模型（any-to-any model）**就是 AI 的等价物：它们接受任意模态组合作为输入，产生任意模态组合作为输出。

• CoDi（Composable Diffusion，可组合扩散）通过训练模态专属的扩散模型，再通过一个共享的条件化机制对齐它们的潜在空间，实现了任意到任意生成。每种模态都有自己的扩散过程（回忆本章第 04 节的扩散模型），但噪声预测网络受一个联合交叉注意力层条件化，该层同时看到所有输入模态的嵌入。这让 CoDi 能从一段文本提示一次性生成，比如说，一张图像和匹配的音频。

• NExT-GPT 采取了不同的架构思路。它通过轻量级的**投影层（projection layer）**把一个 LLM 骨干（「大脑」）连接到输入侧的模态专属编码器和输出侧的模态专属解码器。输入编码器（例如来自 CLIP 的图像编码器、来自 CLAP 的音频编码器）把每种模态翻译进 LLM 的嵌入空间。LLM 在合并后的词元序列上推理，并发出特殊的「模态信号词元」，把信息路由到合适的解码器（例如用于图像的 Stable Diffusion、用于音频的 AudioLDM）。只有投影层被训练；LLM 和专用编码器/解码器保持冻结。

• Gemini（Google DeepMind）从预训练起就是原生多模态的。与 NExT-GPT 的即插即用方式不同，Gemini 的 transformer 从零开始在文本、图像、音频和视频词元的交错序列上训练。这意味着跨模态注意力模式在预训练过程中有机地发展，而不是事后拼凑上去。模型用 SentencePiece 分词器处理文本，并学习一个类似本章第 03 节讨论的 VQ 方法的视觉分词器。

• GPT-4o（「o」代表「omni」，全能）代表了又一种范式：一个端到端模型，所有模态共享同一个 transformer 和同一个下一个词元预测目标。音频输入作为频谱词元处理，图像作为图块词元，文本作为子词词元，全部送入一个单一序列。模型生成输出词元，再由模态专属的头解码。关键创新是通过移除早期系统（如 GPT-4V）所依赖的 ASR、LLM、TTS 串联级联，实现了低延迟。

• 这些模型处于一个集成深度的谱系上：

  • 浅层集成（shallow integration）（NExT-GPT）：用训练过的适配器连接冻结的专用模块。构建快，跨模态推理能力有限。
  • 中层集成（medium integration）（CoDi）：跨模态专属生成器共享条件化。对齐更好，仍是模块化的。
  • 深层集成（deep integration）（Gemini、GPT-4o）：单一模型在所有模态上端到端训练。跨模态推理最丰富，训练最昂贵。

带共享骨干的模态专属编码器与解码器

• 把它想象成一个工厂，只有一条装配线（共享骨干），但有针对不同原材料的不同卸货码头（编码器），以及针对不同成品的不同发货部门（解码器）。每个码头都为它的货物做了专门化，但一旦进入工厂，所有东西都沿着同一条传送带流动。

• 统一模型的主流架构范式采用这种三段式结构：

  • 模态编码器 E_m，把模态 m 的原始输入转换成一串嵌入向量 \mathbf{h}_1^m, \mathbf{h}_2^m, \ldots, \mathbf{h}_{n_m}^m，每个维度为 d。
  • 共享 transformer 骨干 T_\theta，用自注意力处理来自所有输入模态的拼接或交错嵌入。
  • 模态解码器 D_m，把骨干的输出嵌入转换回模态 m 的原生格式（文本词元、图像像素、音频波形）。

• 对于文本，编码器通常是一个嵌入查找表 E_\text{text}(w) = \mathbf{W}_e[w]，其中 w 是词元索引，与你在第 7 章 transformer 中看到的完全相同。对于图像，编码器通常是视觉 transformer（Vision Transformer，ViT），它把图像切成图块并线性投影每个图块，如第 8 章所述。对于音频，编码器计算梅尔频谱图，并用卷积前端或音频频谱 transformer（AST）处理，如第 9 章讨论。

• 共享骨干是一个标准 transformer，对所有模态词元做自注意力。给定一个拼接的输入序列 \mathbf{H} = [\mathbf{h}_1^{m_1}, \ldots, \mathbf{h}_{n_1}^{m_1}, \mathbf{h}_1^{m_2}, \ldots, \mathbf{h}_{n_2}^{m_2}]，自注意力允许每个词元关注所有其他词元，无论属于哪种模态：

• 这与第 7 章的注意力公式相同，但现在 \mathbf{Q}、\mathbf{K}、\mathbf{V} 包含来自多种模态的词元。一个图像图块词元可以关注一个文本词元，从而无需任何独立的交叉注意力模块就能实现跨模态推理。

• **模态嵌入（modality embedding）**被加到每个词元上，让骨干知道一个词元来自哪种模态。这类似于位置嵌入，但编码的是模态身份而非序列位置。一个可学习的向量 \mathbf{e}_m \in \mathbb{R}^d 被加到来自模态 m 的每个词元上：

• 其中 \mathbf{p}_i 是位置 i 的位置嵌入。

多模态分词

• 想象你正在写一封信，信里既有英文文字，也有手绘的草图。你可能写一句话、画一张示意图、再写一句引用这张示意图的话、然后贴上一段乐谱。这封信是一条单一的线性流，交错着不同的「模态」。**多模态分词（multimodal tokenisation）**做的正是这件事：它把文本、图像、音频和视频转换成一条扁平的词元序列，由 transformer 从左到右处理。

• 对于文本，分词已经成熟：**字节对编码（byte-pair encoding，BPE）**或 SentencePiece 产生子词词元的词表，如第 7 章所述。挑战在于把这个想法扩展到连续模态。

• 对于图像，有两大类方法。离散方法用 VQ-VAE 或 VQ-GAN（本章第 03 节详述）把每张图像映射成一串码本索引。如果码本有 |\mathcal{C}| 项、一张图像被编码为 n 个码，那么这张图像就变成了从大小为 |\mathcal{C}| 的词表中抽取的 n 个离散词元，可直接与文本词表兼容。连续方法用 ViT 或 CNN 编码器产生 n 个连续嵌入向量，再线性投影到 transformer 的嵌入维度。Gemini 和 GPT-4o 使用连续方法的变体；Parti 和 LlamaGen 等自回归图像生成器则偏好离散路线。

• 对于音频，信号通常被转换为梅尔频谱图，然后要么用神经音频编解码器（例如 EnCodec、SoundStream，产生层次化离散词元）离散化，要么通过一个学到的编码器连续投影。例如 AudioLM 把音频表示为来自多个码本层级的离散词元序列，再自回归地建模它们。

• 对于视频，分词建立在图像分词之上，但还必须压缩时间维度。一种常见策略用 3D VQ-VAE（如本章第 03 节的 VideoGPT 或 Cosmos 分词器），把时空图块量化成离散词元。时间压缩因子至关重要：24 fps 的原始视频如果不做激进的时间下采样，每秒会产生过多词元。

• 一旦所有模态都被分词，它们就被**交错（interleave）**进一个单一序列，用特殊的分隔词元标记模态边界。一种典型格式如下：

[TEXT] The cat sits on a mat [/TEXT] [IMAGE] <img_tok_1> <img_tok_2> ... <img_tok_n> [/IMAGE] [AUDIO] <aud_tok_1> ... <aud_tok_m> [/AUDIO]
​

• 然后 transformer 用它标准的因果（或双向）注意力机制处理这整段混合序列。模态分隔词元身兼两职：它既告诉模型模态边界，又充当「池化点」，其表征概括了每段模态片段。

• 一个关键设计选择是词元预算（token budget）。一张图像分词成 256 个词元、一段 50 词元的文本说明，意味着图像占用了 5 倍的上下文窗口。模型必须在分辨率（更多词元 = 更多细节）和上下文长度（更多词元 = 更高内存和计算成本）之间取得平衡。词元合并（token merging）（逐步合并相似词元）和自适应分词（adaptive tokenisation）（简单区域用更少词元、复杂区域用更多词元）等技术有助于管理这一权衡。

训练配方：分阶段预训练与联合微调

• 你不会在教孩子算术之前就教他微积分。同样地，你无法从随机初始化开始在所有模态上同时训练一个统一多模态模型并指望它收敛得好。主流方法是分阶段训练（staged training），模型在精心排序的阶段中逐步学习更复杂的跨模态能力。

• 阶段 1：单模态预训练。 每个模态编码器独立地在大型单模态数据集上训练。文本骨干用标准的语言建模目标（下一个词元预测）在数万亿文本词元上预训练，与第 7 章完全相同。视觉编码器在图像分类或自监督目标（MAE、DINO）上预训练，如第 8 章。音频编码器在语音识别或音频分类数据上预训练，如第 9 章。这一阶段产出强大的单模态特征抽取器。

• 阶段 2：跨模态对齐。 把预训练的编码器连接到共享骨干，模型在配对的多模态数据（图文对、音文转录对）上用对比或生成目标训练。在此阶段，编码器权重可以冻结（以保留单模态知识），只更新投影层和骨干。这一阶段把 CLIP 式对齐（本章第 01 节）整合进统一模型。

• 阶段 3：联合多模态预训练。 所有参数（或大部分）解冻，模型在单模态和多模态数据的混合上训练，用单一的下一个词元预测目标覆盖所有模态词元。损失函数为：

• 其中 x_t 可以是文本词元、图像词元或音频词元。模型必须学会无论模态都预测下一个词元，这迫使它发展出真正的跨模态理解。

• 阶段 4：指令微调与对齐。 预训练模型在精选的指令跟随数据集上微调，其中包含多模态指令（例如「详细描述这张图」「这段视频发出什么声音？」「生成一张 X 的图像」）。这一阶段常用**基于人类反馈的强化学习（reinforcement learning from human feedback，RLHF）**或直接偏好优化（direct preference optimisation，DPO）来把模型输出与人类偏好对齐。

• **模态专属预热（modality-specific warm-up）**是阶段内用于防止模态坍缩的技术。如果某种模态（通常是文本，训练数据最多）主导了梯度信号，模型可能「遗忘」较弱的模态。预热策略包括：

  • 梯度平衡（gradient balancing）：缩放每种模态的梯度，使其对参数更新贡献相等。
  • 数据比例调度（data ratio scheduling）：逐步增加多模态数据相对单模态数据的比例。
  • 损失加权（loss weighting）：赋予模态专属权重 \lambda_m，使总损失为 \mathcal{L} = \sum_m \lambda_m \mathcal{L}_m，通过调 \lambda_m 平衡各模态间的学习率。

• 为什么不能跳过阶段？ 从零开始联合训练一切很诱人，但实践中会失败，原因有几个。第一，模型必须同时学习低层特征（边缘检测、音素识别）和高层跨模态推理，这两者学习动态非常不同。第二，模态间的数据分布极度不平衡（数万亿文本词元 vs. 数十亿图像词元 vs. 数亿音频片段）。第三，优化景观高度非凸，分阶段训练提供了一种课程学习（第 6 章的思想），把模型引导到更好的盆地。

多模态思维链推理

• 当你解一道几何题时，你可能会画一张图、标出角度、列出方程，然后一步步求解。你不会直接从题目跳到答案。**多模态思维链（multimodal chain-of-thought，CoT）**推理使模型能做同样的事：在得出最终答案之前，生成可能涉及文本、视觉标注甚至生成图表的中间推理步骤。

• 在纯文本 CoT 中（如第 7 章对提示策略的讨论），模型用自然语言生成一串推理步骤。多模态 CoT 把它扩展为允许中间步骤引用或生成视觉内容。例如，给定一张图表图像和问题「哪一年销售额最高？」，一个多模态 CoT 模型可能先描述图表（「这张图展示了 2018 到 2023 年的销售额……」），再识别相关的视觉特征（「最高的柱子出现在 2021 年……」），最后输出答案（「2021」）。

• 形式上，令 \mathbf{x} 为多模态输入，y 为目标答案。标准预测直接建模 p(y \mid \mathbf{x})。思维链引入中间推理 \mathbf{r} = (r_1, r_2, \ldots, r_L)，把预测因式分解为：

• 实践中，这个求和用对推理链的贪心或束搜索解码来近似。推理步骤 r_i 可以是文本词元、对图像区域的引用，甚至是生成的视觉词元（例如叠加在输入图像上的边界框标注）。

• 训练多模态 CoT 通常涉及精选数据集，由人类标注者提供逐步的多模态推理轨迹，然后在这些轨迹上微调模型。一些方法从更大的教师模型蒸馏 CoT 能力：教师为大型数据集生成推理轨迹，较小的学生模型在输入和教师轨迹上训练。

• 多模态 CoT 对于需要空间推理（例如「红球在蓝立方体的左边吗？」）、基于图表的数学推理（例如几何题）以及多步视觉问答（答案依赖于结合图像多个区域的信息）的任务特别强大。

多模态智能体

• 想象厨房里的一位机器人厨师。它看着台面上的食材（视觉）、读平板上的菜谱（文本）、听定时器的蜂鸣（音频），然后实际拿起刀切洋葱（动作）。**多模态智能体（multimodal agent）**就是它的数字版本：一个通过多种模态感知世界、推理该做什么、并采取扎根于其感知的动作的模型。

• 智能体循环遵循经典的**观察–推理–行动（observe-reason-act）**周期：

  1. 观察（observe）：智能体从环境接收多模态输入（一张截图、用户的语音指令、一路视频流）。
  2. 推理（reason）：统一模型处理多模态输入，可能用思维链来规划一系列步骤。
  3. 行动（act）：模型输出一个动作（一段文本回应、一次工具调用、在坐标 (x, y) 处的一次鼠标点击、一条机器人电机指令）。

• **工具使用（tool use）**是多模态智能体的关键能力。模型被训练成识别自己何时无法直接回答问题，必须调用一个外部工具：计算器、代码解释器、网络浏览器或搜索引擎。模型作为其输出词元序列的一部分，生成一个结构化的工具调用（例如 search("current weather in London")），系统执行调用，结果作为额外的输入词元反馈给模型处理。

• **视觉定位（visual grounding）**把语言连接到图像或视频中的具体区域。当智能体说「点击右上角的蓝色按钮」时，它必须把短语「右上角的蓝色按钮」扎根到像素坐标。在架构上，这通过训练模型把边界框坐标作为特殊词元输出，或让模型在图像上产生一个热图来指示被指代区域来实现。这把本章第 02 节（视觉语言模型）讨论的定位与指代工作扩展到了行动领域。

• **网络智能体（web agent）**如 WebVoyager 和 SeeAct 展示了多模态智能体浏览网站。智能体接收一张网页截图，识别可交互元素（按钮、文本框、链接），并输出动作（点击、输入、滚动）以完成用户指定的目标。关键挑战是巨大的动作空间：一个典型网页有数百个可能的点击目标。

• **具身智能体（embodied agent）**把它扩展到物理环境。一个带摄像头和麦克风的机器人接收视觉和音频输入，通过统一模型处理，输出电机指令。像 PaLM-E（Google）这样的项目把机器人传感器数据直接嵌入到语言模型的词元序列中，让机器人能跟随「拿起碗附近那块绿色的积木」这样的指令——通过把指令扎根到它的视觉观察中，并生成一串电机动作。

• 智能体的训练配方在标准分阶段预训练之上增加一个**强化学习（reinforcement learning，RL）**阶段。智能体与环境（一个模拟桌面、一个网络浏览器、一个机器人模拟器）交互，因任务完成获得奖励，用 PPO 或 REINFORCE 等算法更新策略。奖励信号通常是稀疏的（任务成功为 1，否则为 0），使这种优化极具挑战性，严重依赖多模态预训练带来的强先验。

基准与评估

• 评估一个能看、能听、能读、能行动的模型需要一套多样化的基准。没有任何单一指标能涵盖多模态能力，因此这个领域依赖一组专门的评估。

• MMLU（Massive Multitask Language Understanding，大规模多任务语言理解）测试 57 个学术学科的知识。虽然最初是纯文本的，但它充当一个基线：一个统一的多模态模型在获得视觉能力后不应丢失纯文本性能。多模态训练后 MMLU 下降是灾难性遗忘的信号。

• MMBench 在 20 个细粒度能力维度上评估视觉–语言理解，包括属性识别、空间关系理解和 OCR。每道题给出一幅图像和一道选择题。这个基准系统地测试模型到底是真正理解了图像，还是在依赖纯文本的捷径。

• SEED-Bench 提供 19,000 道选择题，跨越图像和视频理解的 12 个评估维度。它专门测试时间理解（某给定帧之前/之后发生了什么）和组合推理（组合多个视觉属性）。

• MM-Vet 通过要求模型同时使用多种技能来评估整合的多模态能力：识别、OCR、空间感知、语言生成和知识检索，全部在一道题中。

• MathVista 测试在视觉输入上的数学推理：几何图、统计图、函数图和科学插图。这个基准专门针对多模态思维链能力。

• 视听基准如 AVQA（Audio-Visual Question Answering，视听问答）测试模型能否推理所见与所听之间的关系。例如：「说话的人在左边还是右边？」

• 智能体基准如 WebArena、OSWorld 和 SWE-bench 评估在交互环境中的任务完成情况。指标通常是成功率：智能体正确完成了多少比例的任务？这些基准特别有挑战性，因为它们需要长程规划和错误恢复。

• **整体评估（holistic evaluation）**框架如 LMSYS Chatbot Arena 在一对一形式中使用人类偏好判断。两个模型被展示相同的多模态输入，人类评审选择哪个回答更好。从数千次这样的比较中计算 Elo 评分，提供一个与整体模型质量高度相关的单一标量。

• 多模态评估中一个持续的挑战是数据污染（data contamination）：因为这些模型在互联网规模的数据上训练，基准图像和问题可能出现在训练集中。仔细的去重和留出测试集的创建是必不可少但并不完美的保障。

世界模型

• 想象闭上眼睛，想象把一只玻璃杯从桌边推下去会发生什么。你「看到」它落下、「听到」它碎裂、「感觉」到这是个坏主意。你的大脑在运行一个世界模型（world model）：对环境物理与因果结构的内部模拟，能跨多种模态预测未来状态。

• 在 AI 语境下，世界模型是一个学到的函数，在给定当前状态和一个动作的情况下预测世界的下一个状态：

• 其中 s_t 是当前状态表征（可能包括视觉、听觉和本体感觉信息），a_t 是一个动作，\hat{s}_{t+1} 是预测的下一个状态。状态 s_t 生活在学到的潜在空间而非原始像素空间，使预测问题变得可处理。

• 视频预测模型如 Sora（OpenAI）和 Genie（Google DeepMind）代表了迈向世界模型的重要一步。它们学习在文本提示和/或动作序列条件下生成时间连贯的视频帧。虽然它们常被当作视频生成器讨论，但底层能力更接近世界模拟：模型已经内化了足够的物理（重力、碰撞、遮挡、流体动力学），从而能渲染出合理的未来。

• 它与多模态架构的联系很深。一个只预测像素的世界模型是受限的；一个真正有用的世界模型跨模态预测。如果你推那只杯子，世界模型应预测视觉轨迹（杯子落下）、听觉事件（杯子碎裂）以及语义后果（你现在地板上有碎玻璃了）。统一多模态架构是世界模型的天然候选者，因为它们已经把所有模态表示在一个共享空间中。

• 形式上，多模态世界模型优化：

• 其中 s_{t+1}^m 是模态 m 中真实下一个状态的表征，g_\phi^m 是世界模型的模态专属预测头。共享的潜在动力学 g_\phi 在联合多模态空间中运作，而模态专属的头把预测解码为每种模态的原生格式。

• JEPA（Joint Embedding Predictive Architecture，联合嵌入预测架构）由 Yann LeCun 提出，为世界模型提供了一个避免像素级预测陷阱的框架。它不再预测原始像素（这会把容量浪费在无关细节上，比如精确的纹理），而是在嵌入空间中预测。模型学习一个把观察映射到嵌入的编码器，以及一个预测未来嵌入的预测器：

• 损失比较的是嵌入而非原始观察，这更鲁棒于感知歧义（许多不同的像素配置可能表示相同的语义状态）。这种方法对多模态世界模型尤其有前景，因为它天然地运作在统一架构已经提供的共享嵌入空间中。

• 世界模型除了学术兴趣外还有实际应用。在基于模型的强化学习中，智能体用它的世界模型在采取动作之前「想象」后果，从而大幅减少所需的真实世界交互次数（回忆第 11 章对基于模型 RL 的讨论）。在自动驾驶中，世界模型预测在不同转向决策下未来几秒场景将如何演变。在机器人学中，世界模型让机器人在执行操作序列之前先在脑中排练一遍。

• 世界模型研究的前沿正走向交互式世界模型（interactive world model）——能实时运行并对任意用户动作做出响应，本质上成为完全从数据中学习的通用模拟器。Genie 2（Google DeepMind）在 3D 环境中展示了这一点：给定单张图像，它生成一个可交互、可控的 3D 世界，用户可以探索。世界模型与统一多模态架构的融合，预示着一个未来：单一模型能够跨所有模态感知、预测、模拟和行动。

编程练习（使用 CoLab 或 notebook）

任务 1：构建一个最小化的多模态词元交错器

• 写一个函数，接收一个文本字符串和一个虚拟「图像」（一个小型二维数组），把它们的分词表示连同模态嵌入交错成一条扁平序列。

import jax
import jax.numpy as jnp

# 模拟多模态分词：文本词元 +「图像图块」词元
def interleave_modalities(text_tokens, image_patches, embed_dim=32, key=jax.random.PRNGKey(0)):
    """把文本和图像词元连同可学习的模态嵌入交错在一起。"""
    k1, k2, k3 = jax.random.split(key, 3)
    n_text = text_tokens.shape[0]
    n_img = image_patches.shape[0]
    # 随机投影矩阵（真实编码器的替身）
    W_text = jax.random.normal(k1, (text_tokens.shape[-1], embed_dim)) * 0.02
    W_img = jax.random.normal(k2, (image_patches.shape[-1], embed_dim)) * 0.02
    # 模态嵌入：一个给文本、一个给图像
    mod_emb = jax.random.normal(k3, (2, embed_dim)) * 0.02
    text_embs = text_tokens @ W_text + mod_emb[0]  # (n_text, embed_dim)
    img_embs = image_patches @ W_img + mod_emb[1]   # (n_img, embed_dim)
    # 交错：先 [IMG] 词元，再 [TEXT] 词元（像 LLaVA 那样）
    combined = jnp.concatenate([img_embs, text_embs], axis=0)
    print(f"Combined sequence: {n_img} image + {n_text} text = {combined.shape[0]} tokens")
    return combined

# 试一试：5 个文本词元（维度 16）和 4 个图像图块（维度 64）
text = jax.random.normal(jax.random.PRNGKey(1), (5, 16))
image = jax.random.normal(jax.random.PRNGKey(2), (4, 64))
seq = interleave_modalities(text, image)
# 实验：改变 embed_dim，交换交错顺序，加入第三种模态
​

任务 2：可视化跨模态注意力模式

• 创建一个合成的多模态序列并计算自注意力分数，看图像词元如何关注文本词元，反之亦然。

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

def cross_modal_attention(n_text=6, n_img=4, d=32, key=jax.random.PRNGKey(42)):
    """计算并可视化文本和图像词元之间的注意力。"""
    k1, k2, k3 = jax.random.split(key, 3)
    # 模拟两种模态的词元嵌入
    text_embs = jax.random.normal(k1, (n_text, d))
    img_embs = jax.random.normal(k2, (n_img, d))
    seq = jnp.concatenate([img_embs, text_embs], axis=0)  # (n_img+n_text, d)
    # 可学习的 Q、K 投影
    Wq = jax.random.normal(k3, (d, d)) * 0.1
    Wk = jax.random.normal(jax.random.PRNGKey(99), (d, d)) * 0.1
    Q, K = seq @ Wq, seq @ Wk
    scores = Q @ K.T / jnp.sqrt(d)
    attn = jax.nn.softmax(scores, axis=-1)
    # 绘图
    labels = [f"img_{i}" for i in range(n_img)] + [f"txt_{i}" for i in range(n_text)]
    fig, ax = plt.subplots(figsize=(7, 6))
    ax.imshow(attn, cmap="viridis")
    ax.set_xticks(range(len(labels))); ax.set_xticklabels(labels, rotation=45, fontsize=8)
    ax.set_yticks(range(len(labels))); ax.set_yticklabels(labels, fontsize=8)
    ax.set_xlabel("Key (attended to)"); ax.set_ylabel("Query (attending from)")
    ax.set_title("Cross-modal self-attention map")
    plt.colorbar(ax.images[0], ax=ax, shrink=0.8)
    plt.tight_layout(); plt.show()

cross_modal_attention()
# 实验：增大 d，加入因果掩码，观察注意力模式如何变化
​

任务 3：用模态专属损失加权模拟分阶段训练

• 演示模态专属损失权重如何影响一个玩具级多模态训练循环。观察平衡损失如何防止某一种模态占据主导。

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

def staged_training_sim(steps=200, key=jax.random.PRNGKey(7)):
    """用可调的模态损失权重模拟多模态训练。"""
    # 两种「模态」，损失尺度不同（文本损失约比图像损失大 10 倍）
    losses_text, losses_img = [], []
    param = jnp.array([0.0, 0.0])  # 由两种模态损失共同更新的共享参数
    lr = 0.05
    # 试着改变这些权重来看对收敛平衡的影响
    lambda_text, lambda_img = 1.0, 5.0  # 给较弱的模态加权
    for step in range(steps):
        k1, k2, key = jax.random.split(key, 3)
        noise_t = jax.random.normal(k1, ()) * 0.3
        noise_i = jax.random.normal(k2, ()) * 0.1
        loss_t = (param[0] - 3.0) ** 2 + noise_t  # 文本目标 = 3.0
        loss_i = 0.1 * (param[1] - 1.0) ** 2 + noise_i  # 图像目标 = 1.0（尺度更小）
        # 加权的组合梯度
        grad_t = lambda_text * 2 * (param[0] - 3.0)
        grad_i = lambda_img * 0.2 * (param[1] - 1.0)
        param = param - lr * jnp.array([grad_t, grad_i])
        losses_text.append(float(loss_t)); losses_img.append(float(loss_i))

    fig, ax = plt.subplots(figsize=(8, 4))
    ax.plot(losses_text, label=f"Text loss (weight={lambda_text})", alpha=0.7)
    ax.plot(losses_img, label=f"Image loss (weight={lambda_img})", alpha=0.7)
    ax.set_xlabel("Training step"); ax.set_ylabel("Loss"); ax.legend()
    ax.set_title("Modality loss balancing during staged training")
    plt.tight_layout(); plt.show()

staged_training_sim()
# 实验：把 lambda_img 设为 1.0，观察图像损失收敛慢得多
​