Geometric Reasoning in the Embedding Space

Abstract

While neural networks can solve complex geometric problems, as demonstrated by systems like AlphaGeometry, we have limited understanding of how they internally represent and reason about spatial relationships. In this work, we investigate how neural networks develop internal spatial understanding by training Graph Neural Networks and Transformers to predict point positions on a discrete 2D grid from geometric constraints that describe hidden figures. We show that both models develop interpretable internal representations that mirror the geometric structure of the problems they solve. Specifically, we observe that point embeddings self-organize into 2D grid structures during training, and during inference, the models iteratively construct the hidden geometric figures within their embedding spaces. Our analysis reveals how reasoning complexity correlates with prediction accuracy, and shows that models solve constraints through an iterative refinement process, which might resemble continuous optimization. We also find that Graph Neural Networks prove more suitable than Transformers for this type of structured constraint reasoning and scale more effectively to larger problems. These findings provide initial insights into how neural networks can develop structured understanding and contribute to their interpretability.

1. Introduction

While neural networks can solve complex spatial reasoning problems, we have limited understanding of the internal mechanisms they use to represent and manipulate geometric relationships. Many papers have already demonstrated that autoregressive Transformers, for example, Refs. [1,2,3], and Graph Neural Networks (GNNs), for example, Refs. [4,5], are able to learn to solve such problems, but it is not clear what mechanisms they discovered.

A notable example is the work by [6] who developed a system named AlphaGeometry, which was able to solve problems that appeared on the International Mathematical Olympiad. The system contains an autoregressive Transformer that takes as input a sequence of tokens describing the problem in the language of geometric relations and is trained to predict auxiliary points useful for finding a proof for the given statement. AlphaGeometry achieves impressive performance, but it operates as a black box, providing no insights into how the model internally represents geometric relationships or constructs spatial understanding.

A human dealing with such a problem would try to form a mental image of the construction used in the problem (most likely by first drawing it). We can ask a natural question whether NNs could also form such a “mental image”, which would reflect the spatial configuration of points described in the problem. We can also ask whether an autoregressive Transformer is a suitable model for such a problem and whether a GNN, which eliminates a lot of symmetries (variable renaming and constraint reordering), could be easier to train and be more scalable.

In this contribution, we take a closer look at these questions using a simplified approach. We will focus on purpose-built geometric constraint satisfaction problems (CSPs). This setting allows for detailed analysis of neural geometric reasoning that is difficult to achieve in complex systems like AlphaGeometry. We create a simple CSP language with several geometric constraints (relations) for which the domain is a set of points in a discrete 2D grid of certain size. Each instance of this CSP uniquely describes a hidden figure whose points are the solution of the instance. As the discrete grid is finite, we can assign a token/class (when using a Transformer, each point is represented by a token; when using a GNN, each point corresponds to a class) to each point and train the model to predict the points in the hidden figure with a cross-entropy loss.

Analysis shows that both models develop interpretable internal representations that mirror the geometric structure of the problems they solve. After visualizing a low-dimensional projection of the embeddings corresponding to individual points, we can see that they organized themselves in a 2D grid which they represent. We also show that during inference, the embeddings of unknown variables evolve into a configuration that reflects the hidden figure described by the problem. Additionally, we find that GNNs prove more suitable than Transformers for structured constraint reasoning tasks and scale more effectively to larger problems.

This controlled approach provides solid mechanistic insights into how neural networks develop spatial understanding and contributes to a broader understanding of structured reasoning in neural systems.

The rest of the paper has the following structure. In Section 2, we review related work. Section 3 describes our experimental setup, including the two types of architectures and the process for generating problems. In Section 4, we present the results of our experiments for both models, analyze embedding structure emergence, examine the iterative solution process, and provide analysis of failure modes for different problem complexities. We discuss the results and limitations of our work in Section 5. Additional interesting and supporting results can be found in the Appendix A, Appendix B, Appendix C, Appendix D, Appendix E, Appendix F, Appendix G, Appendix H, Appendix I and Appendix J.

2. Related Work

2.1. Reasoning About Geometry

Geometric reasoning has been a focus of research for decades [7], with Wu’s method [8] considered state-of-the-art until recently. In 2024, AlphaGeometry [6] surpassed Wu’s method on International Mathematical Olympiad problems, utilizing a combination of symbolic deduction and neural language model which predicts auxiliary constructions based on problem statement which is described using geometric relations. Our work is partially inspired by AlphaGeometry, but we study a much easier setup in which the model predicts points in a 2D grid that satisfy the relations used to describe the construction.

Other neural approaches include visual reasoning methods that extract geometric primitives from diagrams using deep learning for construction problems [9]. Our work differs by focusing on abstract constraint satisfaction rather than visual diagram processing, enabling detailed analysis of internal spatial representations.

We were also partially inspired by the work of [10], which provides an explanation for the process by which GNNs can learn to solve Boolean CSPs, concretely deciding the satisfiability of Boolean formulas. We augment their setup for the domain of geometric CSPs. In their work, they show empirical evidence that the GNN learns to act as a first-order solver of a semidefinite programming (SDP) relaxation of MAX-SAT. MAX-SAT is an optimization version of SAT, and in SDP relaxation, vectors assigned to variables are optimized in order to minimize an objective which captures how many constraints are satisfied.

In the domain of geometric CSPs, ref. [11] developed an algorithm for the automatic construction of diagrams for geometry problems. The algorithm works by optimizing coordinates of points by gradient descent to maximize an objective function which captures how well the geometric constraints of the problem are satisfied. We emphasize that their algorithm was manually constructed and not learned.

2.2. Spatial Reasoning in Neural Networks

Another related area of research is focused on the ability of NNs to conduct spatial reasoning in a broader sense, such as through spatial navigation, planning and interaction with physical objects.

Several papers explore spatial reasoning in the context of Large Language Models (LLMs) such as ChatGPT (gpt-4-turbo-2024-04-09) or LLama (Llama 3-8B) [12]. In [1,3], the authors point out various failure modes in the planning and navigation capability of a large pool of LLMs. In [2], the authors developed a prompting technique called Visualization of Thought (VoT), which significantly improved LLM performance in visual navigation tasks. The study most similar to our work is [13], in which the authors probe small Transformer models trained for maze navigation in order to understand the representation it uses to solve the task. Unlike the previously mentioned works, ref. [14] uses reinforcement learning to train a language model based on a recurrent NN (RNN) for text-based spatial reasoning.

Apart from language models, other papers also explore spatial reasoning in the context of GNNs. In [4], the authors developed a GNN for spatial navigation, which utilizes a memory mechanism for long-range information storage, while ref. [5] introduced the SpatialSim benchmark together with a GNN that is trained to recognize spatial configurations. Our work contributes to this area by partially elucidating the mechanism by which NNs are able to reason about geometric relations and by showing that GNNs which use a graph of relations/constraints and variables are much easier to train and scale.

3. Problems and Methods

To test the ability to reason about geometric relations, we develop a simple generative procedure of problems with solutions that allows us to control the difficulty and size of the problem. We test two types of models, GNNs and Transformers, on data generated with various generator settings.

3.1. Geometric Problem Generation

Our goal was to create a simple synthetic dataset that would enable us to study the mechanism which NNs use to reason about geometric relations. Specifically, we wanted to investigate whether neural networks develop internal representations that reflect the geometric structure of the problems they solve. In the task of auxiliary point prediction for which AlphaGeometry was using a language model, it is, for example, not clear whether the model took the geometric nature of the constraints into account or treated the language of the problem abstractly without any metric interpretation/grounding.

Therefore, we designed a generator for CSPs whose solutions can be interpreted as geometric figures described by the input constraints. The task of the model is to predict the positions of the points contained in the constraints. We simplify the problem so that the geometric figures lie on a discrete grid, and we can therefore treat the prediction of the positions as a classification task. We note that one can also treat the task as a regression and train the model to process figures with arbitrary 2D coordinates.

Our CSP uses four types of constraints: M, R, S, and T. The constraint $M(A,B,C)$ says that the point B is the midpoint of the line segment given by A and C, and $R(A,B,C,D)$ says that points A, B give an axis of symmetry around which C and D reflect each other. The constraint $S(A,B,C,D)$ says that points A, B, C, D form a square and $T(A,B,C,D)$ says that the vector $D-C$ is a translation of the vector $B-A$. In order to have unique solutions, we also need to fix several points to concrete positions. We can achieve that with the additional constraint P, for which $P(a,[1,1\left]\right)$ signifies that point A lies at position $[1,1]$. For each constraint, a certain number of variables are required to be known so that the constraint can be uniquely resolved; here, we call them determining variables. Once these variables are fixed, the remaining variables are uniquely determined. These are called dependent variables. For the constraint M, two variables are determining and one is dependent, i.e., if we know the positions of two variables, the last one can be resolved. For the constraint R, three variables are determining and one is dependent. For the constraint S, two variables are determining and two are dependent (i.e., we assume a fixed spatial ordering of the square vertices). For the constraint T, three variables are determining and one is dependent. In the Section 4.5 the visualization of the following geometric CSP instance is showed:

$$\begin{array}{cccccccccccccccccccc}\hfill & T(F,D,E,G)\hfill & \hfill & \wedge \hfill & \hfill & T(C,G,A,H)\hfill & \hfill & \wedge \hfill & \hfill & S(H,I,G,J)\hfill & & \wedge & & S(B,J,L,K)& & \wedge & & P(A,[6,15\left]\right)& & \wedge \\ \hfill & P(B,[12,16\left]\right)\hfill & \hfill & \wedge \hfill & \hfill & P(C,[12,14\left]\right)\hfill & \hfill & \wedge \hfill & \hfill & P(D,[14,12\left]\right)\hfill & & \wedge & & P(E,[12,14\left]\right)& & \wedge & & P(F,[12,9\left]\right)& & ,\end{array}$$

where $A-F$ are known points and $G-L$ are unknown.

Our generator creates problems that require deductive reasoning by ensuring the constraints form a dependency structure where some must be resolved before others. Attempting to solve all constraints simultaneously is not feasible because the dependent variables would be incorrectly determined. For example, a network might successfully create a square from four unknown points, but if some of those points are actually determined by previous translation or reflection constraints, the resulting square would be geometrically valid but positioned incorrectly within the overall figure. This dependency structure forces models to discover the proper reasoning sequence rather than solving constraints in isolation.

To create a single problem, we sample a random directed acyclic graph (DAG) with nodes corresponding to types of constraints and edges corresponding to dependencies. Then we add variables to the constraints so that each constraint within the graph can be uniquely resolved given that the parent constraints are already resolved (the parent constraints therefore need to contain the determining variables). To produce a problem with a unique solution, we fix the required number of variables appearing in each of the root constraints which have no parents (i.e., arbitrary two variables in the S constraint, three variables for the T constraint, etc.) We use the Z3 solver [15] to obtain the assignment for all other points. If the resulting figure is larger than the size of the grid, we reject the problem, otherwise we place it in a random position within the grid. The constraints with the fixed points are given as an input to the model, and the positions of the remaining variables are given as labels.

By varying the number and types of constraints, the generator produces problems with different complexity levels and reasoning depths. This allows us to investigate whether models can solve problems harder than those seen during training by leveraging test-time scaling approaches such as increased iterations or resampling.

3.2. Models

In our experiments, we tested two types of architectures: GNNs and autoregressive Transformers.

3.2.1. Graph Neural Network

The architecture of the GNN is inspired by the model presented in [10] and originates in [16]. Similarly to works about Boolean satisfiability, it applies the same update rules repeatedly. It is a GNN operating on a bipartite graph with n variables and m constraints defined by Equations (1) and (2) below which are applied recurrently for several rounds. Equation (2) updates the embeddings of variables which are, after the last update, used to classify a given variable to a point within the grid. This is achieved by applying a linear layer $L(·)$ on each variable embedding. As mentioned in Section 3.1, some variables need to be fixed to concrete points so that the problem has a unique solution.

The embeddings of known variables are initialized using the embedding layer that shares weights with the final classification layer $L(·)$, i.e., both use the same trainable weight matrix $\mathbf{W}\in {\mathbb{R}}^{N\times d}$, where N is the number of grid positions and d is the embedding dimension. This design ensures consistent representation of grid positions throughout the network (similarly to token embeddings sharing weights with the classification head in language models). They are not updated during message passing. The embeddings of unknown variables are initialized randomly, i.e., they are sampled from a unit hypercube ${[0,1]}^d$ following the uniform probability distribution of each coordinate. It enables resampling during inference by using different random initializations.

The embeddings are updated after each message-passing iteration. Let n be the number of variables and m be the number of constraints. The update equations have the following form:

$$\begin{array}{ccc}\hfill Z_c^{t+1}& =U_c\left(\mathrm{cat}\left(A_c{X}^t\right),Z_c^t,{\mathrm{\Phi }}_c\right)\hfill & \hfill \forall c\in C,\end{array}$$

(1)

$$\begin{array}{cc}\hfill X^{t+1}& =U_X\left(A_X{Z}^{t+1},X^t,{\mathrm{\Phi }}_X\right),\hfill \end{array}$$

(2)

where $X^t\in {\mathbb{R}}^{n\times d}$ is the matrix of stacked variable embeddings, $Z^t\in {\mathbb{R}}^{m\times d}$ is the matrix of stacked constraint embeddings, $A_X\in {\mathbb{R}}^{n\times m}$ is the adjacency matrix that matches constraints to variables, $A_c\in {\mathbb{R}}^{4m\times n}$ is the adjacency matrix that selects variables for each constraint ($A_c{X}^t$ will have $4m$ rows because there are four variables in each constraint; for constraints that have only three variables, we append a zero vector to preserve the shape), taking the constraint type and argument position into account, and $\mathrm{cat}\left(A_c{X}^t\right)\in {\mathbb{R}}^{m\times 4d}$ denotes the operation that extracts and concatenates the embeddings of the four variables involved in each constraint in the correct order. $C=\{M,R,S,T\}$ is the set of constraint types. Each problem instance is represented by a bipartite graph containing the denotation of the constraint types, which is fully encoded by the matrix $A_c$, from which the matrix $A_X$ can be easily determined.

Equation (2) updates the embeddings of variables and Equation (1) updates the embeddings of the constraints. The latter equation is further indexed by c which reflects the fact that we are using multiple types of constraints which require distinct update functions. The update functions $U_x$ and $U_c$ ($\forall c\in C$) are in our case realized as LSTMs [17] parameterized by ${\mathrm{\Phi }}_X$ and ${\mathrm{\Phi }}_c$, respectively. Technically, each LSTM also updates a cell state which is not mentioned in the equations. These LSTMs could be replaced by simple RNNs but we found the LSTMs easier to optimize and RNNs do not work well with our data; see Appendix A.4. An example of the simplified schema of the GNN processing can be found in Figure 1.

The update of each embedding happens independently of the other embeddings, i.e., in Equation (2), each embedding of a variable (row in the matrix $X^t$) is updated independently by the same LSTM which takes as input the aggregated message from the constraints containing this particular variable. The aggregated message is simply the sum of the relevant constraint embeddings and it is realized by multiplying the constraint embedding matrix $Z^t$ by the matrix $A_X$. The LSTM $U_c$ which updates the constraint embeddings differs by the fact that the aggregated message is not obtained as a sum but as a concatenation of variable embeddings. The embeddings are concatenated in the order in which the variables appeared within the constraints. For example, for the constraint $S(A,B,C,D)$, we concatenate the embeddings of variables $A,B,C,D$ in that order. One can intuitively view the embeddings of variables as if they are representing the values of these variables (points) and embeddings of constraints as if they represent the information of what is needed to satisfy the constraint. For this reason, the function which updates the constraint embeddings cannot be permutation-invariant because different orders of the determining variables result in different values for the dependent variables.

The embedding dimension d and number of iterations are configurable parameters that we adjust based on problem complexity. Importantly, since the same update rules are applied recurrently, the number of iterations can be modified at test time without changing the model parameters. This architectural property leads to the option of allocating more computation to solve harder problems during inference. Other hyperparameters for training the GNN can be found in

Table A1. Model hyperparameters and training settings for both architectures.

Parameter	GNN	Transformer
Model Architecture
Embedding dimension	128	256
Model iterations/layers	15 ± 10 (diminishing)	6
Number of heads	–	6
Dropout rate	0.1	–
Positional embeddings	–	RoPE
Training Configuration
Optimizer	AdamW	AdamW
Learning rate	${10}^{-3}$	$5\times {10}^{-4}$
Weight decay	${10}^{-3}$	–
Batch size	32	512
Epochs	200	200
Warmup steps	–	200
Learning Rate Schedule
Scheduler	Cosine annealing	Linear
Cycle length	15 epochs	–
Peak decay factor	0.9	–
Min LR factor	0.1	–
Regularization
EMA decay	0.99	–
Gradient clipping	0.65	1.0
Special Configuration
Special tokens	–	[SEP], [UNK], [PAD], [MASK]

.

For clarity, one illustrative example of how a GNN processes and solves a specific problem will be presented later in Section 3.3.

3.2.2. Autoregressive Transformer

The autoregressive Transformer model is based on the GPT-2 architecture developed by [18] with rotary embeddings [19]. It takes as input a sequence of tokens representing the problem together with a query for a variable we want to predict. For example, the input with just one constraint could have the following form: S ([0, 0] [0, 1] C D) ? D. It is querying the position for variable D within a square with two known points. The model reads this sequence of tokens (where the positions of points such as $[0,0]$ correspond to one token) and is trained to predict the token which corresponds to the the correct position of variable D ($[1,0]$ in this case). This means that from each CSP produced by the generator described in Section 3.1, we extract one sequence for each unknown point. After experimenting with the hyperparameters of the model, we set the number of layers to 6, the number of heads to 6, and the embedding dimension to 256. We also experimented with a recurrent application of one layer, as in [20], to more closely mimic the GNN, but having separate weights for each layer produced better results. Other hyperparameters of the model and for training can be found in Table A1.

3.3. GNN: Illustrative Example

To illustrate how the GNN works, we present a concrete example showing how it processes geometric constraints on a $N=20\times 20$ grid with an embedding dimension $d=128$ and iteration count $t=15$. Consider the following constraint satisfaction problem:

$$\begin{array}{cc}\hfill & T(A,B,C,D)\wedge S(B,D,E,F)\wedge S(B,F,G,H)\wedge \hfill \\ \hfill & P(A,[5,2\left]\right)\wedge P(B,[2,2\left]\right)\wedge P(C,[4,3\left]\right)\hfill \end{array}$$

where $A,B,C$ are known points fixed by position constraints P, and $D,E,F,G,H$ are unknown points to be predicted.

3.3.1. Problem Encoding and Dependencies

Each grid position $(x,y)$ maps to a unique index using $\mathrm{index}=x+y\times 20$. Therefore, the known points correspond to indices: $A\to 45$, $B\to 42$, $C\to 64$. The problem forms a bipartite graph with eight point nodes and three constraint nodes, where the dependency structure requires sequential resolution: constraint $T_1$ determines point D, which enables constraint $S_1$ to determine points E and F, which finally allows constraint $S_2$ to determine points G and H.

The point constraints P are not realized as separate constraint nodes in the bipartite graph, but rather as fixed embeddings taken directly from the embedding layer. Figure 2 illustrates this encoding and the target solution.

3.3.2. Embedding Initialization

The model uses a shared embedding matrix $\mathbf{W}\in {\mathbb{R}}^{400\times 128}$ for both the embedding layer and final classifier. Known points are initialized with the fixed embeddings ${\mathbf{W}}_{45}$, ${\mathbf{W}}_{42}$, and ${\mathbf{W}}_{64}\in {\mathbb{R}}^{128}$ corresponding to points A, B, and C, respectively, where ${\mathbf{W}}_i$ denotes the i-th row of matrix $\mathbf{W}$. These embeddings remain unchanged throughout inference:

$${\mathbf{x}}_A^{\left(t\right)}={\mathbf{W}}_{45},{\mathbf{x}}_B^{\left(t\right)}={\mathbf{W}}_{42},{\mathbf{x}}_C^{\left(t\right)}={\mathbf{W}}_{64}\mathrm{for}\mathrm{all}t.$$

Unknown points are initialized with random vectors ${\mathbf{r}}_D,{\mathbf{r}}_E,{\mathbf{r}}_F,{\mathbf{r}}_G,{\mathbf{r}}_H\sim \mathcal{U}{(0,1)}^{128}$, sampled uniformly from the unit hypercube, and constraint embeddings are initialized as random vectors ${\mathbf{c}}_1,{\mathbf{c}}_2,{\mathbf{c}}_3\sim \mathcal{U}{(0,1)}^{128}$. These random vectors are used as the initial hidden states for the LSTM-based message passing. Specifically, for each unknown point i, we set the initial embedding as the hidden state ${\mathbf{x}}_i^{\left(0\right)}:={\mathbf{r}}_i$, and similarly for each constraint c, we set ${\mathbf{z}}_c^{\left(0\right)}:={\mathbf{c}}_c$. The LSTM cell states that ${\mathbf{h}}_i^{\left(0\right)}$ and ${\mathbf{h}}_c^{\left(0\right)}$ are initialized to zero. During inference, the hidden states ${\mathbf{x}}_i^{\left(t\right)}$ and ${\mathbf{z}}_c^{\left(t\right)}$ evolve over time according to the update equations.

3.3.3. Message Passing Dynamics

The update equations are repeated for 15 iterations. Our LSTM-based architecture maintains both hidden states and cell states, where the hidden states serve as the embeddings that flow through the message-passing network, while cell states maintain internal memory. The LSTM also produces a cell output at each iteration, but this output is discarded and only the hidden and cell states are retained.

For constraint updates (Equation (1)), each constraint type uses a specialized LSTM update function. The constraint embedding ${\mathbf{z}}_c^{\left(t\right)}$ (LSTM hidden state) and cell state ${\mathbf{h}}_c^{\left(t\right)}$ are updated as:

$${\mathbf{z}}_c^{(t+1)},{\mathbf{h}}_c^{(t+1)}=U_c^{\left(\mathrm{type}\right)}({\mathbf{m}}_c^{\left(t\right)},{\mathbf{z}}_c^{\left(t\right)},{\mathbf{h}}_c^{\left(t\right)})$$

where ${\mathbf{m}}_c^{\left(t\right)}$ is the concatenated message from variables participating in constraint c, and each constraint type has its own LSTM parameters: $U_c^{\left(T\right)}$ for Translation, $U_c^{\left(S\right)}$ for Square, $U_c^{\left(R\right)}$ for Reflection, and $U_c^{\left(M\right)}$ for Midpoint constraints. The LSTM returns both an output (which is discarded) and a tuple containing the new hidden state ${\mathbf{z}}_c^{(t+1)}$ and cell state ${\mathbf{h}}_c^{(t+1)}$.

For the above example, the input message ${\mathbf{m}}_{T_1}^{\left(t\right)}$ for constraint $T_1$ is the concatenated vector $[{\mathbf{W}}_{45};{\mathbf{W}}_{42};{\mathbf{W}}_{64};{\mathbf{x}}_D^{\left(t\right)}]\in {\mathbb{R}}^{512}$, where ${\mathbf{W}}_{45}$, ${\mathbf{W}}_{42}$, and ${\mathbf{W}}_{64}$ are the fixed embeddings of the known points A, B, and C, respectively, and ${\mathbf{x}}_D^{\left(t\right)}$ is the current hidden state of point D. Similarly, the messages ${\mathbf{m}}_{S_1}^{\left(t\right)}$ and ${\mathbf{m}}_{S_2}^{\left(t\right)}$ for constraints $S_1$ and $S_2$ are given by $[{\mathbf{W}}_{42};{\mathbf{x}}_D^{\left(t\right)};{\mathbf{x}}_E^{\left(t\right)};{\mathbf{x}}_F^{\left(t\right)}]$ and $[{\mathbf{W}}_{42};{\mathbf{x}}_F^{\left(t\right)};{\mathbf{x}}_G^{\left(t\right)};{\mathbf{x}}_H^{\left(t\right)}]$, respectively. Each of these message vectors is passed to the corresponding constraint-specific LSTM update function $U_c^{\left(\mathrm{type}\right)}$.

For variable updates (Equation (2)), unknown point embeddings are updated using messages aggregated from all connected constraints. Each unknown variable i aggregates the sum of constraint embeddings from all connected constraints and passes this sum to the shared LSTM update function $U_X$:

$${\mathbf{x}}_i^{(t+1)},{\mathbf{h}}_i^{(t+1)}=U_X\left(\sum _{c\in \mathcal{N}\left(i\right)}{\mathbf{z}}_c^{(t+1)},{\mathbf{x}}_i^{\left(t\right)},{\mathbf{h}}_i^{\left(t\right)}\right)$$

Here, $\mathcal{N}\left(i\right)$ denotes the set of constraint nodes connected to variable i, ${\mathbf{x}}_i^{\left(t\right)}$ is the hidden state (embedding) of point i at time step t, and ${\mathbf{h}}_i^{\left(t\right)}$ is its corresponding LSTM cell state.

In the illustrative example, point D receives messages from constraints $T_1$ and $S_1$, resulting in an update input of ${\mathbf{z}}_{T_1}^{(t+1)}+{\mathbf{z}}_{S_1}^{(t+1)}$. Point F receives messages from $S_1$ and $S_2$. Points E, G, and H each receive a single constraint message: from $S_1$ for E, and from $S_2$ for G and H.

These aggregated messages are passed into the $U_X$ LSTM, which outputs the updated hidden state ${\mathbf{x}}_i^{(t+1)}$ that becomes the new embedding of variable i.

3.3.4. Solution Recovery

After t iterations, the model predicts the final grid position of each unknown variable using a standard classification head. The final embedding ${\mathbf{x}}_i^{\left(T\right)}$ is projected using the shared embedding matrix $\mathbf{W}\in {\mathbb{R}}^{400\times 128}$ (no bias term) to produce logits:

$${\ell }_i=\mathbf{W}{\mathbf{x}}_i^{\left(T\right)}\in {\mathbb{R}}^{400}$$

The predicted index is obtained as ${\widehat{k}}_i=\arg \max \left({\ell }_i\right)$, which is then mapped back to coordinates via $\widehat{x}={\widehat{k}}_i\mathrm{mod}20$, $\widehat{y}=\lfloor {\widehat{k}}_i/20\rfloor$.

During training, we apply a cross-entropy loss over the logits ${\ell }_i$, comparing predicted and ground truth indices for all unknown variables. The embedding matrix $\mathbf{W}$ is shared with the embedding layer used for known variables, ensuring consistency between input encoding and prediction.

In our example, the correct predictions for points D–H are $[1,3]$, $[0,2]$, $[1,1]$, $[2,0]$, and $[3,1]$, satisfying all constraints.

4. Experiments and Results

We evaluate both architectures on geometric constraint problems. We examine prediction accuracy, dataset complexity effects, embedding structure emergence, initialization strategies, solution dynamics, and failure modes.

4.1. Comparison of the Two Architectures

To compare the two architectures, we measure the accuracy of individual point prediction (point accuracy). In later experiments, we also report the accuracy in terms of correctly assigning all variables within the problem (complete accuracy).

The experiments on the synthetically generated data show that the GNN performs significantly better than the Transformer, as expected. In Appendix G, we show that we were able to train the GNN on grid sizes up to $80\times 80$ points to a validation accuracy larger than $90\%$. In comparison, the Transformer achieved an accuracy of $90\%$ only if we trained it on a grid size $10\times 10$ and limited the number of constraints to 2. If we trained the model on a more complex setting with the grid size $20\times 20$ and up to 6 constraints, it reached an accuracy of approximately $30\%$.

For completeness, we also trained the Transformer model to find the assignments to variables using Chain-of-Thought (CoT) [21] by imitating a log from a simple solver. The solver resolves the constraints one by one in the topological order of the DAG of constraints mentioned in Section 3.1. Using this method of training, the Transformer learns to predict the positions of variables within a $20\times 20$ grid with approximately 50% point accuracy. We did not explore this direction further as reasoning with a CoT is orthogonal to reasoning in the embedding space on which we focus in this work. More details about the CoT experiment can be found in Appendix E.

Given the substantially better results achieved by the GNN on this problem setting, we conduct the main analysis of the embeddings with the GNN. Appendix D contains similar analysis and visualization for the Transformer (see Figure A8).

We also note that the model accuracy can be further improved by leveraging multiple prediction attempts since incorrect predictions often fall close to their true positions (as shown in Appendix C) and according to preliminary tests, sampling multiple different initial embeddings increases the chance of predicting the correct assignment.

4.2. GNN Performance Analysis

Given the GNN’s clear advantage over the Transformer for our CSPs, representable by bipartite graphs, we now analyze the GNN’s behavior in detail across multiple dimensions. We examine how the model performs under distribution shift, visualize the internal representations it develops, and identify failure modes related to problem complexity. The analysis uses our best-performing GNN configuration with an embedding dimension of 128, 15 message-passing iterations, and training procedures detailed in Appendix A. This systematic evaluation reveals both the capabilities and limitations of our geometric constraint solving approach on discrete grids.

4.2.1. Dataset Design and Complexity Distribution

The training dataset contains problems with 1–16 geometric constraints (mean: $4.0$), 3–34 points (mean: $8.7$), and reasoning depths of 1–12 (mean: $2.9$). The test dataset was specifically designed to be more challenging, containing problems with 8–26 geometric constraints (mean: $14.2$), 12–47 points (mean: $23.8$), and reasoning depths of 3–14 (mean: $7.0$). Figure 3 illustrates these distribution differences across key complexity metrics.

Our generator creates problems for specific grid sizes, and we focus on a $20\times 20$ grid as it provides a good balance. It is large enough to accommodate interesting geometric figures while remaining computationally tractable for detailed analysis, which required multiple re-runs. Our constraint set consists of four types, Square (S), Reflection (R), Midpoint (M), and Translation (T), with their relative distribution in the training dataset shown in Figure 4.

This choice of difficulty distributions allows us to study model behavior when problem complexity is increased beyond what the model encountered during training. The high validation accuracy achieved on the training distribution ($99.5$% point accuracy, $98.1$% complete problem accuracy using 10% of training data as validation) demonstrates that the GNN can effectively solve geometric constraint problems when sufficient training data is available. We used $\mathrm{300,000}$ ($\mathrm{270,000}$ without validation) training examples. The size was determined through the scaling analysis outlined in Appendix G, which shows the sample complexity requirements for different grid sizes.

The deliberately harder test distribution provides a challenging evaluation setting, where the model achieves lower baseline accuracy, creating opportunities to study failure modes, reasoning depth effects, and test-time scaling approaches.

4.2.2. Performance Under Distribution Shift

Table 1. Performance of the model. The first row reports validation accuracy (on a separate dataset, using 1 resample). The remaining rows show results on a harder test set drawn from a different distribution, across different numbers of inference iterations (15, 23, Best) and resample counts (1 and 10). We report point-wise accuracy and complete accuracy.

	1 Resample		10 Resamples
Setting	Point Acc.	Complete Acc.	Point Acc.	Complete Acc.
Validation (15 iters)	99.55%	98.93%	–	–
Test (15 iters)	96.07%	76.74%	97.89%	80.93%
Test (23 iters)	98.51%	93.62%	99.02%	95.37%
Test (Best)	99.04%	96.14%	99.44%	97.46%

presents comprehensive performance results on our datasets. The model maintains good performance even under distribution shift, achieving $96.07$% point accuracy and $76.74$% complete accuracy on the harder test set with standard inference (15 iterations, 1 resample).

The results demonstrate clear benefits from test-time scaling approaches. Increasing inference iterations from 15 to 23 improves complete accuracy from $76.74$% to $93.62$%. Using 10 resamples with different random initializations further boosts performance to $95.37$% complete accuracy. We determined the optimal iteration count of 23 through the analysis shown in Figure A10, which also revealed a slight performance increase with multiple resamples.

The “Best” configuration could be explained as an oracle that knows the optimal number of iterations for each individual problem to achieve the most solved points in the fewest iterations (with a maximum of 50 iterations). This approach achieves a complete accuracy of $96.14$% with single resampling and $97.46$% with 10 resamples. This model could in theory get very close to its performance on the validation set.

For the resampling strategy, we tested various counts but found the most significant improvement occurred when going from 1 to 5 resamples. We chose 10 resamples as a conservative buffer, since additional resamples beyond this point showed diminishing returns while increasing computational cost.

These findings suggest several important properties of the learned model. The improvement from additional iterations indicates that the GNN employs an iterative refinement process similar to continuous optimization methods. The benefits of resampling show that different random initializations of unknown variables can lead to different solution paths, with some being more effective for particular problem instances.

The model’s ability to generalize to problems with increased complexity beyond the training distribution are consistent with previous work that applied a very similar GNN architecture to SAT problems [22].

4.3. Visualizing the Embeddings of Individual Points

In the following text, we use the terms static embeddings and dynamic embeddings. By static embeddings we mean the embeddings of points of the grid which are used for initialization of known points and are shared with the classification head of both models. By dynamic embeddings we mean the embeddings of the unknown variables which are updated throughout the forward process.

Both types of models are trained to predict positions of unknown points within the instance. This is achieved by a linear layer which computes the logits for each point. As already mentioned, the weights of this layer are shared with the embedding layer, representing the position of known points.

When visualizing low-dimensional projection of the static embeddings corresponding to individual points, we found that they organize themselves into a 2D grid they represent. In Figure 5, we show how this organization emerges during training and how it is connected to the precision of the prediction (three-dimensional projection is depicted in Appendix B).

We stress that the existence of the grid structure of the data domain and the existence of geometric figures related to constraints is given to a model only indirectly through the constraints and their solutions. The whole training dataset can be thought of as a set of constraints written in an unknown language, and the goal of training is to find a model for this language which will enable correct prediction. As the prediction is based on the similarity (given by the inner product) of the dynamic embeddings of variables and static embeddings of points, it is not that unexpected that the discovered model reflects the geometry behind this language.

Here we focus on point embeddings; for completeness, we provide some results about constraint embeddings in Appendix I. We show that constraint embeddings encode constraint types, geometric properties, satisfaction status, and temporal information to some extent.

4.4. Network Initialization with Grid Structure

We investigated whether providing the model with an initial understanding of grid structure accelerates training convergence. Rather than initializing the shared embedding and classifier weights randomly, we can initialize them to reflect the geometric structure of the $20\times 20$ grid.

To provide the model with an initial geometric inductive bias, we initialize the embedding matrix $\mathbf{W}\in {\mathbb{R}}^{N\times d}$ (used for both known point embeddings and the final classification layer) using a rotated grid construction. We begin by assigning each of the N grid positions ($N=400$ for a $20\times 20$ grid) its exact 2D coordinates $(x_i,y_i)$ and construct a matrix $\mathbf{V}\in {\mathbb{R}}^{N\times d}$ where ${\mathbf{V}}_{i,0}=x_i$, ${\mathbf{V}}_{i,1}=y_i$, and all other dimensions are zero. We then generate a random orthogonal matrix $\mathbf{Q}\in {\mathbb{R}}^{d\times d}$ via QR decomposition of a matrix with entries sampled from a standard normal distribution, and define the initial embedding matrix as:

$$\mathbf{W}=\mathbf{V}\mathbf{Q}$$

This rotation spreads the spatial coordinate information across all d dimensions while avoiding alignment with any particular coordinate axis. Following this initialization, training proceeds identically to the random case, with the weights in $\mathbf{W}$ remaining fully trainable.

Figure 6 demonstrates that grid-structured initialization provides substantial training acceleration compared to random initialization. The model with grid initialization converges to high validation accuracy within the first few epochs, while random initialization requires significantly more training time. This suggests that providing geometric inductive bias through weight initialization helps the model quickly discover the spatial relationships it needs to solve geometric constraint problems.

The initialization method shares weights between the embedding layer (used for known points) and classifier head (used for predictions), ensuring consistent geometric representations throughout the network. This architectural choice proves beneficial as the model can leverage the geometric structure for both encoding known positions and predicting unknown ones.

4.5. Solution Process

In order to find the solution, the GNN model iteratively moves embeddings of unknown points in a high-dimensional space. During the forward pass, the same update rule is applied over and over again. Hence, it is possible to extract the embeddings in each message-passing iteration to obtain a better understanding of the underlying process. Similar examples are provided in Appendix H.

Figure 7 shows an example of the solution process for a random problem given to the GNN. After each update, the closest static point embedding is taken for each variable of a problem and is visualized as the corresponding grid position.

Note that in this example, the GNN first finds a close approximation to the hidden configuration and then refines it. Here, the squares are first “approximated” by quadrilaterals which then converge to exact squares. It can also be observed that one square constraint is approximately satisfied earlier than the other, which reflects the fact that the other constraint can be resolved only after the first constraint is resolved (as explained in the next section). For an extended discussion and visualization of the reasoning dynamics as revealed by UMAP, see Appendix J.

4.6. Analysis of Incorrectly Classified Points

To understand the failure modes of the GNN, we analyze predictions on our $\mathrm{30,000}$ test instances using the basic model setup (15 iterations, 1 resample) from Table 1. We examine how prediction accuracy relates to problem complexity along multiple dimensions.

As explained in Section 3.1, constraints form a dependency structure where some must be resolved before others. For any given point, we can trace the dependency chain of constraints that must be resolved to determine its position. We define point depth as the number of constraint resolution steps required for a specific point. The problems themselves vary in complexity through both the total number of constraints and the maximum depth of dependency chains within the problem.

Figure 8 (left) shows that prediction accuracy correlates negatively with point depth. Points requiring deeper reasoning chains have significantly higher failure rates. The relationship is nearly monotonic, with success rates declining from over $90\%$ for points at depths of 2–3 to below $20\%$ for points requiring 10 or more resolution steps. This pattern aligns with the iterative solution process observed in Section 4.5 where the model progressively constructs solutions by first resolving simpler constraints.

Figure 8 (right) demonstrates that problems with more constraints also show lower complete success rates. This reflects both the increased likelihood of errors accumulating across multiple predictions and the tendency for larger problems to contain deeper dependency structures.

These findings suggest that the GNN’s reasoning capability is fundamentally limited by the length of dependency chains it can reliably process. Importantly, when predictions do fail, the incorrectly classified points tend to be positioned close to their correct target locations, as evidenced by the distance histograms in Appendix C.

5. Limitations and Discussion

We note that we study a much simpler setup than was studied in AlphaGeometry. Predicting positions of points satisfying a set of geometric constraints is straightforward in comparison with predicting useful auxiliary points. Also, in our case, the unknown points need to lie on a discrete 2D grid so that we can train the model with a classification loss. The simplified setup turned out to be sufficient to study how NNs can reason about geometric constraints. Training the model for predicting auxiliary points requires a large computational budget and, moreover, the authors did not release the training dataset, which was also expensive to create. We also mention that we use only four types of constraints, but the whole setup could be easily extended to different sets of constraints by adding more update functions in Equation (1). These simplifications enabled us to provide partial understanding of the process by which the tested models are able to solve individual problems and to demonstrate the benefits of using a GNN instead of a Transformer. We note that variable permutation invariance is not considered in our constraint definitions, as a fixed variable ordering is required to ensure unique solutions.

Compared to the GNN which Hula et al. [10] used for Boolean satisfiability, our GNN deals with fewer constraints and points. In their case, the GNN was able to handle hundreds of constraints, whereas our GNN deals with an order of magnitude fewer constraints and still needs more training samples than in the case of Boolean satisfiability. This could be caused by the fact that the domain of our constraints is much larger, but it could also be the case that different architecture of the GNN (with different update functions) could scale better to larger problems.

It is also possible that a different training procedure might result in significantly easier training. The inference process depicted in Figure 7 and the findings of [10] suggest that the GNN could implicitly optimize an unknown objective during the iterative forward pass. Finding an expression for this objective could be interesting and used to train the network in an unsupervised way. This should be much easier than training it with the classification loss. Another obvious direction is to train the model as a diffusion model which “denoises” randomly assigned points to points of the solution.

Lastly, we mention that our experimental setup can be extended to more complex CSPs which could contain temporal relations and relations between various entities. We believe that such CSPs could be very useful for studying how NNs can generalize to unseen situations [23] and how they can discover models of the world without grounding [24]. Studying geometric CSPs has the advantage that the domain has an obvious geometric interpretation which is visible also in the embedding space.

6. Future Work

While we provided initial insights into geometric reasoning in neural networks, several directions remain for deeper understanding. Future work should investigate the expressiveness requirements of different update mechanisms, building on our finding that LSTMs significantly outperform RNNs for constraint updates. This includes determining the minimal architectural complexity needed to solve geometric constraints and whether these mechanisms can recover explicit mathematical formulas underlying constraints rather than learning approximations. The iterative refinement process that we observed suggests connections to continuous optimization methods. Future work should formalize this relationship and investigate whether training procedures based on explicit optimization objectives could improve efficiency. Additionally, extending beyond our four constraint types while addressing variable permutation invariance would broaden insights. Moving from discrete grid classification to continuous coordinate regression would enable analysis of geometric reasoning in unrestricted spatial domains.

7. Conclusions

We have shown that GNNs as well as Transformers can learn to solve geometric CSPs and provided several insights into the process by which they find the solution. The visualizations show that when processing the problem, models form the hidden spatial configurations described by the problem in the embedding space. During training, the static embeddings of individual points in the 2D grid organize themselves within a 2D subspace and reflect the neighborhood structure of the grid. We showed that the occurrence of errors depends on problem complexity, like the number of steps required to resolve all variables. The models solve constraints through an iterative refinement process and exhibit test-time scaling capabilities, where allocating additional computational resources during inference improves performance on harder (out of training distribution) problem instances. Lastly, we showed that GNNs are much easier to train and can be scaled to significantly larger grids than Transformers.