Comprendre les xarxes neuronals mitjançant circuits escassos
Hem entrenat models perquè pensin en passos més simples i més traçables, per poder entendre millor com funcionen.
Les xarxes neuronals impulsen els sistemes d’IA més capaços d’avui, però continuen sent difícils d’entendre. No escrivim aquests models amb instruccions explícites pas a pas. En canvi, aprenen ajustant milers de milions de connexions internes, o «pesos», fins que dominen una tasca. Dissenyem les regles de l’entrenament, però no els comportaments específics que emergeixen, i el resultat és una xarxa densa de connexions que cap persona no pot desxifrar fàcilment.
A mesura que els sistemes d’IA esdevenen més capaços i tenen un impacte real en decisions dins de la ciència, l’educació i la salut, entendre com funcionen és essencial. La interpretabilitat fa referència als mètodes que ens ajuden a entendre per què un model ha produït una sortida determinada. Hi ha moltes maneres d’aconseguir-ho.
Per exemple, els models de raonament reben incentius perquè expliquin el seu procés fins a arribar a una resposta final. La interpretabilitat de la cadena de pensament aprofita aquestes explicacions per supervisar el comportament del model. Això és útil de seguida: les cadenes de pensament dels models de raonament actuals semblen ser informatives respecte de comportaments preocupants com ara l’engany. Tanmateix, confiar completament en aquesta propietat és una estratègia fràgil, i això es pot degradar amb el temps.
D’altra banda, la interpretabilitat mecanicista, que és l’objecte d’aquest treball, pretén invertir completament l’enginyeria dels càlculs d’un model. Fins ara ha estat menys útil de manera immediata, però en principi podria oferir una explicació més completa del comportament del model. En intentar explicar el comportament del model al nivell més granular, la interpretabilitat mecanicista pot fer menys supòsits i donar-nos més confiança. Però el camí que va dels detalls de baix nivell a les explicacions de comportaments complexos és molt més llarg i difícil.
La interpretabilitat dona suport a diversos objectius clau, per exemple permetent una millor supervisió i proporcionant senyals d’alerta primerenca de comportaments insegurs o estratègicament desalineats. També complementa els nostres altres esforços de seguretat, com ara la supervisió escalable, l’entrenament adversari i el red-teaming.
En aquest treball, mostrem que sovint podem entrenar models de maneres que els facin més fàcils d’interpretar. Considerem la nostra feina un complement prometedor a l’anàlisi a posteriori de xarxes denses.
Aquesta és una aposta molt ambiciosa; hi ha un llarg camí entre el nostre treball i la comprensió completa dels comportaments complexos dels nostres models més potents. Tot i així, per a comportaments simples, trobem que els models escassos entrenats amb el nostre mètode contenen circuits petits i desentrellaçats que són alhora comprensibles i suficients per executar el comportament. Això suggereix que podria haver-hi un camí viable cap a l’entrenament de sistemes més grans amb mecanismes que puguem entendre.
Els treballs anteriors d’interpretabilitat mecanicista han partit de xarxes denses i embullades i han intentat desembullar-les. En aquestes xarxes, cada neurona individual està connectada a milers d’altres neurones. La majoria de neurones semblen dur a terme moltes funcions diferents, cosa que fa que sembli impossible entendre-les.
Però què passaria si entrenéssim xarxes neuronals desentrellaçades, amb moltes més neurones, però en què cada neurona només tingués unes poques dotzenes de connexions? Aleshores, potser la xarxa resultant seria més simple i més fàcil d’entendre. Aquesta és l’aposta central de recerca del nostre treball.
Amb aquest principi al cap, vam entrenar models de llenguatge amb una arquitectura molt semblant a la dels models de llenguatge existents com GPT‑2, amb una petita modificació: forcem que la gran majoria dels pesos del model siguin zeros. Això limitava el model a utilitzar molt poques de les connexions possibles entre les seves neurones. És un canvi senzill que, segons defensem, desentrellaça substancialment els càlculs interns del model.
A les xarxes neuronals denses normals, cada neurona està connectada a totes les neurones de la capa següent. Als nostres models escassos, cada neurona només es connecta a unes poques neurones de la capa següent. Esperem que això faci que les neurones, i la xarxa en conjunt, siguin més fàcils d’entendre.
Volem mesurar fins a quin punt els càlculs dels nostres models escassos estan desentrellaçats. Vam considerar diversos comportaments simples del model i vam comprovar si podíem aïllar les parts del model responsables de cada comportament —que anomenem circuits.
Vam seleccionar manualment un conjunt de tasques algorítmiques simples. Per a cadascuna, vam podar el model fins al circuit més petit que encara podia executar la tasca i vam examinar fins a quin punt era simple aquest circuit. (Per a més detalls, vegeu el nostre article(s'obre en una finestra nova).) Vam observar que, entrenant models més grans i més escassos, podíem produir models cada cop més capaços amb circuits cada cop més simples.
Representem la interpretabilitat en relació amb la capacitat entre models (la part inferior esquerra és millor). Per a una mida fixa de model escàs, augmentar l’escassetat —posar més pesos a zero— redueix la capacitat però augmenta la interpretabilitat. Augmentar la mida del model desplaça aquesta frontera cap enfora, cosa que suggereix que podem construir models més grans que siguin alhora capaços i interpretables.
Per concretar-ho, considerem una tasca en què un model entrenat amb codi Python ha de completar una cadena amb el tipus correcte de cometa. En Python, ‘hello’ ha d’acabar amb cometa simple, i “hello” ha d’acabar amb cometa doble. El model pot resoldre-ho recordant quin tipus de cometa obria la cadena i reproduint-lo al final.
Els nostres models més interpretables semblen contenir circuits desentrellaçats que implementen exactament aquest algorisme.
Exemple de circuit en un transformador escàs que prediu si una cadena s’ha d’acabar amb cometa simple o doble. Aquest circuit utilitza només cinc canals residuals (línies grises verticals), dues neurones MLP a la capa 0, un canal de consulta-clau d’atenció i un canal de valor a la capa 10. El model (1) codifica les cometes simples en un canal residual i les dobles en un altre; (2) utilitza una capa MLP per convertir-ho en un canal que detecta qualsevol cometa i un altre que classifica entre cometes simples i dobles; (3) utilitza una operació d’atenció per ignorar els segments intermedis, trobar la cometa anterior i copiar-ne el tipus al segment final; i (4) prediu la cometa de tancament corresponent.
Segons la nostra definició, les connexions exactes mostrades més amunt són suficients per executar la tasca: si eliminem la resta del model, aquest petit circuit continua funcionant. També són necessàries: si eliminem aquestes poques arestes, el model falla.
També vam analitzar alguns comportaments més complicats. Els nostres circuits per a aquests comportaments (per exemple, l’enllaç de variables que es mostra a continuació) són més difícils d’explicar completament. Fins i tot així, encara podem aconseguir explicacions parcials relativament simples que són predictives del comportament del model.
Un altre exemple de circuit, amb menys detall. Per determinar el tipus d’una variable anomenada current, una operació d’atenció copia el nom de la variable al segment set() quan es defineix, i una altra operació posterior copia el tipus des del segment set() a un ús posterior de la variable, cosa que permet al model inferir el segment següent correcte.
Aquest treball és un primer pas cap a un objectiu més ampli: fer que els càlculs dels models siguin més fàcils d’entendre. Però encara queda molt camí per recórrer. Els nostres models escassos són molt més petits que els models d’avantguarda, i grans parts del seu càlcul continuen sense interpretar-se.
A continuació, esperem escalar les nostres tècniques a models més grans i explicar més del comportament dels models. Si enumerem els motius de circuit que subjauen al raonament més complex en models escassos capaços, podríem desenvolupar una comprensió que ens ajudés a orientar millor les investigacions sobre models d’avantguarda.
Per superar la ineficiència de l’entrenament de models escassos, veiem dos camins possibles. Un és extreure circuits escassos de models densos existents, en lloc d’entrenar models escassos des de zero. Els models densos són fonamentalment més eficients de desplegar que els models escassos. L’altre camí és desenvolupar tècniques més eficients per entrenar models amb vista a la interpretabilitat, cosa que podria ser més fàcil de portar a producció.
Cal tenir en compte que aquestes troballes no garanteixen de cap manera que aquest enfocament s’estengui a sistemes més capaços, però aquests primers resultats són prometedors. El nostre objectiu és ampliar gradualment quina part d’un model podem interpretar de manera fiable i construir eines que facin que els sistemes futurs siguin més fàcils d’analitzar, depurar i avaluar.
Autors
Leo Gao, Achyuta Rajaram, Jacob Coxon, Soham V. Govande, Bowen Baker i Dan Mossing
