CoSTA comprime os tokens de fala do Step-Audio 2 mini com BPE: de 10K a 320K unidades adicionadas

Modelos de diálogo por voz geram a resposta como longas sequências de tokens de fala discretos, e cada token custa um passo de autorregressão. Pesquisadores do Instituto de Automação da Academia Chinesa de Ciências e da Escola de Inteligência Artificial da Universidade da CAS propuseram o CoSTA, uma forma de reduzir o número desses passos sem mexer no tokenizador de fala original nem no decodificador que converte os tokens em áudio. Para isso, eles adicionam ao vocabulário do Step-Audio 2 mini de 10K a 320K unidades BPE, e cada uma substitui um par ou uma cadeia frequente de tokens de fala vizinhos.

O modelo é treinado para emitir essas unidades comprimidas e, antes da síntese, elas são expandidas de forma determinística de volta aos tokens originais.

Por que os tokens de fala comprimem bem

Na maioria dos modelos de diálogo por voz, o mesmo espaço de tokens serve tanto para a geração quanto para a síntese. Por isso as sequências ficam longas: as respostas saem mais devagar, o alinhamento entre texto e fala fica mais complexo e o limite de geração se esgota mais rápido. Isso restringe a duração da conversa e da resposta falada.

Uma alternativa são os tokenizadores com menor frequência de tokens. Mas reduzir demais a frequência piora a qualidade, e trocar o tokenizador de um modelo pronto pode exigir retreinar o tokenizador, o decodificador pareado e o próprio modelo. Além disso, geração e síntese continuam presas a uma única representação.

Os autores observaram que os tokens de fala contêm muitos padrões locais repetidos, e por isso o BPE (byte pair encoding) se aplica a eles: pares vizinhos frequentes são fundidos em novas unidades. A compressão é reversível: a expansão devolve a sequência original sem perdas. O efeito foi testado em oito configurações de tokenizadores, com fala de diálogo em chinês. O vocabulário BPE foi treinado na parte de treino e fixado para o conjunto reservado de avaliação. Quanto mais unidades são adicionadas, mais curta fica a sequência, e a relação é aproximadamente log-linear. Ela se mantém em diferentes fontes de dados e idiomas.

Como o modelo pronto é adaptado

O CoSTA consiste em três etapas de pós-treinamento. Primeiro, o vocabulário é ampliado, e os embeddings e as linhas da LM head das novas unidades são inicializados com a média das linhas dos tokens originais que compõem cada unidade. Isso dá um ponto de partida mais estável do que a inicialização aleatória, e os parâmetros do vocabulário original permanecem inalterados.

Em seguida vem a adaptação da pronúncia em dados de TTS comprimidos. O Step-Audio 2 mini alterna tokens de texto e de fala, então o BPE é aplicado apenas aos de fala, e os de texto não são tocados. O treino é feito em etapas: na primeira, com 50% dos dados, só os embeddings e a LM head são ajustados, com os demais parâmetros congelados, o que calibra as novas linhas do vocabulário. Na segunda etapa, a adaptação se estende ao restante do modelo. O último passo é a adaptação ao diálogo, com dados comprimidos no formato «fala — fala».

Na inferência, apenas uma expansão BPE leve e determinística é inserida antes do decodificador de fala, que continua o mesmo. O tokenizador e o decodificador originais permanecem como estavam, e por isso a compatibilidade com o pipeline existente é preservada.

O que mostraram os experimentos

O método foi testado em duas arquiteturas de modelos de diálogo por voz. No Step-Audio 2 mini, o CoSTA aumenta a velocidade de geração dos tokens de fala expandidos em chinês e inglês, mantendo a qualidade das respostas em avaliações automáticas e humanas. Isso resulta em menor latência da resposta completa, maior vazão e interações mais longas com um orçamento de geração limitado. Os fatores exatos de aceleração não aparecem no trecho de texto disponível.

O tamanho do vocabulário define um compromisso prático entre grau de compressão, eficiência de decodificação e qualidade, portanto precisa ser ajustado. Os experimentos em outra arquitetura mostraram que a abordagem é transferível e compatível com a tokenização de baixa frequência: os dois métodos de encurtar sequências podem ser combinados.

O código foi publicado no GitHub, no repositório yingzhuodeng/CoSTA.