Discrete Neural Algorithmic Reasoning

We thank the reviewer for their review and constructive comments!

the paper could be strengthened by testing on a wider range of algorithms

We agree that a more general architecture is of great interest for future work. In its current form, the proposed model is not capable of executing some algorithms from the CLRS-30 benchmark, however, simple modifications can enhance its expressivity, e.g., by supporting more complex aggregations with top-k attention (with fixed k). Also, as we mention in the paper, the model can be extended with additional modifications, such as edge-based reasoning. For example, edge-based top-2 attention (where each edge chooses the top 2 adjacent edges to receive the message from) can implement the sort of triplet reasoning.

In our work, we aim to describe the key principles that help to build perfectly generalizable models and do not focus on a general architecture capable of executing a wide range of algorithms.

Another example of simple modification which extends the list of algorithms the model can be tested on, is supporting multiple different scalars in ScalarUpdater, which can allow the model to execute several sorting algorithms. While this does not fully address the mentioned limitation, please let us know if you consider such an experiment important and we will add this during the discussion period.

It is unclear how this model’s interpretability … differs from other NAR approaches

The main difference is that with discrete states we “fully observe” the inner computations, while simply looking at the intermediate hint predictions might be misleading.

For example (this particular example is motivated by our analysis of the no-hint model, please note the general response), let us consider the model for the BFS task, which internally uses K discrete node states as distances from the starting node and able to correctly predict all intermediate hints when evaluated on graphs with diameter less than K. If we try to interpret such model by intermediate updates of hints on graphs with small diameter we will simply observe the perfect mimicking of the ground truth algorithm without understanding how exactly the model works: does it uses the notion of visited nodes directly or implements the “distance concept” (or any more complex logic).

In other words, if there is some information flow beyond the predicted hints (which is true for all covered baselines) inside the model, then simple probing of hints might not give a full picture about the underlying computations inside the model.

For the multi-task experiments, do the results in Table 2 correspond to these settings?

For the multitask experiments, we train all 6 algorithms from Table 1 simultaneously. The motivation for the multitask experiment is to check if the proposed processor is capable of multitask learning without losing the generalization capabilities. We did not test the baselines in the multitask experiment. We will add more details regarding the multitask experiment to the text and a separate column to Table 1 to highlight the perfect scores in the multitask setup too.

Without hints, DNAR struggles to generalize. Could you clarify the reverse-engineering issue described? Was any intermediate approach, such as noisy teacher forcing or teacher forcing decay, tested to examine the impact of hints on DNAR?

Please see the second part of the general response for the reverse-engineering details.

We indeed tested some additional approaches, but all of them are more about different hint usage rather than about finding a good intermediate approach: e.g. we checked if it is possible to learn the proposed models without teacher forcing at all and got perfect scores for all problems besides DFS. We did not test teacher forcing decay for the DFS problem, instead, we achieved perfect scores for DFS without teacher forcing by weighting the hint losses along the execution trajectory (by annealing the hint losses along the processor steps) to enforce the sequential learning of the problem.

An interesting intermediate approach that leaves the freedom for the model to learn something different is to supervise the model only on a part of the hints, but we did not test that.

How were GIN and PGN used as baselines?

We reused the implementation from the SALSA-CLRS, we refer to the corresponding paper and public source code for the details. In short, they add a GRU cell after the message-passing step of GIN (GINe is used for the problems with the edge weights), the baselines trained with hint supervision without teacher forcing. They only decode the hint predictions to calculate the loss. The hints of type pointer are predicted from the pairs of node features and node masks are predicted from the node features. Please let us know if some particular details remain unclear.

We hope that our response addresses your concerns and will be happy to answer any additional questions during the discussion period if needed.

Thank you for the detailed review! We addressed the questions and concerns below.

If the claims of this paper are all correct, then we just obtain a recipe for learning a neural reasoner to perfectly mimic a known algorithm. … This is unsatisfying since the goal of neural algorithmic reasoning is to learn correct algorithms from data without knowing the algorithms

First, let us note that the goal of NAR is to learn the network to execute algorithms in the latent space and is not restricted to the setup where the ground truth algorithm is unknown. For example, see the blueprint of NAR (Veličković & Blundell, 2021): “An algorithmic reasoner is trained to imitate A, optimising it to be close to ground-truth abstract outputs, A(x). P is a processor network operating in a high-dimensional latent space, which, if trained correctly, will be able to imitate the individual steps of A”.

Also, we note that the ground truth hints (and different forms of supervision on them during training) are the key component of the CLRS-30 benchmark.

Moreover, for the current state of the field, learning with hints is an important and unsolved problem. E.g., the large body of research (Section 2.1 of the paper), including the state-of-the-art approaches (Bevilacqua et al., 2023; Bohde et al., 2024) heavily rely on different forms of carefully designed step-by-step hints and are not applicable for no-hint learning without additional modifications.

However, we fully agree that learning without hints is an important and challenging problem for further developments of neural algorithmic reasoning.

How do you confirm that the attention block indeed operates as “select_best” selector? I neither see any empirical investigation nor theoretical proof on this point.

In short, we confirm this empirically by directly computing attention weights. As we described, the unnormalized attention weight between nodes U and V depends only on the states of U and V and the indicator if V has the smallest scalar value among the neighbors of U. Thus, the “select_best” property follows from the fact AttentionWeight(U_state, V_state, is_best=True) > AttentionWeight(U_state, V_state, is_best=False) for all possible combinations of (U_state, V_state).

Even if the attention block operates as “select_best” selector, why does this condition lead to the above claim? You should at least demonstrate your detailed logic using an example algorithm like DFS.

Please, see part 1 of our general response where we describe the detailed logic. We will also extend the description in the paper. Please let us know if some details remain unclear.

I think the size of the discretized states would matter a lot to the expressivity of the proposed neural reasoner,

We agree that the model with a greater states count can express strictly more functions. However, for models trained with hints, the states count is directly specified with the hints structure. There is the possibility to design alternative hints to utilize more states, but this requires describing the new hints structure. For no-hint models, states count is a simple hyperparameter. However, as we preliminarily investigated in Section 6, this count should be as small as possible. For example, one can think about extreme cases where the states count is close to the number of training data points; the resulting model can demonstrate perfect training loss but zero generalization capabilities.

An interesting experiment is to check if there is some sort of phase transition in terms of problem size. Specifically, I would imagine the phenomenon of perfect fitting to the correct algorithm would disappear if we decreased the problem size in training

Intuitively, from the perspective of state updates, such phase transition should occur if the training size/distribution is not enough to cover all possible state combinations/subtasks.

For example, for the BFS problem, it is enough to use graphs with only 2 nodes to observe that a not_visited node becomes visited or not depending on the received message. However, the subtask of selecting the parent from the multiple visited neighbors requires at least 4 nodes (where the minimum sufficient example is the complete bipartite graph K(2, 2)).

We conducted additional experiments to empirically find if the described transition occurs and what is the smallest training size for perfect fitting of each covered algorithm. We train the DNAR models for each problem on ER(n, 0.5) graphs for different n keeping the remaining hyperparameters the same as in the main experiments, and test the resulting models on the graphs with 160 nodes (same as in Table 1 of the paper).

Node level scores on graphs with 160 nodes for different training sizes:

Note that the empirical bound is around 4-5 nodes.

The details of how to train task-dependent encoders and decoders are not provided, and how they affect the processor's training is not discussed at all. I would imagine the quality of the encoded embedding is quite important to the success of learning the processor, even with teacher forcing

In our work, encoders and decoders are simple linear layers trained simultaneously with the processor (which is similar to prior work mentioned in Section 2.1).

In section 3.3, I get the high-level idea that the scalars (continuous input) only affect the computation of attention weights and do not affect the node or edge states. However, the description of the idea in the 2nd paragraph of Section 3.3 is not so clear that I could not figure out how exactly the computation is designed. It would be great to either write the equations and or illustrate the computational graph.

We allow the scalars to affect only the Key vector in the attention mechanism:

For each node $U$, we group the neighbors of the node $U$ by their state, and for each neighbor $V$ decide if the edge $(V, U)$ has the smallest scalar value among all neighbors of $U$ with the same state as $V$.

Then, we consider the obtained 0-1 indicator for each $(V, U)$ edge in the attention block as a part of the $(V, U)$ edge state when we compute the Key vector for the $(V, U)$ edge.

You can find the computation of the described 0-1 indicators (variable best_in_group) in lines 30-38 here: https://anonymous.4open.science/r/F4CA/processors.py In short, the mentioned code doubles the amount of used states by splitting each state depending on the discussed indicator and the Key vector of each edge is obtained from these states.

We will add an illustration to the text.

We hope that our response addresses your concerns and we are open to further discussions.

Thank you for your review of our paper! We address your questions and concerns below.

I found one such paper [1], used a text-based version of the CLRS-3o, called TextCLRS-text. Does it make sense to use that dataset here? Or to compare to the proposed TransNAR architecture?

There is some connection between our work and the work by Bounsi et al. (2024). In short, TransNAR discusses a method to enhance reasoning capabilities of language models with the task-specific GNN-based NAR model, where any NAR model can be used as an “internal tool”. Thus, replacing the baselines GNN with the proposed DNAR model improves the quality of the tool that the language model can use and the overall performance will be limited only by the correctness of the “tool usage”, and not by the inaccuracies of the tool itself. However, there are some difficulties in measuring the direct effect of including the proposed DNAR model in the TransNAR pipeline, as the source code for TransNAR is not yet publicly available.

It would be useful to ablate the different algorithmic decisions made in this paper. In its current form, I can't ascertain which of the architectural decisions is responsible for which gains on these datasets.

Please note the ablation experiments in Appendix A. Recall that we have three key components of our contribution: feature discretization, hard attention, and separating discrete and continuous data flows.

In short, we demonstrate that removing each of these components yields the model without provable guarantees of perfect generalization. However, these components differ in terms of the impact on the performance. E.g., using regular attention instead of hard attention yields perfect test scores for given datasets, but it is possible to construct adversarial examples with large neighborhood sizes where performance drops. On the other hand, removing discretization from the scalar updater significantly affects the performance even on the small test graphs.

If you have any other questions or concerns, we are happy to discuss them further.

Dear reviewer, please let us know if you find the mentioned ablation experiments useful and if you have any other concerns. We also note that we added an experiment to investigate if the proposed ScalarUpdate module can be extended to support more complex continuous manipulations with scalars (please see Appendix C of the paper). In short, we find our results positive and consider the proposed discretization as a promising way to overcome the difficulties of value generalization in an interpretable way.

We hope that the revised version of our work with additional details from the general responses is now more clear to readers.

We will be happy to answer any additional questions if needed.

Could you provide a pseudocode behind Discretize_nodes and Discretize_edges functions shown on page 4?

Let $X_i$ be the $i$-th node features ($X_i \in R^{h}$). First, we apply a linear layer to get the unnormalized logits of $X_i$ representing each of possible states: $L_{proj}(X_i) \in R^{k}$. Then, we select the state from these logits (this step is different depending on training/training_with_teacher_forcing/inference) and then we use the (learnable) embedding of the corresponding state.

During training with hints, we use ground truth states (we optimize the logits with CE loss and use teacher forcing) and for no-hint learning we use Gumbel softmax with annealing temperature. At the inference we always use argmax.

def discretize(node_features):
	states_logits = LinearProjection(node_features) # shapes are: (num_nodes, hidden_size) -> (num_nodes, num_states)
	
	if training_without_teacher_forcing:
		states_one_hot = gumbel_softmax(states_logits)
	if training_with_teacher_forcing:
		states_one_hot = ground_truth_states
	if not training:
		states_one_hot = argmax(states_logits, dim=1)
	node_features = states_embeddings(states_one_hot)
	return node_features

Work mentions modifying hints from the benchmark, could you clarify if this was included for tested baselines as well as how the hints are modified?

We only modify hints for the DFS problem, keeping the original execution flow. In particular, the original hints for the DFS problem (Veličković P. et al., 2022) use several additional concepts such as the global time counter and node’s discovery/finish times, which are not directly used in the DFS problem itself, but needed for other DFS-based problems from the CLRS-30 benchmarks (e.g., Bridges and SCC).

As graph-level features are not needed directly for any of the covered problems, we keep the model without the graph-level features and remove graph-level hints for the DFS problem.

These changes were not included for the tested baselines.

Could you further elaborate on the addition of a virtual node in your model (section 4.3)?

The virtual node is needed to select the unique node from the whole graph (for example, to select the current node at each step of the MST algorithm). In prior works, this type of computation is done by using graph-level features (which are updated depending on all node features and used in all node features updates) or using complete graphs for message passing. We use a simpler alternative that does not maintain the graph-level features: the virtual node is not updated from the node features and is only capable of selecting important nodes via the attention mechanism. In other words, the “unique” node becomes the one from which the virtual node received the message.

Could you provide results of the hyperparameter search mentioned in section 6? Could you provide further detail on reverse-engineering (perhaps a toy example) mentioned in section 6?

Please see the second part of the general response, where we give a more detailed description of our analysis. Please let us know if some details remain unclear.

We hope that our response addresses your concerns and will be happy to answer any additional questions during the discussion period if needed.

We thank the reviewer for a thoughtful review! Let us address the raised concerns and questions.

Could you elaborate on reported results and 0 variance between your seeds given that this is not the case in majority of comparable baselines or ablation experiments?

The proposed model indeed converges to the perfect mimicking of the ground truth algorithms and demonstrates perfect scores across different seeds. The main reason why this is not the case for the baselines and ablation experiments is discretization: due to the design of the proposed model it needs only to learn a finite amount of discrete state transitions and (as we discuss in the paper and in the first part of the general response) the same state transitions are required to execute the algorithm on any test data. On the other hand, the baselines and the models from the ablation experiments do not force the node and edge features to be from the fixed sets and usually encounter OOD values when given the larger test graphs. Performance on such data can vary across neural networks initializations and other randomness from the training process.

Note that while we report the perfect scores for the covered tasks, we cannot theoretically guarantee that the training will converge to the correct model for any initialization/training data distribution.

Could you please report variance on results in table 1?

Standard deviation of the node-level scores for the baselines across 5 different seeds (test size is 160 nodes) (we will add the full version of this table to the paper):

Could you please explain the differences in computed discrete states used for SALSA-CLRS and CLRS-30 tasks as well as how the model performs on out-of-distribution data?

Discrete states are the same between the benchmarks as SALSA-CLRS uses the CLRS-30 hints for the problems in CLRS-30 (BFS, DFS, Dijkstra, MST). The main difference is the training and test data: as described in Sections 4.1 and 4.2, the test data for SALSA-CLRS contains more sparse graphs.

As described in Section 5, the performance of the proposed model does not depend on the size and distribution of the test data (please see the first part of the general response for more details).

On the other hand, the performance of the baselines may vary depending on particular test data distributions, and achieving a good generalization for every target distribution is considered challenging in prior work (Georgiev et al., 2023).

Thank you for your response! Let us reply to the raised concerns.

I would recommend further complementing work with more detailed theoretical portion, particularly in the area of training convergence as it is mentioned that this is challenging to guarantee

Let us highlight the main part of our contribution: we consider the proposed model in its current form as a potential answer on where perfect generalization might come from. The proposed architecture allows us to fully interpret the trained model independently from the graph sizes/distributions and, if learned correctly, the model can perfectly imitate algorithms for any test data. To the best of our knowledge, none of the models in the field possesses this property (neither theoretical guarantees of training convergence).

I would also recommend searching for potential counter-examples (perhaps beyond used benchmarks) which could help test limits of the proposed direction

As we mentioned in the paper, one important limitation of our work is the reduced expressiveness of the proposed architecture. For example, the hard-attention mechanism cannot compute averaging over neighbors in a single message-passing step. In this sense, finding potential counter-examples is trivial. On the other hand, we also mentioned that there are several ways to improve expressiveness without losing the generalization and, importantly, to improve expressiveness by reducing the generalization. Thus, we believe that the discrete building blocks of the proposed model might arise in different forms as building blocks in a more general architecture.

For example, our ablation experiments (Appendix A) demonstrate that several components (e.g. hard attention) can be removed to relax strong generalization guarantees while yielding good performance. Also, we conduct an additional experiment to investigate if the proposed ScalarUpdate module can be extended to support more complex continuous manipulations with scalars (please see Appendix C of the paper). While learning to find correct manipulations from the discrete set might be challenging for particular data distributions or updates sets (e.g., such decomposition might not be unique), we find our results positive and consider the proposed discretization as a promising way to overcome the difficulties of value generalization in an interpretable way.

We hope that our response addresses your concerns.

The constraints improve generalization but reduce expressiveness. Is there a way to balance this trade-off by selectively relaxing some constraints?

While there are several ways to improve expressiveness without losing the generalization at all, we also can improve expressiveness by reducing the generalization:

1.Removing hard attention: as we demonstrate in our ablation experiments, using regular attention instead of hard attention yields perfect test scores for the BFS problem, but it is possible to construct adversarial examples with large neighborhood sizes where performance drops. While for more complex attention patterns (besides strictly attending to the single node) the OOD performance might be less robust, the expressivity gain is significant.
2. Removing feature discretization, but updating scalars with discrete operations: as shown by prior work and our ablation experiments, learning precise continuous manipulations is non-trivial and small inaccuracies in such manipulations can significantly affect the overall performance of the NARs. Thus, we can use the proposed separation between the discrete and continuous data flows and do not discretize node/edge features at all (and use discretization only in the ScalarUpdater). However, for non-attention-based models one needs to come up with how scalars will affect the discrete flow.

Can the proposed separation between discrete and continuous data flows be applied to other neural network architectures beyond attention-based models? What challenges might arise in such adaptations?

The general idea is to update scalars with predefined operations and let them affect the other flow only in a discrete manner. The simplest implementation of the latter is to update node/edge features depending on simple comparison/selections of scalars, as we did in the paper. When not aiming for the generalization guarantees, we can simply use such discrete updates of the features in the discrete flow and use any GNN as a processor. The node/edge features in the resulting model will still be discrete (e.g., could express complex substructures of the graph).

We guess that the main challenge arises when we try to answer the question of where perfect generalization might come from and how different GNN models perform on their limits of generalization. E.g., for attention-based models, this leads to the notion of attention weights annealing. We guess that such challenges are architecture-specific. For example, we guess that the potential analysis of the MPNN with MAX aggregation (and predefined fixed sets for node/edge features) is somewhat similar to the one of the hard-attention model (as it is invariant to the exact number of neighbors of each state, only presence of each state is important), but analyzing the possibility of perfect generalization with another aggregation functions might be more challenging.

We thank the reviewer for thoughtful questions and we are happy to discuss further.

We thank the reviewer for the constructive review and positive feedback! We address the questions below.

How does the proposed method handle algorithms that require more complex continuous manipulations or aggregate functions beyond simple increments and selections? Can the model be extended to support such algorithms without losing the benefits of discretization and interpretability?

First, we note that simple manipulations with scalars cover a significant part of the classical algorithms. In our work, we use the minimum set of the required functions, but it can be directly extended by other functions such as pow/sin/cos/exp. However, this complicates the optimization problem of selecting the correct operations/operands from the operations results (e.g., such decomposition might not be unique). Importantly, as ScalarUpdater can be viewed as a separate module, we can separately check if it is possible to train ScalarUpdater with any given set of predefined manipulations for any problem only with MSE on the results. Please let us know if you are interested in a particular type of scalar manipulations in a particular algorithm, we are ready to conduct an additional experiment during the discussion period.

The proposed selection with hard attention indeed significantly reduces expressivity. The simplest way to support more complex aggregations in a single processor step is to use top-k attention (with a fixed k). Also, as we mention in the paper, the model can be extended with additional architectural modifications, such as edge-based reasoning. For example, edge-based attention with top-2 attention (where each edge chooses the top 2 adjacent edges to receive the message from) can implement the sort of triplet reasoning.

Are there potential strategies to improve training without hint supervision? How does the model perform on tasks without hints when compared to continuous models?

We perform a hyperparameter search only for the BFS problem, achieving up to 79 node-level accuracy when tested on the graphs with 64 nodes, please see the second part of the general response for more details. Without hyperparameter search, discrete models are only slightly improved over the untrained models.

As we demonstrate in our no-hint experiment, the proposed discrete model is potentially capable of capturing the desired dynamic of the algorithm (e.g., using the concept of distance for the BFS task) when its architecture is well aligned with the underlying problem. However, the direct regularization of the simplicity of the resulting model (such as reducing the states count) performs poorly as it significantly changes the optimization problem. We guess that some sort of iterative self-distillation might be useful in this setting for BFS (and similar problems):

1. Train a no-hint model with good validation performance (similar to our experiments): such a model can capture the desired (but not the simplest) dynamics of the problem;
2. Train a new model with less states count supervised only on the intermediate updates of the outputs of the model from the first step (e.g., predicted pointers after each step).

However, for sequential problems, such as DFS, obtaining a good model for stage 1 can be difficult. One possible way to overcome this limitation is to design a curriculum learning setup.

Another intermediate approach (similar to step 2) is to supervise the model only on a part of the hints (e.g., pointers) of the ground truth algorithm.

Thank you for being involved in the discussion!

We conduct an additional experiment to investigate if the proposed ScalarUpdate module can be extended to support more complex continuous manipulations with scalars (please see Appendix C of the paper). While learning to find correct manipulations from the discrete set might be challenging for particular data distributions or updates sets, we find our results positive and consider the proposed discretization as a promising way to overcome the difficulties of value generalization in an interpretable way.

We are open to further discussion!