Generación de embeddings de audio a partir de LLM multimodales

jina-ai/audio-embedding-kickstarter

GitHubjina-ai

Google lanzó recientemente Gemini Embedding 2, su primer modelo de embeddings multimodal nativo. Texto, imágenes, vídeo, audio, documentos, todos mapeados en un único espacio vectorial de 3072 dimensiones. Esto forma parte de una tendencia más amplia hacia los modelos de embeddings omni: modelos unificados que manejan todas las modalidades en una sola arquitectura, desde jina-embeddings-v4 hasta Omni-Embed-Nemotron o Omni-5.

Lo que nos llamó la atención fue el audio. La mayoría de la gente escucha "embedding multimodal" y piensa en imágenes, tal vez en vídeo. El audio es la modalidad olvidada: más difícil de recolectar, más difícil de etiquetar y con menos personas trabajando en ella. En Jina AI habíamos explorado exactamente este problema, construyendo modelos de embeddings de audio pequeños (de menos de 1.2B parámetros) como parte de nuestro trabajo hacia los embeddings omni. El lanzamiento de Gemini Embedding 2 es un buen momento para compartir lo que aprendimos en el camino.

Embeddings de audio

Un embedding de audio es una representación vectorial de longitud fija de un clip de audio. Dada una forma de onda bruta, el modelo genera un vector denso (típicamente de 768 a 3072 dimensiones) que captura el contenido semántico del sonido. Dos clips con significados similares producen embeddings similares, y un clip de audio se sitúa cerca de su descripción de texto en el espacio de embeddings compartido. Esta es una pieza del rompecabezas de los embeddings omni: una vez que puedes incrustar audio junto con texto e imágenes en el mismo espacio vectorial, desbloqueas la recuperación cross-modal en todas las modalidades.

El enfoque dominante desde 2022 es el Contrastive Language-Audio Pretraining (CLAP), que extiende CLIP al audio. LAION-CLAP escaló esto con 630K pares y fusión de características. La variante más potente (Elizalde et al., 2023) se entrenó con 4.6M de pares utilizando un codificador de audio a través de 22 tareas de audio diversas emparejadas con un decodificador autorregresivo, logrando un 42.0 cvR@5 en AudioCaps con 250M de parámetros.

Nosotros tomamos una ruta diferente: convertir los LLM multimodales que ya entienden audio en modelos de embeddings.

Arquitectura

Primero, qué entra y qué sale. La entrada es audio bruto: una forma de onda decodificada de cualquier formato estándar (WAV, MP3, FLAC) y remuestreada a 16kHz mono. El codificador de audio convierte esta forma de onda en un espectrograma log-mel de 128 contenedores, luego lo procesa en una secuencia de tokens de características a razón de aproximadamente 150 tokens por segundo. Un clip de 10 segundos se convierte en unos 1,500 tokens. La longitud máxima de entrada es de 30 segundos; audios más largos deben ser fragmentados. La salida es un único vector denso (el embedding), típicamente de 896 a 3584 dimensiones, dependiendo del tamaño del backbone del LLM.

Partimos de Qwen2.5-Omni, un LLM multimodal con comprensión de audio nativa. Tres componentes: un codificador de audio (~0.6-0.8B parámetros) que convierte las formas de onda en vectores de características mediante una proyección lineal de ~4.5M, un backbone de LLM (0.5-7B parámetros) que procesa tanto las características de audio como los tokens de texto con cada capa de transformer añadiendo ~0.2B parámetros, y una capa de pooling que realiza un agrupamiento por media (mean-pooling) del último estado oculto en un único vector de embedding. Ambas modalidades comparten el mismo backbone de LLM, por lo que ya están aproximadamente alineadas desde el preentrenamiento.

El objetivo de entrenamiento es la pérdida contrastiva InfoNCE. Cada modalidad se codifica de forma independiente, la pérdida se calcula en ambas direcciones y se promedia:

def training_step(audio_batch, text_batch):
    audio_embeds = model.encode_audio(audio_batch)  # [B, D]
    text_embeds = model.encode_text(text_batch)      # [B, D]
    
    audio_embeds = F.normalize(audio_embeds, dim=-1)
    text_embeds = F.normalize(text_embeds, dim=-1)
    
    sim = audio_embeds @ text_embeds.T / temperature  # [B, B]
    labels = torch.arange(len(sim), device=sim.device)
    loss = (F.cross_entropy(sim, labels) + 
            F.cross_entropy(sim.T, labels)) / 2
    return loss
Copiar

Datos de entrenamiento

Cinco conjuntos de datos de pares audio-texto, 181K muestras en total:

CLAP utilizó el AudioSet completo (más de 2M de audios) además de otras fuentes, sumando 4.6M de pares. Nosotros usamos solo AudioSetStrong (~100K). Partir de un MLLM preentrenado reduce drásticamente la cantidad de datos necesarios.

def load_sample(audio_path, caption):
    waveform, sr = torchaudio.load(audio_path)
    waveform = torchaudio.transforms.Resample(sr, 16000)(waveform)
    audio_inputs = processor.feature_extractor(
        waveform, sampling_rate=16000, return_tensors="pt"
    )
    text_inputs = processor.tokenizer(caption, padding=True, return_tensors="pt")
    return audio_inputs, text_inputs
Copiar

Cuatro enfoques

Objetivo: un modelo de embeddings de audio de menos de 1.2B que supere a CLAP.

Ajuste fino del modelo completo. Qwen2.5-Omni-7B en pares de audio-texto: AudioCaps T2A cvR@5 = 63.2, Clotho T2A = 39.2. Este es el límite superior, pero un 7B no es fácil de desplegar. Tevatron 2.0 realizó un ajuste fino similar solo en AudioCaps (61.2 pero solo 11.9 en Clotho, mostrando una generalización deficiente por el entrenamiento en un solo conjunto de datos). ColQwen-Omni se ajustó en tareas de documentos visuales sin datos de audio, logrando 37.4 mediante transferencia cross-modal.

Poda de capas. Eliminar capas de transformer del modelo 7B. Cada capa es de ~0.2B, por lo que un modelo de 10 capas tiene ~3.5B en total.

El tamaño del lote (32, 64, 128) no supuso una diferencia significativa. Los lotes más grandes ayudan inicialmente pero pueden degradar el rendimiento después: el lote 128 alcanzó 31.3 NDCG en 2K pasos en Clotho, pero cayó a 29.3 en 10K pasos.

Transferencia de modalidad solo de texto. Ajuste fino solo en pares de texto (MultiNLI, SNLI, FEVER, SciFact), confiando en la alineación cross-modal preentrenada. Funcionó en el modelo 7B completo (AudioCaps 46.1, superando los 42.0 de CLAP), pero falló completamente en el modelo podado de 10 capas (cvR@5 = 5.9). El cableado cross-modal está distribuido por toda la red y no sobrevive a la poda.

Combinación de módulos. El avance: tomar un codificador de audio de un modelo y un LLM pequeño de otro, incluso entre diferentes familias de modelos. Qwen2.5-Omni se entrena en tres etapas: (1) codificadores de audio/visión con LLM congelado, (2) todos los parámetros descongelados, (3) contexto de 32K. Combinamos módulos de diferentes etapas:

Detalle de implementación: Qwen3-Omni utiliza Qwen3OmniMoePreTrainedModel mientras que Qwen3 independiente utiliza Qwen3ForCausalLM. Inicializamos un caparazón de modelo Omni con dimensiones coincidentes y copiamos los pesos en las ubicaciones correspondientes.

Evaluación

Evaluar los embeddings de audio consiste fundamentalmente en la calidad de la recuperación: dada una consulta de texto, ¿puede el modelo encontrar el clip de audio correcto? El desafío clave es que lo "correcto" depende del conjunto de datos. AudioCaps tiene descripciones concretas ("un hombre hablando seguido por una puerta cerrándose"), mientras que Clotho tiene subtítulos abstractos ("una atmósfera tranquila con un estruendo distante"). Un modelo que memoriza características de audio superficiales funcionará bien en AudioCaps pero tendrá dificultades en Clotho. Lo que más nos importa es la generalización entre diferentes estilos de descripción.

CV-Recall@5 (cvR@5): para cada consulta de texto, comprueba si algún clip de audio correcto aparece en los 5 primeros resultados. Puntuación binaria promediada sobre todas las consultas. Métrica estándar en la recuperación de audio de MTEB.

def evaluate_cvr_at_k(model, dataset, k=5):
    audio_embeds = model.encode_audio(dataset.audio_clips)
    text_embeds = model.encode_text(dataset.text_queries)
    sim = F.normalize(audio_embeds) @ F.normalize(text_embeds).T
    
    hits = 0
    for i in range(len(dataset.text_queries)):
        top_k = sim[:, i].argsort(descending=True)[:k]
        if dataset.ground_truth[i] in top_k:
            hits += 1
    return hits / len(dataset.text_queries)
Copiar

Tres conjuntos de datos de evaluación de MTEB: AudioCaps (derivado de video, descripciones humanas), AudioSetStrong (etiquetado temporalmente, descripciones de GPT), Clotho (descripciones diversas y abstractas). CLAP utilizó el AudioSet completo (más de 2 millones), mientras que nosotros utilizamos AudioSetStrong (aprox. 100 000), lo que explica en parte la ventaja de CLAP en ese benchmark.

Aplicaciones

Los embeddings de audio están ganando relevancia más allá de la recuperación tradicional. En los sistemas agénticos, los embeddings de audio permiten el enrutamiento de intenciones: un agente que recibe una entrada de voz puede generar el embedding del audio y dirigirlo a la herramienta o sub-agente adecuado basándose en la similitud semántica, sin esperar a la transcripción completa. La clasificación de eventos de sonido impulsa el monitoreo en tiempo real en entornos industriales, la automatización del hogar inteligente y los sistemas de seguridad. En los flujos de trabajo de agentes multimodales, los embeddings de audio permiten a los agentes buscar, comparar y razonar sobre el contenido de audio de la misma manera que ya lo hacen con el texto y las imágenes. Las aplicaciones de música y medios los utilizan para la búsqueda por similitud, la detección de derechos de autor y la recomendación de contenido. A medida que las interfaces de voz se convierten en el modo de interacción por defecto para los agentes de IA, los embeddings de audio compactos que se ejecutan localmente (on-device) se vuelven críticos para aplicaciones de baja latencia que preservan la privacidad.

Conclusiones

Partir de un MLLM preentrenado es la palanca más importante. Proporciona alineación intermodal, un codificador de texto robusto y un codificador de audio capaz en un solo paquete. La combinación de módulos es la dirección más prometedora: mezclar codificadores de audio y LLM de diferentes modelos y etapas de entrenamiento abre un espacio de diseño apenas explorado. Nuestros modelos dominan en AudioCaps pero solo igualan a CLAP en Clotho, cuyas descripciones abstractas exponen debilidades que AudioCaps no detecta. La transferencia intermodal no sobrevive a la compresión del modelo.

Este trabajo es un paso hacia el modelo de embedding «omni»: un único modelo que integra texto, imágenes, audio, video y documentos en un espacio de recuperación unificado. El enfoque de combinación de módulos demuestra que se pueden impulsar nuevas modalidades de manera eficiente reutilizando componentes preentrenados. Los siguientes pasos incluyen arquitecturas MoE con menos de 500 millones de parámetros de activación, la combinación de la unión de módulos con la transferencia de modalidad y el escalado de datos con WavCaps, MusicCaps y conjuntos de datos de voz.