Mini-Omni 语言模型在流式传输中边思考边听说应用

引入简介

Mini-Omni 是一个开源的多模态大语言模型，能够在思考的同时进行听觉和语言交流。它具有实时端到端语音输入和流媒体音频输出的对话能力。

语言模型的最新进展取得了显著突破。GPT-4o 作为一个新的里程碑，实现了与人类的实时对话，展示了接近人类的自然流畅度。为了实现这种人机交互，模型需要直接通过音频模态进行推理，并在流式传输中生成输出。然而，目前的学术模型通常依赖额外的文本转语音（TTS）系统进行语音合成，从而引入了不必要的延迟。本文介绍了 Mini-Omni，一种基于音频的端到端对话模型，能够实现实时语音交互。为实现这一目标，我们提出了一种文本指导的语音生成方法，并在推理过程中采用了批处理并行策略，以进一步提升性能。我们的方法不仅有助于在最小化退化的同时保留原始模型的语言能力，还使得其他工作能够建立实时交互能力。我们将这种训练方法称为“Any Model Can Talk”。此外，我们引入了 VoiceAssistant-400K 数据集，用于微调针对语音输出优化的模型。Mini-Omni 是第一个完全端到端的开源实时语音交互模型，为未来的研究提供了宝贵的潜力。

图1：Mini-Omni模型架构。

随着大型语言模型的快速发展，模型的能力不断增强，例子包括 Llama 3.1 [Meta, 2024]、Mix [Mix, 2024]、Qwen-2 [Yang et al., 2024a] 以及著名的 GPT-4。这些模型不仅在文本处理上表现出色，还开始掌握对其他模态的理解，如 LLaVA [Liu et al., 2024]、Qwen2-Audio [Chu et al., 2024] 和视频美洲驼 [Zhang et al., 2023b]。尽管如此，一个关键挑战仍然阻碍了它们在日常应用中的广泛集成，那就是实时语音交互能力。

GPT-4o [OpenAI, 2024] 是第一个具备实时多模态语音交互功能的模型，它能够处理视觉、音频和文本信息，并实现实时语音对话，尽管它仍为闭源代码。其他模型通常采用两种方法来实现语音能力：一种是级联方法，其中语言模型生成文本，随后由文本到语音（TTS）模型进行音频合成。

1. 介绍了 Mini-Omni，这是第一个具备音频输入和流输出功能的开源端到端多模态大型模型。
2.  提出了独特的文本指示并行生成方法，使语音推理输出与文本功能对齐，实现了最小数据需求。
3.  增强了延迟并行性，加速了音频推理速度。
4.  引入了“任何模型都可以说话”的创新方法，通过训练和推理提升性能，无需改变大型模型架构。5) 采用了包括退火和SFT在内的三阶段训练过程，对语音到文本和文本到语音适配器进行优化。
5. 发现了现有开源QA数据集的不足，提出了用于语音模型输出的专用数据集——语音助手-400k，以支持音频助手的微调。

二、音频语言建模

1、原文这样子介绍的：

2、解码策略

解码策略音频生成与文本指令。语言模型经历了实质性的进步，在文本模式中展示了特殊的推理能力。作为回应，Mini-Omni进行了重组，通过文本-音频并行解码方法将这些推理能力转移到流音频输出。该方法同时输出音频和文本令牌，以及通过文本到语音合成生成的音频，确保实时交付，同时利用基于文本的推理优势。为了与大型模型的输入对齐，在生成下一个标记之前，将对并行生成的所有序列进行求和，如图1所示。这种方法使模型能够在聊天场景中以最小的第一令牌延迟实现实时语音输出。文本延迟并行解码。并行生成首先由MusicGen [Copet等人，2024]引入，以加速音乐生成过程，我们已经将这种方法集成到文本模式中，以增强推理能力。并行译码是可行的，因为音频标记码本是在语言模型训练中使用的典型方法。

图2： Mini-Omni在批处理并行生成技术中还结合了文本指导机制。

批处理并行解码。虽然之前引入的并行生成方法有效地将推理能力从文本模态转移到音频模态，但我们的实验表明，该模型的推理性能在文本和音频任务之间仍然存在差异，音频响应倾向于更简单。我们假设这是由于模型容量的限制或音频数据的不足。为了解决这个问题，并进一步增强模型在对话过程中的推理能力，最大限度地转移其基于文本的能力，我们实验性地采用了一种批处理方法。考虑到模型在文本模态中的较强性能，我们将单个输入的推理任务扩展为2：一个样本同时需要文本和音频响应，如前所述，而另一个样本只需要文本响应，专注于基于文本的音频合成。但是，第一个样本输出的文本标记被丢弃，第二个样本输出的文本被嵌入到第一个样本对应的相应文本标记位置中。

3、任何模型都可以说话

三阶段培训。我们的训练方法分为三个不同的阶段：

(1)模态对齐。这一阶段的目标是增强文本模型理解和生成语音的能力。Mini-Omni的核心模型完全冻结，只能使用两个适配器。在这一阶段，我们使用来自语音识别和语音合成的数据来训练模型的5个。

(2)适应能力的培训。一旦新的模态与文本模型的输入对齐，适配器就会被冻结。在这一阶段，我们只关注于在给定的音频输入时训练模型的文本能力，因为音频输出是简单地从文本合成的。该模型使用来自语音识别、语音问题回答和文本回答任务的数据进行训练。

(3)多模态的微调。在最后一个阶段，使用综合数据对整个模型进行了微调。此时，所有的模型权值都被解冻和训练。因为主要的模态对齐任务是在适配器传输期间处理的模型输入ID。给定8个并行输出序列，输入还需要8个序列，这导致了显著的复杂性。因此，我们在这里简要地概述了模型输入的组织。该模型可以接受文本或音频输入，它们被放置在相应的模态序列中。对于音频输入，输入令牌和耳语特征通过适配器转换为相同维度的张量，然后连接。根据任务的不同，我们将<回答>特殊标记放置在不同的位置来指导模型的输出，实现多模态输出。一些任务的组织方式如图4所示。在被输入模型之前，所有的序列都被求和并平均以整合特征。

在本节中，我们将介绍我们的培训方法。我们的方法旨在尽可能多地保留原始模型的功能。这首先是因为我们的基础模型的强大性能，其次是因为我们的方法可以应用于其他擅长文本输出但缺乏健壮的语音交互能力的工作。音频编码：音频输入主要集中于从输入的音频中提取特征，与选项包括休伯特或一个单独的预先训练的音频编码器。鉴于我们对语音输入的关注，低语[雷德福等人，2023]和qwen2-音频[Chu等人，2024]也证明了对一般音频任务的有效性能。对于音频输出，使用多码本方法选择音频令牌可以更好地捕获音频细节。我们尝试了音频标记的扁平化建模，但它导致了过长的标记，这不利于流媒体，并导致不稳定的学习。相反，受MusicGen启发[Copet et al.，2024]的并行解码，采用了与文本条件相结合的延迟模式，如图2所示。

4、数据集

为了建立基本的语音能力，我们使用三个语音识别数据集对模型进行训练，总计约8000小时，重点是语音理解和合成。对于文本模式，我们整合了来自开放-Orca[OpenOrca]数据集的200万个数据点，并将它们与其他模式集成，以保持文本的准确性。Moss的SFT数据集[Sun等人，2024]利用零镜头TTS合成了150万对语音QA对。为了避免不合适的代码和符号输出，我们使用GPT-4o创建了Voice助手-400k数据集。数据集的详细信息见表1。阶段1涉及到训练语音适配器的ASR数据。阶段2使用TextQA和AudioQA用于音频/文本输入和文本响应训练。第三阶段侧重于使用AudioQA的音频模态的多模态交互。最后阶段的培训包括退火和使用语音QA的微调。

4.2训练参数我们的模型是在8A100gpu上进行训练的，使用一个余弦退火学习速率调度器，最小学习速率为4e-6，最大学习速率为4e-4。每个训练阶段由40,000步组成，每步的批大小为192。基本语言模型采用了Qwen2-0.5B [Yang等人，2024a]，一个具有24个块、内部维度为896的变压器架构。语音编码器使用耳语小编码器，ASR适配器通过两层MLP连接，而TTS适配器通过添加6个额外的变压器块，扩展了原始模型。在微调期间，我们使用从4e-6到5e-5的学习率。

4.3实验结果我们首先评估了该模型在ASR任务上的表现，以评估其语音理解能力。语音识别能力的基本实验是使用了列表语言[Panayotov等人，2015]的四个测试集进行的：测试清洁、测试其他、开发清洁和开发其他。结果如表2所示，我们比较了我们采用的语音识别系统wavv2vec2[Baevski等人，2020]和低语，以及VITA [Fu等人，2024]的准确性。

4.4在这里，我们提出了几个案例来证明Mini-Omni在语音理解和推理方面的能力。这些例子表明，与基于文本的推理相比，基于语音的推理稍弱，突出了批生成的必要性。

这些特性描述了一个先进的语音识别与合成系统，具备以下能力：

1. 实时语音到语音的会话能力：该系统能够实时地将输入的语音转换为另一种语言的语音输出，而无需额外的自动语音识别（ASR）或文本到语音（TTS）模型。这意味着用户可以直接进行跨语言对话，系统能够即时理解和翻译，并以语音形式反馈，极大地促进了不同语言背景人群之间的即时交流。

2. 边思考边说话的能力：系统具备同时生成文本和音频的能力。这意味着在处理用户输入的语音时，系统能够几乎同时地将其转换为文本（用于内部处理或显示）和语音（用于即时反馈或输出到其他系统）。这种特性对于需要快速响应的交互式应用场景尤为重要。

3. 流式音频输出能力：系统支持流式音频输出，即能够连续不断地处理并输出音频流，而不是等待整个音频文件处理完成后再输出。这对于实时通信和流媒体应用至关重要，因为它能够确保音频传输的连续性和实时性。

4. “音频到文本”和“音频到音频”批量推理功能：除了实时处理能力外，该系统还提供了批量处理模式，支持对大量音频数据进行“音频到文本”或“音频到音频”的转换。这种批量推理功能可以进一步提高处理效率，适用于需要对大量录音进行转录或翻译的场景。通过批量处理，用户可以在较短时间内获得大量数据的处理结果，从而加速工作流程并提高效率。

NOTE: need to unmute first.
 

主要代码：

import os
import lightning as L
import torch
import time
from snac import SNAC
from litgpt import Tokenizer
from litgpt.utils import (num_parameters,
)
from litgpt.generate.base import (generate_AA,generate_ASR,generate_TA,generate_TT,generate_AT,generate_TA_BATCH,next_token_batch
)
import soundfile as sf
from litgpt.model import GPT, Config
from lightning.fabric.utilities.load import _lazy_load as lazy_load
from utils.snac_utils import layershift, reconscruct_snac, reconstruct_tensors, ge