In AI, the dominant approach to develop neural nets is to first pretrain a large model on a big general dataset and then fine-tune it on specific tasks. In vision, there are, for example, Vision Transformers (ViTs) pretrained on ImageNet and fine-tuned on downstream tasks (e.g., classification of satellite images, textures, etc.). In language modeling, there are, for example, large language models (LLMs) specialized in Math, Coding or specific languages (e.g., French). Even though the original pretrained model can also perform well on many tasks, fine-tuning remains essential to achieve optimal performanceTraining-free adaptation methods such as in-context learning, prompt tuning and others often show competitive results on some tasks and are less costly. However, direct fine-tuning of model parameters is often more effective and applicable to a wider range of tasks, especially when a lot of domain knowledge is necessary requiring a lot of parameter updates (e.g., some medical tasks or rare languages).. As a result, there are hundreds or thousands of models fine-tuned for various tasks.

As AI tackles increasingly more tasks, fine-tuning costs and storage requirements add up. Consider a situation when we first collected a lot of Math data and fine-tuned an LLM on it, then some other lab released a model excelling at Code and another lab released a model excelling at French. First of all, we now have tripled the number of models to store and maintain. Secondly, despite having three models, we cannot effectively solve mixed tasks such as a Math task in French, since there may be no single model trained on such a mixture. To solve the issues in this example, we can collect a desired mix of data and fine-tune the model on it, but this is costly and time-consuming especially if we need to change the mixing ratio. We can also build some smart ensemble with routing conditioned on the input, but doing this effectively is not trivial and still requires some training of the router and storing an entire ensemble.

Model merging is a recent training-free approach to combine multiple models into one. As many papers and our experiments show, model merging can mitigate limitations of naive fine-tuning in a simple yet effective way. Surprisingly, the resulting merged models perform well on individual tasks and on the mixed tasks they often approach fine-tuning performance, but without any training or data collection.

In this post, we introduce a new model merging method called MetaMerge.

MetaMerge uses a pretrained graph neural network (GNN), that takes weights of multiple models as input and produces the weights for a single merged model as output without any training or finetuning.

Before introducing MetaMerge and why it is called “Meta”, let us briefly describe a common baseline approach and several key concepts.

TL;DR

• We introduce MetaMerge, a new model merging method that uses a pretrained graph neural network (GNN).
• The “Meta” in MetaMerge comes from the fact that both input and output of the pretrained GNN are model weights.
• MetaMerge does not require any training/finetuning.
• MetaMerge can be competitive to weight averaging in some cases even though the metamodel (GNN) was not trained for merging, and more critically, the specific model we use for MetaMerge only observed tiny GPT2 style transformers (with <=1.6M parameters) during training.

Simple Merge

In general, model merging can be defined as some aggregation function $f$ that takes weights of $N$ models as input and produces the weights of a single merged model as output:

where \(\mathbf{W}_1, \mathbf{W}_2, \ldots, \mathbf{W}_N\) are the weights of the individual models, and \(\mathbf{W}_{\text{merged}}\) are the weights of the merged model.

Most of the merging methods implement $f$ directly in the weight space. For example, the simplest way to merge models is to average their weights:

More advanced methods include task arithmetic, ties-merge, and othersTools like https://github.com/arcee-ai/mergekit implement many merging methods.. The averaging method works well only when the models are trained from the same initialization, usually from the same pretrained model. And if they are trained for too long, the performance of the merged model can degrade significantly. In case of different initializations or fine-tuning for too long, merging is still possible by using neuron alignment and other techniques, but this is out of scope of this post.

Meta Networks

The key concept of MetaMerge is Meta Networks. A high level idea of Meta Networks is to have a learnable model that can take weights of other models as input to produce some representation or modified weights as output. Hence, the term “Meta” is used. While most merging methods implement $f$ in the weight space directly, some recent methods implement $f$ in the low-dimensional space produced by singular value decomposition (SVD). MetaMerge can be viewed as a more general weight-space transformation learned by a meta-network.

NiNo

Among various meta networks, we focus on NiNo. To the best of our knowledge, a pretrained NiNo checkpoint is the only meta network that can potentially be used for merging models in a realistic setup out-of-the-box, since it can both encode and decode models and it supports diverse architectures, including modern transformers. NiNo is a meta network that was trained to accelerate training. It takes past $c$ checkpoints as input and predicts future (at step $t+K$) parameters as:

This step is applied every 1,000 steps, while Adam or SGD is used for the rest of the steps. So, the original NiNo that we are going to use is not trained for any model merging objective, but for accelerating training. However, it does not mean it cannot be used for merging as we show.

NiNo for Merging

The only challenge when adapting NiNo to the merging setup is the preparation of the input. Specifically, in the equation above NiNo expects $c$ input models in the order corresponding to an optimization trajectory (where $c=5$ in pretrained NiNo models). One straightforward way to address this issue is to create the “trajectory” as [pretrained model, model fine-tuned for task 1, model fine-tuned for task 2, ...]. Since tasks do not have a particular order and there may be fewer/more than $c$ tasks, we need certain heuristics. Specifically, for the setup with two tasks, we found the following heuristic to work well in practice:

For the setup with four tasks, the equation is more straightforward since we have exactly 5 models to feed into NiNo:

In principle, the order of the tasks can be optimized, but we fixed it after some trial and error.

Experiments

We evaluate MetaMerge on two vision and two language setups. For vision, we use ViT-B-16 pretrained on ImageNet and fine-tuned on the DTD, RESISC45, MNIST and SVHN datasets. This is a subset from standard 8 tasks used in model merging papers. For language, we created a similar setup by fine-tuning Qwen3-0.6 and Qwen3-1.7B models on the GSM8K and French datasetsWe also released Qwen3-4B models fine-tuned on these datasets and models fine-tuned on the code dataset..

ViT

For ViT-B-16, we use Task Arithmetic code and their checkpoints for evaluation, which contains trained classification heads for all the task, so merging is done only for the backboneCombining classification heads with different backbones in the Task Vectors code allows for flexible evaluation (e.g., a model fine-tuned on task B can be evaluated on task A by attaching a respective head).. We run experiments for 2 and 4 tasks. “N models” in Tables denote the number of models required to obtain the results in a given row. For example, “Finetuned (per task)” requires a fine-tuned model for each task, which is often considered as an upper bound for merging methods. The code to merge ViTs using MetaMerge is provided at merge_vit.py.

Table 1. Results using ViT-B-16 on 2 tasks.

Table 2. Results using ViT-B-16 on 4 tasks.

Surprisingly, on 2 tasks (Table 1) the ablated NiNo variant (without the GNN part) actually outperforms the full NiNo even though the full one excelled at accelerating training as shown in the NiNo paper’s ablations. One logical explanation is that merging is quite different from predicting future parameters, so good results on merging are not expected in the first place. Also, the MLP version of NiNo has a simplicity inductive bias so its predictions can be more generic (kind of trend prediction which may be closer to weight averaging) and less overfitted to the training objective of accelerating training. On 4 tasks (Table 2), the ablated variant performs much worse. In future work, it would be interesting to push NiNo’s performance specifically for merging, however, defining a proper objective for that is not trivial.

Below we visualize detailed results from Table 2 for all the tasks and models.

Qwen3

In the language setup, we first fine-tuned base Qwen3 models on the train split of GSM8K and the subset of luth-sft. These are Math and French datasets, respectively. We chose to fine-tune our own models instead of using existing checkpoints to have a more controlled consistent setup for different model sizesOur fine-tuning is likely to be far from compute and performance efficient, but our goal is not to provide SOTA models.. Then we evaluate on the test split of GSM8K, FrenchBench and GSM8K-Fr (GSM8K translated to French). See details about the evaluation pipeline and model checkpoints at 🤗SamsungSAILMontreal/qwen3-small. We report Qwen3-0.6B and Qwen3-1.7B results in Tables 3 and 4, respectively. However, we also have Qwen3-4B results and checkpoints on the provided link.

In Qwen3 experiments, we do not have a model fine-tuned on GSM8K-Fr to showcase a common scenario when a training dataset for the mixed task is unavailable, which is prevailing for rare mixes. Therefore, Qwen3 (fine-tuned per task) only contains two models. Models with “Math-Fr” in their name are obtained by merging the Math and French models using either simple averaging or MetaMerge. The code to merge Qwen models using MetaMerge is provided at merge_qwen.py.

Table 3. Results using Qwen3-0.6B on 3 tasks.

Table 4. Results using Qwen3-1.7B on 3 tasks.

For Qwen3-0.6B, the simple averaging method outperforms MetaMerge and as in ViT (on 2 tasks), the MLP version of NiNo outperforms the full NiNo. One possible explanation is that the pretrained NiNo has only seen tiny GPT2 style transformers (with <=1.6M parameters) during training. So making predictions for a Qwen architecture with 600M parameters is an extreme out-of-distribution scenario for NiNo. As in ViT experiments, the MLP version of NiNo performs slightly better than the full model potentially due to its simplicity inductive bias. Qwen3-1.7B is even a more severe out-of-distribution scenario for NiNo, so it is not surprising that MetaMerge does not perform as well as simple averaging. Despite the average benchmark scores being worse or comparable, an interesting question for future studies is whether MetaMerge produces functionally different models than simple averaging or the ablated NiNo variant.

Visualization of MetaMerge

In this demo, we compare weight averaging to MetaMerge on ViT using 4 tasks. We sampled 100 weight positions per layer (same positions for all models) to enable efficient visualization. For each position, parameters are ordered in the way we pass them to NiNo (i.e., pretrained, dtd, resisc45, mnist, svhn). Since NiNo was trained to predict future parameters, we show a MetaMerge prediction as the last point on the trajectory. The weight averaging baseline is shown as a horizontal line for comparison.

MetaMerge is used for all layers in the visual backbone (starting with “visual.”) except for conv1 as explained in our merge_vit.py code and layers not starting with “visual.”. For those we use weight averaging, so in the visualization the “meta-merge” point is the same as the “average” in such cases. Overall, the visualization shows that there is no simple relationship between the input models and the MetaMerge prediction. But the visualization may give some insights on problematic behavior (e.g., for biases the MetaMerge predictions tend to be too far from the overall trajectory).

Conclusion

• MetaMerge shows competitive performance to weight averaging in some experiments even though the metamodel (GNN) was not trained for merging.
• In some cases, MetaMerge underperforms like due to the distribution shifts and a different objective. So we believe there is a lot of potential for improving MetaMerge by training a metamodel with an objective more aligned with merging.
• Besides introducing MetaMerge, we also released fine-tuned and merged Qwen3 models obtained in a controllable way that could be used for merging experiments.

License

Diagrams and text are licensed under Creative Commons Attribution CC-BY 4.0, unless noted otherwise.

Citation

@inproceedings{camacho2026meta,
  title={Meta-Merging by Checkpoint Nowcasting},
  author={Camacho, Albert Manuel Orozco and Knyazev, Boris and Belilovsky, Eugene and Wolf, Guy},
  booktitle={ICML 2026 Workshop on Weight-Space Symmetries: from Foundations to Practical Applications},
  year={2026},
  url={https://openreview.net/forum?id=whl40fLEnM}
}

When you deploy a very large model, the first challenge you will face would likely be the memory demand. For example, one of the largest and performant Qwen3 MoE models (Qwen3-235B-A22B) requires around 235B params $\times$ 2 bytes/param = 470GB of VRAM (e.g., 6$\times$GPUs each having 80GB of VRAM such as NVIDIA H100) to merely load all its parameters on GPUs in 16 bit precision (see the sidenote on the right). Traditional approaches to deal with this challenge include pruning and quantization (and of course getting more GPUs 💰 💰 💰).

I explore an alternative approach specific to MoE, which is reducing the number of experts by merging groups of experts, which I call REAM.

REAM stands for Router-weighted Expert Activation Merging motivated by the Router-weighted Expert Activation Pruning (REAP) method.

The core idea is that merging can be more effective than pruning as it aims to preserve functions by leveraging all experts instead of discarding some of them.

But before going into details, let me first describe the MoE architecture and its benefits. Then, I will present my results of merging experts in Qwen3-30B-A3B, Qwen3-235B-A22B and Qwen3-Next-80B-A3BTheir Instruct versions., reducing the number of experts by 25%From 128 to 96 for Qwen3-30B-A3B and Qwen3-235B-A22B, and from 512 to 384 for Qwen3-Next., thereby reducing memory demands by around 25%Since most (>98% in Qwen3) of parameters belong to MoE layers, reducing the number of experts is roughly equivalent to reducing the number of parameters (and hence, memory) in the overall model., while maintaining high performance across multiple benchmarks.

TL;DR

• I propose REAM, a retraining-free expert merging method for MoE LLMs.
• REAM is a sequential merging algorithm using REAP-based saliency.
• It reduces experts by 25% with strong results on Qwen3 MoE LLMs.
• REAM outperforms REAP and HC-SMoE on multiple tasks.
• Compressed Qwen3 models are released on huggingface🤗.

Introducing MoE

Mixture-of-Experts (MoE) is an old idea developed in 1989-1991 (e.g., see Robert A. Jacobs et al.). MoEs were later proved useful in deep learning, including for language models in 2017 by Noam Shazeer et al., then scaled to 1T parameter models in 2022 by William Fedus et al..

As illustrated in the Figure above and formally defined below, MoE is a special neural network layer consisting of multiple expert networks ($E$)I drop the layer index for simplicity as the same equation is used for all layers (with trainable parameters in all layers being different unless shared). and a gating network ($g$) that selects which experts to use for $d$-dimensional input $\mathbf{x} \in \mathbb{R}^{n \times d}$ with $n$ tokens:

where each $E_i$ has the same architecture as an FFNFFN stands for "feed-forward network", which is usually implemented as a simple MLP or a Gated Linear Unit (GLU) in Transformers. layer in Transformers. The gating network is typically implemented asThe order of the Softmax and TopKMask operations can vary in implementations.:

where $\mathbf{W}_g \in \mathbb{R}^{N \times d}$ are the gating weights, and TopKMask is an operation that retains only the TopK values per row (i.e. per token) and sets the rest to zero. Softmax is applied along each row to produce a probability distribution over the $N$ experts.

Expert Network Architecture

Each expert network $E_i$ is an FFN layer that processes the input $\mathbf{x}$ independently of other experts. Recent FFN layers in both original (non-MoE or “dense”) and MoE Transformers are based on Gated Linear Units (GLU) instead of simple MLPs:

where act is SiLU/Swish or GELU, $\odot$ is element-wise multiplication; \(\mathbf{W}_{i,\text{gate}}\) and \(\mathbf{W}_{i,\text{up}}\)
project \(\mathbf{x}\) to a lower dimensional space (e.g., 2048 to 768 in Qwen3), \(\mathbf{W}_{i,\text{down}}\) projects it back to the high dimensional space.

Key Concepts of MoE

Having defined an MoE layer above, let’s highlight the following key MoE concepts to better understand expert pruning and merging techniques:

1. MoE layers replace regular FFN (MLP) layers in Transformers (see the first Figure on the right). So for example in a Qwen3 model with 96 Transformer layers, there are 96 MoE layers (that are interleaved with self-attention layers within each Transformer layer).
2. Each MoE layer has $N$ experts (e.g., 128 in Qwen3 MoE models) with a Gated Linear Unit architecture described above. Having many FFNs massively increases the total number of parameters compared to non-MoE (dense) models (with $N=1$).
3. The gating network $g$ selects TopK experts (e.g., 8 in Qwen3 MoE models) for each input token based on its representation (row in $\mathbf{x}$) in the current Transformer layer. So for example, some experts in early layers learn to specialize on punctuation or number tokens, while in later layers they can specialize on more abstract concepts.
4. During inference, only the selected experts are used in the forward pass, making MoEs efficient (in terms of FLOPs) despite their large number of parameters. However, to make inference efficient in practice, expert weights need to be stored on GPUs to avoid data transfer overheads creating the memory demand.

Further technical details of MoE layers can be found in papers and a more animated and educational explanation in Maarten Grootendorst’s blog post.

Strengths of MoE

Although dense Transformers can learn from diverse data well, MoEs more explicitly promote distributed specialized knowledge. This is achieved by having multiple experts that can each specialize in different aspects of the data with the help of the gating mechanism that promotes sparsity.

It has been also shown that MoEs can scale well given a fixed FLOPs budget at inference (see the Figures below). Specifically, Tay et al. in 2023 showed that MoEs (or Switch Transformers back then) were one of the few architectures that demonstrated strong scaling trends in both upstream (pretraining) and downstream (fine-tuning) cases.

As figures above show, Switch Transformers still underperformed dense Transformers at the largest scale. However, the design choices and training of MoEs have improved since then (e.g., see), leading to even better scaling trends in recent MoE LLMs. Moreover, recently Kimi MoE LLMs demonstrated strong scaling trends with increasing MoE sparsity (see the Figure below). They define sparsity as the total number of experts ($N$) divided by TopK showing that for a fixed TopK increasing $N$ and, hence, sparsity consistently improves results.

These findings explain increasing the number of experts to 384 in Kimi K2 compared to 256 in DeepSeek and 128 in Qwen3 (more recent Qwen3-Next-80B-A3B also increased $N$ to 512). Even though the total number of parameter is soared to 1T, the inference remains efficient (with only TopK=8 experts activated). However, the memory demand to store all expert weights on GPUs becomes a bottleneck for deploying such large MoE LLMs.

Background: MoE Compression

Let me briefly describe some key expert pruning and merging techniques below before presenting my approach and results. I follow an assumption standard in the MoE compression literature that there is a pretrained MoE model with $N$ experts per MoE layer and the goal is to reduce it to $k < N$ experts. I do not consider retraining/fine-tuning (or any further compression) after pruning/merging experts, which could be a complementary step.

Pruning Experts

Given hundreds of experts in each MoE layer, it is tempting to assume there is some redundancy in them and aim to reduce their number without sacrificing performance too much. One straightforward approach to do so is to prune some experts based on their activation frequencies. These frequencies are computed based on gate logits $g(\mathbf{x})$ by counting how often each of the $N$ experts is among TopK experts. For that purpose, we can forward pass some calibration data through the MoE layer and compute the activation frequencies (or saliency scores, $A$) for each expert. Given $A$, we keep only $ k < N $ experts with the highest scores.

The following is my pseudocode of computing activation frequencies $A$.

# Pseudocode for computing activation frequencies
N = 128  # number of experts
n = 1024  # number of tokens in calibration data
topk = 8  # TopK experts activated per token (based on the model configuration)

g = torch.randn(n, N)  # assume router logits for n tokens and N experts after some layer
g = torch.softmax(g, dim=-1) # apply softmax according to Equation above
topk_indices = torch.topk(g, k=topk, dim=-1).indices  # (n, topk)
# count how often each expert is selected
A = torch.zeros(N)
for i in range(N):
    A[i] = (topk_indices == i).sum().item()
# keep k < N experts with highest A

Recent Router-weighted Expert Activation Pruning (REAP) method significantly improves simple frequency counting. REAP defines $i$-th expert’s importance ($S_i$) based on the expert output’s norm weighted by the gate output:

where \(\mathbf{x}^{(i)}\) is the subset of \(\mathbf{x}\) corresponding to the tokens that activated expert $ i $. Compared to activation frequencies ($A$), scores $S$ measure more accurately the contribution of each expert to the final output of the MoE layer.

The following is my pseudocode of computing REAP scores based on the equation aboveMy implementation may differ from the official REAP code https://github.com/CerebrasResearch/reap. I implement a simple version aligned with the equation..

# Pseudocode for computing REAP scores
N = 128  # number of experts
n = 1024  # number of tokens in calibration data
topk = 8  # TopK experts activated per token (based on the model configuration)
d = 2048 # expert output dimension

expert_activations = torch.randn(N, n, d)  # assume expert outputs for n tokens and N experts after some layer

g = torch.randn(n, N)  # assume router logits for n tokens and N experts after some layer
g = torch.softmax(g, dim=-1) # apply softmax according to Equation above
topk_values, topk_indices = torch.topk(g, k=topk, dim=-1)  # (n, topk)

# compute REAP score for each expert
S = torch.zeros(N)
for i in range(N):
    # get tokens routed to expert i
    top_x, idx = torch.where(topk_indices == i)
    # top_x - indices of tokens activating expert i (values between 0 and n-1)
    # idx - indices of experts in topk (values between 0 and topk-1)
    if len(idx) == 0:
        continue  # no tokens activated expert i
    expert_state = expert_activations[i, top_x]  # (selected tokens, d)
    S[i] = (expert_state.norm(dim=-1) * topk_values[top_x, idx]).mean().item()
# keep k < N experts with highest S

Merging Experts

Compared to pruning experts, several works have explored merging them instead (e.g., ). The algorithms to merge experts typically have two steps applied to each MoE layer:

1. Grouping or clustering similar experts together based on some similarity metric.
2. Combining the weights of experts within each group to form a single expert.

Step 1 (grouping) is usually the key difference across the methods. For example, MC-SMoE finds $k$ cluster centroids first. To do so, the activation frequencies ($A$), defined above, are used as in pruning methods. However, in contrast to simply removing low-score experts as in pruning, MC-SMoE assigns them to the centroids based on the similarity of expert representations (e.g., expert weights or expert/gate logits given some calibration data). In HC-SMoE, the grouping step is based on hierarchical clustering of experts based on their outputs (also given some calibration data). HC-SMoE outperformed MC-SMoE, therefore I use HC-SMoE as a baseline in my experiments below.

Step 2 (combining) in both MC-SMoE and HC-SMoE involves weighted averaging of expert weights within each group, where the weighting coefficients are computed as in pruning methods (i.e., based on normalized activation frequencies $A$). In the next subsection, I describe an important nuance of this step.

Combining the Weights of Experts: Permutation Alignment

Computing the average of neural network weights can lead to poor results in certain casesIt is actually very surprising that averaging of weights can work at all in some cases, so papers like Model soups were very impactful.. Among such cases is when the networks are trained from different initializations and different data orderings/subsets. The difficulty of averaging weights in this case arises because of the different symmetries in neural network weights, primarily neuron permutation symmetry.

Neuron permutation symmetry states that the order of neurons in adjacent layers of a neural network can be permuted in certain ways without affecting the overall function of the network.

For example, in a 2 layer MLP, we can randomly permute rows of the first layer weights and permute in the same order the corresponding columns of the second layer weights without changing the function computed by the MLP.

Neuron permutation symmetry and other symmetries are studied in the Weight Space Learning (WSL) area of machine learning. One of the key findings in this area is that permutation alignment of the weights of one network w.r.t. another network prior to averaging can significantly improve the results. To model neuron permutation symmetry better, we proposed neural graphs and recently showed that we can train powerful graph neural networks based on neural graphs.

When training MoE LLMs, each expert network $E_i$ is initialized randomly and trained on different data subsets (because of the gating mechanism). To address the permutation issue, MC-SMoE and HC-SMoE align expert weights by either expert weights or hidden activations. This can be done, for example, using the scipy.optimize.linear_sum_assignment function that implements the Hungarian algorithm to solve the assignment problem. I show a pseudocode for expert merging with improved permutation alignment in the next section.

REAM

My approach is based on REAP with merging instead of pruning, hence it is called Router-weighted Expert Activation Merging or REAM. It can also be seen as a modification of MC-SMoE with REAP scores and additional tricks as described below.

Algorithm

For each MoE layer with $N$ experts originally (and $k$ experts as a result):

1. Compute saliency scores $S$ for each expert as in REAP.
2. Pick $k$ experts with the highest $S$ and label them as cluster centroids similarly to MC-SMoE.
3. Group experts, but in contrast to MC-SMoE’s grouping, REAM uses pseudo-pruningThis is called pseudo-pruning, because most of the scores $S$ of the non-centroid experts will be low, so the weighted average (in step 5) will be dominated by the centroid expert.. The idea behind pseudo-pruning is to have a few big clusters and many singletons with unchanged weights. The grouping starts with the centroid having the highest $S$ and assigning $C$ most similar experts to it (that are not assigned yet). This step is repeated for all centroids until all $N$ experts are assigned ($C=16$ in my experiments).
4. The expert similarity metric, used in the previous step, is improved compared to prior work based on two key tricks. First, REAM uses an average of cosine similarity of expert outputs and cosine similarity of gate logits. Second, gated similarity is computed by multiplying expert outputs by gate logits motivated by the REAP equation.
5. Merge the weights in each group accounting for permutation alignment of weights based on both hidden activations and weights of $E$ as shown in the pseudocode below.
6. Use the merged MoE layer to compute the inputs for the next MoE layer. In contrast, previous works perform merging/pruning at each layer given the original inputs (i.e., computed using the original uncompressed model) to that layer. It means they can do pruning/merging in an arbitrary order of layers since the inputs/outputs are precomputed based on the original model. In REAM, the process is inherently sequentialTo implement sequential merging, calibration data has to be propagated through each MoE layer twice: first time to get necessary outputs to perform merging, second time to get the inputs for the next layer using the merged experts..

The highlighted steps are ablated in the Experiments section below.

The following pseudocode illustrates step 5 of merging a group of experts with permutation alignment based on both hidden activations and weights.

# Pseudo-code for merging a group of experts with logits+weights permutation alignment
G = 16 #  number of experts in a group (cluster) to be merged found in steps 3-4
d = 768 # bottleneck dimension in the expert network 
d_model = 2048 # model dimension

expert_hidden = torch.randn(G, d, n)  # hidden activations of G experts for n tokens in a MoE layer
expert_weights = torch.randn(G, d, d_model*3)   # concatenated expert weights in a group (3 weight matrices per expert)

# optional: perform dimensionality reduction and normalization steps on expert_weights

S = torch.rand(G)  # REAP scores for G experts in a group (from step 1)
S_norm = S / S.sum()  # normalized scores (sum=1)

avg_weights = S_norm[0] * expert_weights[0].clone()

# compute permutation cost matrices w.r.t. expert 0 
for i in range(1, G):
    cost1 = torch.cdist(expert_hidden[0], expert_hidden[i])    # (d, d)
    cost2 = torch.cdist(expert_weights[0], expert_weights[i])  # (d_model*3, d_model*3)
    row_ind, col_ind = scipy.optimize.linear_sum_assignment(cost1 + cost2)
    perm = col_ind  # permutation of expert i w.r.t. expert 0
    expert_weights[i] = expert_weights[i][perm]  # permute expert i hidden weights accordingly
    avg_weights += S_norm[i] * expert_weights[i]  # weighted average

# Result: G experts merged into one expert with avg_weights

Calibration Data

The design of calibration data is critical in expert pruning/merging methods. In expert merging and early pruning papers, the c4 dataset is often used. REAP uses c4 and/or evol-codealpaca depending on the downstream task. In my experiments, to make evaluation consistent while the model performant, a fixed mix of c4, math and coding data is used.

The calibration data consists of 2048 sequences from diverse sources. It is fixed in all my experiments for all the methods unless explicitly mentioned otherwise (like in ablations).

For math data, the AI-MO/NuminaMath-1.5 dataset is used, namely subsets ‘cn_k12’ and ‘olympiads’, with the idea to avoid the overlap with benchmark datasets. For coding data, bigcode/the-stack-smol is used. A different number of samples (sequences) and max tokens are used per dataset to balance their proportions as shown below.

Table 1. Calibration data used across all the experiments.

Math and coding data are dominant as they are often shown important in adapting LLMs for reasoning and coding tasks. Also, math and coding benchmarks dominate the evaluation in practice. However, the exact proportion in REAM may be suboptimal and can be further explored in future work.

Gate Weight Adjustment

Finally, after experts are merged, the gate ($g$) weights have to be adjusted to account for the reduced number of experts. In expert merging implementations, merged experts are tied in memory instead of actually keeping only $k$ experts, while gate weights remain unchanged. This leads to gate logits corresponding to the experts within a merged group being summed during inference (see the REAP paper for additional analysis). In contrast, REAM follows REAP and simply removes the weights of the non-centroid experts from the gate weights.

Experiments

Experiments are run for three models: Qwen3-30B-A3B-Instruct-2507, Qwen3-235B-A22B-Instruct-2507 and Qwen3-Next-80B-A3B-Instruct. These are MoE LLMs with 30B, 235B and 80B parameters, respectively. Qwen3-30B-A3B and Qwen3-235B-A22B have $N$=128 experts and TopK=8The key differences between the two models are the number of Transformer layers (48 vs 94) and hidden size (2048 vs 4096).. Qwen3-Next-80B-A3B has $N$=512 experts and TopK=10. They are versatile models fine-tuned for instruction-following tasks heavily used in practice.

The experiments aim to reduce the number of experts by 25% (i.e., from 128 to 96 or from 512 to 384), while maintaining high performance across multiple benchmarks. The proposed REAM method is evaluated compared to the leading expert pruning method (REAP) and leading expert merging method (HC-SMoE).

🤗Qwen3 models compressed with HC-SMoE/REAP/REAM are released at huggingface-ream. Details about the evaluation are provided at huggingface-qwen3-ream.

Benchmarks

Compressed models are evaluated on two sets of benchmarks: multi-choice question answering (MC) and long context reasoning tasks (GEN) tasks. The MC tasks are often evaluated by loglikelihood of the correct choice among multiple choices, so the model does not generate free-form text answers. In contrast, in long context reasoning tasks (GEN), the model need to generate free-form answers, which is often more challenging and useful in practice.

The MC set includes 8 tasks used in the HC-SMoE paper: Winogrande, ARC-C, ARC-E, BoolQ, HellaSwag, MMLU, OpenBookQA and RTE.

Table 2. GEN benchmarks used in the experiments.

Qwen3-30B-A3B

Qwen3-30B-A3B-Instruct-2507 was used to tune the REAM hyperparameters (like calibration data mix, number of experts per group $C$, etc.). Also, several ablations of REAM are performed using this model to highlight the importance of REAM components. The results are shown in the Figure below with the average scores across MC ($x$ axis) and GEN tasks ($y$ axis).

{
  "data": [
    {
      "name": "Original",
      "type": "scatter",
      "mode": "markers",
      "x": [69.7],
      "y": [70.9],
      "marker": { "size": 12, "symbol": "circle" }
    },
    {
      "name": "HC-SMoE (c4)",
      "type": "scatter",
      "mode": "markers",
      "x": [63.3],
      "y": [65.2],
      "marker": { "size": 12, "symbol": "square" }
    },
    {
      "name": "HC-SMoE",
      "type": "scatter",
      "mode": "markers",
      "x": [64.9],
      "y": [66.3],
      "marker": { "size": 12, "symbol": "diamond" }
    },
    {
      "name": "REAP",
      "type": "scatter",
      "mode": "markers",
      "x": [65.0],
      "y": [64.2],
      "marker": { "size": 12, "symbol": "cross" }
    },
    {
      "name": "REAM",
      "type": "scatter",
      "mode": "markers",
      "x": [65.8],
      "y": [67.7],
      "marker": { "size": 15, "symbol": "star" }
    },
    {
      "name": "REAM (c4)",
      "type": "scatter",
      "mode": "markers",
      "x": [69.3],
      "y": [35.3],
      "marker": { "size": 10, "symbol": "x" }
    },
    {
      "name": "REAM (no REAP)",
      "type": "scatter",
      "mode": "markers",
      "x": [52.3],
      "y": [64.1],
      "marker": { "size": 10, "symbol": "triangle-up" }
    },
    {
      "name": "REAM (no gated sim)",
      "type": "scatter",
      "mode": "markers",
      "x": [64.6],
      "y": [51.3],
      "marker": { "size": 10, "symbol": "triangle-down" }
    },
    {
      "name": "REAM (no pseudo-pruning)",
      "type": "scatter",
      "mode": "markers",
      "x": [63.5],
      "y": [61.3],
      "marker": { "size": 10, "symbol": "triangle-right" }
    },
    {
      "name": "REAM (no seq)",
      "type": "scatter",
      "mode": "markers",
      "x": [65.5],
      "y": [66.0],
      "marker": { "size": 10, "symbol": "triangle-left" }
    }    
  ],
  "layout": {
    "title": { "text": "Pareto Tradeoff Between MC and Generative tasks", "size": 20 },
    "xaxis": {
      "title": { "text": "MC tasks score (%)" },
      "range": [45, 77]
    },
    "yaxis": {
      "title": { "text": "GEN tasks score (%)" },
      "range": [34, 77]
    },
    "legend": {
      "orientation": "v"
    }
  }
}

REAM obtains 65.8% and 67.7% average scores on MC and GEN tasks, respectively, while the original model achieves 69.7% and 70.9%. HC-SMoE and REAP achieve 64.8%/66.3% and 65.2%/64.1%, respectively, with REAP being better on MC tasks and HC-SMoE being better on GEN tasks. The proposed REAM method outperforms both of them.

The ablations highlight the importance of REAM components. Notably, using only c4 as calibration data significantly degrades GEN tasks performance, however it makes MC tasks performance close to the original model. On the contrary, HC-SMoE is not as sensitive to the calibration data choiceThis result confirms the stability of HC-SMoE regardless data as shown in their paper..

Another important component of REAM is using REAP scores for selecting centroids. When simple frequency-based scores are used instead, the performance (especially MC) drops significantly. Using proposed gated similarity, pseudo-pruning and sequential merging also improve the results with gated similarity being the most important among them.

Interestingly, the results reveal an inherent tradeoff between MC and GEN tasks performance when compressing MoE LLMs. Achieving Pareto optimal results seems challenging, however REAM shows promising results in this direction.

Detailed results per task are shown below.

Table 3. MC results for Qwen3-30B-A3B-Instruct-2507.

Table 4. GEN results for Qwen3-30B-A3B-Instruct-2507.

Qwen3-235B-A22B

Results on Qwen3-235B-A22B-Instruct-2507 are obtained by running REAM and baselines with the same hyperparameters and calibration data as for Qwen3-30B-A3B-Instruct-2507. So no tuning specific to this model is performed.

Evaluation on a model of this size is very challenging and computationally expensive. Specifically, 8xH100 GPUs were required to fit the model in memory and perform evaluationIn comparison, for Qwen3-30B, 1-4 GPUs were sufficient depending on the task.. Each task took from a few minutes to 1-2 hours to evaluate. So evaluation is costly, especially in academic settings. The results are shown in the Table below.

Table 5. GEN results for Qwen3-235B-A22B-Instruct-2507.

On Qwen3-235B-A22B-Instruct-2507, REAP slightly outperforms REAM on 4 out of 6 GEN tasks and on average. However, both achieve the results surprisingly close to the original model (even outperforming the original model on some tasks!). HC-SMoE lags behind both methods significantly, which is different from the results on Qwen3-30B-A3B-Instruct-2507. It is possible that hyperparameters of REAM (as well as the baselines) are suboptimal for this model, so further tuning may improve the results.

Qwen3-Next-80B-A3B

Qwen3-Next-80B-A3B-Instruct is a newer MoE model with a similar architecture as Qwen3-30B-A3B, but with 512 experts and TopK=10 instead of 128 and TopK=8, respectively. Even though the total number of experts is increased 4x, the number of activated parameters is mainly affected by TopK. And although TopK is also increased in this model, to keep the number of activated parameters around 3B, the bottleneck dimension in each expert is reduced from 768 to 512. The results of compressing this model from 512 to 384 experts are shown below. Other than the number of experts and TopK, the same hyperparameters and calibration data as for Qwen3-30B-A3B-Instruct-2507 are used in all methods.

Table 6. GEN results for Qwen3-Next-80B-A3B-Instruct.

On Qwen3-Next-80B-A3B-Instruct, REAM outperforms REAP on 3 out of 6 GEN tasks and on average. HC-SMoE was not evaluated on this model, since it underperformed in previous experiments. Notably, the original Qwen3-Next model has almost the same GEN performance as its previous much larger (Qwen3-235B-A22B) variant (see Table 5 above), while having only 3B activated parameters instead of 22B. This further confirms the effectiveness of increasing MoE sparsity.

Conclusion

Merging vs Pruning. REAM shows promising results by reducing the memory requirement by around 25% while maintaining high performance across multiple benchmarks. There is still a room for improvement, especially in merging methods to better combine the weights of multiple experts. So it is possible that pruning methods (REAP) remain strong because the merging methods are still quite suboptimal. Improving merging methods may allow for further reducing the number of experts (e.g., by 50%), while maintaining high performance.

MC vs GEN Tradeoff. The tradeoff between MC and GEN tasks performance is surprising, since MC tasks are assumed to be easier than GEN tasks. So I mistakenly expected that improving GEN tasks performance would also lead to better MC results.

High Engineering and Computation Barrier. Another observation is that evaluation of LLMs is way more challenging than expected, in some cases taking more time to set up properly than doing research. Besides the high engineering and computation barrier, some benchmarks such as AIME25 are quite noisy because of the small number of samples (30 questions). So evaluation should be done on more diverse and larger benchmarks, further increasing the engineering and computational cost.

Weight Space Learning. Finally, relying on calibration data is generally undesirable, because the results may be sensitive to the data choice. However, only using the weights of experts have led to poor results in my experiments and in ablation studies in the literature. Future work may explore how to use Weight Space Learning (WSL) and specifically methods based on neural graphs to better perform merging without relying on calibration data.

More REAM Models and Code. We continue to release more REAM-compressed models such as Qwen3-Coder-Next and GLM-4.5-Air.

🔥 Our code and paper with all the details have been released!

License

Diagrams and text are licensed under Creative Commons Attribution CC-BY 4.0, unless noted otherwise. The figures that have been reused from other sources do not fall under this license and can be recognized by a note in their caption: “Figure from …”. The license of the code and paper is specified in their respective repositories.

Acknowledgements

Minjoong Lee and Hoshik Lee from Samsung DS provided initial project recommendations, feedback and infrastructure support. The experiments were in part enabled by computational resources provided by Calcul Québec and Compute Canada.

Citation

Paper:

@article{jha2026ream,
  title={REAM: Merging Improves Pruning of Experts in LLMs},
  author={Jha, Saurav and Hashemzadeh, Maryam and Pasand, Ali Saheb and Parviz, Ali and Lee, Min-Joong and Knyazev, Boris},
  journal={arXiv preprint arXiv:2604.04356},
  year={2026}
}

This blog post:

@misc{knyazev2026compressing,
  title={REAM: Compressing Mixture-of-Experts LLMs},
  author={Boris Knyazev},
  year={2026},
  url={https://bknyaz.github.io/blog/2026/moe/}
}

Training large neural networks is famously slow and expensive. In our paper, Accelerating Training with Neuron Interaction and Nowcasting Networks, presented at ICLR 2025 in Singapore, we introduced a new way to speed things up. We treat a neural network as a graph of interacting neurons, or neural graph, and train a graph neural network (GNN) to predict how the parameters of the network will evolve during training.

From Adam to NiNo

Due to massive training costs, recent years have seen a surge of faster adaptive optimizers: Shampoo, SOAP, Muon, and others. Each improves on Adam by smarter scaling and/or orthogonalization of parameter updates for a given weight matrix, often inspired by second-order methods or preconditioning. But they work at the level of parameters or layers, not the network as a whole. Moreover, they do not learn from the previous optimization runs, i.e., the optimization algorithms are based on manually-designed gradient-descent rules. Our method, NiNo (Neuron Interaction and Nowcasting), is different. Motivated by Kofinas et al., we model neurons as nodes and weights as edges, and use a graph neural network (GNN) to predict how weights will evolve. This lets us “nowcast” future parameters and reduce the number of steps required to reach the same performance metric.

NiNo is a GNN-based model that takes a history of past parameter values along the optimization trajectory (obtained with Adam or another optimizer) and makes a jump by predicting future parameter values. After making the prediction, optimization is continued with Adam, then followed by NiNo’s another prediction and so on.

This periodic nowcasting idea is borrowed from Weight Nowcaster Network (WNN) that revealed predictable patterns in optimization trajectories, but neural graphs synergized with GNNs make it work really well. To give some context in terms of optimization runs, below is Figure 1 reproduced from the NiNo paper.

{
  "data": [
    {
      "type": "scatter",
      "name": "Adam",
      "x": [
        645.4,
        1098.0,
        1559.6,
        2021.3,
        2482.9,
        2944.5,
        3406.2,
        3472.6,
        3867.8,
        4329.4,
        4791.1,
        5252.7,
        5714.3,
        6176.0,
        6637.6,
        6769.6,
        7099.3,
        7560.9,
        8022.5,
        8484.2,
        8945.8,
        9407.4,
        9869.1,
        10066.6,
        10330.7,
        10792.3,
        11254.0,
        11715.6,
        12177.2,
        12638.9,
        13100.5,
        13363.6
      ],
      "y": [
        1290.42,
        654.89,
        491.01,
        415.95,
        367.21,
        333.12,
        306.04,
        302.24,
        282.43,
        264.16,
        246.73,
        233.03,
        221.53,
        211.24,
        202.02,
        199.9,
        194.86,
        187.62,
        181.82,
        176.28,
        171.77,
        167.21,
        163.92,
        162.06,
        160.25,
        157.63,
        153.97,
        151.23,
        149.08,
        146.99,
        144.91,
        144.02
      ],
      "mode": "lines",
      "line": {
        "color": "#1f77b4",
        "width": 2.5
      }
    },
    {
      "type": "scatter",
      "name": "WNN (Jang et al., 2023)",
      "x": [
        654.8,
        1098.0,
        1559.6,
        2021.3,
        2482.9,
        2944.5,
        3406.2,
        3472.6,
        3867.8,
        4329.4,
        4791.1,
        5252.7,
        5714.3,
        6176.0,
        6637.6,
        6769.6,
        7099.3,
        7560.9,
        8022.5,
        8484.2,
        8945.8,
        9407.4,
        9869.1,
        10066.6,
        10330.7,
        10792.3,
        11254.0,
        11715.6,
        12177.2,
        12638.9,
        13100.5,
        13363.6
      ],
      "y": [
        1290.42,
        653.73,
        444.74,
        387.85,
        326.93,
        303.61,
        266.84,
        264.77,
        251.27,
        227.12,
        216.34,
        199.39,
        192.44,
        180.71,
        175.41,
        170.13,
        167.42,
        163.54,
        157.7,
        154.93,
        150.43,
        148.24,
        144.45,
        143.78,
        143.02,
        140.42,
        138.9,
        136.51,
        135.4,
        133.53,
        132.68,
        131.58
      ],
      "mode": "lines",
      "line": {
        "color": "#2ca02c",
        "width": 3,
        "dash": "dash"
      }
    },
    {
      "type": "scatter",
      "name": "NiNo (ours)",
      "x": [
        648.7,
        1098.0,
        1559.6,
        2021.3,
        2482.9,
        2944.5,
        3406.2,
        3472.6,
        3867.8,
        4329.4,
        4791.1,
        5252.7,
        5714.3,
        6176.0,
        6637.6,
        6769.6,
        7099.3,
        7560.9,
        8022.5,
        8484.2,
        8945.8,
        9407.4,
        9869.1,
        10066.6,
        10330.7,
        10792.3,
        11254.0,
        11715.6,
        12177.2,
        12638.9,
        13100.5,
        13363.6
      ],
      "y": [
        1290.42,
        652.54,
        425.26,
        371.58,
        294.78,
        271.73,
        231.12,
        228.8,
        217.43,
        193.29,
        184.29,
        169.48,
        163.7,
        154.62,
        150.46,
        149.22,
        144.61,
        141.77,
        138.42,
        136.37,
        134.1,
        132.88,
        131.09,
        130.71,
        130.05,
        128.99,
        128.08,
        127.1,
        126.53,
        125.55,
        124.76,
        124.48
      ],
      "mode": "lines",
      "line": {
        "color": "#D44841",
        "width": 5
      }
    },
    {
      "type": "scatter",
      "name": "target perplexity",
      "x": [
        636.4,
        13363.6
      ],
      "y": [
        147.0,
        147.0
      ],
      "mode": "lines",
      "line": {
        "color": "#7f7f7f",
        "width": 2,
        "dash": "dot"
      }
    }
  ],
  "layout": {
    "title": {
      "text": "NiNo achieves ~2× speedup compared to Adam.",
      "font": {
        "size": 20
      }
    },
    "xaxis": {
      "title": "training iteration",
      "range": [
        0,
        14000
      ],
      "tickvals": [
        0,
        2000,
        4000,
        6000,
        8000,
        10000,
        12000,
        14000
      ]
    },
    "yaxis": {
      "title": "validation perplexity",
      "type": "log",
      "range": [
        2,
        3
      ],
      "tickvals": [
        100,
        200,
        300,
        400,
        600,
        1000
      ],
      "ticktext": [
        "10\u00b2",
        "2\u00d710\u00b2",
        "3\u00d710\u00b2",
        "4\u00d710\u00b2",
        "6\u00d710\u00b2",
        "10\u00b3"
      ]
    },
    "legend": {
      "x": 1.0,
      "y": 1.0,
      "xanchor": "right",
      "yanchor": "top"
    },
    "margin": {
      "l": 70,
      "r": 20,
      "t": 60,
      "b": 60
    },
    "annotations": [
      {
        "x": 300,
        "xref": "x",
        "y": 0.60,
        "yref": "paper",
        "text": "Nowcast",
        "showarrow": false,
        "textangle": -45,
        "font": {
          "size": 10
        },
        "xanchor": "left",
        "yanchor": "bottom"
      },
      {
        "x": 1300,
        "xref": "x",
        "y": 0.40,
        "yref": "paper",
        "text": "Nowcast",
        "showarrow": false,
        "textangle": -45,
        "font": {
          "size": 10
        },
        "xanchor": "left",
        "yanchor": "bottom"
      },
      {
        "x": 10300,
        "xref": "x",
        "y": -0.03,
        "yref": "paper",
        "text": "Nowcast",
        "showarrow": false,
        "textangle": -45,
        "font": {
          "size": 10
        },
        "xanchor": "left",
        "yanchor": "bottom"
      }
    ]
  }
}

The figure shows the results of Adam without and with nowcasting using our NiNo. “Nowcast” steps are shown at step 1000, 2000 and 11,000 for visualization purposes, but this step is applied every 1000 steps in our experiments. Note the ~2× reduction of the number of steps required by NiNo to achieve the same validation perplexity as by Adam. The optimization task in this example is autoregressive next-token prediction on the WikiText103 dataset that NiNo has not seen during its training.

NiNo Details

As NiNo is a neural network (namely, a GNN), it needs to be trained first before we can use it to speed up optimization. To do so, we collected and publicly released a 🤗dataset of checkpoints with optimization trajectories on 2 vision and 2 language tasks. Even though collecting these checkpoints and training NiNo is computationally expensive, this one-time cost is amortized, meaning the same trained NiNo can potentially be used across many tasks, ultimately reducing total training cost.

Neuron Permutation Symmetry

Developing a strong parameter nowcasting model requires many specific design choices and, perhaps most critically, accurate modeling of neuron permutation symmetryModeling neuron permutation symmetry imposes a strong inductive bias similarly to using convolution for images. So our model can be used in more diverse tasks and should be able to learn parameter prediction rules that generalize better with fewer samples compared to models that do not explicitly take this symmetry into account.. Neuron permutation symmetry states that the order of neurons in adjacent layers of a neural network can be permuted in certain ways without affecting the overall function of the network. To better understand this symmetry, let me introduce a simple example based on a two layer neural network with weights \(\mathbf{W}_1\), \(\mathbf{W}_2\) and an element-wise activation function σ. Given input $\mathbf{x}$, the output of such a network can be expressed as:

We can permute the neurons in the hidden layer by applying a permutation matrix \(\mathbf{P}\) to the weights in layer 1 (permuting rows) and the inverse permutation to the weights in layer 2 (permuting columns), resulting in:

Given the orthogonality property of permutation matrices, \(\mathbf{P}^{-1} = \mathbf{P}^T\), and that the activation function σ is element-wise, we can see that the output of the network remains unchanged:

Let me now visualize this symmetry using a simple demo, where the output is dynamically computed given the current order of neurons for some fixed input. You can click the “Swap Random Pair” button to randomly swap two (just for simplicity) hidden neurons and see how the network diagram and weight matrices change, while the computed output should (hopefully!) remain the same. You can also toggle the “Correct Permutation” option off to turn off the permutation of the weights in the second layer, in which case the output changes. You can reset the demo to the original state by clicking the “Reset” button.

Neuron Permutation Symmetry in Transformers

To make NiNo work well for LLMs and Transformers in general, it was critical to carefully construct the neural graph for multi-head self-attention, making it stand out compared to WNN (which ignores the neural network structure). This is a tricky part as the illustration below shows, but we implemented it for many different Transformer layers.

Visualizing the permutation symmetry in Transformers using a demo is also possible, but I leave it for future work (or please submit a PR to this blog post).

Results that Stand Out

• NiNo achieves roughly a 2× speedup compared to Adam, which is rarely achieved by manually designed optimizers in practiceWhich are considered remarkable if a 1.2–1.3× speedup is achieved..
• NiNo speeds up optimization across 9 vision and language tasks that we systematically evaluated, achieving the same validation performance (accuracy or perplexity) as Adam in about half the steps on average across these tasks (see Table 2 in the paper).
• NiNo continues to speed up training for larger models beyond our main 9 tasks that are much larger than the models in the training checkpoints used to train NiNo (see Table 4 in the paper).

In addition, as shown below, applying NiNo is straightforward:

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM
from optim import NiNo

model = AutoModelForCausalLM.from_config(...)  # some model

# NiNo is implemented as a wrapper around the base optimizer
# any optimizer other than Adam should also be possible to use with NiNo
opt = NiNo(base_opt=torch.optim.AdamW(model.parameters(), 
           lr=1e-3, 
           weight_decay=1e-2),
           ckpt='checkpoints/nino.pt',
           subgraph=False, # can be set to True for larger models (see Llama 3.2 example below)
           edge_sample_ratio=0,  # can be set to a small positive number for larger models (see Llama 3.2 example below)
           model=model,
           period=1000,
           max_train_steps=10000)
for step in range(10000):
    if opt.need_grads:  # True/False based on the step number and period
        opt.zero_grad()  # zero out gradients
        data, targets = ...  # get some batch of data
        # base optimizer step (majority of the time)
        outputs = model(data)  # forward pass
        loss = F.cross_entropy(outputs, targets)  # compute some loss
        loss.backward()  # only compute gradients for the base optimizer            
    opt.step()  # base_opt step or nowcast params every 1000 steps using NiNo    

Learning to Optimize, Revisited

Our work connects to the broader “learning to optimize” literature, such as VeLO. Unlike many learned optimizers that struggle with cost, stability or show only a ~1.2–1.3× speedupWhile a ~1.2–1.3× speedup is remarkable, in practice it usually does not justify the immense amount of extra work (e.g. efficient distributed implementation, tuning, potential instabilities especially in mixed/low-bit precision, investigating unexpected side effects like overfitting or poor generalization) that is required for actual large-scale usefulness., NiNo is conceptually lightweight and stable, because it is only applied every 1,000 steps (by default), while for all the other steps any base optimizer, such as Adam, can be applied to allow for stable convergence. At the same time, recent learned optimizers such as our recent Celo and μLO that will be presented at ICLR 2026, make a significant step in improving learned optimizersIn particular, Celo and μLO make learned optimizers more cost-effective and stable (i.e. without big loss spikes) to train and use..

Conclusion

Even though NiNo shows great speedups, it requires further work, for example:

• Making NiNo’s step more scalable (especially memory-efficient) to larger models - currently the message passing step in its GNN is a bottleneck;
• Improving speedups on larger tasks - for tasks with >1B parameter models we currently observe no speedup;
• Adding automatic ways to construct neural graphs and verify their correctness;
• Combining NiNo with learned optimizers to get the merits of both;
• Showing theoretical guarantees for convergence.

We have open-sourced our code and pretrained NiNo checkpoints at github.com/SamsungSAILMontreal/nino under MIT License and welcome contributions.

License

Diagrams and text are licensed under Creative Commons Attribution CC-BY 4.0, unless noted otherwise.

Citation

@inproceedings{knyazev2024accelerating,
  title={Accelerating Training with Neuron Interaction and Nowcasting Networks}, 
  author={Boris Knyazev and Abhinav Moudgil and Guillaume Lajoie and Eugene Belilovsky and Simon Lacoste-Julien},  
  booktitle={International Conference on Learning Representations},
  year={2025},
}

The British Machine Vision Conference (BMVC), finished about two weeks ago in Cardiff, UK, is one of the top conferences in computer vision & pattern recognition with a competitive acceptance rate of 28%. Compared to others, it’s a small event, so you have plenty of time to walk around posters and talk to presenters one-on-one, which I found really nice.

BMVC 2019 on Twitter

Paper decisions were released yesterday. Congratulations to all of you! In total, we received 1008 submissions, of which 815 were valid. Of these, a total of 231 papers were accepted (38 as Oral Presentations, 193 as Poster Presentations). This amounts to a 28% acceptance rate.

I presented a poster on Image Classification with Hierarchical Multigraph Networks on which I mainly worked during my internship at SRI International under the supervision of Xiao Lin, Mohamed Amer (homepage) and my PhD advisor Graham Taylor.

In the paper, we basically try to answer the question “Can we do better than Convolutional Neural Networks?”. Here I discuss this question and support my arguments by results. I also walk you through the forward pass of the whole pipeline for a single image from PASCAL VOC 2012 using PyTorch.

The complete code for this post is in my notebook on Github. It should be easy to adapt it for training and validating on the whole PASCAL dataset.

So, why do we want to do better than ConvNets? Haven’t they outperformed humans in many tasks?

For example, you could say that image classification is a solved task. Well, in terms of ImageNet, yes. But despite great contribution of ImageNet, it is a weird task. Why would you want to discriminate between hundreds of dog breeds? So, in result we have models succeeding in that, but failing to discriminate between slightly rotated dogs and cats. Fortunately, we now have ImageNet-C and other similar benchmarks showing that we are nowhere close to solving it.

Another open problem arising in related tasks, like object detection, is training on really large images (e.g., 4000×3000), which is addressed, for example, by Katharopoulos & Fleuret (ICML, 2019) and Ramapuram et al. (BMVC, 2019). Thanks to the latter I now know that if the background of a poster is black, then it’s likely from Apple. I should reserve some color too!

So, maybe we need something different than a Convolutional Neural Network? Instead of constantly patching its bugs, maybe we should use a model that has nicer properties from the beginning?

We argue that such a model can be a Graph Neural Network (GNN) — a neural network that can learn from graph-structured data. GNNs have some appealing properties. For example, compared to ConvNets, GNNs are inherently rotation and translation invariant, because there is simply no notion of rotation or translation in graphs, i.e. there is no left and right, there are only “neighbors” in some sense (Khasanova & Frossard, ICML, 2017). So, the problem of making a ConvNet generalize better to different rotations, that people have been trying to solve for years, is solved automatically with GNNs!

Regarding learning from large images, how about extracting superpixels from images and feeding a much lower dimensional input to a GNN instead of feeding a downsampled (e.g. 224×224) image to a ConvNet? Superpixels seem to be a much better way to downsample an image compared to, say, bilinear interpolation, because they often preserve a lot of semantics by keeping the boundaries between objects. With a ConvNet we cannot directly learn from this kind of an input, however, there are some nice works proposing to leverage them (Kwak et al., AAAI, 2017).

So, a GNN sounds wonderful! Let’s see how it performs in practice.

Oh no! Our baseline GNN based on (Kipf & Welling, ICLR, 2017) achieves merely 19.2% (mean average precision or mAP) on PASCAL, compared to 32.7% of a ConvNet with the same number of layers and filters in each layer.

We propose several improvements that eventually beat the ConvNet!

1. Hierarchical Graph

In ConvNets, the hierarchical structure of images is implicitly modeled by pooling layers. In GNNs, you can achieve this in at least two ways. First, you can use pooling similar to ConvNets, but for graphs, defining a fast and good pooling method is really challenging. Instead, we can compute superpixels at multiple scales and pool superpixels by their correspondence to a larger parent superpixel. However, for some reasons this kind of pooling didn’t work well in our case (I still think it should work well). So, instead we model a hierarchy at the input level. In particular, we combine superpixels of all scales into a single set and compute hierarchical relations based on intersection over union (IoU), commonly used in semantic segmentation.

Based on that principle, I build the hierarchical graph in the code below. I also build a multiscale version of the spatial graph, but it encodes only spatial relationships, while IoU should better encode hierarchical ones. For example, using IoU we can create shortcuts between remote child nodes, i.e. connect two small superpixels (e.g., wheels) that are far away spatially, but belong to the same parent node (e.g., a car) as shown on the image above.

And indeed, the hierarchical graph boosts mAP to 31.7%, making it just 1% lower than a ConvNet while having 4 times fewer trainable parameters! If we use only the spatial multiscale graph, the results are much worse as explored in the paper.

Great! What else can we do to further improve results?

2. Learnable relations

So far, if we visualize our filters, they will look very primitive (as Gaussians). See my tutorial on GNNs for more details. We want to learn some edge detectors similar to ConvNets, because they work so well. But it turned out to be very challenging to learn them with GNNs. To do that, we basically need to generate edges between superpixels depending on the difference between coordinates. By doing so, we will endow our GNN with the ability to understand the coordinate system (rotation, translation). We will use a 2 layer neural network defined in PyTorch like this:

pred_edge = nn.Sequential(nn.Linear(2, 32),
                          nn.ReLU(True),
                          nn.Linear(32, L))

where L is the number of predicted edges or the number of filters, such as 4 in the visualization below.

We restrict the filter to learn edges only based on the absolute difference between coordinates, |(x₁,y₁) - (x₂,y₂)|, instead of raw values, so that the filters become symmetric. This limits the capacity of filters, but it is still much better than a simple Gaussian filter used by our baseline GCN.

In my Jupyter notebook, I created a class LearnableGraph that implements the logic to predict edges given node coordinates (or any other features) and the spatial graph. The latter is used to define a small local neighborhood around each node to avoid predicting edges for all possible node pairs, because it’s expensive and doesn’t make much sense to connect very remote superpixels.

Below, I visualize the trained pred_edge function. To do that, I assume that the current node with index 1, where we apply the convolution, is in the center of a coordinate system, (x₁,y₁)=0. Then I simply sample coordinates of other nodes, (x₂,y₂), and feed them to pred_edge. The color shows the strength of an edge depending on the distance from a center node.

The learned graph is also very powerful, but at a larger computational cost, which is negligible if we generate a very sparse graph. The result of 32.3% is just 0.4% lower than a ConvNet and can be easily improved if we generate more filters!

3. Multiscale GNN

We now have three graphs: spatial, hierarchical and learned. A single graph convolutional layer with the spatial or hierarchical graph permits feature propagation only within the “first neighbors”. Neighbors are soft in our case, since we use a Gaussian to define the spatial graph and IoU for the hierarchical one. Defferrard et al. (NIPS, 2016) proposed a multiscale (multihop) graph convolution, which aggregates features within a K-hop neighborhood and approximates spectral graph convolution. See my other post for an extensive explanation of this method. For our spatial graph, it essentially corresponds to using multiple Gaussians of different width. For the hierarchical graph, this way we can create K-hop shortcuts between remote child nodes. For the learned graph, this method will create multiple scales of the learned filters visualized above.

Using multiscale graph convolution, implemented in my GraphLayerMultiscale class, turned out to be extremely important allowing us to outperform the baseline ConvNet by 0.3%!

4. Improving fusion of relation types at low cost

So far, to learn from our three graphs, we have used a standard concatenation method. This method, however, has a couple of problems. First, the number of trainable parameters of such a fusion operator is linear w.r.t. the input and output feature dimensionalities, scale (K) and number of relation types, so it can really grow fast if we increase two or more of these parameters at once. Second, the relation types we try to fuse can have very different natures and occupy very different subspaces of a manifold. To solve both problems at the same time, we propose learnable projections similar to (Knyazev et al., NeurIPS-W, 2018). This way we decouple the linear dependency reducing the number of parameters by a factor of 2–3 compared to concatenation. In addition, learnable projections transform multirelational features so that they should occupy nearby subspaces of the manifold, facilitating the propagation of information from one relationship to another.

By using the proposed fusion method, implemented in the GraphLayerFusion class below, we achieve 34.5% beating the ConvNet by 1.8%, while having 2 times fewer parameters! Quite impressive for the model that initially didn’t know anything about the spatial structure of images, except for information encoded in superpixels. It would be interesting to explore other fusion methods, like this one, to get even better results.

Conclusion

It turned out that with a multirelational graph network and some tricks, we can do better than a Convolutional Neural Network!

Unfortunately, during our process of improving the GNN we slowly lost its invariance properties. For example, the shape of superpixels might change after rotating the image, and superpixel coordinates that we use for node features to improve the model also make it less robust.

Nevertheless, our work is a small step towards a better image reasoning model and we show that GNNs can pave a promising direction.

See my notebook on Github for implementation details.

I also highly recommend Matthias Fey’s Master’s thesis with the code on a very related topic.

Find me on Github, LinkedIn and Twitter. My homepage.

<hr><p>Can we do better than Convolutional Neural Networks? was originally published in TDS Archive on Medium, where people are continuing the conversation by highlighting and responding to this story.</p>

First, let’s recall what is a graph. A graph G is a set of nodes (vertices) connected by directed/undirected edges. In this post, I will assume an undirected graph G with N nodes. Each node in this graph has a C-dimensional feature vector, and features of all nodes are represented as an N×C dimensional matrix X⁽ˡ⁾. Edges of a graph are represented as an N×N matrix A, where the entry Aᵢⱼ indicates if node i is connected (adjacent) to node j. This matrix is called an adjacency matrix.

Spectral analysis of graphs (see lecture notes here and earlier work here) has been useful for graph clustering, community discovery and other mainly unsupervised learning tasks. In this post, I basically describe the work of Bruna et al., 2014, ICLR 2014 who combined spectral analysis with convolutional neural networks (ConvNets) giving rise to spectral graph convolutional networks that can be trained in a supervised way, for example for the graph classification task.

Despite that spectral graph convolution is currently less commonly used compared to spatial graph convolution methods, knowing how spectral convolution works is still helpful to understand and avoid potential problems with other methods. Plus, in the conclusion I refer to some recent exciting works making spectral graph convolution more competitive.

1. Graph Laplacian and a little bit of physics

While “spectral” may sound complicated, for our purpose it’s enough to understand that it simply means decomposing a signal/audio/image/graph into a combination (usually, a sum) of simple elements (wavelets, graphlets). To have some nice properties of such a decomposition, these simple elements are usually orthogonal, i.e. mutually linearly independent, and therefore form a basis.

When we talk about “spectral” in signal/image processing, we imply the Fourier Transform, which offers us a particular basis (DFT matrix, e.g. scipy.linalg.dft in Python) of elementary sine and cosine waves of different frequencies, so that we can represent our signal/image as a sum of these waves. But when we talk about graphs and graph neural networks (GNNs), “spectral” implies eigen-decomposition of the graph Laplacian L. You can think of the the graph Laplacian L as an adjacency matrix A normalized in a special way, whereas eigen-decomposition is a way to find those elementary orthogonal components that make up our graph.

Intuitively, the graph Laplacian shows in what directions and how smoothly the “energy” will diffuse over a graph if we put some “potential” in node i. A typical use-case of Laplacian in mathematics and physics is to solve how a signal (wave) propagates in a dynamic system. Diffusion is smooth when there is no sudden changes of values between neighbors as in the animation below.

In the rest of the post, I’m going to assume “symmetric normalized Laplacian”, which is often used in graph neural networks, because it is normalized so that when you stack many graph layers, the node features propagate in a more smooth way without explosion or vanishing of feature values or gradients. It is computed based only on an adjacency matrix A of a graph, which can be done in a few lines of Python code as follows:

# Computing the graph Laplacian
# A is an adjacency matrix of some graph G
import numpy as np

N = A.shape[0] # number of nodes in a graph
D = np.sum(A, 0) # node degrees
D_hat = np.diag((D + 1e-5)**(-0.5)) # normalized node degrees
L = np.identity(N) — np.dot(D_hat, A).dot(D_hat) # Laplacian

Here, we assume that A is symmetric, i.e. A = Aᵀ and our graph is undirected, otherwise node degrees are not well-defined and some assumptions must be made to compute the Laplacian. An interesting property of an adjacency matrix A is that Aⁿ (matrix product taken n times) exposes n-hop connections between nodes (see here for more details).

Let’s generate three graphs and visualize their adjacency matrices and Laplacians as well as their powers.

For example, imagine that the star graph above in the middle is made from metal, so that it transfers heat well. Then, if we start to heat up node 0 (dark blue), this heat will propagate to other nodes in a way defined by the Laplacian. In the particular case of a star graph with all edges equal, heat will spread uniformly to all other nodes, which is not true for other graphs due to their structure.

In the context of computer vision and machine learning, the graph Laplacian defines how node features will be updated if we stack several graph neural layers. Similarly to the first part of my tutorial, to understand spectral graph convolution from the computer vision perspective, I’m going to use the MNIST dataset, which defines images on a 28×28 regular grid graph.

2. Convolution

In signal processing, it can be shown that convolution in the spatial domain is multiplication in the frequency domain (a.k.a. convolution theorem). The same theorem can be applied to graphs. In signal processing, to transform a signal to the frequency domain, we use the Discrete Fourier Transform, which is basically matrix multiplication of a signal with a special matrix (basis, DFT matrix). This basis assumes a regular grid, so we cannot use it for irregular graphs, which is a typical case. Instead, we use a more general basis, which is eigenvectors V of the graph Laplacian L, which can be found by eigen-decomposition: L=VΛVᵀ, where Λ are eigenvalues of L.

PCA vs eigen-decomposition of the graph Laplacian. To compute spectral graph convolution in practice, it’s enough to use a few eigenvectors corresponding to the smallest eigenvalues. At first glance, it seems to be an opposite strategy compared to frequently used in computer vision Principal component analysis (PCA), where we are more interested in the eigenvectors corresponding to the largest eigenvalues. However, this difference is simply due to the negation used to compute the Laplacian above, therefore eigenvalues computed using PCA are inversely proportional to eigenvalues of the graph Laplcacian (see this paper for a formal analysis). Note also that PCA is applied to the covariance matrix of a dataset for the purpose to extract the largest factors of variation, i.e. the dimensions along which data vary the most, like in Eigenfaces. This variation is measured by eigenvalues, so that the smallest eigenvalues essentially correspond to noisy or “spurious” features, which are assumed to be useless or even harmful in practice.

Eigen-decomposition of the graph Laplacian is applied to a single graph for the purpose to extract subgraphs or clusters (communities) of nodes, and eigenvalues tell us a lot about graph connectivity. I will use eigenvectors corresponding to the 20 smallest eigenvalues in our examples below, assuming that 20 is much smaller than the number of nodes N (N=784 in case of MNIST). To find eigenvalues and eigenvectors below on the left, I use a 28×28 regular graph, whereas on the right I follow the experiment of Bruna et al. and construct an irregular graph by sampling 400 random locations on a 28×28 regular grid (see their paper for more details about this experiment).

So, given graph Laplacian L, node features X and filters W_spectral, in Python spectral convolution on graphs looks very simple:

# Spectral convolution on graphs
# X is an N×1 matrix of 1-dimensional node features
# L is an N×N graph Laplacian computed above
# W_spectral are N×F weights (filters) that we want to train
from scipy.sparse.linalg import eigsh # assumes L to be symmetric

Λ,V = eigsh(L,k=20,which=’SM’) # eigen-decomposition (i.e. find Λ,V)
X_hat = V.T.dot(X) # 20×1 node features in the "spectral" domain
W_hat = V.T.dot(W_spectral)  # 20×F filters in the "spectral" domain
Y = V.dot(X_hat * W_hat)  # N×F result of convolution

Formally:

where we assume that our node features X⁽ˡ⁾ are 1-dimensional, e.g. MNIST pixels, but it can be extended to a C-dimensional case: we will just need to repeat this convolution for each channel and then sum over C as in signal/image convolution.

Formula (3) is essentially the same as spectral convolution of signals on regular grids using the Fourier Transform, and so creates a few problems for machine learning:

• the dimensionality of trainable weights (filters) W_spectral depends on the number of nodes N in a graph;
• W_spectral also depends on the graph structure encoded in eigenvectors V.

These issues prevent scaling to datasets with large graphs of variable structure. Further efforts, summarized below, were focused on resolving these and other issues.

3. “Smoothing” in the spectral domain

Bruna et al. were one of the first to apply spectral graph analysis to learn convolutional filters for the graph classification problem. The filters learned using formula (3) above act on the entire graph, i.e. they have global support. In the computer vision context, this would be the same as training convolutional filters of size 28×28 pixels on MNIST, i.e. filters have the same size as the input (note that we would still slide a filter, but over a zero-padded image). While for MNIST we can actually train such filters, the common wisdom suggests to avoid that, as it makes training much harder due to the potential explosion of the number of parameters and difficulty of training large filters that can capture useful features shared across different images.

I actually successfully trained such a model using PyTorch and this code from my GitHub. You should run it using mnist_fc.py --model conv. After training for 100 epochs, the filters look like mixtures of digits:

To reiterate, we generally want to make filters smaller and more local (which is not exactly the same as I’ll note below).

To enforce that implicitly, they proposed to smooth filters in the spectral domain, which makes them more local in the spatial domain according to the spectral theory. The idea is that you can represent our filter W_spectral from formula (3) as a sum of 𝐾 predefined functions, such as splines, and instead of learning N values of W, we learn K coefficients α of this sum:

While the dimensionality of fk does depend on the number of nodes N, these functions are fixed, so we don’t learn them. The only thing we learn are coefficients α, and so W_spectral is no longer dependent on N. Neat, right?

To make our approximation in formula (4) reasonable, we want K<<N to reduce the number of trainable parameters from N to K and, more importantly, make it independent of N, so that our GNN can digest graphs of any size. We can use different bases to perform this “expansion”, depending on which properties we need. For instance, cubic splines shown above are known as very smooth functions (i.e. you cannot see knots, i.e. where the pieces of the piecewise spline polynomial meet). The Chebyshev polynomial, which I discuss in my another post, has the minimum 𝑙∞ distance between the approximating function. The Fourier basis is the one that preserves most of the signal energy after transformation. Most bases are orthogonal, because it would be redundant to have terms that can be expressed by each other.

Note that filters W_spectral are still as large as the input, but their effective width is small. In case of MNIST images, we would have 28×28 filters, in which only a small fraction of values would have an absolute magnitude larger than 0 and all of them should be located close to each other, i.e. the filter would be local and effectively small, something like the one below (second from the left):

To summarize, smoothing in the spectral domain allowed Bruna et al. to learn more local filters. The model with such filters can achieve similar results as the model without smoothing (i.e. using our formula (3)), but with much fewer trainable parameters, because the filter size is independent of the input graph size, which is important to scale the model to datasets with larger graphs. However, learned filters W_spectral still depend on eigenvectors V, which makes it challenging to apply this model to datasets with variable graph structures.

Conclusion

Despite the drawbacks of the original spectral graph convolution method, it has been developed a lot and has remained a quite competitive method in some applications, because spectral filters can better capture global complex patterns in graphs, which local methods like GCN (Kipf & Welling, ICLR, 2017) cannot unless stacked in a deep network. For example, two ICLR 2019 papers, of Liao et al. on “LanczosNet” and Xu et al. on “Graph Wavelet Neural Network”, address some shortcomings of spectral graph convolution and show great results in predicting molecule properties and node classification. Another interesting work of Levie et al., 2018 on “CayleyNets” showed strong performance in node classification, matrix completion (recommender systems) and community detection. So, depending on your application and infrastructure, spectral graph convolution can be a good choice.

In another part of my Tutorial on Graph Neural Networks for Computer Vision and Beyond I explain Chebyshev spectral graph convolution introduced by Defferrard et al. in 2016, which is still a very strong baseline that has some nice properties and is easy to implement as I demonstrate using PyTorch.

Acknowledgement: A large portion of this tutorial was prepared during my internship at SRI International under the supervision of Mohamed Amer (homepage) and my PhD advisor Graham Taylor (homepage). I also thank Carolyn Augusta for useful feedback.

Find me on Github, LinkedIn and Twitter. My homepage.

If you want to cite this blog post in your paper, please use:
@misc{knyazev2019tutorial,
title={Tutorial on Graph Neural Networks for Computer Vision and Beyond},
author={Knyazev, Boris and Taylor, Graham W and Amer, Mohamed R},
year={2019}
}

<hr><p>Spectral Graph Convolution Explained and Implemented Step By Step was originally published in TDS Archive on Medium, where people are continuing the conversation by highlighting and responding to this story.</p>

I’m presenting an overview of important Graph Neural Network works, by distilling key ideas and explaining simple intuition behind milestone methods using Python and PyTorch. This post continues the first part of my tutorial.

20+ years of Graph Neural Networks

In the “Graph of Graph Neural Network (GNN) and related works” above, I added papers on graphs that I have come across in the last year. In this graph, a directed edge between two works denotes that one paper is based on the other (while not necessary citing it) and a color of the work denotes:

• Red — spectral methods (require eigen-decomposition of the graph Laplacian, which will be explained below)
• Green — methods that work in the spatial domain (do not require eigen-decomposition of the graph Laplacian)
• Blue — equivalent to spectral methods, but do not require eigen-decomposition (so, effectively, spatial methods)
• Black — are methods complementary to GNNs and agnostic to the choice of a GNN itself (i.e. pooling, attention).

Note, that some other important works and edges are not shown to avoid further clutter, and only a tiny fraction of works, highlighted in bold boxes, will be covered in this post. Disclaimer: I still found room to squeeze our own recent works there 😊.

Most of the important methods are covered in this non-exhaustive list of works:

• Nicket et al., 2015, A Review of Relational Machine Learning for Knowledge Graphs
• Bronstein et al., 2016, Geometric deep learning: going beyond Euclidean data
• Hamilton et al., 2017, Representation Learning on Graphs: Methods and Applications
• Kipf et al., 2018, Structured deep models: Deep learning on graphs and beyond, presentation slides.
• Battaglia et al., 2018, Relational inductive biases, deep learning, and graph networks
• Zhang et al., 2018 Deep Learning on Graphs: A Survey
• Zhou et al., 2018, Graph Neural Networks: A Review of Methods and Applications
• Wu et al., 2019, A Comprehensive Survey on Graph Neural
  Networks
• Petar Veličković, 2019, The resurgence of structure in deep neural networks, PhD Thesis.
• NIPS and CVPR video tutorials

The first work where graphs were classified using a neural network seems to be a 1997 paper by Alessandro Sperduti and Antonina Starita on “Supervised Neural Networks for the Classification of Structures”.

Sperduti & Starita, 1997: “Until now neural networks have been used for classifying unstructured patterns and sequences. However, standard neural networks and statistical methods are usually believed to be inadequate when dealing with complex structures because of their feature-based approach.”

From 1997, the body of works on learning from graphs has grown so much and in so many diverse directions that it is very hard to keep track without some smart automated system. I believe we are converging to using methods based on neural networks (based on our formula (2) explained in the first part of my tutorial), or some combination of neural networks and other methods.

To recap the notation we used in the first part, we have some undirected graph G with N nodes. Each node in this graph has a C-dimensional feature vector, and features of all nodes are represented as an N×C dimensional matrix X⁽ˡ⁾. In a typical graph network, such as GCN (Kipf & Welling, ICLR, 2017), we feed these features X⁽ˡ⁾ to a graph neural layer with C×F dimensional trainable weights W⁽ˡ⁾ , so that the output of this layer is an N×F matrix X⁽ˡ⁺¹⁾ encoding updated (and hopefully better in some sense) node features. 𝓐 is an N×N matrix, where the entry 𝓐ᵢⱼ indicates if node i is connected (adjacent) to node j. This matrix is called an adjacency matrix. I use 𝓐 instead of plain A to highlight that this matrix can be normalized in a way to facilitate feature propagation in a deep network. For the purpose of this tutorial, we can assume that 𝓐=A, i.e. each i-th row of the matrix product 𝓐X⁽ˡ⁾ will contain a sum of features of node i neighbors.

In the rest of this part of the tutorial, I’ll briefly explain works of my choice showed in bold boxes in the overview graph. I recommend Bronstein et al.’s review for a more comprehensive and formal analysis.

Note that even though I dive into some technical details of spectral graph convolution below, many recent works (e.g., GIN in Xu et al., ICLR, 2019) are built without spectral convolution and show great results in some tasks. However, knowing how spectral convolution works is still helpful to understand and avoid potential problems with other methods.

1. Spectral graph convolution

Bruna et al., 2014, ICLR 2014

I explain spectral graph convolution in detail in my another post.

I’ll briefly summarize it here for the purpose of this part of the tutorial. A formal definition of spectral graph convolution, which is very similar to the convolution theorem in signal/image processing, can be written as:

where V are eigenvectors and Λ are eigenvalues of the graph Laplacian L, which can be found by eigen-decomposition: L=VΛVᵀ; W_spectral are filters. Throughout this tutorial I’m going to assume “symmetric normalized Laplacian”. It is computed based only on an adjacency matrix A of a graph, which can be done in a few lines of Python code as follows:

# Computing the graph Laplacian
# A is an adjacency matrix
import numpy as np

N = A.shape[0] # number of nodes in a graph
D = np.sum(A, 0) # node degrees
D_hat = np.diag((D + 1e-5)**(-0.5)) # normalized node degrees
L = np.identity(N) — np.dot(D_hat, A).dot(D_hat) # Laplacian

Here, we assume that A is symmetric, i.e. A = Aᵀ and our graph is undirected, otherwise node degrees are not well-defined and some assumptions must be made to compute the Laplacian. In the context of computer vision and machine learning, the graph Laplacian defines how node features will be updated if we stack several graph neural layers in the form of formula (2).

So, given graph Laplacian L, node features X and filters W_spectral, in Python spectral convolution on graphs looks very simple:

# Spectral convolution on graphs
# X is an N×1 matrix of 1-dimensional node features
# L is an N×N graph Laplacian computed above
# W_spectral are N×F weights (filters) that we want to train
from scipy.sparse.linalg import eigsh # assumes L to be symmetric

Λ,V = eigsh(L,k=20,which=’SM’) # eigen-decomposition (i.e. find Λ,V)
X_hat = V.T.dot(X) # 20×1 node features in the "spectral" domain
W_hat = V.T.dot(W_spectral)  # 20×F filters in the "spectral" domain
Y = V.dot(X_hat * W_hat)  # N×F result of convolution

where we assume that our node features X⁽ˡ⁾ are 1-dimensional, e.g. MNIST pixels, but it can be extended to a C-dimensional case: we will just need to repeat this convolution for each channel and then sum over C as in signal/image convolution.

Formula (3) is essentially the same as spectral convolution of signals on regular grids using the Fourier Transform, and so creates a few problems for machine learning:

1. the dimensionality of trainable weights (filters) W_spectral depends on the number of nodes N in a graph;
2. W_spectral also depends on the graph structure encoded in eigenvectors V.

These issues prevent scaling to datasets with large graphs of variable structure.

To solve the first issue, Bruna et al. proposed to smooth filters in the spectral domain, which makes them more local in the spatial domain according to the spectral theory. The idea is that you can represent our filter W_spectral from formula (3) as a sum of 𝐾 predefined functions, such as splines, and instead of learning N values of W, we learn K coefficients α of this sum:

While the dimensionality of fk does depend on the number of nodes N, these functions are fixed, so we don’t learn them. The only thing we learn are coefficients α, and so W_spectral is no longer dependent on N. To make our approximation in formula (4) reasonable, we want K<<N to reduce the number of trainable parameters from N to K and, more importantly, make it independent of N, so that our GNN can digest graphs of any size.

While solves the first issue, this smoothing method does not address the second issue.

2. Chebyshev graph convolution

Defferrard et al., NeurIPS, 2016

The main drawback of spectral convolution and its smooth version above is that it still requires eigen-decomposition of an N×N dimensional graph Laplacian L, which creates two main problems:

1. 🙁 The complexity of eigen-decomposition is huge, O(N³). Moreover in case of large graphs, keeping the graph Laplacian in a dense format in RAM is infeasible. One solution is to use sparse matrices and find eigenvectors using scipy.sparse.linalg.eigs in Python. Additionally, you may preprocess all training graphs on a dedicated server with a lot of RAM and CPU cores. In many applications, your test graphs can also be preprocessed in advance, but if you have a constant influx of new large graphs, eigen-decomposition will make you sad.
2. 🙁 Another problem is that the model you train ends up being closely related to the eigenvectors V of the graph. This can be a big problem if your training and test graphs have very different structures (numbers of nodes and edges). Otherwise, if all graphs are very similar, it is less of a problem. Moreover, if you use some smoothing of filters in the frequency domain like splines discussed above, then your filters become more localized and the problem of adapting to new graphs seems to be even less noticeable. However, the models will still be quite limited.

Now, what does Chebyshev graph convolution have to do with all that?

It turns out that it solves both problems at the same time! 😃

That is, it avoids computing costly eigen-decomposition and the filters are no longer “attached” to eigenvectors (yet they still are functions of eigenvalues Λ). Moreover, it has a very useful parameter, usually denoted as K having a similar intuition as K in our formula (4) above, determining the locality of filters. Informally: for K=1, we feed just node features X⁽ˡ⁾ to our GNN; for K=2, we feed X⁽ˡ⁾ and 𝓐X⁽ˡ⁾; for K=3, we feed X⁽ˡ⁾, 𝓐X⁽ˡ⁾ and 𝓐²X⁽ˡ⁾; and so forth for larger K (I hope you’ve noticed the pattern). See more accurate and formal definition in Defferrard et al. and my code below, plus additional analysis is given in (Knyazev et al., NeurIPS-W, 2018).

Due to the power property of adjacency matrices, when we perform 𝓐²X⁽ˡ⁾ we actually average (or sum depending on how 𝓐 is normalized) over 2-hop neighbors, and analogously for any n in 𝓐ⁿX⁽ˡ⁾ as illustrated below, where we average over n-hop neighbors.

Note that to satisfy the orthogonality of the Chebyshev basis, 𝓐 assumes no loops in the graph, so that in each i-th row of matrix product 𝓐X⁽ˡ⁾ we will have features of the neighbors of node i, but not the features of node i itself. Features of node i will be fed separately as a matrix X⁽ˡ⁾.

If K equals the number of nodes N, the Chebyshev convolution closely approximates a spectral convolution, so that the receptive field of filters will be the entire graph. But, as in the case of convolutional networks, we don’t want our filters to be as big as the input images for a number of reasons that I already discussed, so in practice, K takes reasonably small values.

In my experience, this is one of the most powerful GNNs, achieving great results in a very wide range of graph tasks. The main downside is the necessity to loop over K in the forward/backward pass (since Chebyshev polynomials are recursive, so it’s not possible to parallelize them), which slows down the model.

Same as with Splines discussed above, instead of training filters, we train coefficients, but this time, of the Chebyshev polynomial.

To generate the Chebyshev basis, you can use the following Python code:

# Set K to some integer > 0, like 4 or 5 in our plots above
# Set n_points to a number of points on a curve (we set to 100)
import numpy as np

x = np.linspace(-1, 1, n_points)
T = np.zeros((K, len(x)))
T[0,:] = 1
T[1,:] = x
for n in range(1, K-1):
    T[n+1, :] = 2*x*T[n, :] - T[n-1, :] # recursive computation   
return T

The full code to generate spline and Chebyshev bases is in my github repo.

To illustrate how a Chebyshev filter can look on a irregular grid, I follow the experiment from Bruna et al. again and sample 400 random points from the MNIST grid in the same way as I did to show eigenvectors of the graph Laplacian. I trained a Chebyshev graph convolution model on the MNIST images sampled from these 400 locations (same irregular grid is used for all images) and one of the filter for K=1 and K=20 is visualized below.

3. GCN

Kipf & Welling, ICLR, 2017

As you may have noticed, if you increase K of the Chebyshev convolution, it increases the total number of trainable parameters. For example, for K=2, our weights W⁽ˡ⁾ will be 2C×F instead of just C×F. This is because we concatenate features X⁽ˡ⁾ and 𝓐X⁽ˡ⁾ into a single N×2C matrix. More training parameters means the model is more difficult to train and more data must be labeled for training. Graph datasets are often extremely small. Whereas in computer vision, MNIST is considered a tiny dataset, because images are just 28×28 dimensional and there are only 60k training images, in terms of graph networks MNIST is quite large, because each graph would have N=784 nodes and 60k is a large number of training graphs. In contrast to computer vision tasks, many graph datasets have only around 20–100 nodes and 200–1000 training examples. These graphs can represent certain small molecules and labeling chemical/biological data is usually more expensive than labeling images. Therefore, training Chebyshev convolution models can lead to severe overfitting of the training set (i.e. the model will have the training loss close to 0 yet will have a large validation or test error). So, GCN of Kipf & Welling essentially “merged” matrices of node features X⁽ˡ⁾ and 𝓐X⁽ˡ⁾ into a single N×C matrix. As a result, the model has two times fewer parameters to train compared to Chebyshev convolution with K=2, yet has the same receptive field of 1 hop. The main trick involves adding “self-loops” to your graph by adding an identity matrix I to 𝓐 and normalizing it in a particular way, so now in each i-th row of matrix product 𝓐X⁽ˡ⁾ we will have features of the neighbors of node i, as well as features of node i.

This model seems to be a standard baseline choice well-suited for many application due to its lightweight, good performance and scalability to larger graphs.

3.1. GCN vs Chebyshev layer

The difference between GCN and Chebyshev convolution is illustrated below.

The code above follows the same structure as in the first part of my tutorial, where I compared classical NN and GNN. One of the main steps both in GCN and Chebyshev convolution is computation of the rescaled graph Laplacian L. This rescaling is done to make eigenvalues in the range [-1,1] to facilitate training (this might be not a very important step in practice as weights can adapt during training). In GCN, self-loops are added to the graph by adding an identity matrix before computing the Laplacian as discussed above. The main difference between the two methods is that in the Chebyshev convolution we recursively loop over K to capture features in the K-hop neighborhood. We can stack such GCN or Chebyshev layers interleaved with nonlinearities to build a Graph Neural Network.

Now, let me politely interrupt 😃 our spectral discussion and give a general idea behind two other exciting methods: Edge-conditioned filters by Simonovsky & Komodakis, CVPR, 2017 and MoNet by Monti et al., CVPR, 2017, which share some similar concepts.

4. Edge-conditioned filters

Simonovsky & Komodakis, CVPR, 2017

As you know, in ConvNets we learn the weights (filters) by optimizing some loss like Cross Entropy. In the same way, we learn our W⁽ˡ⁾ in GNNs. Imagine that instead of learning these weights, you have another network that predicts them. So during training, we learn the weights of that auxiliary network, which takes an image or a graph as an input and returns weights W⁽ˡ⁾ (Θ in their work) as the output. The idea is based on Dynamic Filter Networks (Brabandere et al., NIPS, 2016), where “dynamic” means that filters W⁽ˡ⁾ will be different depending on the input as opposed to standard models in which filters are fixed (or static) after training.

This is a very general form of convolution that, besides images, can be easily applied to graphs or point clouds as they did in their CVPR paper and got excellent results. However, there is no “free lunch”, and training such models is quite challenging, because the regular grid constraint is now relaxed and the scope of solutions increases dramatically. This is especially true for larger graphs with many edges or for convolution in deeper layers, which often have hundreds of channels (number of features, C), so you might end up generating thousands of numbers in total for each input! In this regard, standard ConvNets are so good, because we don’t waste the model’s capacity on training to predict these weights, instead we directly enforce that the filters should be the same for all inputs. But this prior makes ConvNets limited and we cannot directly apply them to graphs or point clouds. So, as always, there’s some trade-off between flexibility and performance in a particular task.

When applied to images, like MNIST, the Edge-conditioned model can learn to predict anisotropic filters — filters that are sensitive to orientation, such as edge detectors. Compared to Gaussian filters discussed in the first part of my tutorial, these filters are able to better capture certain patterns in images, such as strokes in digits.

I want to highlight one more time that whenever we have a complicated model with auxiliary networks, it becomes a chicken-or-the-egg problem in some sense. To solve it, one of the networks — the auxiliary or the main one — should receive a very strong signal, so that it can implicitly supervise another network. In our BMVC paper, which is similar to Simonovsky & Komodakis’s work, we apply additional constraints on the edge-generating network to facilitate training. I will describe our work in detail in later posts.

5. MoNet

Monti et al., CVPR, 2017

MoNet is different from other works discussed in this post, as it assumes to have the notion of node coordinates, and therefore is more suited for geometric tasks such as 3D mesh analysis or image/video reasoning. It is somewhat similar to edge-conditioned filters of Simonovsky & Komodakis, since they also introduce an auxiliary learnable function 𝐷(𝑤, 𝜃, ρ) that predicts weights. The difference is that these weights depend on the node polar coordinates (angle 𝜃 and radius ρ); and trainable parameters 𝑤 of that function are constrained to be means and variances of Gaussians, so that instead of learning N×N matrices, we only learn fixed-size vectors (means and variances) independently of the graph size N. In terms of standard ConvNets, it would be the same as learning only 2 values (the mean and variance of a Gaussian) for each filter instead of learning 9, 25 or 121 values for 3×3, 5×5 or 11×11 dimensional filters respectively. This parameterization would greatly reduce the number of parameters in a ConvNet, but the filters would be very limited in their power to capture image features.

Monti et al. train 𝐽 means and variances of Gaussians and the process of transforming node coordinates is similar to fitting them into a Gaussian Mixture Model. The model is quite computationally intensive to train if we want our filters to be global enough, but it can be a good choice for visual tasks (see our BMVC paper for comparison), yet it is often worse than simple GCN on non-visual tasks (Knyazev et al., NeurIPS-W, 2018). Since function D depends on coordinates, the generated filters are also anisotropic and have a shape of oriented and elongated Gaussians as illustrated below.

Pseudo-code of the MoNet layer using PyTorch

# assume X to be input N×C node features
# coord are N×N×2 node coordinate differences between all pairs of nodes (node degrees for non-geometric tasks)
# coord can be viewed as angular and radial edges between nodes

1. Generate J Gaussian-shaped filters based on coordinates of nodes    using some trainable function D
   weights = D(coord)  # weights: J×N×N
2. Multiply node features X by these weights
   X = torch.bmm(weights, X.expand(J, N, C))  # X: J×N×C
3. Project features by a learnable linear transformation
   X = fc(X.permute(1, 2, 0).view(N, J*C))  # X: N×F
4. Feed X to the next layer

Conclusion

Despite a lengthy discussion, we have only scratched the surface. Applications of graph neural networks are expanding far beyond typical graph reasoning tasks, like molecule classification. The number of different graph neural layers is increasing very quickly, similar to how it was for convolutional networks a few years ago, so it’s hard to keep track of them. On that note, PyTorch Geometric (PyG) — a nice toolbox to learn from graphs — frequently populates its collection with novel layers and tricks.

Acknowledgement: A large portion of this tutorial was prepared during my internship at SRI International under the supervision of Mohamed Amer (homepage) and my PhD advisor Graham Taylor (homepage). I also thank Carolyn Augusta for useful feedback.

Find me on Github, LinkedIn and Twitter. My homepage.

If you want to cite this blog post in your paper, please use:
@misc{knyazev2019tutorial,
title={Tutorial on Graph Neural Networks for Computer Vision and Beyond},
author={Knyazev, Boris and Taylor, Graham W and Amer, Mohamed R},
year={2019}
}

<hr><p>Anisotropic, Dynamic, Spectral and Multiscale Filters Defined on Graphs was originally published in TDS Archive on Medium, where people are continuing the conversation by highlighting and responding to this story.</p>

I’m answering questions that AI/ML/CV people not familiar with graphs or graph neural networks typically ask. I provide PyTorch examples to clarify the idea behind this relatively new and exciting kind of model.

The questions addressed in this part of my tutorial are:

1. Why are graphs useful?
2. Why is it difficult to define convolution on graphs?
3. What makes a neural network a graph neural network?

To answer them, I’ll provide motivating examples, papers and Python code making it a tutorial on Graph Neural Networks (GNNs). Some basic knowledge of machine learning and computer vision is expected, however, I’ll provide some background and intuitive explanation as we go.

First of all, let’s briefly recall what is a graph? A graph G is a set of nodes (vertices) connected by directed/undirected edges. Nodes and edges typically come from some expert knowledge or intuition about the problem. So, it can be atoms in molecules, users in a social network, cities in a transportation system, players in team sport, neurons in the brain, interacting objects in a dynamic physical system, pixels, bounding boxes or segmentation masks in images. In other words, in many practical cases, it is actually you who gets to decide what are the nodes and edges in a graph.

In many practical cases, it is actually you who gets to decide what are the nodes and edges in a graph.

This is a very flexible data structure that generalizes many other data structures. For example, if there are no edges, then it becomes a set; if there are only “vertical” edges and any two nodes are connected by exactly one path, then we have a tree. Such flexibility is both good and bad as I’ll discuss in this tutorial.

1. Why graphs can be useful?

In the context of computer vision (CV) and machine learning (ML), studying graphs and the models to learn from them can give us at least four benefits:

1. We can become closer to solving important problems that previously were too challenging, such as: drug discovery for cancer (Veselkov et al., Nature, 2019); better understanding of the human brain connectome (Diez & Sepulcre, Nature Communications, 2019); materials discovery for energy and environmental challenges (Xie et al., Nature Communications, 2019).
2. In most CV/ML applications, data can be actually viewed as graphs even though you used to represent them as another data structure. Representing your data as graph(s) gives you a lot of flexibility and can give you a very different and interesting perspective on your problem. For instance, instead of learning from image pixels you can learn from “superpixels” as in (Liang et al., ECCV, 2016) and in our forthcoming BMVC paper. Graphs also let you impose a relational inductive bias in data — some prior knowledge you have about the problem. For instance, if you want to reason about a human pose, your relational bias can be a graph of skeleton joints of a human body (Yan et al., AAAI, 2018); or if you want to reason about videos, your relational bias can be a graph of moving bounding boxes (Wang & Gupta, ECCV, 2018). Another example can be representing facial landmarks as a graph (Antonakos et al., CVPR, 2015) to make reasoning about facial attributes and identity.
3. Your favourite neural network itself can be viewed as a graph, where nodes are neurons and edges are weights, or where nodes are layers and edges denote flow of forward/backward pass (in which case we are talking about a computational graph used in TensorFlow, PyTorch and other DL frameworks). An application can be optimization of a computational graph, neural architecture search, analyzing training behavior, accelerating training, etc.
4. Finally, you can solve many problems, where data can be more naturally represented as graphs, more effectively. This includes, but is not limited to, molecule and social network classification (Knyazev et al., NeurIPS-W, 2018) and generation (Simonovsky & Komodakis, ICANN, 2018), 3D Mesh classification and correspondence (Fey et al., CVPR, 2018) and generation (Wang et al., ECCV, 2018), modeling behavior of dynamic interacting objects (Kipf et al., ICML, 2018), visual scene graph modeling (see the upcoming ICCV Workshop) and question answering (Narasimhan, NeurIPS, 2018), program synthesis (Allamanis et al., ICLR, 2018), different reinforcement learning tasks (Bapst et al., ICML, 2019) and many other exciting problems.

As my previous research was related to recognizing and analyzing faces and emotions, I particularly like this figure below.

2. Why is it difficult to define convolution on graphs?

To answer this question, I first give some motivation for using convolution in general and then describe “convolution on images” using the graph terminology which should make the transition to “convolution on graphs” more smooth.

2.1. Why is convolution useful?

Let’s understand why we care about convolution so much and why we want to use it for graphs. Compared to fully-connected neural networks (a.k.a. NNs or MLPs), convolutional networks (a.k.a. CNNs or ConvNets) have certain advantages explained below based on the image of a nice old Chevy.

First, ConvNets exploit a natural prior in images, more formally described in (Bronstein et al., 2016), such as:

1. Shift-invariance — if we translate the car on the image above to the left/right/up/down, we still should be able to detect and recognize it as a car. This is exploited by sharing filters across all locations, i.e. applying convolution.
2. Locality — nearby pixels are closely related and often represent some semantic concept, such as a wheel or a window. This is exploited by using relatively large filters, which can capture image features in a local spatial neighborhood.
3. Compositionality (or hierarchy)— a larger region in the image is often a semantic parent of smaller regions it contains. For example, a car is a parent of doors, windows, wheels, driver, etc. And a driver is a parent of head, arms, etc. This is implicitly exploited by stacking convolutional layers and applying pooling.

Second, the number of trainable parameters (i.e. filters) in convolutional layers does not depend on the input dimensionality, so technically we can train exactly the same model on 28×28 and 512×512 images. In other words, the model is parametric.

Ideally, our goal is to develop a model that is as flexible as Graph Neural Nets and can digest and learn from any data, but at the same time we want to control (regularize) factors of this flexibility by turning on/off certain priors.

All these nice properties make ConvNets less prone to overfitting (high accuracy on the training set and low accuracy on the validation/test set), more accurate in different visual tasks, and easily scalable to large images and datasets. So, when we want to solve important tasks where input data are graph-structured, it is appealing to transfer all these properties to graph neural networks (GNNs) to regularize their flexibility and make them scalable. Ideally, our goal is to develop a model that is as flexible as GNNs and can digest and learn from any data, but at the same time we want to control (regularize) factors of this flexibility by turning on/off certain priors. This can open research in many interesting directions. However, controlling of this trade-off is challenging.

2.2. Convolution on images in terms of graphs

Let’s consider an undirected graph G with N nodes. Edges E represent undirected connections between nodes. Nodes and edges typically come from your intuition about the problem. Our intuition in the case of images is that nodes are pixels or superpixels (a group of pixels of weird shape) and edges are spatial distances between them. For example, the MNIST image below on the left is typically represented as an 28×28 dimensional matrix. We can also represent it as a set of N=28*28=784 pixels. So, our graph G is going to have N=784 nodes and edges will have large values (thicker edges in the Figure below) for closely located pixels and small values (thinner edges) for remote pixels.

When we train our neural networks or ConvNets on images, we implicitly define images on a graph — a regular two-dimensional grid as the one on the figure below. Since this grid is the same for all training and test images and is regular, i.e. all pixels of the grid are connected to each other in exactly the same way across all images (i.e. have the same number of neighbors, length of edges, etc.), this regular grid graph has no information that will help us to tell one image from another. Below I visualize some 2D and 3D regular grids, where the order of nodes is color-coded. By the way, I’m using NetworkX in Python to do that, e.g. G = networkx.grid_graph([4, 4]).

Given this 4×4 regular grid, let’s briefly look at how 2D convolution works to understand why it’s difficult to transfer this operator to graphs. A filter on a regular grid has the same order of nodes, but modern convolutional nets typically have small filters, such as 3×3 in the example below. This filter has 9 values: W₁,W₂,…, W₉, which is what we are updating during training using backprop to minimize the loss and solve the downstream task. In our example below, we just heuristically initialize this filter to be an edge detector (see other possible filters here):

When we perform convolution, we slide this filter in both directions: to the right and to the bottom, but nothing prevents us from starting in the bottom corner — the important thing is to slide over all possible locations. At each location, we compute the dot product between the values on the grid (let’s denote them as X) and the values of filters, W: X₁W₁+X₂W₂+…+X₉W₉, and store the result in the output image. In our visualization, we change the color of nodes during sliding to match the colors of nodes in the grid. In a regular grid, we always can match a node of the filter with a node of the grid. Unfortunately, this is not true for graphs as I’ll explain later below.

The dot product used above is one of so called “aggregator operators”. Broadly speaking, the goal of an aggregator operator is to summarize data to a reduced form. In our example above, the dot product summarizes a 3×3 matrix to a single value. Another example is pooling in ConvNets. Keep in mind, that such methods as max or sum pooling are permutation-invariant, i.e. they will pool the same value from a spatial region even if you randomly shuffle all pixels inside that region. To make it clear, the dot product is not permutation-invariant simply because in general: X₁W₁+X₂W₂ ≠X₂W₁+X₁W₂.

Now let’s use our MNIST image and illustrate the meaning of a regular grid, a filter and convolution. Keeping in mind our graph terminology, this regular 28×28 grid will be our graph G, so that every cell in this grid is a node, and node features are an actual image X, i.e. every node will have just a single feature — pixel intensity from 0 (black) to 1 (white).

Next, we define a filter and let it be a famous Gabor filter with some (almost) arbitrary parameters. Once we have an image and a filter, we can perform convolution by sliding the filter over that image (of digit 7 in our case) and putting the result of the dot product to the output matrix after each step.

This is all cool, but as I mentioned before, it becomes tricky when you try to generalize convolution to graphs.

Nodes are a set, and any permutation of this set does not change it. Therefore, the aggregator operator that people apply should be permutation-invariant.

As I have already mentioned, the dot product used above to compute convolution at each step is sensitive to the order. This sensitivity permits us to learn edge detectors similar to Gabor filters important to capture image features. The problem is that in graphs there is no well-defined order of nodes unless you learn to order them, or come up with some heuristic that will result in a consistent (canonical) order from graph to graph. In short, nodes are a set, and any permutation of this set does not change it. Therefore, the aggregator operator that people apply should be permutation-invariant. The most popular choices are averaging (GCN, Kipf & Welling, ICLR, 2017) and summation (GIN, Xu et al., ICLR, 2019) of all neighbors, i.e. sum or mean pooling, followed by projection by a trainable vector W. See Hamilton et al., NIPS, 2017 for some other aggregators.

For example, for the graph above on the left, the output of the summation aggregator for node 1 will be X₁=(X₁+X₂+X₃+X₄)W₁, for node 2: X₂=(X₁+X₂+X₃+X₅)W₁ and so forth for nodes 3, 4 and 5, i.e. we need to apply this aggregator for all nodes. In result, we will have the graph with the same structure, but node features will now contain features of neighbors. We can process the graph on the right using the same idea.

Colloquially, people call this averaging or summation “convolution”, since we also “slide” from one node to another and apply an aggregator operator in each step. However, it’s important to keep in mind that this is a very specific form of convolution, where filters don’t have a sense of orientation. Below I’ll show how those filters look like and give an idea how to make them better.

3. What makes a neural network a graph neural network?

You know how a classical neural network works, right? We have some C-dimensional features X as the input to the net. Using our running MNIST example, X will be our C=784 dimensional pixel features (i.e. a “flattened” image). These features get multiplied by C×F dimensional weights W that we update during training to get the output closer to what we expect. The result can be directly used to solve the task (e.g. in case of regression) or can be further fed to some nonlinearity (activation), like ReLU, or other differentiable (or more precisely, sub-differentiable) functions to form a multi-layer network. In general, the output of some layer l is:

The signal in MNIST is so strong, that you can get an accuracy of 91% by just using the formula above and the Cross Entropy loss without any nonlinearities and other tricks (I used a slightly modified PyTorch example to do that). Such model is called multinomial (or multiclass, since we have 10 classes of digits) logistic regression.

Now, how do we transform our vanilla neural network to a graph neural network? As you already know, the core idea behind GNNs is aggregation over “neighbors”. Here, it is important to understand that in many cases, it is actually you who specifies “neighbors”.

Let’s consider a simple case first, when you are given some graph. For example, this can be a fragment (subgraph) of a social network with 5 persons and an edge between a pair of nodes denotes if two people are friends (or at least one of them think so). An adjacency matrix (usually denoted as A) in the figure below on the right is a way to represent these edges in a matrix form, convenient for our deep learning frameworks. Yellow cells in the matrix represent the edge and blue — the absence of the edge.

Now, let’s create an adjacency matrix A for our MNIST example based on coordinates of pixels (complete code is provided in the end of the post):

import numpy as np
from scipy.spatial.distance import cdist

img_size = 28  # MNIST image width and height
col, row = np.meshgrid(np.arange(img_size), np.arange(img_size))
coord = np.stack((col, row), axis=2).reshape(-1, 2) / img_size
dist = cdist(coord, coord)  # see figure below on the left
sigma = 0.2 * np.pi  # width of a Gaussian
A = np.exp(- dist ** 2 / sigma ** 2)  # see figure below in the middle

This is a typical, but not the only, way to define an adjacency matrix for visual tasks (Defferrard et al., NIPS, 2016, Bronstein et al., 2016). This adjacency matrix is our prior, or our inductive bias, we impose on the model based on our intuition that nearby pixels should be connected and remote pixels shouldn’t or should have very thin edge (edge of a small value). This is motivated by observations that in natural images nearby pixels often correspond to the same object or objects that interact frequently (the locality principle we mentioned in Section 2.1.), so it makes a lot of sense to connect such pixels.

So, now instead of having just features X we have some fancy matrix A with values in the range [0,1]. It’s important to note that once we know that our input is a graph, we assume that there is no canonical order of nodes that will be consistent across all other graphs in the dataset. In terms of images, it means that pixels are assumed to be randomly shuffled. Finding the canonical order of nodes is combinatorially unsolvable in practice. Even though for MNIST we technically can cheat by knowing this order (because data are originally from a regular grid), it’s not going to work on actual graph datasets.

Remember that our matrix of features X has 𝑁 rows and C columns. So, in terms of graphs, each row corresponds to one node and C is the dimensionality of node features. But now the problem is that we don’t know the order of nodes, so we don’t know in which row to put features of a particular node. If we just pretend to ignore this problem and feed X directly to an MLP as we did before, the effect will be the same as feeding images with randomly shuffled pixels with independent (yet the same for each epoch) shuffling for each image! Surprisingly, a neural network can in principle still fit such random data (Zhang et al., ICLR, 2017), however test performance will be close to random prediction. One of the solutions is to simply use the adjacency matrix A, we created before, in the following way:

We just need to make sure that row i in A corresponds to features of node in row i of X. Here, I’m using 𝓐 instead of plain A, because often you want to normalize A. If 𝓐=A, the matrix multiplication 𝓐X⁽ˡ⁾ will be equivalent to summing features of neighbors, which turned out to be useful in many tasks (Xu et al., ICLR, 2019). Most commonly, you normalize it so that 𝓐X⁽ˡ⁾ averages features of neighbors, i.e. 𝓐=A/ΣᵢAᵢ. A better way to normalize matrix A can be found in (Kipf & Welling, ICLR, 2017).

Below is the comparison of NN and GNN in terms of PyTorch code:

And HERE is the full PyTorch code to train two models above: python mnist_fc.py --model fc to train the NN case; python mnist_fc.py --model graph to train the GNN case. As an exercise, try to randomly shuffle pixels in code in the --model graph case (don’t forget to shuffle A in the same way) and make sure that it will not affect the result. Is it going to be true for the --model fc case?

Here is the full PyTorch code to train two models.

After running the code, you may notice that the classification accuracy is actually about the same. What’s the problem? Aren’t graph networks supposed to work better? Well, they are, in many cases. But not in this one, because the 𝓐X⁽ˡ⁾ operator we added is actually nothing else, but a Gaussian filter:

So, our graph neural network turned out to be equivalent to a convolutional neural network with a single Gaussian filter, that we never update during training, followed by the fully-connected layer. This filter basically blurs/smooths the image, which is not a particularly useful thing to do (see the image above on the right). However, this is the simplest variant of a graph neural network, which nevertheless works great on graph-structured data. To make GNNs work better on regular graphs, like images, we need to apply a bunch of tricks. For example, instead of using a predefined Gaussian filter, we can learn to predict an edge between any pair of pixels by using a differentiable function like this:

import torch.nn as nn  # using PyTorch

nn.Sequential(nn.Linear(4, 64),  # map coordinates to a hidden layer
              nn.ReLU(),         # nonlinearity
              nn.Linear(64, 1),  # map hidden representation to edge
              nn.Tanh())         # squash edge values to [-1, 1]

To make GNNs work better on regular graphs, like images, we need to apply a bunch of tricks. For example, instead of using a predefined Gaussian filter, we can learn to predict an edge between any pair of pixels.

This idea is similar to Dynamic Filter Networks (Brabander et al., NIPS, 2016), Edge-conditioned Graph Networks (ECC, Simonovsky & Komodakis, CVPR, 2017) and (Knyazev et al., NeurIPS-W, 2018). To try it using my code, you just need to add the --pred_edge flag, so the entire command is python mnist_fc.py --model graph --pred_edge. Below I show the animation of the predefined Gaussian and learned filters. You may notice that the filter we just learned (in the middle) looks weird. That’s because the task is quite complicated since we optimize two models at the same time: the model that predicts edges and the model that predicts a digit class. To learn better filters (like the one on the right), we need to apply some other tricks from our BMVC paper, which is beyond the scope of this part of the tutorial.

The code to generate these GIFs is quite simple:

I’m also sharing an IPython notebook showing 2D convolution of an image with a Gabor filter in terms of graphs (using an adjacency matrix) compared to using circulant matrices, which is often used in signal processing.

In the next part of the tutorial, I’ll tell you about more advanced graph layers that can lead to better filters on graphs.

Update:

Throught this blog post and in the code the dist variable should have been squared to make it a Gaussian. Thanks Alfredo Canziani for spotting that. All figures and results were generated without squaring it. If you observe very different results after squaring it, I suggest to tune sigma.

Conclusion

Graph Neural Networks are a very flexible and interesting family of neural networks that can be applied to really complex data. As always, such flexibility must come at a certain cost. In case of GNNs it is the difficulty of regularizing the model by defining such operators as convolution. Research in that direction is advancing quite fast, so that GNNs will see application in increasingly wider areas of machine learning and computer vision.

See another nice blog post about GNNs from Neptune.ai.

Acknowledgement: A large portion of this tutorial was prepared during my internship at SRI International under the supervision of Mohamed Amer (homepage) and my PhD advisor Graham Taylor (homepage).

Find me on Github, LinkedIn and Twitter. My homepage.

If you want to cite this tutorial in your paper, please use:
@misc{knyazev2019tutorial,
title={Tutorial on Graph Neural Networks for Computer Vision and Beyond},
author={Knyazev, Boris and Taylor, Graham W and Amer, Mohamed R},
year={2019}
}