Skip to main content

t-SNE implementation in TensorFlow

t-Distributed Stochastic Neighbor Embedding (t-SNE) is widely used for visualizing high-dimensional data by reducing it to 2 or 3 dimensions. In this article, we will implement t-SNE using TensorFlow, a popular deep learning library.

1. Introduction​

TensorFlow provides powerful tools for building and training machine learning models, but it doesn't include a direct implementation of t-SNE. However, with TensorFlow, you can manually compute the necessary steps for t-SNE, giving you more control and flexibility over the process.

In this example, we will apply t-SNE to the MNIST dataset of handwritten digits to visualize the data in 2D space.

2. Importing necessary libraries​

We start by importing TensorFlow and other necessary libraries.

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.preprocessing import StandardScaler

Explanation:​

  • tensorflow: TensorFlow is used for numerical computations and creating the t-SNE implementation.
  • sklearn.datasets: Provides the MNIST dataset.
  • matplotlib.pyplot: Used for plotting the results.

3. Loading and preprocessing the dataset​

We'll use the MNIST dataset, which contains 8x8 pixel images of handwritten digits.

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

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

Explanation:​

  • load_digits: Loads the MNIST dataset.
  • StandardScaler: Standardizes the data, which is important for the performance of t-SNE.

4. Implementing t-SNE in TensorFlow​

Since TensorFlow doesn’t have a built-in t-SNE function, we will manually implement the key steps.

Step 1: Compute pairwise distances​

First, we compute the pairwise distances between all points in the dataset.

# Compute pairwise squared Euclidean distances
def pairwise_distances(X):
    sum_X = tf.reduce_sum(tf.square(X), 1)
    distances = tf.add(tf.transpose(sum_X), sum_X) - 2 * tf.matmul(X, tf.transpose(X))
    return distances

pairwise_dist = pairwise_distances(X_scaled)

Explanation:​

  • pairwise_distances: This function computes the squared Euclidean distances between all points in the dataset.

Step 2: Compute joint probabilities​

Next, we compute the joint probabilities of the high-dimensional data points.

def joint_probabilities(distances, sigma=1.0):
    P = tf.exp(-distances / (2 * tf.square(sigma)))
    sum_P = tf.reduce_sum(P)
    return P / sum_P

P = joint_probabilities(pairwise_dist)

Explanation:​

  • joint_probabilities: Converts distances into probabilities using a Gaussian kernel.

Step 3: Initialize the low-dimensional map​

We initialize a random 2D map for our low-dimensional representation.

# Initialize low-dimensional map
Y = tf.Variable(tf.random.normal([X_scaled.shape[0], 2]))

Explanation:​

  • tf.random.normal: Initializes the low-dimensional representation randomly.

Step 4: Optimize the t-SNE objective function​

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

optimizer = tf.keras.optimizers.Adam(learning_rate=200)

def tsne_step():
    with tf.GradientTape() as tape:
        low_dim_distances = pairwise_distances(Y)
        Q = joint_probabilities(low_dim_distances, sigma=1.0)
        kl_divergence = tf.reduce_sum(P * tf.math.log(P / (Q + 1e-10)))

    gradients = tape.gradient(kl_divergence, [Y])
    optimizer.apply_gradients(zip(gradients, [Y]))
    return kl_divergence

# Perform the 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.numpy()}")

Explanation:​

  • tsne_step: Implements one step of the t-SNE optimization, minimizing the Kullback-Leibler divergence between the high-dimensional and low-dimensional joint probabilities.
  • optimizer.apply_gradients: Applies the computed gradients to update the low-dimensional representation.

5. Visualizing the results​

After the optimization, we can visualize the resulting 2D map.

# Plot the t-SNE result
Y_np = Y.numpy()

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 (TensorFlow)')
plt.xlabel('t-SNE Dimension 1')
plt.ylabel('t-SNE Dimension 2')
plt.grid(True)
plt.show()

Explanation:​

  • Y_np: Converts the TensorFlow variable to a NumPy array for plotting.
  • plt.scatter: Plots the 2D t-SNE result.

6. Conclusion​

In this article, we manually implemented t-SNE using TensorFlow and applied it to the MNIST dataset. While TensorFlow doesn’t provide a built-in t-SNE function, our implementation demonstrates how you can leverage TensorFlow's flexibility to build custom machine learning models. By understanding the t-SNE algorithm in depth, you can apply it to a wide range of high-dimensional datasets.