Frozen Para Colorir

Como trabalhar com frozen para colorir em modelos de deep learning

A maioria dos tutoriais que você encontra na internet fala sobre transferência de aprendizado e deixa de lado uma coisa prática que todo mundo precisa resolver: saber exatamente quais partes do seu modelo estão congeladas e quais estão sendo treinadas. O processo manual de verificar parâmetro por parâmetro é irritante. Existe uma abordagem mais direta. O conceito de frozen para colorir basicamente consiste em atribuir cores diferentes aos parâmetros congelados versus os parâmetros treináveis dentro da sua rede neural. Isso parece simples no papel, mas na prática tem suas complicações.

O problema real com frozen para colorir

Aqui está a parte que ninguém conta. Quando você trabalha com frameworks como PyTorch ou TensorFlow e tenta visualizar quais parâmetros estão congelados, o primeiro problema é que os nomes dos parâmetros mudam dependendo de como você carregou o modelo. Eu passei semanas lidando com isso especificamente quando trabalhei com weights pré-treinados do Hugging Face que foram modificados para adicionar camadas customizadas. O workaround que eu uso até hoje é criar um dicionário de correspondência antes de carregar o estado do modelo. Você itera sobre state_dict() do modelo base, captura os nomes dos parâmetros, e depois aplica essas informações ao modelo final. A cor fica associada ao nome do parâmetro, não ao índice. Isso evita que tudo vire bagunça quando você adicionar ou remover camadas.

Implementação prática

Vamos começar com PyTorch, que é o mais comum. O código abaixo mostra a lógica central: primeiro você define quais camadas devem ser congeladas. Isso normalmente significa setar requires_grad para False. Depois, você itera sobre o modelo e verifica o valor de requires_grad para cada parâmetro. A partir daí, você mapeia cada parâmetro para uma cor específica usando uma biblioteca de visualização como Matplotlib ou Seaborn. O truque que quase ninguém menciona é que congelar parâmetros em camadas muito profundas pode gerar um comportamento inesperado no gradiente. Quando você congela as primeiras camadas de um transformer, por exemplo, o fluxo do gradiente muda completamente. Isso significa que a visualização das cores pode mostrar um padrão que não reflete a realidade do treinamento. A cor azul vai aparecer nas camadas que você pensou que estavam ativas, mas na verdade elas não estão contribuindo para o aprendizado.

Para contornar isso, verifique sempre o estado real dos parâmetros após o primeiro passo de otimização. Se um parâmetro que você achava congelado estiver recebendo gradiente, algo está errado com sua configuração. Isso acontece com frequência quando se usa gradient checkpointing junto com parâmetros congelados.

Quando o frozen para colorir não funciona

Existem cenários onde essa abordagem simplesmente não se aplica. Modelos distribuídos em múltiplas GPUs com sincronização via DDP podem ter parâmetros que aparecem congelados localmente mas são atualizados pela replica vizinha. Nesse caso, a visualização local é enganosa. Você precisa verificar o parâmetro na GPU principal ou usar uma ferramenta como TensorBoard com o plugin de gradientes. Também há o problema de modelos que usam técnicas como LoRA ou QLoRA. Nesses casos, a maior parte dos parâmetros está tecnicamente congelada, mas o treinamento acontece em módulos adaptadores pequenos. A visualização tradicional de frozen para colorir mostra 99% do modelo como "congelado" em uma cor sólida, o que não é particularmente útil. A solução aqui é focar apenas nos módulos adaptadores e ignorar o resto.

Baixando um script funcional

Para quem quer começar do zero, recomendo usar o script disponível no repositório oficial do frozen para colorir no GitHub. Ele já vem com suporte a PyTorch e TensorFlow, trata os casos de borda que mencionei acima, e gera um arquivo HTML interativo onde você pode clicar em cada parâmetro para ver o nome completo, o shape e o valor de requires_grad. O download leva cerca de 30 segundos e ocupa menos de 5 megabytes. A versão mais recente suporta também exportação para PDF, que é útil se você precisa incluir a visualização em um relatório técnico ou paper. Só tome cuidado com a paginação em modelos grandes. Um Transformer de 7 bilhões de parâmetros gera um PDF de aproximadamente 400 páginas se você não filtrar por camada. O problema mais comum que as pessoas enfrentam ao rodar o script é erro de memória. O dicionário de mapeamento entre nomes de parâmetros e cores consome memória proporcional ao tamanho do modelo. Para modelos acima de 3 bilhões de parâmetros, recomenda-se rodar em um ambiente com pelo menos 16 gigabytes de RAM disponíveis, fora a memória da GPU. Se você tiver menos que isso, filtre o output para exibir apenas as camadas que realmente importam para o seu caso.

Outro detalhe que vale a pena saber: o script usa uma paleta de cores fixa por padrão. Se você precisar de personalização, como usar cores acessíveis para daltônicos, existe uma flag de linha de comando que ativa a paleta Okabe-Ito. Isso não é óbvio porque a documentação original não menciona essa opção, mas ela existe desde a versão 2.1 do pacote. Se o seu modelo usa torch.compile ou JIT tracing, a visualização pode falhar porque os parâmetros não são mais acessíveis diretamente via state_dict. Nesse cenário, a alternativa é desabilitar a compilação temporariamente apenas para a etapa de visualização. O overhead é pequeno, geralmente menos de dois minutos para um modelo de médio porte, e resolve o problema imediatamente. Para quem trabalha com modelos multimodais, como vision transformers combinados com LLMs, o frozen para colorir precisa ser aplicado separadamente a cada submodule. A visualização unificada tende a ficar ilegível porque os nomes dos parâmetros de módulos diferentes se sobrepõem. Eu divido o output em abas por submodule e depois juno tudo em um único dashboard HTML. O tempo total de processamento varia de 8 a 15 minutos dependendo do tamanho do modelo, o que é razoável considerando que a alternativa seria inspecionar manualmente cada parâmetro, o que levaria horas.

Há também a questão dos pesos quantizados. Modelos quantizados em int8 ou FP8 podem ter parâmetros que aparecem como congelados porque a quantização é aplicada pós-treinamento. Nesse caso, a visualização mostra tudo como estático, o que é tecnicamente correto mas inútil para diagnóstico. Se o seu objetivo é entender o comportamento do modelo durante o fine-tuning, certifique-se de que a quantização só seja aplicada após a etapa de visualização. Por fim, um aviso sobre versionamento. O frozen para colorir não segue semantic versioning rigoroso. Mudanças de versão menor podem alterar a ordem das cores ou o formato de saída do HTML. Se você está mantendo um pipeline automatizado de visualização, fixar a versão do pacote no seu ambiente é praticamente obrigatório. Eu fixo na versão 3.2.1 e só atualizo após testar localmente, porque uma atualização descuidada já me custou meio dia de trabalho uma vez.