rwkv-quant: uma matriz em bf16 elimina 39% das perdas na quantização do RWKV-7 2.9B

A quantização de modelos da família RWKV, com sua estrutura recorrente linear, ainda é menos estudada do que a de transformers. O autor do rwkv-quant, ferramenta que comprime o RWKV-7 com dois presets (cerca de 2.1x e cerca de 3.1x), notou uma anomalia. O COMPRESSION agressivo resulta em +4.04% de perplexidade no 2.9B, mas só +2.85% no 1.5B. A análise das causas mostrou que boa parte das perdas está em uma única matriz de 2560×2560. Armazená-la em bf16 custa 0.5% do arquivo e elimina 39% do KL.

Dois presets e seis tamanhos

O estudo examina a linha RWKV-7 em seis tamanhos. Três pares permitem separar a influência da largura e da profundidade: 0.4B/1.5B (mesma profundidade, largura duas vezes maior), 2.9B/7.2B (profundidade 32, largura 2560 contra 4096) e 7.2B/13.3B (largura 4096, profundidade 32 contra 61).

• 0.1B: 12 camadas, largura 768
• 0.4B: 24 camadas, largura 1024
• 1.5B: 24 camadas, largura 2048
• 2.9B: 32 camadas, largura 2560
• 7.2B: 32 camadas, largura 4096
• 13.3B: 61 camadas, largura 4096

O preset REDUCTION é quase sem perdas: projeções em 8 bits com layout simétrico e grupo de 16 pesos, cmix (FFN) em 6 bits, embedding e head em 8 bits. O COMPRESSION é mais agressivo: projeções em 4 bits com layout assimétrico em superblocos e grupo de 32, cmix em 4 bits (ffn.key em 5), head em 5 bits, embedding em 6 bits. Nos dois presets, os ramos LoRA ficam em 6 bits (g em 8). Em ambos, a escala e o deslocamento dos grupos são escolhidos por busca, ponderada pela média do quadrado da ativação de entrada do canal, E[x²].

Onde o 2.9B quebra

A decomposição do KL por grupos de matrizes, pelo método leave-one-out, mostrou que no 2.9B o o_proj responde por 46% de todo o KL e o embedding por 17%. No 1.5B, essas parcelas são de 16% e 3%, e a maior contribuição vem do cmix.value (37%). Dentro do o_proj, 85–86% do dano vem de uma única matriz: a projeção de saída do bloco zero.

Um experimento direto no 2.9B (caminho real, MLX) confirmou isso:

• preset base: KL 0.052048, +4.04% de ppl;
• o_proj da camada 0 em bf16: KL 0.031966, +2.43% de ppl, arquivo 9.4 MB maior;
• o_proj da camada 31 em bf16: sem efeito (+4.07%);
• todo o o_proj em 5 bits: +3.54%, +26.2 MB;
• todo o o_proj em 6 bits: +2.89%, +52.4 MB.

Uma única matriz em bf16 funciona melhor do que passar todo o o_proj para 6 bits, e custa mais de cinco vezes menos espaço. Em outros tamanhos o efeito é mais fraco: o KL cai 7.6% no 0.1B, 1.5% no 0.4B, 0.4% no 1.5B e 3.9% no 7.2B. Só o 0.1B e o 2.9B são significativos.

A entrada dessa matriz carrega um "canal massivo". No 2.9B, oito canais concentram 98% da energia, e o número do canal, 517, coincide em duas séries de checkpoints; no 1.5B, é o canal 1202. Ou seja, ele é determinado pela largura e pela arquitetura, e não por um treinamento específico. Mas a intensidade do canal não prevê o dano: o 7.2B tem a maior concentração (99%) e o menor KL base (0.023 contra 0.078 do 2.9B). Sem a ponderação AW, 4 bits no o_proj destroem o 2.9B: KL 0.340, +37.2% de ppl.

O autor mostra que três métricas locais de erro, todas naturais, não preveem o dano. A que funcionou foi o erro no fluxo residual multiplicado pelo ganho medido.

Vazamento de calibração e medições honestas

O primeiro conjunto de avaliação (38 janelas de 512 tokens, en/ru/sr) estava inteiramente dentro do corpus de calibração: 100% dos 16-gramas eram comuns. A escala do vazamento foi medida no GPTQ. No 1.5B, o ganho de KL foi de −61.8% no conjunto antigo e de −27.1% no limpo, ou seja, o vazamento inflou o resultado em 2.3 vezes. O ajuste automático de bitwidth não se beneficiou do vazamento: −16.1% contra −17.6%.

Por isso, as seções 7–10 do trabalho usam apenas conjuntos limpos: Wikipédia en/ru/sr (36 janelas), um conjunto fechado de código Python/Swift (24 janelas) e um aberto (48 janelas). Os intervalos são construídos por bootstrap pareado por janelas, com 20 000 repetições; a meia-largura do intervalo de 95% para o KL equivale a 7–12% do valor.

Além disso, o trabalho descreve o ajuste automático de bitwidth pela regra de "dano por byte", que funciona nos dois sentidos, e o GPTQ na grade do preset, sem mudar o formato, os kernels nem o tamanho do arquivo. Esses métodos se somam de forma quase multiplicativa. Há ainda uma comparação separada com GGUF e MLX, com arquivos do mesmo tamanho.

Na prática, a conclusão é simples. Se você quantiza o RWKV-7 para execução local, principalmente em Mac via MLX, mantenha o o_proj da primeira camada em bf16: são 6–10 MB nos modelos de 1.5B a 2.9B. O embedding deve ficar em pelo menos 6 bits. Antes de confiar nos números de um método calibrado, verifique se a interseção de 16-gramas entre o texto de calibração e o de avaliação é zero. E é melhor medir a contribuição de cada grupo de matrizes por leave-one-out do que supor que, em um modelo grande, as perdas se distribuam de maneira uniforme.