How to Calculate KL Divergence Loss of Neural Networks in PyTorch?

PyTorch allows users to build neural networks and evaluate their performance using different loss functions like MAE, MSE, and KL Divergence. The Kullback-Leibler (KL) Divergence loss measures the distance between the probability distributions of correct and incorrect predictions. PyTorch provides torch.nn.KLDivLoss() for this purpose, while Keras provides the equivalent kl_divergence loss argument in compile().

Sequential vs MLP Neural Network for KL Divergence

This guide covers two model architectures for computing KL Divergence loss. Here’s a quick comparison before diving in:

Feature	Sequential Model	MLP Model (nn.Module)
Framework	Keras/TensorFlow	PyTorch (torch.nn)
Loss function	compile(loss='kl_divergence')	nn.KLDivLoss()
Architecture style	Linear stack — one input, one output per step	Fully connected layers with feedforward pass
Best for	Streaming data, time series, quick prototyping	Image data, custom architectures, full PyTorch pipelines
Optimizer used	Adam via string argument	torch.optim.Adam()

Model 1: Using Sequential Neural Network

The Sequential model processes data in a strict sequence — one input produces one output. It works efficiently for streaming data like time series and strings.

Step 1: Access Python Notebook

Open a Python environment using Jupyter or Google Colab. This guide uses Google Colaboratory:

Step 2: Import Libraries

import torch
from keras.layers import Dense
from sklearn.datasets import make_circles
from keras.models import Sequential
from matplotlib import pyplot
from numpy import where

• torch — deep learning library with methods for model training and loss computation.
• Keras — Google-built API for building and stacking neural network layers.
• sklearn — Python machine learning library for dataset generation and splitting.
• Matplotlib — visualization library for plotting loss and accuracy curves.
• NumPy — array library used here to build conditional logic on the dataset.

Step 3: Build the Dataset

a, b = make_circles(n_samples=1000, noise=0.1, random_state=1)
for i in range(2):
    samples_ix = where(b == i)
    pyplot.scatter(a[samples_ix, 0], a[samples_ix, 1])
pyplot.show()

Step 4: Split into Training and Testing Samples

n_test = 500
traina = a[:n_test]
testa = a[n_test:]
trainb = b[:n_test]
testb = b[n_test:]

Step 5: Build the Sequential Model

model = Sequential()
model.add(Dense(100, input_shape=(2,), activation='relu'))
model.add(Dense(1, activation='sigmoid'))
model.compile(loss='kl_divergence', optimizer='adam', metrics=['accuracy'])

• A 100-neuron hidden layer with ReLU activation introduces non-linearity.
• The output layer uses Sigmoid to produce a probability in the range [0, 1].
• The kl_divergence loss measures how far predictions diverge from true labels.
• Adam optimizer adapts the learning rate dynamically during training.

Step 6: Fit and Evaluate the Model

result = model.fit(traina, trainb, validation_data=(testa, testb), epochs=30, verbose=0)

_, train_acc = model.evaluate(traina, trainb, verbose=1)
_, test_acc = model.evaluate(testa, testb, verbose=1)

Similar training and test loss values close to zero indicate a well-balanced, accurate model — both key indicators of good generalization.

Step 7: Plot the Loss Curve

pyplot.subplot(211)
pyplot.title('Loss')
pyplot.plot(result.history['loss'], label='train')
pyplot.plot(result.history['val_loss'], label='test')
pyplot.legend()
pyplot.show()

Model 2: Using MLP Neural Network

The MultiLayer Perceptron (MLP) model uses fully connected neurons across all layers. Each layer’s output feeds as the next layer’s input in a feedforward approach.

Step 1: Import Libraries

import torch
from torch import nn
from torchvision.datasets import FakeData
from torch.utils.data import DataLoader
from torchvision import transforms

Step 2: Configure the MLP Model

class MLP(nn.Module):

  def __init__(self):
    super().__init__()
    self.layers = nn.Sequential(
      nn.Flatten(),
      nn.Linear(28 * 28 * 3, 64),
      nn.ReLU(),
      nn.Linear(64, 32),
      nn.ReLU(),
      nn.Linear(32, 1),
      nn.Sigmoid()
    )

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

Step 3: Build and Load the Dataset

if __name__ == '__main__':
    torch.manual_seed(42)
    dataset = FakeData(size=15000, image_size=(3, 28, 28), num_classes=2, transform=transforms.ToTensor())
    trainloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True)

Step 4: Initialize the KLDivLoss Function

mlp = MLP()
kl = nn.KLDivLoss()
optimizer = torch.optim.Adam(mlp.parameters(), lr=1e-4)

Step 5: Train the Model

for epoch in range(0, 3):
    print(f'Starting epoch {epoch+1}')
    current_loss = 0

    for i, data in enumerate(trainloader, 0):
      inputs, targets = data
      targets = targets \
                .type(torch.FloatTensor) \
                .reshape((targets.shape[0], 1))

      optimizer.zero_grad()
      outputs = mlp(inputs)
      loss = kl(outputs, targets)
      loss.backward()
      optimizer.step()
      current_loss += loss.item()
      if i % 10 == 0:
          print('Loss after mini-batch %5d: %.3f' %
                (i + 1, current_loss / 500))
          current_loss = 0.0

print('\n Training process has finished')

Loss values near zero confirm that the MLP model’s predictions closely match the true probability distribution.

Conclusion

To calculate KL Divergence loss of neural networks in PyTorch, use nn.KLDivLoss() with an MLP model or Keras’ compile(loss='kl_divergence') for a Sequential model. Both approaches evaluate how far the model’s predicted probability distribution deviates from the true distribution — loss values close to zero indicate accurate, well-generalized predictions. Use the Sequential model for quick prototyping with Keras, and the MLP with nn.KLDivLoss() for full PyTorch control over training loops and architectures.