Kernel Renormalization in Bayesian Deep Neural Networks: the Equivalent Wishart Ansatz in the Proportional Regime

Titolo: Rinormalizzazione del Kernel nelle Reti Neurali Profonde Bayesiane: l'Ansatz Wishart Equivalente nel Regime Proporzionale

Enunciazione del Problema
Le reti neurali profonde (DNN) sono centrali nell'apprendimento automatico moderno, tuttavia la loro comprensione teorica rimane impegnativa, in particolare per le architetture multistrato e non lineari. Mentre il limite "pigro" di larghezza infinita (Processo Gaussiano delle Reti Neurali o NNGP) fornisce descrizioni esatte per le reti larghe, non riesce a catturare gli effetti di larghezza finita responsabili delle prestazioni superiori delle DNN rispetto ai metodi basati su kernel. Al contrario, il regime di apprendimento "ricco", in cui le caratteristiche si adattano significativamente, richiede spesso parametri d'ordine ad alta dimensionalità (ad esempio, matrici di kernel $P \times P$) che sono computazionalmente intrattabili.

Il divario specifico affrontato da questo lavoro è la mancanza di una teoria efficace a bassa dimensionalità per le DNN bayesiane nel regime proporzionale, in cui il numero di campioni di addestramento $P$ e la larghezza della rete $N$ scalano proporzionalmente all'infinito ($P, N \to \infty$ con $\alpha = P/N$ fisso). In questo regime, le fluttuazioni di larghezza finita non sono trascurabili e le teorie esistenti si basano o su correzioni perturbative (che falliscono per le reti profonde) o richiedono la risoluzione di equazioni accoppiate ad alta dimensionalità che non offrono alcuna riduzione della dimensionalità.

Metodologia: l'Ansatz Wishart Equivalente (EWA)
Gli autori propongono una teoria approssimata ma efficace basata sull'Ansatz Wishart Equivalente (EWA). L'idea centrale è modellare le fluttuazioni stocastiche dei kernel empirici gerarchici in una rete non lineare profonda come se fossero distribuite secondo Wishart, simile al comportamento esatto trovato nelle reti lineari profonde.

1. Gerarchia dei Kernel: La priori dell'output della rete è descritta da una gerarchia di kernel empirici $K^{(\ell)}_E$ a ogni strato $\ell$. Nel limite proporzionale, gli autori assumono che la distribuzione condizionata di un kernel allo strato $\ell$, dato il kernel allo strato $\ell-1$, si comporti come una distribuzione Wishart:
   $$K^{(\ell)}_E | K^{(\ell-1)}_E \sim \mathcal{W}_P\left( \frac{\Theta^{L-\ell+1}(K^{(\ell-1)}_E)}{N_\ell}, N_\ell \right)$$
   dove $\Theta$ è la funzione del kernel NNGP e $N_\ell$ è la larghezza.

2. Riduzione della Dimensionalità: Sfruttando le proprietà delle matrici Wishart (in particolare la proprietà di contrazione per cui le forme quadratiche seguono una distribuzione $\chi^2$), gli autori riducono l'integrale complesso sullo spazio dei parametri a un integrale a bassa dimensionalità su $L$ parametri d'ordine scalari $q_\ell$ (o matrici $q_\ell$ per le CNN).

   • Per le reti Fully Connected (FC), l'azione efficace dipende da un prodotto scalare $Q = \prod q_\ell$.
   • Per le Reti Neurali Convoluzionali (CNN), viene introdotto un Ansatz Wishart Equivalente Impilato, che porta a un parametro d'ordine matriciale a bassa dimensionalità $Q$ che cattura le correlazioni tra patch.

3. Azione Efficace: La funzione di partizione è espressa tramite un principio di grande deviazione, producendo un'azione efficace $S(q)$ che dipende solo da questi parametri d'ordine e dal dataset. Il kernel rinormalizzato è dato da $K^{(R)}_Q = Q \Theta^L(C)$, dove $C$ è la matrice di Gram degli input.

4. Gestione dei Medie Non Nulle: Gli autori dimostrano che per il regime proporzionale, la distinzione tra distribuzioni Wishart centrali e non centrali (che derivano da funzioni di attivazione con medie non nulle come ReLU) è asintoticamente trascurabile. L'EWA a media zero rimane una descrizione efficace valida.

Contributi Chiave

• Teoria a Bassa Dimensionalità per DNN Non Lineari: Il documento fornisce la prima descrizione approssimata a bassa dimensionalità dei perceptron multistrato (MLP) bayesiani e delle CNN nel regime proporzionale, catturando effetti di larghezza finita non perturbativi senza ricorrere a equazioni di auto-consistenza $P \times P$.
• Ansatz Wishart Equivalente: Introduce l'EWA come meccanismo per catturare le fluttuazioni stocastiche dominanti dei kernel empirici gerarchici, riducendo il problema alla ricerca di $L$ (o di una piccola matrice) parametri d'ordine.
• Estensione alle CNN: Il framework è esteso alle architetture convoluzionali, identificando un meccanismo gerarchico di rinormalizzazione del kernel locale che tiene conto degli effetti di larghezza finita nei dati correlati spazialmente.
• Validazione nel Regime Proporzionale: La teoria è testata contro estesi esperimenti numerici (campionamento MCMC) su dataset reali (MNIST, CIFAR-10) e dati sintetici, coprendo profondità fino a $L \sim 10$ e varie funzioni di attivazione (Erf, ReLU).

Risultati

• Accordo Quantitativo: Le previsioni dell'EWA per la perdita di addestramento e l'errore di generalizzazione mostrano un "accordo complessivo molto buono" con gli esperimenti di campionamento numerico per reti con profondità fino a $L \sim O(10)$ e varie dimensioni del dataset. La teoria traccia con successo il comportamento delle reti a larghezza finita dove il limite NNGP a larghezza infinita fallisce.
• Deviazioni Sistemiche: Vengono osservate due tipi di deviazioni quando sia la profondità $L$ che il carico $\alpha$ sono grandi:
  1. Deriva Graduale: Negli esperimenti con dati gaussiani ad alto $\alpha$, l'errore dell'EWA cresce gradualmente con la profondità, probabilmente a causa di errori cumulativi nella larghezza della funzione di tasso o di eterogeneità direzionali trascurate.
  2. Metastabilità: Per le reti ReLU su CIFAR-10 con $L > 5$ e $\alpha \approx 2.5$, si osserva un calo netto e improvviso dell'errore di addestramento e di test nelle simulazioni MCMC. Questo fenomeno, interpretato come una transizione metastabile nel posterior, non è catturato dalla previsione liscia dell'EWA, suggerendo un collasso dell'ansatz in regimi estremi.
• Parametrizzazione Mean-Field vs. Standard: La teoria predice con successo le prestazioni di generalizzazione sia nella parametrizzazione standard (SP) che in quella mean-field ($\mu$P). Nel regime di apprendimento ricco ($\mu$P), il miglioramento nella generalizzazione risulta essere guidato principalmente dalla soppressione della varianza del predittore piuttosto che da una forte adattazione del bias del predittore, coerente con la visione del kernel rinormalizzato dell'EWA.

Significato e Affermazioni
Gli autori affermano che l'EWA fornisce una teoria efficace trattabile che colma il divario tra il limite pigro di larghezza infinita e il complesso regime di apprendimento delle caratteristiche ad alta dimensionalità. Riducendo la descrizione dell'apprendimento a larghezza finita a pochi parametri d'ordine scalari (o matrici a rango basso), il documento sostiene che la scala relativa della dimensione dei dati e della larghezza ($\alpha = P/N$) è un principio organizzativo più fondamentale per l'apprendimento a larghezza finita rispetto alla tradizionale dicotomia "pigro vs. ricco".

Il lavoro suggerisce che, per profondità e carichi moderati, la complessa adattazione delle caratteristiche nelle DNN non lineari può essere efficacemente catturata da una semplice rinormalizzazione globale (o locale, per le CNN) del kernel. Ciò offre una linea di base per comprendere l'apprendimento delle caratteristiche, indicando che il grande divario tra le previsioni a larghezza infinita e la fenomenologia a larghezza finita è in gran parte catturato da questa rinormalizzazione omogenea direzionalmente. Il documento conclude che, sebbene l'EWA sia un'approssimazione, il suo successo nel predire le curve di apprendimento empiriche attraverso diverse architetture e dataset lo stabilisce come uno strumento potente per analizzare l'apprendimento profondo bayesiano nel regime proporzionale.