FlexAct: Why Learn when you can Pick?

Resumo Técnico: Flex-Act

Declaração do Problema
Apesar do papel crítico que as funções de ativação desempenham na propagação de sinais, treinabilidade e generalização em redes neurais profundas, elas são tipicamente tratadas como hiperparâmetros fixos (ex: ReLU, GELU, Swish) aplicados uniformemente a todas as camadas e tarefas. Este paradigma estático assume que uma única não-linearidade é ideal independentemente da profundidade da camada ou do regime de dados, o que pode levar a um viés indutivo subótimo, gradientes fragmentados (shattered gradients) ou problemas de saturação. Embora existam aproximações de funções parametrizadas (ex: PReLU, SPLASH) e aproximadores de funções aprendíveis, elas frequentemente introduzem alta complexidade de otimização, redução de interpretabilidade ou exigem o compromisso com uma família funcional específica. Os autores argumentam que nenhuma não-linearidade é universalmente ótima e que redes profundas se beneficiariam da capacidade de selecionar dinamicamente a função de ativação mais apropriada para cada camada sem ajuste manual de hiperparâmetros.

Metodologia
O artigo apresenta o Flex-Act, um framework que permite a seleção discreta, camada a camada, de funções de ativação a partir de um conjunto predefinido durante o treinamento. O mecanismo central baseia-se no truque Gumbel-Softmax para criar uma relaxação diferenciável de escolhas discretas.

1. Roteamento Discreto via Gumbel-Softmax:
   Para cada camada $i$, o modelo computa uma combinação convexa de funções de ativação candidatas $\{\sigma^{(1)}, \dots, \sigma^{(p)}\}$ aplicadas à transformação afim $h_i$. As probabilidades de seleção $p_i$ são derivadas de logits treináveis $\pi_i$ usando a distribuição Gumbel-Softmax:
   $$X_i = \sum_{j=1}^{p} p_i^{(j)} \cdot \sigma^{(j)}(h_i)$$
   onde $p_i \sim \text{GumbelSoftmax}(\log \pi_i, \tau)$. À medida que a temperatura $\tau$ é reduzida (annealed), o modelo transita da exploração de ativações mistas para o compromisso com uma função discreta específica por camada. Isso permite que o processo de seleção seja otimizado de ponta a ponta via gradiente descendente.

2. Correção de Viés via Normalização de Gradiente:
   Os autores identificam uma falha crítica no roteamento ingênuo de Gumbel-Softmax: um viés de escala em direção a funções não limitadas (como ReLU). Como a retropropagação escala os gradientes pela magnitude da ativação, funções com saídas maiores dominam o processo de seleção, independentemente de sua adequação para os dados.
   Para combater isso, o Flex-Act introduz um regularizador baseado na norma do gradiente. Ele computa as normas de gradiente de cada ativação candidata em relação à entrada e as converte nessas pseudo-probabilidades. A função de perda total combina a perda da tarefa primária ($L_{task}$) com um termo de divergência de Kullback-Leibler ($L_{KL}$) que alinha as probabilidades de roteamento aprendidas com esses pseudo-rótulos baseados em gradiente:
   $$L = L_{task} + \alpha L_{KL}$$
   Esta regularização impede que o modelo recorra por padrão a funções não limitadas apenas devido a diferenças de magnitude, garantindo que a seleção seja impulsionada pela adequação funcional.

Principais Contribuições

• Roteamento de Ativação Discreta: Um novo mecanismo usando reparametrização Gumbel-Softmax para permitir a otimização de ponta a ponta das escolhas de funções de ativação discretas por camada.
• Derivações Analíticas de Gradiente: Derivação formal das equações de retropropagação para redes que empregam roteamento de ativação discreta, resolvendo a não-diferenciabilidade do processo de seleção.
• Regularização de Norma de Gradiente: Uma técnica específica para corrigir o viés inerente a ativações não limitadas no roteamento discreto, melhorando significativamente a precisão da seleção.
• Interpretabilidade e Modularidade: Ao tratar a seleção de ativação como uma variável discreta e aprendível, o framework fornece insights sobre quais não-linearidades são utilizadas em diferentes profundidades, oferecendo uma alternativa modular às famílias parametrizadas.

Resultados Experimentais
O artigo avalia o Flex-Act tanto em tarefas de regressão sintética quanto em benchmarks do mundo real:

• Regressão Sintética: Em tarefas onde a função de ativação de verdade fundamental (ground-truth) é conhecida (ReLU, Sigmoid, Tanh, LeakyReLU, Identidade), o Flex-Act converge consistentemente para a função correta.
  • Sem regularização ($\alpha=0$), o modelo falha em selecionar funções limitadas (Sigmoid/Tanh), recorrendo à ReLU devido ao viés de magnitude de gradiente.
  • Com regularização ($\alpha=0.3$), o Flex-Act atinge um Erro Quadrático Médio (MSE) próximo de zero em todos os cenários de verdade fundamental, igualando ou superando baselines de ativação fixa sem exigir múltiplas execuções de treinamento ou varreduras de ativação.
• Benchmarks do Mundo Real:
  • Classificação de Imagens (CIFAR-10/100): Integrado ao ResNet-18 e ResNet-34, o Flex-Act (aplicado à penúltima camada e a todas as camadas) alcança melhorias estatisticamente significativas sobre baselines de ReLU fixos e outros métodos parametrizados (PReLU, Swish, AReLU). Por exemplo, o ResNet-34 com Flex-Act (todas as camadas) atinge 95,75% de acurácia no CIFAR-10 e 79,42% no CIFAR-100, superando o baseline de ReLU.
  • Classificação de Texto: O framework também foi testado em benchmarks GLUE usando modelos BERT-base, demonstrando compatibilidade com arquiteturas Transformer.

Significância e Alegações
O artigo afirma que o Flex-Act oferece uma solução "plug-and-play" para o aprendizado profundo adaptativo. Argumenta-se que, ao permitir que as redes recuperem dinamicamente a não-linearidade apropriada para uma tarefa ou camada específica, os praticantes podem evitar o custo computacional de varreduras manuais de hiperparâmetros ou o treinamento de múltiplos modelos com diferentes ativações fixas. O framework preenche a lacuna entre flexibilidade teórica e utilidade prática, forneando uma abordagem robusta, interpretável e modular para o design de não-linearidade que escala de tarefas sintéticas rasas para arquiteturas profundas do mundo real. Os autores enfatizam que seu método alcança esses ganhos sem aumentar significativamente o número de parâmetros, baseando-se, em vez disso, na seleção inteligente de formas funcionais existentes e bem caracterizadas.