PQuantML: A Tool for End-to-End Hardware-aware Model Compression

O PQuantML é uma biblioteca de código aberto para compressão de modelos de redes neurais consciente de hardware que simplifica o fluxo de trabalho de ponta a ponta, permitindo a aplicação conjunta ou individual de poda e quantização de ponto fixo para alcançar reduções significativas de parâmetros e largura de bits mantendo a precisão em tarefas como a classificação de subestrutura de jatos no LHC.

O essencial

Imagine que você tem um carro de Fórmula 1 (um modelo de Inteligência Artificial super inteligente) que é incrível em pistas de corrida, mas é impossível de dirigir no trânsito da sua cidade porque é muito grande, consome muita gasolina e demora para acelerar.

O PQuantML é como uma oficina de engenharia mágica que pega esse carro gigante e o transforma em um carro compacto, econômico e rápido o suficiente para andar na cidade, sem perder a capacidade de chegar ao destino.

Aqui está a explicação do que o artigo diz, usando analogias do dia a dia:

1. O Problema: O "Trânsito" de Dados

No CERN (o laboratório de física onde se estuda o universo), eles têm uma máquina gigante chamada Grande Colisor de Hádrons (LHC). Ela bate partículas a uma velocidade insana, gerando uma quantidade de dados tão enorme que seria impossível guardar tudo em um disco rígido.

Para resolver isso, eles usam um sistema de "tráfego" (chamado trigger) que decide, em microssegundos (milésimos de milésimo de segundo), quais dados são importantes e quais devem ser descartados. É como um porteiro de balada que precisa decidir quem entra em 0,000001 segundos.

• O Desafio: Os modelos de Inteligência Artificial (IA) são ótimos para identificar padrões, mas geralmente são "gordos" e lentos. Eles não cabem nos chips eletrônicos (FPGAs) usados nessas máquinas porque ocupam muita memória e demoram demais para pensar.

2. A Solução: O PQuantML (O "Personal Trainer" da IA)

O PQuantML é uma nova ferramenta de código aberto criada por cientistas para "emagrecer" essas IAs antes de colocá-las nos chips. Ela faz isso de duas formas principais, como se estivesse fazendo uma dieta e uma ginástica ao mesmo tempo:

A. Poda (Pruning) – "Cortar o que não serve"

Imagine que você tem um escritório cheio de funcionários. Alguns trabalham muito, outros apenas tomam café e olham pela janela.

• O que o PQuantML faz: Ele identifica os "funcionários" (números dentro do modelo) que não estão ajudando de verdade e os demite.
• O resultado: A equipe fica menor, mas continua fazendo o mesmo trabalho. Isso economiza espaço e energia. O PQuantML sabe cortar de formas diferentes: pode demitir funcionários inteiros (estruturado) ou apenas tirar a cadeira de quem está sentado (não estruturado), adaptando-se ao que o chip de hardware precisa.

B. Quantização – "Trocar a calculadora de luxo por uma de bolso"

Normalmente, essas IAs pensam usando números com muitos dígitos decimais (como 3,14159265...), o que exige calculadoras super precisas e caras.

• O que o PQuantML faz: Ele ensina a IA a pensar com números mais simples (como 3,14 ou até apenas 3).
• A analogia: É como trocar uma calculadora de engenharia de R$ 500,00 por uma calculadora de bolso de R$ 20,00. A precisão é um pouco menor, mas para a maioria das tarefas, é suficiente, e o custo e o espaço são drasticamente reduzidos. O PQuantML faz isso enquanto a IA está aprendendo, para que ela não se "confunda" com os números mais simples.

3. Como Funciona na Prática?

Antes, os cientistas tinham que usar várias ferramentas separadas: uma para cortar, outra para simplificar os números e outra para testar. Era como tentar montar um móvel com instruções em três idiomas diferentes.

O PQuantML junta tudo em um único manual de instruções:

1. Configuração Fácil: Você diz ao programa o que quer (ex: "quero cortar 50% dos dados e usar números simples").
2. Treinamento Inteligente: A IA aprende a ser pequena e rápida durante o treinamento, não depois.
3. Pronto para o Chip: No final, o modelo é transformado em código que os chips eletrônicos (FPGAs) conseguem entender e executar instantaneamente.

4. Os Resultados: O Carro Compacto

Os autores testaram essa ferramenta em um problema real de física: identificar o que compõe um "jato" de partículas (como distinguir se uma partícula veio de um quark ou de um glúon).

• Comparação: Eles compararam o PQuantML com outras ferramentas famosas (como QKeras e HGQ).
• O Veredito: O PQuantML conseguiu criar modelos que eram muito menores e mais rápidos (ocupando menos espaço no chip e respondendo mais rápido), mantendo a mesma precisão de acerto. Em alguns casos, reduziu o uso de recursos do chip em mais de 90%, sem perder a capacidade de identificar as partículas corretamente.

Resumo Final

O PQuantML é uma ferramenta que permite que cientistas peguem modelos de Inteligência Artificial complexos e os "compactem" para caber em equipamentos eletrônicos rápidos e limitados, como os usados no CERN.

É como pegar um elefante e transformá-lo em um coelho que ainda consegue correr tão rápido quanto o elefante original, permitindo que a ciência do futuro (como a descoberta de novas partículas) aconteça em tempo real, sem que os dados se percam no trânsito.

Resumo técnico

1. O Problema

Os experimentos de Física de Alta Energia (HEP) no Grande Colisor de Hádrons (LHC), como ATLAS e CMS, enfrentam um desafio extremo de dados: colisões próton-próton ocorrem a 40 MHz, gerando centenas de terabytes por segundo. Para gerenciar isso, os sistemas de gatilho (trigger) utilizam uma abordagem em dois níveis:

• Nível 1 (L1T): Baseado em hardware (FPGAs), deve tomar decisões em microssegundos.
• Nível Alto (HLT): Baseado em software, opera em escala de segundos.

O uso de modelos de Aprendizado de Máquina (ML) no L1T é promissor para aproximar algoritmos complexos com alta velocidade, mas a implantação direta é difícil. Os FPGAs possuem recursos limitados (DSPs, BRAM, LUTs) e exigem latência determinística. Técnicas tradicionais de ML (como superparametrização) são ineficientes nesses ambientes. Além disso, as ferramentas existentes de compressão de modelos (como QKeras e HGQ) focam principalmente em quantização, exigindo que os usuários implementem estratégias de poda (pruning) separadamente, o que fragmenta o fluxo de trabalho e dificulta a otimização conjunta.

2. Metodologia

O PQuantML é uma biblioteca de código aberto projetada para fornecer um fluxo de trabalho unificado e consciente de hardware para a compressão de modelos de redes neurais. A metodologia baseia-se em três pilares principais:

• Compressão Consciente de Hardware (Hardware-Aware): O treinamento não é apenas uma etapa de pós-processamento. O PQuantML integra poda e quantização diretamente no ciclo de treinamento (Quantization-Aware Training - QAT e Pruning-Aware Training).
• Arquitetura Unificada:
  • Interface Flexível: Suporta tanto PyTorch quanto TensorFlow/Keras.
  • Definição de Modelo: Permite a criação de modelos através de camadas nativas do PQuantML ou através de uma função de substituição automática de camadas em arquiteturas existentes.
  • Configuração Declarativa: Todos os parâmetros de compressão, treinamento e otimização são definidos em arquivos YAML validados via Pydantic, facilitando a reprodutibilidade.
• Técnicas de Compressão Implementadas:
  • Quantização: Suporta quantização de ponto fixo (com controle granular de bits inteiros e fracionários) e Quantização de Alta Granularidade (HGQ), que aprende bit-widths individuais para pesos e ativações via otimização baseada em gradiente, minimizando o custo de operações efetivas (EBOPs).
  • Poda (Pruning): Implementa múltiplos métodos com diferentes granularidades:
    • Estruturada: Remove canais inteiros ou filtros.
    • Não estruturada: Remove pesos individuais.
    • Semi-estruturada (N:M): Padrões regulares (ex: 2:4) que equilibram esparsidade e eficiência de hardware.
    • Métodos específicos incluem: Activation Pruning (AP), MDMM, AutoSparse, Continuous Sparsification (CS), DST, PDP e Wanda.
• Otimização de Hiperparâmetros: Integração nativa com Optuna (busca Bayesiana) e MLflow para rastreamento de experimentos, permitindo a busca automática pelo melhor equilíbrio entre precisão, latência e uso de recursos.
• Fluxo de Trabalho End-to-End: O modelo treinado e comprimido é convertido diretamente para firmware FPGA usando o pacote hls4ml, garantindo que as previsões sejam bit-precisas em relação ao modelo original.

3. Principais Contribuições

1. Primeira Biblioteca Unificada: PQuantML é a primeira ferramenta que integra sistematicamente poda e quantização em um único framework para HEP, eliminando a necessidade de orquestrar ferramentas separadas.
2. Suporte a Múltiplas Granularidades: Oferece uma vasta gama de algoritmos de poda (desde não estruturada até padrões N:M e estruturada) combinados com quantização de ponto fixo e HGQ.
3. Facilidade de Uso e Reprodutibilidade: Abstrai a complexidade da compressão através de configurações YAML e pipelines de otimização automatizados, reduzindo a barreira de entrada para físicos.
4. Compatibilidade com Hardware: Projetado especificamente para atender às restrições de latência e recursos de FPGAs, com integração direta ao ecossistema hls4ml.

4. Resultados

O PQuantML foi avaliado na tarefa de classificação de subestrutura de jatos (jet tagging), um problema crítico para o processamento de dados do LHC.

• Desempenho vs. Recursos: Em comparação com o estado da arte (QKeras e HGQ), o PQuantML alcançou:
  • Redução de Recursos: Diminuição significativa no uso de LUTs (Look-Up Tables) e DSPs. Por exemplo, no conjunto de dados HLF (OpenML), o método DST com PQuantML reduziu o uso de LUTs de 5.504 para 3.895 e a latência de 105 ns para ~47 ns, mantendo uma acurácia de ~76,3% (comparável ao QKeras).
  • Acurácia: Os modelos comprimidos mantiveram a acurácia física necessária, superando ou igualando as ferramentas existentes quando combinam poda e quantização.
  • Comparação com HGQ Nativo: Modelos treinados com PQuantML usando o algoritmo HGQ produziram resultados de acurácia e uso de recursos comparáveis aos treinados diretamente com a biblioteca HGQ nativa, validando a integração.
• Flexibilidade: O uso de FITCompress (um algoritmo integrado) permitiu ajustar bit-widths mistos e níveis de esparsidade para atingir taxas de compressão específicas, embora com pequenas perdas de acurácia em troca de ganhos significativos em DSPs e latência.

5. Significado e Impacto

O PQuantML representa um avanço crucial para a implementação de Inteligência Artificial em sistemas de gatilho em tempo real do LHC.

• Viabilidade de Implantação: Demonstra que é possível reduzir drasticamente a complexidade computacional e o uso de recursos de FPGA sem sacrificar a precisão física, tornando viável a execução de modelos complexos no L1T.
• Padronização: Oferece um padrão unificado para a comunidade HEP, facilitando a comparação de diferentes estratégias de compressão e acelerando o desenvolvimento de soluções otimizadas para hardware.
• Futuro: A ferramenta estabelece uma base para futuras otimizações, como a integração de knowledge distillation e técnicas de fatoração de baixo posto, além de expandir o suporte para mais backends e arquiteturas complexas (como redes com conexões residuais).

Em resumo, o PQuantML preenche uma lacuna crítica entre a pesquisa em compressão de modelos e a engenharia de implantação em hardware, fornecendo uma solução robusta, automatizada e eficiente para os desafios de latência e recursos dos futuros upgrades do LHC.