What Makes a Good Feedforward Computational Graph?

Dear Reviewer QmmN,

Thank you for your careful review! We hope to provide useful clarifications, and that you may reconsider the relevance of our work:

On feedforward networks

Your comments focus on our work not easily representing multi-layer feedforward NNs.

For us, the term “feedforward graph” is defined more narrowly (cf. Section 3): we study feedforward computation over a data axis (ie. we assume the input nodes have a sequential ordering to them) rather than the layer axis! It is hence a spatial constraint rather than one pertaining to multilayer architectures.

Prior work on analysing data-axis feedforward graphs is rare: there was no established theory even for single-layer models over such inputs. The entirety of our paper aims to fill this gap, and study how to integrate information across the data axis with a single graph. This corresponds to studying multi-layer GNNs which use the same feedforward graph in every layer. It was never our intention to study arbitrary multi-layer signal propagation in this paper.

This was sufficient to fill the page limit and pose important directions for future work. Progressing into the general case (where graphs may differ across layers) is a natural next step, but it induces a combinatorial explosion of design choices. We wanted to understand the basics well before inviting such analyses.

This gap risks misleading readers when considering new architecture development.

We recognise that, due to the established meaning of “feedforward”, our wording might pose a risk of misleading. We are happy to commit to changing our title to be more precise, eg.: “What makes a good one-step feedforward computational graph?”

We also commit to detailedly discussing how our setup differs from general multilayer NNs.

On metrics

Our notion of mixing time measures how quickly salient information (ie. input features) can mix within the model. That is, how many feedforward layers (for a given graph) are necessary for a certain level of mixing of salient data to occur.

ResNets or GPT architectures have ±96 layers to enable the learning of "strong" hidden representations without directly propagating information to output neurons.

This point is irrelevant to our metric, as intermediate neurons in a multi-layer network do not hold salient information at the beginning. If every neuron had salient data at the start, then our mixing time would be particularly high; but in the common scenario, our metric’s one-step approach wouldn’t measure mixing in the output neurons only, but at any intermediate point.

On minimax fidelity:

Modern networks routinely prune away large portions of weights or nodes, indicating that focusing on relatively few “important” nodes/edges is advantageous.

The problem is that modern NNs over feedforward data (eg. GPT) are topologically biased towards earlier nodes (proved in Barbero et al., NeurIPS’24), and hence not equally well-prepared when important data is not at the start of the input. This is precisely what we try to quantify with minimax fidelity: “which node is in the worst possible position by this choice of graph structure”?

Fidelity focuses on averaging (rather than pruning) because the mechanisms current models have for pruning, such as the softmax function, will provably fail to prune at OOD problem sizes (proved in Veličković et al., 2024).

In particular, given that the nodes are an ordered set, are nodes allowed to have the same order?

While not explicitly allowed, our framework supports it: simply choose a random order for equal nodes. One component of the FS graph is a block-wise sequence of bipartite communications of this kind.

Further points

Related work [1, 2]

We appreciate these relevant papers and will be sure to discuss them in our revision. While [1] generalises oversmoothing to directed graphs, they are not “feedforward”; backwards edges are allowed, making a spectral approach far easier.

The datasets do not provide meaningful insights into what constitutes a “good” computational graph

We designed our tasks to allow fine-grained control over which (& how many!) nodes are relevant for computing the answer, allowing us to validate our theory’s prescriptions in an unbiased manner.

That said, we agree evaluating real-world setups would make the work more valuable. We have now trained Gemma 2B models – varying the graph used as the attention mask – on the standard Wiki dataset (tensorflow.org/datasets/catalog/wikipedia) containing texts from Wikipedia.

After 3,000 batches of training, our model achieves the following perplexities:

This trend is consistent throughout training, demonstrating the sparse FS graph is competitive with the full feedforward graph, even when nodes are natural language tokens.

Dear Reviewer y3Tf,

We are delighted that you have found our foundations to be strong and our results to be interesting. We hope that our responses will strengthen your view of our contributions even further!

However, Theorem 6.1's proof relies on strong assumptions

Thank you for raising this!

We have now improved the proof of Theorem 6.1 so that all simplifying assumptions have been removed. This turned out to be fairly straightforward, and mostly required the judicious use of ceiling and floor functions; eg. the correct value of the number of levels, $D$, is now:

$D = \left \lfloor \frac{\log n}{\log \lceil \log n \rceil} \right \rfloor,$

but the key arguments of our random walk analysis are unchanged.

The proof is unwieldy to paste here but we will of course update it in our revision, and we are happy to discuss further aspects if you would like to know more!

The experiments use synthetic tasks that lack real-world complexity

We designed our tasks to allow fine-grained control over which (and how many!) nodes are relevant for computing the answer; this allows us to validate our theory’s prescriptions in the most unbiased manner.

That said, we agree that evaluating on real world data would make the work more valuable. We have now trained Gemma 2B models – varying the graph used as the attention mask – on the standard Wiki dataset (https://www.tensorflow.org/datasets/catalog/wikipedia) containing natural language texts scraped from Wikipedia.

After 3,000 batches of training, our model achieves the following perplexities:

This trend is consistent throughout training. This demonstrates that the sparse FS graph can remain competitive with the full feedforward graph, even when nodes are natural language tokens.

Providing more detailed information on the experimental settings

Thank you; we will include more details. We used default hyperparameters in most places, and tuned the order-of-magnitude of the learning rate and the weight decay coefficient using the fully-connected graph experiments only (then reused the tuned parameters everywhere else).

The lack of access to the code as part of the supplementary material limits the ability to verify the experimental results.

We appreciate your remark, and thank you for acknowledging that the paper’s findings should be easier to reproduce since tasks are synthetic. We understand the positive impact of open-sourcing, and would aim to release our code for the synthetic tasks upon acceptance.

Including discussions of random walks on directed acyclic graphs, or connections to causal inference frameworks, could further highlight the novelty of extending these analyses to feedforward graphs.

We agree that these points are useful, and will add further discussions! To give a concrete example, since our original submission, we managed to make our statement in Section 4.4. (mixing time is ‘small’ iff there are ‘lots’ of paths from ‘most’ vertices to $\tau$) precise:

Proposition: Suppose that the outdegree of every vertex other than $\tau$ is at least $2$. Let $t$ be the average mixing time. Then for some $s \leq t$, the average number of paths from vertex $i$ to $\tau$ with length $s$ is at least $(3/4t)2^s$, where the average is taken over all vertices $i$ between $0$ and $n-1$.

We will include this in the next revision – it provides a relevant connection between our metrics and paths on a DAG and its proof is only a few paragraphs long.

Related works ([1, 2])

Thank you for bringing these to our attention. They are highly relevant to discuss and we will incorporate them. The key difference is that both Big Bird and GRASS optimise for bidirectional graphs, where back edges are allowed.

Rigorous experiments and ablation studies, e.g., add/remove self-edges, experimenting with various in-degrees

This is an excellent suggestion! We have now performed both of these ablations, and found that removing self-edges led to significant performance regressions on all tasks – especially on the Parity task where multiple elements need to interact meaningfully for the final answer computation. We also evaluated our graph generators at different levels of in-degree OOM: $O(1)$, $O(\log n)$ and $O(\sqrt{n})$. We find that, as the orders are increased, the models get more performant in-distribution, without a significant effect out-of-distribution. We will incorporate both of these analyses in our revision!

The submission would benefit from a discussion on limitations and future research directions.

We’ll explicitly include the directions you discussed (integration with Transformer variants, dynamically adapting the computation graph) and we’ll also suggest additional ones (connecting the mixing and fidelity metrics and finding further connections between them and spectrally-motivated analyses!).

Dear Reviewer GmC2,

We are very thankful for your kind review and recognising the strengths of our work!

To address your questions:

The paper discusses the asymptotic behaviour of various graphs. why certain structures perform better from an ML perspective.

This is an excellent question. We believe that we can relate at least some of our metrics to established approaches in graph machine learning that study which graphs will have better information propagation from a machine learning perspective.

To name a few examples:

• High mixing time implies that it takes a long time for information to travel to the final vertex (in the worst case – when there is more than one sink – mixing time is infinite). If there are insufficiently many layers in a model to cover these long paths, this directly relates to the well-known under-reaching phenomenon in graph neural networks (Barceló et al., ICLR’20).
• Low fidelity for a particular node implies that this node’s information is not reaching the final node in a way that is not “drowned out” by others, especially as more model layers are added – this effect is well-studied in the GNN literature as the over-smoothing effect, and was recently also studied in the domain of Transformers as well (Barbero et al., NeurIPS’24).
• High fidelity in certain nodes, but low minimax fidelity, implies that, while some nodes are able to sharply send their information into the final node, this is done at the expense of other nodes, which fail to arrive at the sink quickly enough to avoid being overwhelmed. This kind of biased communication is related to the over-squashing effect, which Barbero et al. (NeurIPS’24) have also studied from the perspective of Transformers.

Since our original submission, we have also managed to make our statement in Section 4.4. (that mixing time is ‘small’ if and only if there are ‘lots’ of paths from ‘most’ vertices to $\tau$) precise, as follows:

Proposition: Suppose that the outdegree of every vertex other than $\tau$ is at least $2$. Let $t$ be the average mixing time. Then for some $s \leq t$, the average number of paths from vertex $i$ to $\tau$ with length $s$ is at least $(3/4t)2^s$, where the average is taken over all vertices $i$ between $0$ and $n-1$.

The proof is only a couple of paragraphs long, and it uses the intuition given in Section 4.1. More precisely, $(\mathbf{W}^t)_{\tau i}$ represents the probability of a random walk starting at $i$ and reaching $\tau$ by time $t$. This is equal to the sum, over all $s \leq t$, of the probability that the random walk first reaches $\tau$ at time $s$. For any such $s$, this probability is equal to the sum, over all paths from $i$ to $\tau$ with length $s$ avoiding the loop based at $\tau$, of the probability of taking that path. This latter probability is at most $1/ 2^{s}$ because at each step along the path there were at least two options that could have been taken. So, for some $s \leq t$, the average number of paths from vertex $i$ to $\tau$ with length $s$ is at least $(3/4t)2^s$.

We intend to include this new result in the appendices.

The paper presents empirical results on tasks like maximum retrieval and parity. However, the tasks seem relatively simple. Could the authors expand their evaluation to more complex tasks, such as those involving real-world datasets (eg. natural language processing or graph-based tasks)?

Thank you for your comment!

We have designed our tasks to allow us fine-grained control over which (and how many!) nodes are relevant for computing the final answer, as this would allow us to validate our theory’s prescriptions in the most unbiased manner.

That being said, we fully agree that evaluating our proposal on real world datasets would make the work more valuable. To this end, we have now trained Gemma 2B language models – varying the graph used as the attention mask – on the standard Wiki dataset (https://www.tensorflow.org/datasets/catalog/wikipedia) containing natural language texts scraped from Wikipedia.

After 3,000 batches of training, our model achieves the following perplexities:

This trend is consistent throughout training. This demonstrates that the sparse FS graph can remain competitive with the full feedforward graph, even when nodes are natural language tokens.