An Exploration of Mamba for Speech Self-Supervised Models

1. Problema e Motivação

Os modelos baseados em Transformers e seus mecanismos de atenção multi-head dominaram o campo de processamento de linguagem e fala nos últimos anos. No entanto, eles sofrem de complexidade computacional quadrática em relação ao comprimento da sequência, o que resulta em custos de implantação elevados e limitações de memória ao processar contextos longos (ex: documentos de áudio inteiros).

Como alternativa, a arquitetura Mamba (baseada em Modelos de Espaço de Estado Selecionados - SSMs) oferece complexidade linear no tempo, mantendo a capacidade de seleção de conteúdo. Embora o Mamba tenha mostrado desempenho superior em modelagem de linguagem, sua aplicação em aprendizado auto-supervisionado (SSL) de fala permanece subexplorada. Estudos anteriores limitaram-se a tarefas isoladas ou mostraram desempenho inferior aos Transformers em certas configurações, sem uma avaliação unificada e abrangente.

O objetivo deste trabalho é preencher essa lacuna, explorando sistematicamente modelos HuBERT baseados em Mamba como alternativas viáveis aos Transformers para SSL de fala, focando em eficiência, modelagem de contexto longo e extração de características.

2. Metodologia

Os autores desenvolveram e treinaram modelos Mamba-based HuBERT, substituindo os blocos Transformer do HuBERT original por blocos Mamba, mantendo os módulos de extração de características iniciais (CNN).

• Treinamento: Seguiram o pipeline de duas iterações do HuBERT Base:
  1. Primeira iteração: Treinamento com 250k passos usando características MFCC como alvos.
  2. Segunda iteração: Reinício do treinamento do zero usando as saídas da sexta camada da primeira iteração como alvos (400k passos).
  • Nota: O treinamento foi realizado em uma única GPU NVIDIA V100 com um tamanho de lote ajustado para limitações de memória, mas com configurações idênticas para os modelos Transformer de comparação.
• Configurações Arquiteturais:
  • Causal: Modelos Mamba e Transformers Causais (usando máscaras triangulares inferiores).
  • Bidirecional: Modelos com ExtBiMamba (External Bidirectional Mamba) e InnBiMamba (Inner Bidirectional Mamba), além do Transformer padrão. O ExtBiMamba utiliza duas camadas independentes de Mamba (uma para a ordem original e outra para a ordem reversa), enquanto o InnBiMamba compartilha projeções.
• Avaliação:
  • Eficiência: Medição de MACs (operações de multiplicação-acumulação) e Fator de Tempo Real (RTF) em sequências de 5 segundos a 5 minutos.
  • Ajuste Fino (Fine-tuning): Tarefas de Reconhecimento Automático de Fala (ASR) em contexto longo (TEDLIUM3) e ASR causal/streaming (LibriSpeech).
  • Análise de Representação: Pureza de fonemas (quantização) e Análise de Correlação Canônica (CCA) para características de fonemas, falantes e emoções.
  • Benchmarks: Avaliação no SUPERB (MiniSUPERB) para tarefas de reconhecimento de fonemas (PR), identificação de falante (SID), reconhecimento de emoção (ER) e classificação de intenção (IC).

3. Principais Contribuições

1. Validação da Arquitetura Causal: Demonstraram que a arquitetura causal inerente do Mamba é particularmente adequada para modelos SSL de fala, superando os Transformers causais em tarefas de ASR streaming com menos parâmetros.
2. Eficiência em Contexto Longo: Provaram que os modelos Mamba-based HuBERT permitem o ajuste fino para ASR de contexto longo (processamento de documentos inteiros sem segmentação), onde os Transformers falham devido a erros de memória (OOM).
3. Qualidade de Representação Quantizada: Encontraram que os modelos Mamba produzem representações quantizadas com maior qualidade fonética, o que é benéfico para modelos de linguagem falada que utilizam unidades SSL como entrada.
4. Exploração Abrangente: Este é o primeiro estudo a explorar o Mamba-based HuBERT tanto como modelo fundamental de fala quanto como extrator de características, cobrindo pré-treinamento, ajuste fino e análise de representações.

4. Resultados Chave

Eficiência Computacional

• Escalabilidade Linear: O custo computacional (MACs) e o RTF do Mamba permanecem quase constantes à medida que o comprimento da sequência aumenta (de 5s a 5min).
• Limitação dos Transformers: Os modelos baseados em Transformers mostram um aumento quadrático no custo e falham com erros de memória (OOM) além de 80 segundos de áudio.

Desempenho em ASR (Ajuste Fino)

• ASR de Contexto Longo: No nível de documento, o ExtBiMamba alcançou uma Taxa de Erro de Palavras (WER) de 11,08%, enquanto o Transformer não pôde ser executado. No nível de enunciado, o Mamba teve desempenho ligeiramente inferior (13,37% vs 11,86% do Transformer), mas a vantagem do contexto global foi decisiva.
• ASR Causal/Streaming: O modelo Mamba Base (78M parâmetros) superou o Causal Transformer Base (94M parâmetros) com uma WER de 15,77% contra 16,66%, demonstrando maior eficiência com ~17% menos parâmetros.

Análise de Representação

• Pureza de Fonemas: O Mamba demonstrou maior pureza de fonemas (correlação com rótulos fonéticos) em camadas intermediárias e tardias, especialmente no modo causal, indicando uma codificação fonética superior.
• Características de Falante: O Mamba capturou características relacionadas ao falante de forma mais distinta do que os Transformers, conforme medido pela similaridade CCA com embeddings de falantes.

Benchmarks SUPERB (Probing)

• Configuração Causal: O Mamba-based HuBERT superou consistentemente o Transformer em reconhecimento de fonemas e pontuação geral (SUPERBS), especialmente em modelos de tamanho Base.
• Configuração Bidirecional:
  • Escala Base: O Transformer superou o ExtBiMamba em todas as métricas, indicando desafios de escalabilidade do Mamba bidirecional em tamanhos maiores.
  • Escala Pequena: O ExtBiMamba superou o Transformer na maioria das tarefas (PR, SID, ER), mostrando eficácia em escalas menores.
  • Ablação: O InnBiMamba (com MLP) superou o ExtBiMamba em modelos grandes, sugerindo que a escolha da arquitetura bidirecional é crítica e depende do tamanho do modelo.

5. Significado e Conclusão

Este trabalho estabelece o Mamba-based HuBERT como uma direção promissora e complementar para a modelagem de fala. As principais implicações são:

• Viabilidade para Contexto Longo: O Mamba resolve o gargalo de memória dos Transformers, permitindo o processamento eficiente de áudios longos e documentos completos, essencial para aplicações de transcrição de reuniões ou livros de áudio.
• Eficiência em Tempo Real: Sua natureza causal e complexidade linear o tornam ideal para sistemas de ASR streaming de baixa latência.
• Qualidade de Representação: A capacidade de extrair representações fonéticas e de falante de alta qualidade sugere que o Mamba é um forte candidato para alimentar modelos de linguagem falada (SLMs) e tarefas de extração de unidades de fala.

O estudo conclui que, embora existam desafios de escalabilidade na configuração bidirecional (especialmente em modelos grandes), o Mamba oferece vantagens significativas em cenários causais e de contexto longo, redefinindo o potencial de arquiteturas de espaço de estado na área de processamento de fala. O código é disponibilizado publicamente para reprodutibilidade.