GSoC 24 | Quantum Generative Adversarial Networks for HEP event generation the LHC

How did Gen AI start? well, it’s GANs

I’m Adithya Penagonda, one of the contributors for ML4SCI/QMLHEP for Google Summer of Code 2024. Here is my project page: link This is a mid-term update of my project, all links available as footnotes.

Understanding the Problem: Jet Generation in the LHC

A marvel today in science, the Large Hadron Collider (LHC) smashes protons together almost as quickly as light in order to simulate conditions right after the Big Bang. Jets are among the different particles produced by these high-energy collisions. Particles that are sprayed into space as a result of quark and gluon hadronization are known as jets. Because they shed light on the fundamental particles and forces that comprise our universe, these jets are essential to the study of particle physics.

Why Use GANs for Jet Generation?

Conventional jet simulation techniques require large computational resources and intricate algorithms. Generative Adversarial Networks, or GANs, provide a viable substitute by using data to create realistic jet images. Two neural networks make up a GAN: the discriminator, which assesses the validity of the generated data, and the generator, which produces synthetic data. GANs can generate extremely realistic jet images by iteratively improving, greatly accelerating the simulation process.

The process is similar to that of a police officer trying to catch a fake note and the thief trying to generate better and better notes similar to that of original. A classical GAN looks something like below:

classical GAN architecture

GANs: How They Work

A GAN consists of two main components:

1. Generator: Takes random noise as input and generates synthetic data.
2. Discriminator: Evaluates the generated data against real data, aiming to distinguish between the two.

The discriminator’s goal in the training process is to correctly identify real and fake data, while the generator’s goal is to trick the discriminator in a minimax game. This can be stated mathematically as:

The most important parts of the a classical GAN pipeline are

1. Data: Data is the foundation of any machine learning model, including GANs. For this project, we are using two specific datasets

• Electron-Photon Dataset: This dataset contains images or data points representing electron-photon interactions.
• Quark-Gluon Dataset: Similar to the Electron-Photon dataset, this dataset contains images or data points representing quark-gluon interactions. These were explained in detail by one of the previous contributors Marçal Comajoan Cara, check out their blog for more details about the datasets used.

2. Loss functions: The performance and training dynamics of GANs are greatly influenced by the selection of loss function. We experimented with the following loss functions, we will explore more about them in later sections

• Wasserstein GAN Loss: This loss function stabilizes GAN training by providing a smoother gradient flow. It calculates the difference between the distribution of real and generated data, which encourages the generator to create more realistic images.
• Total Variation: This loss function promotes smoother images by penalizing significant differences between neighboring pixels.
• Perceptual Loss: Perceptual loss employs pre-trained neural networks to compare the similarity of real and generated images in a feature space, rather than pixel-wise.

3. Network Evaluation: Understanding a GAN’s effectiveness requires evaluating its performance. We use the following metrics for network evaluation:

• Cross Entropy: This metric measures the discriminator’s ability to distinguish between real and fake images. It will help to determining how effectively the generator fools the discriminator.
• Fréchet Inception Distance (FID): FID score measures the distance between the distributions of real and generated data. It is a robust metric for evaluating the quality and diversity of generated images. A lower FID score indicates better performance. This was suggested by one of the mentors (Gopal Dahale) and proved to be very effective.

4. Gradient (Optimizers): An optimizer calculates gradients, updates parameters and also controls hyperparameters like learning rate. We experimented with:

• RMSprop Optimizer: RMSprop adjusts the learning rate for each parameter based on the average of recent gradients, helping to stabilize training
• Adam Optimizer: Adam combines the benefits of RMSprop and momentum, often leading to faster convergence and more stable training dynamics.

Previous contributors working on the same project have already worked on classical GANs like Amey Bhatuse, you can look at their work here. Though GANs sound intuitive to implement they come with major issues like mode collapse, Barren plateaus and vanishing gradients, apart from the ones that are common to neural networks like overfitting.

This year, we wanted to test the impact loss function has over training GANs. It’s obvious that every other part of the pipeline would remain same. Log loss, familiar to almost everyone who learned about neural networks would straight away lead to mode collapse. To brief what mode collapse is, it’s when the generator is not able to use a good enough distribution to produce samples that are unique. For jets, mode collapse would typically look something like below:

Mode collapse output

Can you notice how similar almost every output looks. So, after a lot of trial and errors, the following pipeline seems to work well:

Generator Architecture:

class Generator(nn.Module):
    def __init__(self, latent_dim):
        super(Generator, self).__init__()
        self.latent_dim = latent_dim

        self.model = nn.Sequential(
            nn.ConvTranspose2d(latent_dim, 256, 4, 1, 0, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.Dropout(0.3),
            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.Dropout(0.3),
            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.Dropout(0.3),
            nn.ConvTranspose2d(64, 1, 3, 1, 1, bias=False),
            nn.Tanh()
        )

    def forward(self, z):
        return self.model(z)

Discriminator Architecture:

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()

        self.model = nn.Sequential(
            spectral_norm(nn.Conv2d(1, 64, 4, 2, 1, bias=False)),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Dropout(0.3),
            spectral_norm(nn.Conv2d(64, 128, 4, 2, 1, bias=False)),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Dropout(0.3),
            nn.Flatten(),
            spectral_norm(nn.Linear(128*4*4, 1, bias=False))
        )

    def forward(self, x):
        return self.model(x)

Network evaluation: We use FID scores, it works by passing both the real and generated images through a pre-trained Inception network (usually Inception-v3). This network extracts feature representations from images, typically from the layer before the output layer. These are high-dimensional vectors that capture important information about the images.

The FID is computed as the Fréchet distance between the two multivariate Gaussian distributions defined by the means and covariances of the real and generated images. Mathematically,

def calculate_fid(real_images, fake_images, batch_size=128):
    inception_model = models.inception_v3(pretrained=True, transform_input=False).cuda()
    inception_model.eval()

    def get_features(images):
        features = []
        for i in range(0, len(images), batch_size):
            batch = images[i:i+batch_size].cuda()
            batch = F.interpolate(batch, size=(299, 299), mode='bilinear', align_corners=False)
            batch = batch.repeat(1, 3, 1, 1)
            with torch.no_grad():
                pred = inception_model(batch)
            features.append(pred.cpu().numpy())
        return np.concatenate(features, axis=0)

    real_features = get_features(real_images)
    fake_features = get_features(fake_images)

    mu1, sigma1 = real_features.mean(axis=0), np.cov(real_features, rowvar=False)
    mu2, sigma2 = fake_features.mean(axis=0), np.cov(fake_features, rowvar=False)

    if sigma1.ndim == 0:
        sigma1 = np.array([[sigma1]])
    if sigma2.ndim == 0:
        sigma2 = np.array([[sigma2]])

    ssdiff = np.sum((mu1 - mu2) ** 2.0)
    covmean = sqrtm(sigma1.dot(sigma2))
    if np.iscomplexobj(covmean):
        covmean = covmean.real
    fid = ssdiff + np.trace(sigma1 + sigma2 - 2.0 * covmean)
    return fid

We’ve used RMSprop for loss testing, but later also tested another model by using Adam optimizer.

How effective are loss functions?

1. Wasserstein GAN loss: The Wasserstein GAN takes a different approach than the traditional GAN loss. Instead of using the sigmoid cross-entropy loss, it uses the Wasserstein distance to provide a meaningful measure of how far the generated data distribution deviates from the actual data distribution. This loss function is used with a critic (rather than a discriminator) to approximate the Wasserstein distance. To enforce the Lipschitz constraint, the gradient penalty is added.

Wasserstein GAN loss

2. Total Variation loss: Total Variation (TV) loss is used to reduce noise and improve the smoothness of generated images. It penalizes differences between adjacent pixel values.

The TV loss minimizes the sum of absolute differences between neighboring pixel values in the generated image. This encourages the image to have areas with smooth intensity transitions, lowering noise.

Total Variation loss

3. Perceptual Loss: Perceptual loss, also known as content loss or feature loss, is the difference between high-level feature representations of real and generated images extracted by a pre-trained network, most commonly a convolutional neural network (CNN) such as VGG.

The perceptual loss algorithm compares feature maps obtained from intermediate layers of a pre-trained network to both real and generated images. This aids in capturing perceptually relevant differences that are not adequately captured by pixel-wise losses such as L2.

Perceptual loss

Here are the results of training, all of them were trained for 30 epochs:

Wasserstein GAN loss

Total variation loss

Perception loss

We can compare how well each loss function did by seeing the FID scores, lower scores mean they are close to real images.

These findings suggest that loss functions do make a significant difference. Out of our testing, Total Variation Loss appears to work the best, resulting in the lowest FID scores. You can explore the project repository for all the image outputs. You can look at other plots like generator loss and discriminator loss below:

Let's introduce Quantum

We were working with a hybrid model rather than a fully quantum model. A hybrid model means that one of the neural networks runs on a quantum computer while the other runs on a classical one.

We usually have a Quantum generator and Classical Discriminator, and not the other way around because of a phenomenon called no-cloning theorem, which essentially means that it’s impossible to copy a quantum state without losing information.

A typical Hybrid QGAN

Before we understand how this works, it’s important to know what encoding and Ansatz mean.

1. Encoding: Classical data must be encoded into quantum states, for quantum machines to work on. We’ve used Angle Encoding which involves encoding classical data points into the amplitudes of the quantum state through rotation gates.
2. PQC: A parameterized quantum circuit (PQC), also known as a quantum ansatz, is a quantum circuit with gate operations determined by a set of tunable parameters. These parameters are adjusted during training to optimize the quantum circuit’s performance for a specific task.

You can create a circuit by using Pennylane:

n_qubits = 4
latent_dim = n_qubits
dev = qml.device('default.qubit', wires=n_qubits)

def quantum_circuit(params, data, n_qubits=4):
    depth = len(params) // (2 * n_qubits)
    for d in range(depth):
        for i in range(n_qubits):
            qml.RY(data[i], wires=i)
            qml.RY(params[d * 2 * n_qubits + i], wires=i)
        for i in range(n_qubits):
            qml.CNOT(wires=[i, (i + 1) % n_qubits])
        for i in range(n_qubits):
            qml.RZ(params[d * 2 * n_qubits + n_qubits + i], wires=i)
    return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

qnode = qml.QNode(quantum_circuit, dev)

What this code does is:

• The quantum_circuit function defines the structure of the quantum circuit.
• params is an array of parameters used in the rotation gates.
• data is the input data to be encoded into the quantum circuit.
• The circuit consists of RY gates for data and parameter encoding, CNOT gates for entanglement, and RZ gates for additional parameterization.
• The qml.QNode creates a quantum node that can be used to evaluate the circuit on a given quantum device (dev).

Usually since, the number of physically available qubits are comparatively less than dimensions of the image, we use a classical upscaler.

Upscaler for MNIST dataset

class QuantumGenerator(nn.Module):
    def __init__(self, n_qubits=4, depth=3, output_dim=16*16):
        super(QuantumGenerator, self).__init__()
        self.n_qubits = n_qubits
        self.depth = depth
        self.output_dim = output_dim
        self.params = nn.Parameter(torch.randn((depth * 2 * n_qubits,), requires_grad=True))
        self.fc1 = nn.Linear(n_qubits, 128)
        self.fc2 = nn.Linear(128, 256)
        self.fc3 = nn.Linear(256, self.output_dim)
        self.bn1 = nn.BatchNorm1d(128)
        self.bn2 = nn.BatchNorm1d(256)

    def forward(self, x):
        q_out = []
        for i in range(x.shape[0]):
            data = x[i].detach().cpu().numpy()
            result = np.array(qnode(self.params.detach().cpu().numpy(), data))
            q_out.append(result)
        q_out = torch.tensor(q_out, dtype=torch.float32).to(device)
        q_out = F.relu(self.bn1(self.fc1(q_out)))
        q_out = F.relu(self.bn2(self.fc2(q_out)))
        q_out = torch.tanh(self.fc3(q_out))
        return q_out.view(-1, 1, 16, 16)

The idea is to use an encoder that takes input data and applies parameterized rotations and entanglements and outputs expectation values of Pauli-Z operators. then use classical fully connected layers with batch normalization to upscale the image (28x28 for MNIST). The quantum circuit is parameterized with trainable parameters to adjust during training which are optimized through classical backpropagation.

The Generator circuit (CNOT gates add entanglement)

So how does this work?

It’s always a standard to test any pipeline with MNIST dataset to make sure there is no issues with the models. In our case, making sure that no mode collapse occurs.

Hybrid QGAN training on MNIST

Final image after training the model

Looks like a good output, right? not really though the output looks noiseless it’s the loss landscape that is worrisome. Convergence during training is a very important aspect to look at, and this model did not.

A messy loss landscape

On a similar note, the model didn’t perform well on jets images as well.

Remember earlier, we tested the importance of loss function? We put that data to use. Instead of going with random experimentation with the model we started with using Total variation as the loss function.

and well, these are the results:

Classical GANs (loss: Total variation)

Hybrid QGANs (loss: Total variation)

Note that gaussian blur was applied to the output of Hybrid QGANs

The Hybrid model was trained on much lesser data than the classical models. (1:7), but still performed equivalently to the classical model.

Final thoughts and future work

If we thoroughly look at Hybrid QGANs, it’s truly doesn’t explore the quantum advantage. A lot of parts are still done by classical computing including upscaling. It would be interesting to see implementing a fully quantum model. That would be where I would spend time to explore.

I finally thank my mentors Abhay Kamble, Tom Magorsch, Gopal Ramesh Dahale, Rui Zhang for their valuable time and suggestions. A Special mention to Sergei V. Gleyzer, Ph.D., for the way he organizes the program.

Here is a link to my Code: Click here Link to slides: Click here.

You can also look at my proposal at here.