An equivariant graph embedding¶

Permuted adjacency matrices describe the same graph¶

Let us first verify that permuted adjacency matrices really describe one and the same graph. We also gain some useful data generation functions for later.

First we create random adjacency matrices. The entry $a_{ij}$ of this matrix corresponds to the weight of the edge between nodes $i$ and $j$ in the graph. We assume that graphs have no self-loops; instead, the diagonal elements of the adjacency matrix are interpreted as node weights (or “node attributes”).

Taking the example of a Twitter user retweet network, the nodes would be users, edge weights indicate how often two users retweet each other and node attributes could indicate the follower count of a user.

import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

rng = np.random.default_rng(4324234)

def create_data_point(n):
    """
    Returns a random undirected adjacency matrix of dimension (n,n).
    The diagonal elements are interpreted as node attributes.
    """
    mat = rng.random((n, n))
    A = (mat + np.transpose(mat))/2
    return np.round(A, decimals=2)

A = create_data_point(3)
print(A)

[[0.36 0.52 0.27]
 [0.52 0.32 0.8 ]
 [0.27 0.8  0.92]]

Let’s also write a function to generate permuted versions of this adjacency matrix.

def permute(A, permutation):
    """
    Returns a copy of A with rows and columns swapped according to permutation.
    For example, the permutation [1, 2, 0] swaps 0->1, 1->2, 2->0.
    """

    P = np.zeros((len(A), len(A)))
    for i,j in enumerate(permutation):
        P[i,j] = 1

    return P @ A @ np.transpose(P)

A_perm = permute(A, [1, 2, 0])
print(A_perm)

[[0.32 0.8  0.52]
 [0.8  0.92 0.27]
 [0.52 0.27 0.36]]

If we create networkx graphs from both adjacency matrices and plot them, we see that they are identical as claimed.

fig, (ax1, ax2) = plt.subplots(1, 2)

# interpret diagonal of matrix as node attributes
node_labels = {n: A[n,n] for n in range(len(A))}
np.fill_diagonal(A, np.zeros(len(A)))

G1 = nx.Graph(A)
pos1 = nx.spring_layout(G1, seed=1)
nx.draw(G1, pos1, labels=node_labels, ax=ax1, node_size = 800, node_color = "#ACE3FF")
edge_labels = nx.get_edge_attributes(G1,'weight')
nx.draw_networkx_edge_labels(G1,pos1,edge_labels=edge_labels, ax=ax1)

# interpret diagonal of permuted matrix as node attributes
node_labels = {n: A_perm[n,n] for n in range(len(A_perm))}
np.fill_diagonal(A_perm, np.zeros(len(A)))

G2 = nx.Graph(A_perm)
pos2 = nx.spring_layout(G2, seed=1)
nx.draw(G2, pos2, labels=node_labels, ax=ax2, node_size = 800, node_color = "#ACE3FF")
edge_labels = nx.get_edge_attributes(G2,'weight')
nx.draw_networkx_edge_labels(G2,pos2,edge_labels=edge_labels, ax=ax2)

ax1.set_xlim([1.2*x for x in ax1.get_xlim()])
ax2.set_xlim([1.2*x for x in ax2.get_xlim()])
plt.tight_layout()
plt.show()

Permutation equivariant embeddings¶

When we design a machine learning model that takes graph data, the first step is to encode the adjacency matrix into a quantum state using an embedding or quantum feature map $\phi:$

We may want the resulting quantum state to be the same for all adjacency matrices describing the same graph. In mathematical terms, this means that $\phi$ is an invariant embedding with respect to simultaneous row and column permutations $\pi(A)$ of the adjacency matrix:

However, invariance is often too strong a constraint. Think for example of an encoding that associates each node in the graph with a qubit. We might want permutations of the adjacency matrix to lead to the same state up to an equivalent permutation of the qubits $P_{\pi},$ where

The function $\text{perm}_{\pi}$ maps each index to the permuted index according to $\pi.$

This results in an equivariant embedding with respect to permutations of the adjacency matrix:

This is exactly what the following quantum embedding is aiming to do! The mathematical details behind these concepts use group theory and are beautiful, but can be a bit daunting. Have a look at this paper if you want to learn more.

Implementation in PennyLane¶

Let’s get our hands dirty with an example. As mentioned, we will implement the permutation-equivariant embedding suggested in Skolik et al. (2022) which has this structure:

The image can be found in Skolik et al. (2022) and shows one layer of the circuit. The $\epsilon$ are our edge weights while $\alpha$ describe the node weights, and the $\beta,$ $\gamma$ are variational parameters.

In PennyLane this looks as follows:

import pennylane as qml

def perm_equivariant_embedding(A, betas, gammas):
    """
    Ansatz to embedd a graph with node and edge weights into a quantum state.

    The adjacency matrix A contains the edge weights on the off-diagonal,
    as well as the node attributes on the diagonal.

    The embedding contains trainable weights 'betas' and 'gammas'.
    """
    n_nodes = len(A)
    n_layers = len(betas) # infer the number of layers from the parameters

    # initialise in the plus state
    for i in range(n_nodes):
        qml.Hadamard(i)

    for l in range(n_layers):

        for i in range(n_nodes):
            for j in range(i):
                # factor of 2 due to definition of gate
                qml.IsingZZ(2*gammas[l]*A[i,j], wires=[i,j])

        for i in range(n_nodes):
            qml.RX(A[i,i]*betas[l], wires=i)

We can use this ansatz in a circuit.

n_qubits = 5
n_layers = 2

dev = qml.device("lightning.qubit", wires=n_qubits)

@qml.qnode(dev)
def eqc(adjacency_matrix, observable, trainable_betas, trainable_gammas):
    """Circuit that uses the permutation equivariant embedding"""

    perm_equivariant_embedding(adjacency_matrix, trainable_betas, trainable_gammas)
    return qml.expval(observable)


A = create_data_point(n_qubits)
betas = rng.random(n_layers)
gammas = rng.random(n_layers)
observable = qml.PauliX(0) @ qml.PauliX(1) @ qml.PauliX(3)

qml.draw_mpl(eqc, decimals=2)(A, observable, betas, gammas)
plt.show()

Validating the equivariance¶

Let’s now check if the circuit is really equivariant!

This is the expectation value we get using the original adjacency matrix as an input:

result_A = eqc(A, observable, betas, gammas)
print("Model output for A:", result_A)

Model output for A: 0.4293394554691393

If we permute the adjacency matrix, this is what we get:

perm = [2, 3, 0, 1, 4]
A_perm = permute(A, perm)
result_Aperm = eqc(A_perm, observable, betas, gammas)
print("Model output for permutation of A: ", result_Aperm)

Model output for permutation of A:  0.4536992005613013

Why are the two values different? Well, we constructed an equivariant ansatz, not an invariant one! Remember, an invariant ansatz means that embedding a permutation of the adjacency matrix leads to the same state as an embedding of the original matrix. An equivariant ansatz embeds the permuted adjacency matrix into a state where the qubits are permuted as well.

As a result, the final state before measurement is only the same if we permute the qubits in the same manner that we permute the input adjacency matrix. We could insert a permutation operator qml.Permute(perm) to achieve this, or we simply permute the wires of the observables!

observable_perm = qml.PauliX(perm[0]) @ qml.PauliX(perm[1]) @ qml.PauliX(perm[3])

Now everything should work out!

result_Aperm = eqc(A_perm, observable_perm, betas, gammas)
print("Model output for permutation of A, and with permuted observable: ", result_Aperm)

Model output for permutation of A, and with permuted observable:  0.42933945546913965

Et voilà!

Conclusion¶

Equivariant graph embeddings can be combined with other equivariant parts of a quantum machine learning pipeline (like measurements and the cost function). Skolik et al. (2022), for example, use such a pipeline as part of a reinforcement learning scheme that finds heuristic solutions for the traveling salesman problem. Their simulations compare a fully equivariant model to circuits that break permutation equivariance and show that it performs better, confirming that if we know about structure in our data, we should try to use this knowledge in machine learning.

References¶

1. Andrea Skolik, Michele Cattelan, Sheir Yarkoni,Thomas Baeck and Vedran Dunjko (2022). Equivariant quantum circuits for learning on weighted graphs. arXiv:2205.06109

2. Quynh T. Nguyen, Louis Schatzki, Paolo Braccia, Michael Ragone, Patrick J. Coles, Frédéric Sauvage, Martín Larocca and Marco Cerezo (2022). Theory for Equivariant Quantum Neural Networks. arXiv:2210.08566

About the author¶

Total running time of the script: (0 minutes 0.487 seconds)