Het artikel beschrijft een methode om het attention-mechanisme van transformer-gebaseerde taalmodellen (LLM's) te visualiseren. Door attention-gewichten te aggregeren over alle lagen en heads, wordt inzichtelijk welke tokens het meest bijdragen aan de generatie van een nieuw token.
Belangrijkste inzichten uit de visualisatie:
- Kopiëren: De tool laat zien dat LLM's zeer accuraat kunnen kopiëren en plakken omdat ze direct naar specifieke brongegevens in de tekstsequentie kunnen verwijzen, in plaats van enkel te vertrouwen op probabilistische voorspellingen.
- Synthese: Modellen kunnen informatie uit verschillende delen van een prompt combineren om tot één passend woord te komen.
Technische implementatie: De auteur heeft een React-applicatie gebouwd met Transformers.js. Omdat standaard .onnx-bestanden in WASM beperkte toegang bieden tot interne waarden, is het model handmatig aangepast om de benodigde data bloot te leggen. Dit geïnstrumenteerde model is beschikbaar via een eigen Hugging Face-repository.
Visualisatie van het Attention-mechanisme in LLM's
Het blijkt dat we dit mechanisme kunnen visualiseren. Door op gegenereerde tokens te tikken of eroverheen te bewegen, kun je zien welke eerdere tokens de generatie hebben beïnvloed*.
\ "Beïnvloed" is wellicht niet volledig accuraat, aangezien deze visualisatie sterk vereenvoudigd is. Er wordt het attention-gewicht berekend, geschaald naar de grootte van de value-vector, geaggregeerd over alle attention-heads en opgeteld over alle lagen. Dit resultaat wordt gebruikt om de ondoorzichtigheid (opacity) van de vorige tokens te bepalen. De hoogste waarden krijgen altijd een opacity van 1, terwijl de rest wordt geïnterpoleerd.*
Patronen en inzichten
Om de visualisatie te beperken tot één numerieke waarde per vorig token, moest er veel informatie worden weggegooid. Toen ik begon met de implementatie, vreesde ik dat het resultaat onbegrijpelijk zou zijn, maar het produceert juist interessante patronen.
Accurate kopieeracties
In het standaardvoorbeeld "Office Move Summary" kun je over tekst bewegen die letterlijk is gekopieerd, zoals adressen en datums. Je ziet dan dat de originele data sterk naar voren komt, omdat het gegenereerde token veel informatie put uit de brongegevens.
Dit werpt een nieuw licht op een aspect van LLM's dat voorheen onintuïtief leek: als ze werken door het volgende token probabilistisch te voorspellen, waarom zijn ze dan zo goed in het kopiëren en plakken van teksten? Zouden ze niet door puur toeval uiteindelijk een fout maken?
Door dit mechanisme is te zien dat het model de volledige sequentie niet voorspelt vanuit een beperkte interne status. Omdat het toegang heeft tot alle vorige tokens, kan het simpelweg beslissen uit welke tokens het moet putten bij het kopiëren, waardoor de kans op fouten zeer gering is. In het voorbeeld "Debugging an Average Function" is te zien dat een relatief klein model (600 miljoen parameters) moeiteloos een volledige JavaScript-functie kan reproduceren, op de beoogde aanpassing na (hoewel het model de fout niet zelfstandig kon vinden en daarom hints nodig had).
Combineren van informatie
Een ander interessant voorbeeld is het woord "remain" in de zin "Existing access cards and phone numbers remain" uit de "Office Move Summary"-prompt. De visualisatie laat zien dat het model informatie put uit het woord "work" (uit de zin "Existing employee access cards will work") en "stay the same" (uit "company phone numbers will stay the same"). Het combineert dus als het ware de informatie uit beide zinnen, wat een fascinerend proces is.
Implementatie
De visualisatie is een eenvoudige React-applicatie die Transformers.js gebruikt voor het genereren van tekst. Omdat er meer data uit het model gehaald moet worden dan gebruikelijk, kon de reguliere generatie-loop niet worden gebruikt. Ik heb daarom de generatie-loop in de app zelf geschreven, zodat de waarden voor de visualisatie bijgehouden kunnen worden.
Aangezien zelfs een kleiner model honderden megabytes groot is, is het niet praktisch om de gebruiker te laten wachten op de download voordat er iets getoond wordt. Daarom heb ik een reeks prompts vooraf gegenereerd die direct geladen en bekeken kunnen worden.
Een technische uitdaging was dat sommige elementen in de visualisatie niet bedoeld zijn om gelezen te worden en daarom niet als outputs zijn gedefinieerd. In Python ML-bibliotheken is dit eenvoudig toegankelijk, maar Transformers.js maakt gebruik van .onnx-bestanden die de volledige computation graph bevatten. De logica voor het laden en de berekeningen is geïmplementeerd in WASM, waardoor er geen eenvoudige manier is om toegang te krijgen tot andere waarden dan de vooraf gedefinieerde outputs.
Uiteindelijk heb ik een klein script gebruikt om het .onnx-bestand zodanig aan te passen dat deze interne waarden worden blootgesteld. Dit betekent dat ik geen regulier .onnx-model kan gebruiken. Om een browser-gebaseerde generatiefunctie mogelijk te maken, heb ik een apart, geïnstrumenteerd model geüpload naar mijn eigen Hugging Face-repository waar de app naar verwijst.
De code is beschikbaar in de GitHub-repository.
Visualisatie van het Attention-mechanisme in LLM's
Het blijkt dat we dit mechanisme kunnen visualiseren. Door op gegenereerde tokens te tikken of eroverheen te bewegen, kun je zien welke eerdere tokens de generatie hebben beïnvloed*.
\ "Beïnvloed" is wellicht niet volledig accuraat, aangezien deze visualisatie sterk vereenvoudigd is. Er wordt het attention-gewicht berekend, geschaald naar de grootte van de value-vector, geaggregeerd over alle attention-heads en opgeteld over alle lagen. Dit resultaat wordt gebruikt om de ondoorzichtigheid (opacity) van de vorige tokens te bepalen. De hoogste waarden krijgen altijd een opacity van 1, terwijl de rest wordt geïnterpoleerd.*
Patronen en inzichten
Om de visualisatie te beperken tot één numerieke waarde per vorig token, moest er veel informatie worden weggegooid. Toen ik begon met de implementatie, vreesde ik dat het resultaat onbegrijpelijk zou zijn, maar het produceert juist interessante patronen.
Accurate kopieeracties
In het standaardvoorbeeld "Office Move Summary" kun je over tekst bewegen die letterlijk is gekopieerd, zoals adressen en datums. Je ziet dan dat de originele data sterk naar voren komt, omdat het gegenereerde token veel informatie put uit de brongegevens.
Dit werpt een nieuw licht op een aspect van LLM's dat voorheen onintuïtief leek: als ze werken door het volgende token probabilistisch te voorspellen, waarom zijn ze dan zo goed in het kopiëren en plakken van teksten? Zouden ze niet door puur toeval uiteindelijk een fout maken?
Door dit mechanisme is te zien dat het model de volledige sequentie niet voorspelt vanuit een beperkte interne status. Omdat het toegang heeft tot alle vorige tokens, kan het simpelweg beslissen uit welke tokens het moet putten bij het kopiëren, waardoor de kans op fouten zeer gering is. In het voorbeeld "Debugging an Average Function" is te zien dat een relatief klein model (600 miljoen parameters) moeiteloos een volledige JavaScript-functie kan reproduceren, op de beoogde aanpassing na (hoewel het model de fout niet zelfstandig kon vinden en daarom hints nodig had).
Combineren van informatie
Een ander interessant voorbeeld is het woord "remain" in de zin "Existing access cards and phone numbers remain" uit de "Office Move Summary"-prompt. De visualisatie laat zien dat het model informatie put uit het woord "work" (uit de zin "Existing employee access cards will work") en "stay the same" (uit "company phone numbers will stay the same"). Het combineert dus als het ware de informatie uit beide zinnen, wat een fascinerend proces is.
Implementatie
De visualisatie is een eenvoudige React-applicatie die Transformers.js gebruikt voor het genereren van tekst. Omdat er meer data uit het model gehaald moet worden dan gebruikelijk, kon de reguliere generatie-loop niet worden gebruikt. Ik heb daarom de generatie-loop in de app zelf geschreven, zodat de waarden voor de visualisatie bijgehouden kunnen worden.
Aangezien zelfs een kleiner model honderden megabytes groot is, is het niet praktisch om de gebruiker te laten wachten op de download voordat er iets getoond wordt. Daarom heb ik een reeks prompts vooraf gegenereerd die direct geladen en bekeken kunnen worden.
Een technische uitdaging was dat sommige elementen in de visualisatie niet bedoeld zijn om gelezen te worden en daarom niet als outputs zijn gedefinieerd. In Python ML-bibliotheken is dit eenvoudig toegankelijk, maar Transformers.js maakt gebruik van .onnx-bestanden die de volledige computation graph bevatten. De logica voor het laden en de berekeningen is geïmplementeerd in WASM, waardoor er geen eenvoudige manier is om toegang te krijgen tot andere waarden dan de vooraf gedefinieerde outputs.
Uiteindelijk heb ik een klein script gebruikt om het .onnx-bestand zodanig aan te passen dat deze interne waarden worden blootgesteld. Dit betekent dat ik geen regulier .onnx-model kan gebruiken. Om een browser-gebaseerde generatiefunctie mogelijk te maken, heb ik een apart, geïnstrumenteerd model geüpload naar mijn eigen Hugging Face-repository waar de app naar verwijst.
De code is beschikbaar in de GitHub-repository.
Groetjes,