Geometric Reasoning in the Embedding Space

Dit artikel toont aan dat Graph Neural Networks en Transformers in staat zijn om geometrische redenering te leren en verborgen figuren in de inbeddingsruimte te reconstrueren, waarbij Graph Neural Networks significant beter presteren en schaalbaar zijn dan Transformers.

De kern

Hoe neurale netwerken "ruimtelijk denken": Een reis door het binnenste van een AI

Stel je voor dat je een kind leert tekenen. Je geeft het een opdracht: "Teken een vierkant, en zorg dat punt A precies in het midden ligt." Het kind moet niet alleen de regels begrijpen, maar ook een mentaal beeld vormen van hoe de lijnen en hoeken in elkaar zitten.

Deze paper onderzoekt hoe kunstmatige intelligentie (AI) datzelfde doet. Wetenschappers hebben gekeken hoe neurale netwerken (de hersenen van AI) ruimtelijke problemen oplossen en of ze daadwerkelijk een soort "mentaal beeld" van de wereld ontwikkelen, net als mensen.

Hier is een simpele uitleg van wat ze hebben ontdekt, met wat creatieve vergelijkingen.

1. Het Proefje: Een Raadsel in een Rooster

De onderzoekers hebben een soort spelletje bedacht. Stel je een groot, digitaal rooster voor (zoals een schaakbord, maar dan met duizenden vakjes). Ze geven de AI een reeks regels, zoals:

• "Punt B is precies in het midden tussen A en C."
• "Punt D is een spiegelbeeld van C."
• "Punten A, B, C en D vormen een perfect vierkant."

De AI moet nu raden waar de onbekende punten op het rooster liggen. Het is als een raadsel waarbij je de positie van de stukjes moet bepalen op basis van de regels.

2. Twee Spelers: De "Lijst" vs. De "Bouwer"

Ze hebben twee soorten AI-modellen getest om dit raadsel op te lossen:

• De Transformer (zoals ChatGPT): Dit model werkt als iemand die een lange lijst tekst leest. Het probeert de regels te begrijpen door de woorden (tokens) achter elkaar te lezen.
• De GNN (Graph Neural Network): Dit model werkt als een bouwer die direct kijkt naar de connecties. Het ziet de regels en punten als een netwerk van knopen en lijnen, net als een plattegrond.

Het resultaat? De "bouwer" (GNN) was veel beter. De "lijstlezer" (Transformer) raakte snel de weg kwijt als het probleem groter werd. De GNN kon zelfs problemen oplossen die veel groter waren dan die tijdens het leren.

3. De Magie: Het "Mentale Beeld" Ontstaat

Dit is het meest fascinerende deel van de paper. De onderzoekers keken niet alleen naar het antwoord, maar naar wat er in de AI gebeurde terwijl het dacht.

Stel je voor dat elke punt op het rooster een kleurrijke bal is in een grote, onzichtbare kamer (de "embedding space").

• Aan het begin: De onbekende punten zijn als ballen die willekeurig door de kamer worden gegooid. Ze hebben geen vaste plek.
• Tijdens het denken: Terwijl de AI de regels toepast, beginnen deze ballen zich te verplaatsen. Ze trekken naar elkaar toe of duwen uit elkaar, precies alsof ze door onzichtbare veertjes aan elkaar hangen.
• Het eindresultaat: Na een tijdje vormen de ballen vanzelf een perfect roosterpatroon in die kamer!

De analogie: Het is alsof je een doos met losse Lego-blokjes schudt. Als je de AI de regels geeft, "schudt" het model de blokjes in zijn hoofd totdat ze vanzelf een mooi huis vormen. De AI heeft geen kaart van het rooster gekregen; het heeft de structuur van het rooster zelf ontdekt door de regels te volgen.

4. Het Iteratieve Proces: Het "Slepen" naar de oplossing

De AI lost het probleem niet in één flits op. Het is meer als het oplossen van een puzzel waarbij je stukjes eerst grofweg op de juiste plek legt en ze daarna steeds fijner bijstelt.

• De AI begint met een ruwe schets (bijvoorbeeld een vierkant dat eruitziet als een rommelig vierhoekje).
• In elke volgende stap (iteratie) "trekt" het model de lijnen strakker.
• Uiteindelijk wordt het rommelige vierhoekje een perfect vierkant.

Dit lijkt erg op hoe een mens een tekening maakt: eerst een lichte schets, en daarna steeds preciezer de lijnen trekken.

5. Waarom is dit belangrijk?

Vroeger zagen we AI vaak als een "zwarte doos": je geeft een vraag, en je krijgt een antwoord, maar je weet niet hoe het tot dat antwoord kwam.
Deze paper laat zien dat AI's niet zomaar gokken. Ze bouwen een intern, ruimtelijk model van de wereld.

• GNNs zijn beter: Voor ruimtelijke problemen werkt de "bouwer" (GNN) veel beter dan de "lijstlezer" (Transformer).
• Meer tijd = Beter: Als je de AI meer tijd geeft om te "nadenken" (meer iteraties), wordt het antwoord vaak beter, zelfs bij moeilijke problemen.

Conclusie

Kortom: deze AI's hebben niet alleen geleerd om antwoorden te raden. Ze hebben geleerd om een mentale kaart te tekenen in hun eigen hoofd. Ze zien de ruimte niet als een lijst met cijfers, maar als een fysieke structuur waar punten en lijnen met elkaar verbonden zijn. Het is een stap dichter bij het begrijpen van hoe machines werkelijk kunnen "denken" over de wereld om hen heen.

Technische samenvatting

Titel: Geometrisch Redeneren in de Embeddingruimte

Auteurs: David Mojžíšek, Jan Hůla, Jiří Janeček, David Herel, Mikoláš Janota
Publicatie: Machine Learning and Knowledge Extraction (MAKE), 2025

---

1. Probleemstelling en Doel

Hoewel neurale netwerken (zoals AlphaGeometry) complexe geometrische problemen kunnen oplossen, is er beperkt inzicht in de interne mechanismen die ze gebruiken om ruimtelijke relaties te representeren en te manipuleren. Bestaande systemen werken vaak als "black boxes".

Het doel van dit onderzoek is om te begrijpen hoe neurale netwerken interne ruimtelijke begrippen ontwikkelen. De auteurs onderzoeken dit door Graph Neural Networks (GNNs) en Transformers te trainen om de posities van punten op een discreet 2D-rooster te voorspellen, gebaseerd op een reeks geometrische beperkingen (constraints) die een verborgen figuur beschrijven. Het centrale vraagstuk is of deze modellen een interne "mentale afbeelding" vormen die de geometrische structuur van het probleem weerspiegelt.

2. Methodologie

2.1 Geometrische Probleemgeneratie (CSP)

De auteurs hebben een synthetische dataset ontwikkeld voor Constraint Satisfaction Problems (CSPs) op een discreet 2D-rooster (bijv. 20x20).

• Beperkingen (Constraints): Er worden vier soorten relaties gebruikt:
  • M (Midpoint): Punt B is het midden van lijnstuk AC.
  • R (Reflection): A en B vormen een symmetrieas voor C en D.
  • S (Square): A, B, C, D vormen een vierkant.
  • T (Translation): Vector D-C is een translatie van vector B-A.
• Oplossing: Een deel van de punten is vast (bekend), en het model moet de posities van de overige punten (onbekend) voorspellen. De problemen zijn opgebouwd als een Directed Acyclic Graph (DAG) van afhankelijkheden, waarbij sommige beperkingen eerst opgelost moeten worden voordat andere kunnen worden opgelost.

2.2 Architecturen

Twee modellen werden vergeleken:

1. Graph Neural Network (GNN):
   • Gebaseerd op een bipartiete graaf met variabele-nodes (punten) en constraint-nodes.
   • Gebruikt recurrente update-regels (LSTMs) voor zowel variabele- als constraint-embeddings.
   • Bekende punten worden geïnitieerd met vaste embeddings (gedeeld met de classificatielaag), terwijl onbekende punten willekeurig worden geïnitieerd en iteratief worden bijgewerkt.
   • De output is een classificatie over de roosterposities.
2. Autoregressive Transformer:
   • Gebaseerd op GPT-2 met RoPE (Rotary Positional Embeddings).
   • Ontvangt een sequentie van tokens die de beperkingen beschrijven en voorspelt de positie van een specifiek onbekend punt.
   • Werkt als een black-box sequentiemodel zonder expliciete grafische structuur van de beperkingen.

2.3 Training en Evaluatie

• De modellen werden getraind op een dataset van 300.000 voorbeelden.
• Er werd gebruikgemaakt van test-time scaling: het verhogen van het aantal iteraties tijdens de inferentie en het herhalen van de voorspelling met verschillende random initialisaties (resampling).
• De embeddings werden gevisualiseerd via UMAP en PCA om te zien of er een ruimtelijke structuur ontstaat.

3. Belangrijkste Bijdragen en Resultaten

3.1 Vergelijking van Architecturen

• GNN vs. Transformer: De GNN presteert aanzienlijk beter dan de Transformer.
  • De GNN bereikte >90% validatie-accuraatheid op roosters tot 80x80.
  • De Transformer faalde bij complexere problemen (20x20 rooster, >2 beperkingen) en bereikte slechts ~30% accurate. Zelfs met Chain-of-Thought (CoT) training bleef de Transformer beperkt tot ~50%.
• Conclusie: GNNs zijn beter geschikt voor gestructureerd redeneren met beperkingen omdat ze de graafstructuur van de afhankelijkheden expliciet modelleren, terwijl Transformers moeite hebben met de symmetrieën en de complexe afhankelijkheidsketens.

3.2 Ontstaan van Interne Ruimtelijke Representaties

• Self-Organizing Grids: Een van de meest opvallende bevindingen is dat de static embeddings (de weergave van de roosterposities) zich tijdens het trainen zelf organiseren in een 2D-roosterstructuur in de embeddingruimte.
  • Hoewel het model geen expliciete ruimtelijke supervisie krijgt (alleen de beperkingen), leert het model dat punten die ruimtelijk dicht bij elkaar liggen in het rooster, ook dicht bij elkaar liggen in de embeddingruimte.
  • Dit wordt zichtbaar in UMAP-projecties die een duidelijke roosterstructuur tonen, zelfs in 3D-projecties waar een gebogen manifold (een "kom"-vorm) wordt waargenomen.
• Dynamische Evolutie: Tijdens de inferentie bewegen de embeddings van de onbekende punten iteratief door de ruimte totdat ze een configuratie bereiken die de verborgen figuur weerspiegelt. Het model bouwt de figuur stap voor stap op in de embeddingruimte.

3.4 Iteratief Redeneren en Test-Time Scaling

• Het oplossen van beperkingen is een iteratief verfijningsproces. Het model begint met een ruwe benadering en verbetert deze gedurende meerdere iteraties (message passing).
• Test-Time Scaling: Het verhogen van het aantal iteraties tijdens de inferentie (van 15 naar 23) en het gebruik van meerdere resamples (random initialisaties) leidt tot een significante stijging in accurate.
  • Op het zware testset (buiten de trainingsverdeling) steeg de volledige probleem-accuraatheid van 76,7% naar 95,4% met 10 resamples en 23 iteraties.
  • Dit suggereert dat het model een proces volgt dat lijkt op continue optimalisatie.

3.5 Foutanalyse

• De nauwkeurigheid neemt af naarmate de diepte van de afhankelijkheidsketen toeneemt (punten die verder verwijderd zijn van de bekende startpunten).
• Fouten zijn echter vaak "zacht": onjuist voorspelde punten liggen meestal dicht bij de ware positie (kleine Manhattan-afstand), wat aangeeft dat het model de globale structuur wel begrijpt maar de exacte positie soms mist.

4. Significatie en Conclusie

Deze studie biedt waardevolle mechanistische inzichten in hoe neurale netwerken ruimtelijk begrip ontwikkelen:

1. Interpreteerbaarheid: Het paper toont aan dat neurale netwerken geen abstracte symbolen manipuleren, maar een interne geometrische structuur vormen die de fysieke ruimte nabootst. De embeddingruimte wordt een "mentale kaart" van het probleem.
2. Architecturale Voorkeur: Voor gestructureerde redeneertaken met beperkingen zijn GNNs superieur aan Transformers, voornamelijk vanwege hun vermogen om afhankelijkheidsrelaties direct te modelleren en te schalen naar grotere problemen.
3. Redeneermechanisme: Het oplossen van geometrische problemen gebeurt via een iteratief optimalisatieproces in de embeddingruimte, vergelijkbaar met numerieke optimalisatie, waarbij de model iteraties kan toevoegen om moeilijkere problemen op te lossen (test-time scaling).

De auteurs concluderen dat hun vereenvoudigde CSP-omgeving een krachtige testbed is voor het bestuderen van de interne werking van neurale netwerken bij ruimtelijk redeneren, en dat de gevonden patronen (zoals de zelforganiserende roosters) fundamenteel zijn voor het begrijpen van hoe AI ruimtelijke relaties leert.