CoSTA comprime los tokens de voz de Step-Audio 2 mini con BPE: de 10K a 320K unidades añadidas

Los modelos conversacionales de voz generan su respuesta como largas secuencias de tokens de voz discretos, y cada token equivale a un paso de autorregresión. Investigadores del Instituto de Automatización de la Academia China de Ciencias y de la Escuela de Inteligencia Artificial de la Universidad de la CAS han propuesto CoSTA, un método para reducir el número de esos pasos sin tocar ni el tokenizador de voz original ni el decodificador que convierte los tokens en audio. Para ello, añaden al vocabulario de Step-Audio 2 mini entre 10K y 320K unidades BPE, cada una de las cuales sustituye a un par o una cadena frecuente de tokens de voz contiguos.

El modelo se entrena para emitir estas unidades comprimidas, y antes de la síntesis se expanden de forma determinista a los tokens originales.

¿Por qué los tokens de voz se comprimen bien?

En la mayoría de los modelos conversacionales de voz, el mismo espacio de tokens sirve tanto para la generación como para la síntesis. Por eso las secuencias resultan largas: las respuestas son más lentas, se complica la alineación entre texto y voz y el límite de generación se agota antes. Esto restringe la duración de la conversación y de la respuesta hablada.

Una alternativa son los tokenizadores con menor frecuencia de tokens. Pero reducir la frecuencia de forma demasiado agresiva empeora la calidad, y cambiar el tokenizador en un modelo ya entrenado puede obligar a reentrenar el tokenizador, el decodificador asociado y el propio modelo. Además, la generación y la síntesis siguen atadas a una misma representación.

Los autores observaron que los tokens de voz contienen muchos patrones locales repetidos, por lo que se les puede aplicar BPE (byte pair encoding): los pares contiguos frecuentes se fusionan en nuevas unidades. La compresión es reversible: la expansión devuelve la secuencia original sin pérdidas. Comprobaron el efecto en ocho configuraciones de tokenizadores con voz conversacional en chino: el vocabulario BPE se entrenó con la parte de entrenamiento y se fijó para el conjunto de evaluación reservado. Cuantas más unidades se añaden, más corta es la secuencia, y la relación es aproximadamente log-lineal. Se mantiene con distintas fuentes de datos e idiomas.

Cómo se adapta un modelo ya entrenado

CoSTA consta de tres pasos de posentrenamiento. Primero se amplía el vocabulario, y los embeddings y las filas de la cabeza del LM de las nuevas unidades se inicializan con la media de las filas de los tokens originales que componen cada unidad. Así se obtiene un punto de partida más estable que con una inicialización aleatoria, y los parámetros del vocabulario original no cambian.

Después llega la adaptación de la pronunciación con datos TTS comprimidos. Step-Audio 2 mini alterna tokens de texto y de voz, por lo que BPE se aplica solo a los de voz y los de texto no se tocan. El entrenamiento es por etapas: en la primera, con el 50% de los datos, solo se ajustan los embeddings y la cabeza del LM mientras el resto de los parámetros permanece congelado, lo que calibra las nuevas filas del vocabulario. En la segunda etapa, la adaptación se extiende al resto del modelo. El último paso es la adaptación al diálogo con datos comprimidos en formato «voz a voz».

En la inferencia, antes del decodificador de voz, que no cambia, solo se inserta una ligera expansión BPE determinista. El tokenizador y el decodificador originales siguen siendo los mismos, así que se mantiene la compatibilidad con el pipeline existente.

Qué mostraron los experimentos

El método se probó en dos arquitecturas de modelos conversacionales de voz. En Step-Audio 2 mini, CoSTA aumenta la velocidad de generación de los tokens de voz ampliados en chino e inglés y mantiene la calidad de las respuestas según las evaluaciones automáticas y humanas. Esto se traduce en una menor latencia de la respuesta completa, un mayor rendimiento (throughput) y una interacción más prolongada con un presupuesto de generación limitado. El fragmento de texto disponible no incluye los factores exactos de aceleración.

El tamaño del vocabulario fija un compromiso práctico entre el grado de compresión, la eficiencia de la decodificación y la calidad, por lo que hay que elegirlo con cuidado. Los experimentos con otra arquitectura mostraron que el enfoque se traslada y que es compatible con la tokenización de baja frecuencia: ambas formas de acortar las secuencias pueden combinarse.

El código está publicado en GitHub, en el repositorio yingzhuodeng/CoSTA.