Skip to main content

t-SNE implementation in PyTorch

t-Distributed Stochastic Neighbor Embedding (t-SNE) is a popular technique for visualizing high-dimensional data by reducing it to 2 or 3 dimensions. In this article, we will demonstrate how to implement t-SNE using PyTorch, a powerful deep learning framework.

1. Introduction​

PyTorch is widely used for its flexibility and ease of use in building and training deep learning models. While PyTorch does not include a built-in t-SNE function, we can manually implement the key components of t-SNE to visualize high-dimensional datasets effectively.

2. Importing necessary libraries​

We begin by importing PyTorch and other essential libraries.

import torch
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.preprocessing import StandardScaler

Explanation:​

  • torch: The main PyTorch library for building and training neural networks.
  • torch.optim: Provides optimization algorithms for training.
  • sklearn.datasets: Used to load the MNIST dataset.
  • matplotlib.pyplot: For visualizing the results.

3. Loading and preprocessing the dataset​

We will use the MNIST dataset of handwritten digits for our example.

# Load the digits dataset
digits = load_digits()
X = digits.data
y = digits.target

# Standardize the data
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Convert to PyTorch tensors
X_tensor = torch.tensor(X_scaled, dtype=torch.float32)

Explanation:​

  • StandardScaler: Standardizes the data for better performance in the t-SNE algorithm.
  • torch.tensor: Converts the NumPy array into a PyTorch tensor.

4. Implementing t-SNE in PyTorch​

Step 1: Compute pairwise distances​

We first compute the pairwise distances between all data points.

def pairwise_distances(X):
    sum_X = torch.sum(X ** 2, dim=1)
    distances = torch.add(torch.add(-2 * torch.mm(X, X.t()), sum_X).t(), sum_X)
    return distances

pairwise_dist = pairwise_distances(X_tensor)

Explanation:​

  • pairwise_distances: Computes squared Euclidean distances between all points in the dataset.

Step 2: Compute joint probabilities​

Next, we convert the pairwise distances into joint probabilities.

def joint_probabilities(distances, sigma=1.0):
    P = torch.exp(-distances / (2 * sigma ** 2))
    P_sum = torch.sum(P)
    return P / P_sum

P = joint_probabilities(pairwise_dist)

Explanation:​

  • joint_probabilities: Applies a Gaussian kernel to convert distances into probabilities.

Step 3: Initialize the low-dimensional map​

We initialize the low-dimensional map, which will be optimized.

# Initialize low-dimensional map
Y = torch.randn(X_tensor.shape[0], 2, requires_grad=True)

Explanation:​

  • torch.randn: Initializes the low-dimensional map with random values.

Step 4: Optimize the t-SNE objective function​

Finally, we optimize the t-SNE objective function using gradient descent.

optimizer = optim.Adam([Y], lr=200)

def tsne_step():
    optimizer.zero_grad()
    low_dim_distances = pairwise_distances(Y)
    Q = joint_probabilities(low_dim_distances, sigma=1.0)
    kl_divergence = torch.sum(P * torch.log(P / (Q + 1e-10)))
    kl_divergence.backward()
    optimizer.step()
    return kl_divergence

# Perform optimization for a set number of iterations
for i in range(1000):
    loss = tsne_step()
    if i % 100 == 0:
        print(f"Iteration {i}, Loss: {loss.item()}")

Explanation:​

  • tsne_step: Computes the Kullback-Leibler divergence and updates the low-dimensional map using gradient descent.

5. Visualizing the results​

We can visualize the results after the optimization process.

# Convert the tensor to a NumPy array for plotting
Y_np = Y.detach().numpy()

# Plot the t-SNE result
plt.figure(figsize=(10, 7))
scatter = plt.scatter(Y_np[:, 0], Y_np[:, 1], c=y, cmap='viridis', s=50, alpha=0.7)
plt.colorbar(scatter, label='Digit Label')
plt.title('t-SNE Visualization of MNIST Digits (PyTorch)')
plt.xlabel('t-SNE Dimension 1')
plt.ylabel('t-SNE Dimension 2')
plt.grid(True)
plt.show()

Explanation:​

  • detach().numpy(): Converts the PyTorch tensor to a NumPy array for plotting.

6. Conclusion​

This article demonstrated how to implement t-SNE using PyTorch and applied it to the MNIST dataset. The flexibility of PyTorch allows you to build custom machine learning models, and understanding the inner workings of t-SNE can help you apply it to various high-dimensional datasets.