Technische Samenvatting: Hoe Feature Learning de Functieruimte Herstructureert

1. Probleemstelling

Moderne neurale netwerken (NN's) onderscheiden zich van vaste-kernmethoden (zoals Random Feature Models of lazy training-regimes) door hun vermogen om feature learning uit te voeren: het aanpassen van interne representaties aan de datastructuur tijdens het trainen. Hoewel het succes van NN's vaak wordt toegeschreven aan dit mechanisme, blijft de precieze impact van feature learning op de geïnduceerde functieruimte slecht begrepen. Bestaande karakteriseringen van kernen en neurale netwerken zijn grotendeels statisch; ze beschrijven de functieruimte bij initialisatie of in de limiet van oneindige breedte, maar falen in het uitleggen hoe trainingsalgoritmen, specifiek gradient descent, deze ruimte dynamisch herstructureren.

Dit werk adresseert de vraag: Hoe evolueert de kern (en de bijbehorende functieruimte) onder gradient-updates, en welke informatie wordt geleidelijk geleerd? De auteurs focussen op de vroege fasen van het trainen van een tweelaags neurale netwerk op een Gaussisch single-index model, een setting die bekend staat om het vangen van de meest algemene klasse van functies die leerbaar zijn in het hoogdimensionale proportionele regime.

2. Methodologie

De auteurs analyseren de evolutie van de functieruimte na een enkele grote gradient descent-stap op de gewichten van de eerste laag van een tweelaagsnetwerk. De opzet omvat:

• Data Model: Een Gaussisch single-index model $y_i = g(\langle w^*, x_i \rangle) + \epsilon_i$, waarbij $x_i \sim \mathcal{N}(0, I_d)$.
• Netwerkarchitectuur: Een tweelaagsnetwerk $f(x) = \frac{1}{\sqrt{m}} \sum_{j=1}^m a_j \sigma(w_j^\top x)$ met ReLU- of algemene Lipschitz-activaties.
• Trainingsregime: Het "proportionele regime" waarbij het aantal samples $n$, het aantal verborgen eenheden $m$ en de dimensie $d$ naar oneindig groeien bij vaste verhoudingen ($\alpha = n/d, \beta = m/d$). De stapgrootte $\eta$ schaalt als $\Theta(d^\zeta)$ met $\zeta \in [1/2, 1)$, wat aanzienlijke gewichtsupdates toelaat.
• Tweefasige Training:
  1. Feature Learning: Een enkele gradient-stap update de gewichten van de eerste laag $W$ terwijl de tweede laag $a$ vast blijft.
  2. Lineaire Readout: De gewichten van de tweede laag worden geoptimaliseerd via ridge-regressie op een nieuwe set samples.

De kern van de analytische aanpak bestaat uit het karakteriseren van de post-update feature-distributie en het afleiden van de resulterende data-adaptieve kern. De auteurs maken gebruik van random matrixtheorie, spectrale analyse van integraaloperatoren en Taylor-achtige expansies om de bijgewerkte kern te relateren aan de initiële isotrope kern.

3. Belangrijkste Bijdragen en Resultaten

A. Benadering door een Doelafhankelijke Spiked Gaussische Verdeling

Het artikel bewijst dat na een grote gradient-stap de distributie van de geleerde features goed wordt benaderd door een doelafhankelijke spiked Gaussische verdeling. Specifiek transformeert de covariantiematrix van de feature-gewichten van een isotrope vorm naar:
$$\Gamma = A I_d + B w^* (w^*)^\top$$
waarbij $A$ en $B$ constanten zijn die afhankelijk zijn van de stapgrootte, datadimensies en de linkfunctie $g$. Dit leidt tot een data-adaptieve kern $k_1$ die kan worden uitgedrukt als:
$$k_1(x, x') = k_0(\Gamma^{1/2}x, \Gamma^{1/2}x')$$
Deze formulering interpreteert feature learning als een distributieve transformatie, hetzij in de parameterruimte (het verschuiven van de gewichtsverdeling) of in de inputruimte (het transformeren van de data-representatie).

B. Spectrale Expansie en Kernvervorming

De auteurs voeren een Taylor-expansie uit van de nieuwe kern $k_1$ rondom de initiële isotrope kern $k_0$. Ze tonen aan dat de "spike" (de rang-één update uitgelijnd met $w^*$) hogere-orde termen introduceert die verdwijnen naarmate de dimensie groeit, waardoor een dominante eerste-orde perturbatie overblijft.

• Resultaat: Het feature learning-proces legt effectief een extra, doelafhankelijk kerncomponent op.
• Mechanisme: Dit component selecteertief eigenwaarden versterkt die zijn uitgelijnd met de doelrichting $w^*$ en mengt leidende eigenfuncties.

C. Expliciete Karakterisatie voor ReLU-activaties

Voor het specifieke geval van ReLU-activatie bieden de auteurs een nauwkeurige spectrale analyse:

1. Lineaire Eigenruimte: De lineaire functie uitgelijnd met het doel, $\psi^*(x) = \langle x, w^* \rangle$, ontvangt een selectieve boost van zijn eigenwaarde. Lineaire functies die orthogonaal zijn op $w^*$ blijven onaangetast.
2. Mixing van de Top Eigenruimte: De top-eigenfunctie van de bijgewerkte operator is niet puur de constante harmonische (zoals bij de initiële kern), maar wordt een superpositie van de constante radiale modus ($Y_0$) en een doel-uitgelijnde kwadratische harmonische ($Y_2$).
3. Koppeling: Dit onthult dat feature learning de top radiale modus koppelt aan een kwadratische harmonische die is uitgelijnd met het signaal, waardoor de geometrie van de functieruimte effectief wordt herstructureerd om richtingen te prioriteren die het signaal bevatten.

D. Numerieke Validatie

Experimenten op synthetische data bevestigen de theoretische bevindingen:

• Alignering: De leidende eigenvector van de bijgewerkte kern ($k_1$) toont een toenemende alignering met het kwadratische feature $Y_2$ naarmate de spike-strength $B$ toeneemt, terwijl de initiële kern ($k_0$) geen dergelijke alignering vertoont.
• Generalisatie: Neurale netwerken getraind met een enkele gradient-stap convergeren naar de prestaties van de spiked kern $k_1$, en presteren beter dan de vaste kern $k_0$ bij het leren van de doelfunctie, met name wanneer het doel niet-lineaire componenten bevat.

4. Betekenis en Claims

Het artikel biedt een precieze perspectief op de functieruimte voor feature learning in een vroege fase. De primaire claim is dat gradient descent niet zomaar een vaste kern herschaalt; in plaats daarvan induceert het een data-adaptieve vervorming van de functieruimte.

• Mechanisme van Leren: Feature learning wordt gekarakteriseerd als een transformatie die richtingen in de functieruimte die zijn uitgelijnd met het signaal in de data, preferentieel versterkt.
• Afhankelijkheid van Stapgrootte: De grootte van deze transformatie wordt gemoduleerd door de stapgrootte; grotere stappen induceren sterkere vervormingen en verhogen de bijdrage van zowel lineaire als hogere-orde niet-lineaire features.
• Theoretische Brug: De resultaten overbruggen de kloof tussen data-adaptieve kernen en neurale netwerken, en suggereren dat de "magie" van feature learning kan worden begrepen door de lens van hoe de spectrale structuur van de kern wordt gewijzigd om uitgelijnd te raken met de doelfunctie.

De auteurs concluderen dat het begrijpen van deze distributieve verschuivingen een potentiële route biedt om vroege trainingsfasen te "overslaan" door netwerken direct te initialiseren met deze aangepaste kernen, hoewel zij dit noemen als een richting voor toekomstig onderzoek en niet als een gedemonstreerde toepassing in het huidige werk.