t-SNE (t-Distributed Stochastic Neighbor Embedding) › Implementation in Scikit-learn
Machine Learning from ScratchAll chapters t-SNE (t-Distributed Stochastic Neighbor Embedding) › Implementation in Scikit-learnOne section · Sections in this chapter5 / 8 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Scikit-learn provides an implementation of t-SNE through the TSNE class. This implementation is optimized for performance and includes several important features for practical use. Let's walk through a complete example using the digits dataset, which contains images of handwritten digits.
Loading and preparing the data
We'll start by loading the digits dataset and examining its basic properties.
from sklearn.datasets import load_digits
# Load the digits dataset as an example
digits = load_digits()
X_digits = digits.data
y_digits = digits.target
n_samples, n_features = X_digits.shape
n_classes = len(np.unique(y_digits))Dataset: 1797 samples with 64 features Number of classes: 10
The digits dataset contains 1,797 samples, each representing an 8×8 pixel image (64 features total). With 10 digit classes (0-9), this is a perfect test case for t-SNE visualization, as it has enough dimensions to make visualization challenging but not so many points that computation becomes impractical.
Applying t-SNE
Before applying t-SNE, we standardize the data to ensure all features contribute equally to the distance calculations.
# Standardize the data (important for t-SNE)
scaler = StandardScaler()
X_digits_scaled = scaler.fit_transform(X_digits)
# Apply t-SNE with default parameters
start_time = time.time()
tsne = TSNE(n_components=2, random_state=42, perplexity=30, max_iter=1000)
X_digits_tsne = tsne.fit_transform(X_digits_scaled)
end_time = time.time()
computation_time = end_time - start_time
kl_divergence = tsne.kl_divergence_Computation time: 4.18 seconds Final KL divergence: 0.8368
The computation takes a few seconds on this moderately-sized dataset. The KL divergence value indicates how well the low-dimensional representation preserves the high-dimensional similarities. Lower values are better, though there's no universal threshold for what constitutes a "good" value. The key is that the optimization has converged (the value stopped decreasing significantly).
Visualizing the results
Let's visualize the t-SNE embedding to see how well it separates the different digit classes.
# Visualize the t-SNE results
use_mls_style()
fig, ax = plt.subplots(figsize=(6.0, 4.8))
scatter = ax.scatter(
X_digits_tsne[:, 0],
X_digits_tsne[:, 1],
c=y_digits,
cmap=digit_cmap,
norm=digit_norm,
alpha=0.74,
s=15,
edgecolors=PALETTE["paper"],
linewidths=0.2,
)
ax.set_title("t-SNE of Digits Dataset")
ax.set_xlabel("t-SNE 1")
ax.set_ylabel("t-SNE 2")
colorbar = fig.colorbar(scatter, ax=ax, ticks=range(10), pad=0.025)
colorbar.set_label("Digit")
polish_axes(ax)
plt.show()
The visualization shows clear separation between most digit classes, with distinct clusters forming for each number. This demonstrates that t-SNE successfully captured the underlying structure in the 64-dimensional space. Notice that some digits like 0, 1, and 6 form tight, well-separated clusters, while others like 3, 5, and 8 show more overlap.
# Show some sample digits
use_mls_style()
fig, axes = plt.subplots(2, 5, figsize=(6.0, 3.25))
for i in range(10):
# Find a representative sample for each digit
digit_indices = np.where(y_digits == i)[0]
sample_idx = digit_indices[0]
row = i // 5
col = i % 5
axes[row, col].imshow(
digits.images[sample_idx],
cmap="gray",
interpolation="nearest",
)
axes[row, col].set_title(f"Digit {i}")
axes[row, col].axis("off")
fig.suptitle("Sample Digits from Dataset")
plt.show()
Exploring different perplexity values
Perplexity is the most important parameter in t-SNE. Let's see how different values affect the visualization.
# Demonstrate different t-SNE parameters
use_mls_style(fixed_canvas=True)
# Different perplexity values
perplexities = [5, 15, 30]
for perplexity in perplexities:
tsne = TSNE(
n_components=2, random_state=42, perplexity=perplexity, max_iter=1000
)
X_tsne = tsne.fit_transform(X_digits_scaled)
fig, ax = plt.subplots(figsize=(2.7, 2.55))
ax.scatter(
X_tsne[:, 0],
X_tsne[:, 1],
c=y_digits,
cmap=digit_cmap,
norm=digit_norm,
alpha=0.72,
s=10,
edgecolors="none",
)
ax.set_title(f"Perplexity = {perplexity}")
ax.set_xlabel("t-SNE 1")
ax.set_ylabel("t-SNE 2")
polish_axes(ax)
fig.subplots_adjust(left=0.22, right=0.97, bottom=0.20, top=0.84)
plt.show()
use_mls_style(fixed_canvas=True)
for perplexity in [50, 100, 200]:
tsne = TSNE(
n_components=2, random_state=42, perplexity=perplexity, max_iter=1000
)
X_tsne = tsne.fit_transform(X_digits_scaled)
fig, ax = plt.subplots(figsize=(2.7, 2.55))
ax.scatter(
X_tsne[:, 0],
X_tsne[:, 1],
c=y_digits,
cmap=digit_cmap,
norm=digit_norm,
alpha=0.72,
s=10,
edgecolors="none",
)
ax.set_title(f"Perplexity = {perplexity}")
ax.set_xlabel("t-SNE 1")
ax.set_ylabel("t-SNE 2")
polish_axes(ax)
fig.subplots_adjust(left=0.22, right=0.97, bottom=0.20, top=0.84)
plt.show()
The perplexity value dramatically affects the visualization structure. Low values (5, 15) create many small, tight clusters that emphasize very local structure but can fragment natural groups. Medium values (30, 50) typically provide the best balance, revealing both local and global structure. High values (100, 200) create more spread-out visualizations that emphasize global structure but may lose fine-grained details. For this dataset with ~1,800 points, perplexity values between 30 and 50 work best.
Exploring different learning rates
The learning rate controls the step size during optimization. Let's examine its effect.
# Demonstrate the effect of learning rate
use_mls_style(fixed_canvas=True)
learning_rates = [10, 100, 1000]
for lr in learning_rates:
tsne = TSNE(
n_components=2,
random_state=42,
perplexity=30,
learning_rate=lr,
max_iter=1000,
)
X_tsne = tsne.fit_transform(X_digits_scaled)
fig, ax = plt.subplots(figsize=(2.7, 2.55))
ax.scatter(
X_tsne[:, 0],
X_tsne[:, 1],
c=y_digits,
cmap=digit_cmap,
norm=digit_norm,
alpha=0.72,
s=10,
edgecolors="none",
)
ax.set_title(f"Learning Rate = {lr}")
ax.set_xlabel("t-SNE 1")
ax.set_ylabel("t-SNE 2")
polish_axes(ax)
fig.subplots_adjust(left=0.22, right=0.97, bottom=0.20, top=0.84)
plt.show()
The learning rate affects how quickly the optimization converges. Very low values (10) may result in poor optimization, with points not fully separated. Moderate values (100) typically work well. Very high values (1000) can cause the optimization to overshoot, though scikit-learn's default of 200 handles this reasonably. If your visualization looks poor, try adjusting the learning rate before changing other parameters.
Key Parameters
Below are the main parameters that affect how t-SNE works and performs.
-
n_components: Number of dimensions in the embedding (default: 2). Use 2 for visualization, occasionally 3 for interactive 3D plots. Higher values defeat the purpose of visualization. -
perplexity: Controls the effective number of neighbors considered for each point (default: 30). Think of it as a smoothness parameter. Typical range is 5 to 50. Lower values (5-15) for small datasets or fine-grained structure, higher values (30-50) for larger datasets or global structure. Must be less than the number of samples. -
learning_rate: Step size for gradient descent (default: 200). Typical range is 10 to 1000. If clusters look compressed or poorly separated, try values between 100 and 1000. Too low can result in poor optimization, too high can cause instability. -
max_iter: Maximum number of optimization iterations (default: 1000). Increase to 2000 or 5000 if the KL divergence hasn't converged. For large datasets, you may need more iterations. -
random_state: Seed for reproducibility (default: None). Set to an integer to get consistent results across runs. Always use this when comparing different parameter settings. -
init: Initialization method (default: 'pca'). Using PCA initialization often leads to more stable results than random initialization. Can also pass a custom array of initial positions. -
method: Algorithm variant (default: 'barnes_hut'). Barnes-Hut t-SNE is faster (O(n log n)) for large datasets. Use 'exact' for smaller datasets (< 1000 points) when you need precise results. -
n_jobs: Number of CPU cores to use (default: None). Only affects the gradient computation. Set to -1 to use all available cores, though the speedup is limited.
Key Methods
The following are the most commonly used methods for interacting with t-SNE.
-
fit_transform(X): Fits t-SNE on data X and returns the embedded coordinates. This is the primary method you'll use. Returns an array of shape (n_samples, n_components). -
kl_divergence_: Attribute (not a method) that stores the final KL divergence after fitting. Lower values indicate better preservation of similarities. Use this to assess convergence and compare different parameter settings.