🤖 machine learning

BCJR-QAT: A Differentiable Relaxation of Trellis-Coded Weight Quantization

Este artigo apresenta o BCJR-QAT, uma relaxação diferenciável da quantização de pesos codificada em treliça que substitui o argmax não diferenciável do Viterbi por um algoritmo de soma-produto BCJR controlado por temperatura, permitindo treinamento consciente da quantização de ponta a ponta que supera empiricamente a fronteira de PTQ mais avançada em LLMs.

Autores originais: Venugopalan Iyengar

Publicado 2026-05-12

📖 6 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Venugopalan Iyengar

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você tem uma biblioteca massiva de livros (um Modelo de Linguagem de Grande Escala) que deseja reduzir para que caiba em um laptop ou telefone padrão. Para fazer isso, você precisa comprimir as "palavras" (pesos) dentro da biblioteca.

O artigo apresenta um novo método chamado BCJR-QAT para comprimir esses livros ainda mais — até apenas 2 bits por palavra (o que é como comprimir uma foto de alta definição em uma miniatura pequena e granulada).

Aqui está a história de como eles fizeram isso, usando analogias simples:

1. O Problema: A Armadilha da "Rua de Mão Única"

Anteriormente, a melhor maneira de reduzir esses modelos era um método chamado QTIP. Pense no QTIP como um bibliotecário muito inteligente que olha para uma página de texto e diz: "Ok, vou substituir essa frase complexa pela frase simples mais próxima do meu dicionário."

No entanto, havia um limite. Uma vez que o bibliotecário fazia essa escolha, ele não podia voltar atrás e mudar de ideia para fazer o livro inteiro soar melhor. Era como uma "rua de mão única". Se o bibliotecário escolhesse uma frase ligeiramente errada no início, o resto do livro sofria, e ele não podia corrigir porque a decisão era "rígida" e final.

Para corrigir isso, os pesquisadores geralmente usam Treinamento (QAT), onde você deixa o bibliotecário praticar e ajustar suas escolhas. Mas aqui está a pegadinha: o processo de decisão do bibliotecário envolve um labirinto complexo (um "treliça"). Para encontrar o melhor caminho através do labirinto, eles usam uma regra chamada Viterbi, que é como um botão de "escolher a opção absolutamente melhor".

O Problema: Você não pode ensinar um computador a aprender através de um botão de "escolher o melhor", porque esse botão é um beco sem saída matemático. Se você pressionar o botão, o computador não sabe como ajustar suas escolhas para melhorar; ele apenas salta para a opção mais próxima. É como tentar dirigir um carro permitindo apenas que ele pule instantaneamente para a próxima faixa — você não consegue fazer curvas suaves.

2. A Solução: A Decisão "Suave" (BCJR)

Os autores, Venugopalan Iyengar, inventaram uma nova maneira de deixar o bibliotecário praticar. Em vez de forçar uma decisão rígida de "escolher o melhor" imediatamente, eles introduziram um conceito chamado Temperatura ( $T$ ).

Temperatura Alta (Quente): Imagine que o bibliotecário está muito relaxado e indeciso. Em vez de escolher apenas uma frase, ele considera muitas frases ao mesmo tempo, atribuindo uma "probabilidade" a cada uma. É como um dia nebuloso onde você consegue ver vários caminhos, não apenas um. Essa visão "suave" é matematicamente suave, o que significa que o computador pode calcular facilmente como ajustar as escolhas para melhorar o livro.
Temperatura Baixa (Fria): À medida que o treinamento termina, o bibliotecário fica "mais frio" e mais decisivo. A neblina se dissipa, e ele salta para a única melhor frase (a decisão rígida).

Esse método, chamado BCJR, transforma a "rua de mão única" em uma estrada suave e transitável. O computador agora pode aprender exatamente como ajustar os pesos para fazer o livro inteiro soar melhor, não apenas frases individuais.

3. O Erro de "Excesso"

Os autores descobriram uma peculiaridade surpreendente em como essa "temperatura" funciona.

Na física tradicional, quando você tenta encontrar a melhor solução, geralmente começa muito quente (muito indeciso) para explorar todas as possibilidades e, em seguida, esfria lentamente.

A Descoberta do Artigo: Para esse tipo específico de compressão, começar "muito quente" é um desastre. Se o bibliotecário estiver muito indeciso no início, ele se perde em um bairro ruim (uma "bacia de Voronoi pior") e fica preso lá. Quando ele esfria e tenta encontrar o melhor caminho, já está perdido em uma área ruim e não consegue voltar.
A Correção: Eles descobriram que começar em uma temperatura moderada (nem muito quente, nem muito fria) funciona melhor. É como começar uma trilha com um mapa claro em vez de vaguear cegamente na neblina. Ao pular a fase "super quente", eles evitaram se perder e encontraram uma solução melhor.

4. Os Resultados: Uma Vitória para Computadores Pequenos

A equipe testou isso em dois tipos diferentes de modelos:

O Teste "Proxy" (OLMoE): Eles tentaram otimizar a compressão apenas observando o quão bem cada camada do modelo reconstituía seus próprios dados (como verificar se uma fotocópia se parece com o original). Resultado: Falhou. O modelo ficou, na verdade, pior do que o método antigo. Isso ensinou a eles que apenas fazer uma "fotocópia melhor" não significa que o livro fica melhor de ler.
O Teste "Real" (Llama-3.2): Eles otimizaram o modelo para realmente ler e entender texto melhor (usando um método chamado "destilação", onde um modelo professor inteligente guia o aluno). Resultado: Funcionou!
- Em uma camada específica do modelo, seu novo método superou o antigo melhor método por uma margem pequena, mas significativa.
- Quando aplicaram isso a várias camadas, as melhorias se somaram de uma maneira "superaditiva" (o todo tornou-se maior que a soma das partes).

5. O Motor Sob o Capô

Fazer essa matemática geralmente é muito lento e requer supercomputadores caros. Os autores também construíram um "motor" especializado (um kernel Triton) que executa essa matemática complexa incrivelmente rápido em uma única placa gráfica de consumidor (como uma RTX 4080). Eles a tornaram 6,5 vezes mais rápida do que os métodos padrão, provando que essa técnica avançada não precisa de um data center para rodar.

Resumo

O artigo apresenta uma nova maneira de reduzir modelos de IA para caber em dispositivos de consumo.

O Truque: Eles substituíram um processo de decisão rígido e não treinável por um "suave" e treinável que endurece gradualmente.
A Percepção: Você não deve começar o processo de aprendizado muito "quente" (muito aleatório); um início moderado impede que o modelo se perca.
O Resultado: Quando treinado corretamente para entender texto (e não apenas copiar dados), esse método produz resultados melhores do que as técnicas anteriores mais avançadas, e roda rápido o suficiente em um PC gamer comum.

Os autores lançaram seu código e os modelos treinados para que outros possam experimentá-los.

Resumo Técnico: BCJR-QAT

Declaração do Problema

Alcançar 2 bits por peso (bpw) ou menos para Modelos de Linguagem Grande (LLMs) em hardware de consumidor é um gargalo crítico. Embora métodos de Quantização Pós-Treinamento (PTQ) como o QTIP tenham estabelecido uma fronteira forte para quantização de 2 bits usando Quantização Codificada por Treliça (TCQ), eles são limitados pelo "teto do PTQ". Empurrar além disso requer Treinamento Consciente da Quantização (QAT). No entanto, técnicas padrão de QAT falham em códigos de treliça porque o algoritmo de decodificação de Viterbi envolve uma operação argmax não diferenciável sobre caminhos exponencialmente numerosos. Aplicar um Estimador Direto (STE) à treliça destrói o acoplamento entre decisões de pesos adjacentes, que é a principal vantagem da codificação por treliça sobre a quantização escalar.

Metodologia: BCJR-QAT

O artigo introduz o BCJR-QAT, um relaxamento diferenciável da quantização codificada por treliça que substitui o argmax rígido de Viterbi por uma expectativa suave derivada do algoritmo BCJR (Bahl-Cocke-Jelinek-Raviv) de forward-backward.

Mecanismo Central

Analogia da Mecânica Estatística: A treliça é tratada como um grafo de fatores equivalente a uma cadeia de spins unidimensional do tipo Ising. O argmax de Viterbi é visto como o limite de temperatura zero ( $T \to 0$ ) de uma distribuição de Boltzmann sobre caminhos.
Palavra-Código Suave: Em uma temperatura finita $T > 0$ , o método calcula as probabilidades posteriores marginais exatas de cada estado usando o algoritmo BCJR. A saída quantizada é a expectativa de Boltzmann da palavra-código emitida:
$\hat{w}_t(T) = \langle c(s_t) \rangle_{p_T} = \sum_s c(s) p_T(s_t = s | w)$
Esta palavra-código suave é exatamente diferenciável em relação à pré-imagem do peso e aos parâmetros da treliça.
Recozimento: Durante o treinamento, a temperatura $T$ é recozida de um valor inicial $T_0$ para um valor próximo de zero $T_{end}$ . À medida que $T \to 0$ , a palavra-código suave converge para a saída rígida de Viterbi, efetivamente "cristalizando" a otimização contínua em uma solução discreta.

Inovações Técnicas Chave

Diferenciabilidade: A recursão forward-backward do BCJR (complexidade $O(L \cdot S^2)$ ) fornece gradientes exatos através da treliça, preservando o acoplamento inter-estado que métodos de STE escalar perdem.
Kernel Triton: Os autores implementaram um kernel fundido do Triton que combina operações de coleta (gather), log-sum-exp e adição em um único lançamento. Isso reduz a passagem forward/backward de centenas de lançamentos de kernel para um, alcançando uma aceleração de 6,57× sobre uma implementação de referência do PyTorch autograd, mantendo a paridade numérica FP32.
Correção de Cronograma (Skip-High-T): O artigo diagnostica um modo de falha crítico em cronogramas padrão de recozimento simulado (começando em $T_0=1.0$ ). Em temperaturas altas, o gradiente da palavra-código suave está "espalhado" e não informativo, fazendo com que o otimizador derive para uma bacia de Voronoi pior antes que a fase de resfriamento possa recuperar. A correção proposta é pular a fase de alta- $T$ , inicializando em uma temperatura mais baixa (por exemplo, $T_0=0.3$ ) onde os gradientes já são informativos.

Resultados Experimentais

O artigo avalia o BCJR-QAT em duas arquiteturas: OLMoE-1B-7B (MoE) e Llama-3.2-1B (Densa).

1. OLMoE-1B-7B (Resultado Negativo: Lacuna do Proxy)

Configuração: QAT gulosa por camada otimizando para MSE de reconstrução (um objetivo proxy).
Resultado: O BCJR-QAT reduziu com sucesso a MSE de reconstrução por camada em 17–20% em comparação com o QTIP-PTQ. No entanto, essa melhoria não se traduziu em desempenho na tarefa final. Os modelos regrediram em $\sim 1,3$ de perplexidade (PPL) no WikiText-2 em comparação com a linha de base QTIP-PTQ.
Conclusão: Isso confirma uma "lacuna de proxy" onde minimizar o erro de reconstrução local não garante melhor desempenho na tarefa final em 2 bpw.

2. Llama-3.2-1B (Resultado Positivo: KL End-to-End)

Configuração: QAT de camada única otimizando para destilação de forward-KL end-to-end contra um professor FP16.
Resultado:
- Excesso de Cronograma: Usar um cronograma ingênuo de alta- $T$ ( $T_0=1.0$ ) resultou em desempenho ligeiramente pior que o PTQ (+0,005 a +0,022 PPL).
- Correção Skip-High-T: Usar o cronograma corrigido ( $T_0=0.3$ ) eliminou o excesso. Na Camada 4, o modelo alcançou 10,1347 PPL, uma melhoria de −0,084 PPL sobre a linha de base QTIP-PTQ (10,2189).
- Super-Aditividade: Ao combinar camadas treinadas (Camada 4 e Camada 8), o ganho conjunto (−0,077 PPL) excedeu a soma dos ganhos individuais (−0,062 PPL), demonstrando um "excedente de cooperação" onde as palavras-código treinadas pelo BCJR se reforçam mutuamente.

3. Análise Teórica

Teoria do Orçamento de Deriva: Os autores derivam uma condição para escapar da bacia de Voronoi do PTQ: a deriva cumulativa do peso ( $\eta \cdot N \cdot g_{max}$ ) deve exceder o raio da bacia ( $r_{Voronoi} \approx 10^{-3}$ ). Experimentos confirmam que cronogramas que falham nessa condição não podem escapar da solução PTQ, enquanto aqueles que a excedem podem encontrar soluções melhores se o objetivo e o cronograma estiverem corretos.

Significado e Alegações

O artigo alega três contribuições primárias:

Algorítmica: O BCJR-QAT é o primeiro método a tornar a quantização codificada por treliça nativamente diferenciável, habilitando QAT em TCQ sem perder os benefícios do acoplamento de caminhos.
Engenharia: O kernel fundido do Triton torna isso computacionalmente viável em hardware de consumidor (uma única GPU de 12 GB), removendo a barreira de implantação no tempo de inferência para pesquisa.
Insight Empírico:
- Identifica que a MSE de reconstrução por camada é um proxy pobre para QAT de 2 bits, enquanto a destilação de KL end-to-end é eficaz.
- Desafia o conhecimento convencional do recozimento simulado para relaxamentos discretos, mostrando que começar "quente" ( $T=1.0$ ) é prejudicial para o BCJR-QAT porque o gradiente de alta temperatura é dominado por ruído. A estratégia ótima é começar em uma temperatura onde o gradiente já é informativo.

O artigo conclui que, embora os resultados de camada única sejam promissores, o potencial total do BCJR-QAT provavelmente reside no treinamento conjunto end-to-end através de todas as camadas, uma tarefa tornada acessível pelo kernel liberado, mas que requer hardware em escala de nuvem (H100/H200) para modelos maiores que 1 bilhão de parâmetros. Os autores afirmam explicitamente que não realizaram esse treinamento conjunto em escala total neste trabalho.