Suppose I have a large number of distinct classes, some of which are related.
My model has high classification accuracy for some classes, whilst other classes are hard to predict.
How could I group together similar classes such that accuracy improves, whilst maintaining as many distinct classes as possible?
An exhaustive search over the space of combinations is intractable.
To me it sounds like there are two different requirements for the aggregation. The overall task is to aggregate related classes, in a manner that improves accuracy.
Breaking this down into two simpler objectives: we want to aggregate for accuracy (irrespective of class similarity), and at the same time enforce that aggregations are sensible, i.e. that class similarity is respected. Thus we have aggregation of labels, conditioned on features.
One way of approaching this is to map the data into a space that is modelled on the aforementioned requirements. Classes that end up close together in that space should be similar and their proximity should maximise classification accuracy.
I think a variational autoencoder (VAE) could be used to learn such a mapping. The VAE would balance 3 losses:
Having situated the data in this new space, we have a soft grouping of samples where distances are an amalgam of feature similarity and classification optimality. We can then apply hard clustering to group classes that are close together, using cross-validation to determine an appropriate threshold.
I will apply this to the classification task of identifying handwritten digits. The dataset below comprises about 1800 samples, where each sample is a grayscale image of a handwritten digit. There are 10 classes, representing digits the 0 to 9:
Suppose we want to group classes that look similar, in a way that improves classification accuracy. I train a VAE to 96% validation accuracy, and it produces this latent space distribution:
This suggests that '8' and '9' are both similar in terms of how they look, and that classification is optimised by placing them close together. On that basis, we can group these two classes together and expect an improvement in accuracy. A more empirical way of assessing appropriate clusterings is using hierarchical agglomeration:
The dendrogram highlights how we can transition from high granularity (left side), to coarser clusters towards the right. Using a validation fold (or cross-validation), we can track classification performance for different levels of clustering:
9 clusters | class memberships: [0] [5] [8, 9] [4] [6] [2] [3] [7] [1]
...
6 clusters | class memberships: [0, 5] [8, 9] [4, 6] [2, 3] [7] [1]
If we want high accuracy, whilst maximising the number of clusters, then using 9 clusters (grouping '8' and '9') or 6 clusters look like appropriate choices. This is using an MLP classifier; other models would in general have a different accuracy-granularity tradeoff (the example code includes a random forest classifier for reference).
A feature of this method is that classes get sorted on a scale from granular groupings to coarse groupings. At any single threshold, some classes remain independent whilst others get agglomerated at one or multiple levels. Thus, the final grouping is a composite of multiple hierarchies, which is more nuanced than if you were to enforce a uniform depth.
Imports. Load, split, and view data.
import numpy as np
from matplotlib import pyplot as plt
from sklearnex import patch_sklearn
patch_sklearn()
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
#Split the data - just a train and val set for the purposes of this demo
X, y = load_digits(return_X_y=True, as_frame=False)
classes = np.unique(y)
n_classes = classes.size
X_trn, X_val, y_trn, y_val = train_test_split(
X, y, random_state=0, stratify=y, test_size=0.30
)
scaler = StandardScaler().fit(X_trn)
f, axs = plt.subplots(nrows=4, ncols=12, figsize=(8, 3.5))
for i, ax in enumerate(axs.ravel()):
ax.imshow(X_trn[i].reshape(8, 8), cmap='Greys')
ax.axis('off')
ax.set_title(str(y_trn[i]), fontsize=8, weight='bold')
Define the VAE. It comprises ~5200 parameters, and is trained with weight-decayed NAdam to help mitigate over-fitting.
def calc_class_centroids(tfmd, use_y):
class_centroids = np.row_stack(
[tfmd[use_y==clas].mean(axis=0) for clas in classes]
)
return class_centroids
def calc_distance_matrix(class_centroids):
distance_matrix = np.row_stack([
[np.linalg.norm(class_centroids[i] - class_centroids[j]) for j in range(n_classes)]
for i in range(n_classes)
])
#Enforce distance matrix constraints
np.fill_diagonal(distance_matrix, 0)
distance_matrix = (distance_matrix + distance_matrix.T) / 2
return distance_matrix
def calc_distance_matrix_cl(tfmd, use_y):
#Use complete linkage for distance
distance_matrix = np.row_stack([
[np.linalg.norm(tfmd[use_y==i][None, :, :] - tfmd[use_y==j][:, None, :], axis=2).mean() for j in classes]
for i in range(n_classes)
])
np.fill_diagonal(distance_matrix, 0)
return (distance_matrix + distance_matrix.T) / 2
import torch
from torch import nn
from torch.utils.data import DataLoader
X_trn_t, X_val_t = [torch.tensor(scaler.transform(arr)).float() for arr in [X_trn, X_val]]
y_trn_t, y_val_t = [torch.tensor(arr).long() for arr in [y_trn, y_val]]
train_loader = DataLoader(list(zip(X_trn_t, y_trn_t)), shuffle=True, batch_size=16, drop_last=True)
class VAE (nn.Module):
def __init__(self, input_size=64, encoding_size=2, n_classes=10):
super().__init__()
self.encoder_backbone = nn.Sequential(
nn.Linear(input_size, 32),
nn.ReLU(),
nn.BatchNorm1d(32),
nn.Linear(32, 16),
nn.ReLU(),
nn.BatchNorm1d(16),
)
self.mu_layer = nn.Linear(16, encoding_size)
self.logvar_layer = nn.Linear(16, encoding_size)
self.decoder_recon = nn.Sequential(
nn.Linear(encoding_size, 32),
nn.ReLU(),
nn.BatchNorm1d(32),
nn.Linear(32, input_size),
)
self.decoder_clf = nn.Sequential(
nn.Linear(encoding_size, n_classes),
nn.ReLU(),
nn.BatchNorm1d(n_classes),
nn.Linear(n_classes, n_classes),
)
def encoder_forward(self, x):
x = self.encoder_backbone(x)
mu, logvar = [layer(x) for layer in [self.mu_layer, self.logvar_layer]]
return mu, logvar
def sample_from_encoding(self, mu, logvar):
return mu + torch.randn_like(mu) * torch.exp(logvar / 2)
def decoder_forward(self, x):
return self.decoder_recon(x), self.decoder_clf(x)
def forward(self, x):
mu, logvar = self.encoder_forward(x)
sampled = self.sample_from_encoding(mu, logvar)
recon, logits = self.decoder_forward(sampled)
return recon, logits, mu, logvar
print(
'VAE model comprises',
sum(p.numel() for p in VAE().parameters() if p.requires_grad),
'trainable parameters'
)
def kl_loss_fn(mu, logvar):
return (-0.5 * (1 + logvar - torch.exp(logvar) - mu**2)).sum(dim=1).mean()
def calc_losses(recon, logits, mu, logvar, X, y, loss_alphas=[0.1, 1, .01]):
alpha_recon, alpha_clf, alpha_kl = loss_alphas
mse_loss = alpha_recon * nn.MSELoss()(recon, X)
clf_loss = alpha_clf * nn.CrossEntropyLoss()(logits, y)
kl_loss = alpha_kl * kl_loss_fn(mu, logvar)
acc = (logits.argmax(dim=1) == y).float().mean()
return (mse_loss + clf_loss + kl_loss), mse_loss, clf_loss, kl_loss, acc
Training loop. Record and print metrics.
from collections import defaultdict
metrics_dict = defaultdict(list)
torch.manual_seed(0)
model = VAE()
optimiser = torch.optim.NAdam(model.parameters(), weight_decay=1e-3)
for epoch in range(n_epochs := 35):
model.train()
for X_minibatch, y_minibatch in train_loader:
recon, logits, mu, logvar = model(X_minibatch)
loss, mse_loss, ce_loss, kl_loss, _ = calc_losses(
recon, logits, mu, logvar, X_minibatch, y_minibatch,
)
optimiser.zero_grad()
loss.backward()
optimiser.step()
#/end of epoch
time_to_print = (epoch==0) or (epoch+1 == n_epochs) or ((epoch+1) % 5 == 0)
if not time_to_print:
continue
model.eval()
with torch.no_grad():
trn_recon, trn_logits, trn_mu, trn_logvar = model(X_trn_t)
val_recon, val_logits, val_mu, val_logvar = model(X_val_t)
trn_loss, trn_mse, trn_ce, trn_kl, trn_acc = calc_losses(
trn_recon, trn_logits, trn_mu, trn_logvar, X_trn_t, y_trn_t
)
val_loss, val_mse, val_ce, val_kl, val_acc = calc_losses(
val_recon, val_logits, val_mu, val_logvar, X_val_t, y_val_t
)
print(
f'[EP{epoch + 1:>2d}/{n_epochs:>2d}]',
f'TRN L{trn_loss:>5.3f}|mse {trn_mse:>5.3f}|ce {trn_ce:>5.3f}|kl {trn_kl:>4.2f}|acc {trn_acc:>6.2%}] |',
f'VAL L{val_loss:>5.3f}|mse {val_mse:>5.3f}|ce {val_ce:>5.3f}|kl {val_kl:>4.2f}|acc {val_acc:>6.2%}'
)
#Record metrics
metrics_dict['epoch'].append(epoch + 1)
metrics_dict['trn_loss'].append(trn_loss)
metrics_dict['trn_mse'].append(trn_mse)
metrics_dict['trn_ce'].append(trn_ce)
metrics_dict['trn_kl'].append(trn_kl)
metrics_dict['trn_acc'].append(trn_acc * 100)
metrics_dict['val_mse'].append(val_mse)
metrics_dict['val_loss'].append(val_loss)
metrics_dict['val_ce'].append(val_ce)
metrics_dict['val_kl'].append(val_kl)
metrics_dict['val_acc'].append(val_acc * 100)
f, ax = plt.subplots(figsize=(8, 3))
epoch_ax = metrics_dict['epoch']
ax.plot(epoch_ax, metrics_dict['trn_loss'], color='tab:blue', lw=3, label='trn_loss')
ax.plot(epoch_ax, metrics_dict['trn_mse'], marker='.', color='tab:blue', lw=1, label='trn_mse')
ax.plot(epoch_ax, metrics_dict['trn_ce'], marker='.', color='tab:blue', lw=1, ls='--', label='trn_ce')
# ax.plot(epoch_ax, metrics_dict['trn_kl'], marker='.', color='tab:blue', ls=':', label='trn_kl')
ax.plot(epoch_ax, metrics_dict['val_loss'], color='tab:red', lw=3, label='val_loss')
ax.plot(epoch_ax, metrics_dict['val_mse'], marker='.', color='tab:red', lw=1, label='val_mse')
ax.plot(epoch_ax, metrics_dict['val_ce'], marker='.', color='tab:red', lw=1, ls='--', label='val_ce')
# ax.plot(epoch_ax, metrics_dict['val_kl'], marker='.', color='tab:red', ls=':', label='val_kl')
ax.legend(ncols=2)
ax.set(xlabel='epoch', ylabel='loss')
#View recons
ex_per_row = 5
f, axs = plt.subplots(nrows=4, ncols=ex_per_row * 2, figsize=(8, 4))
for row, row_axs in enumerate(axs):
for col, ax in enumerate(row_axs):
linear_ix = row * ex_per_row + col // 2 + 10#start_at
if not col % 2:
x = X_trn_t[linear_ix].reshape(8, 8)
else:
x = trn_recon[linear_ix].reshape(8, 8)
ax.imshow(x, cmap='Greys')#, vmin=-.2, vmax=.2)
ax.axis('off')
ax.set_title('recon' if (col % 2) else str(y_trn[linear_ix]), size=8, weight='bold')
Visualise the encodings:
model.eval()
with torch.no_grad():
trn_mu, trn_logvar = [tens.numpy() for tens in model.encoder_forward(X_trn_t)]
val_mu, val_logvar = [tens.numpy() for tens in model.encoder_forward(X_val_t)]
# mu_proj = sklearn.manifold.MDS(n_components=2, n_jobs=-1).fit_transform(trn_mu)
mu_proj = trn_mu #VAE encoding is in 2D, no need for projection step to 2D
f, ax = plt.subplots(figsize=(9, 4.5))
ax.scatter(
mu_proj[:, 0], mu_proj[:, 1], c=y_trn, s=30,
cmap='tab10', edgecolor='none', alpha=0.3
)
im = ax.scatter(*[[None]*y_trn.size]*2, c=y_trn, cmap='tab10', alpha=0.8) #adjust cbar
f.colorbar(mappable=im, pad=-0.02, label='class')
ax.spines[:].set_visible(False)
ax.set(xlabel='encoding$_0$', ylabel='encoding$_1$')
[getattr(ax, f'ax{hv}line')(0, lw=1, ls=':', color='black') for hv in 'hv']
#Distance matrix
class_centroids = calc_class_centroids(mu_proj, y_trn)
ax.scatter(
class_centroids[:, 0], class_centroids[:, 1], c=classes, cmap='tab10',
marker="v", edgecolor='black', s=60, facecolor='none',
)
for clas, centroid in zip(classes, class_centroids):
ax.text(*centroid + [-0.06, 0.16], str(clas), fontweight='bold', color='black')
distance_matrix = calc_distance_matrix(class_centroids)
distance_matrix_cl = calc_distance_matrix_cl(mu_proj, y_trn)
#Centroid distances underestimate complete linkage distances by about 1% avg
# ((distance_matrix - distance_matrix_cl) / distance_matrix_cl * 100).round(1)
Apply hard clustering to the latent space:
from scipy.cluster.hierarchy import linkage, dendrogram, fcluster
from scipy.spatial.distance import squareform
#Cluster
Z = linkage(squareform(distance_matrix), method='complete')
cut_at = 1.7
dendro = dendrogram(
Z,
orientation='right',
labels=[f'class {i}' for i in classes],
color_threshold=cut_at
)
ax = plt.gca()
ax.axvline(cut_at, lw=6, alpha=0.3, color='blue')
ax.axvline(cut_at, lw=3, alpha=0.2, color='red')
ax.figure.set_size_inches(8, 3.5)
ax.spines[['top', 'right', 'left']].set_visible(False)
ax.set_xlabel('linkage distance')
ax.tick_params(axis='y', labelsize=10, rotation=15, left=True, width=3, color='dimgray')
Using a validation set, visualise and tune for the tradeoff between class granularity and classification accuracy:
from sklearn.ensemble import RandomForestClassifier
from sklearn.neural_network import MLPClassifier
np.random.seed(0)
scores_dict = {}
for linkage_distance in [0] + Z[:-1, 2].tolist():
clusters = fcluster(Z, t=linkage_distance + 1e-6, criterion='distance') - 1
y_trn_clustered = np.empty_like(y_trn)
y_val_clustered = np.empty_like(y_val)
for clas, cluster_id in zip(classes, clusters):
y_trn_clustered[y_trn==clas] = cluster_id
y_val_clustered[y_val==clas] = cluster_id
n_clusters = np.unique(clusters).size
clf = RandomForestClassifier()
clf = MLPClassifier(
[n_classes, n_classes], #Mimic the VAE decoder
learning_rate_init=0.002,
batch_size=16,
max_iter=1000,
random_state=np.random.RandomState(0)
)
scores_dict[n_clusters] = (
clf.fit(X_trn, y_trn_clustered).score(X_val, y_val_clustered) * 100
)
print(
n_clusters, 'clusters | class memberships:',
*[classes[clusters==cid].tolist() for cid in np.unique(clusters)]
)
plt.plot(scores_dict.keys(), scores_dict.values(), marker='o', color='darkolivegreen')
plt.gcf().set_size_inches(5, 2)
plt.gca().spines[['top', 'right']].set_visible(False)
plt.xlabel('number of clusters')
plt.ylabel('validation accuracy (%)')
plt.gca().invert_xaxis()
ax2 = plt.gca().secondary_xaxis(location=-0.4)
ax2.set_xticks([2, n_classes], ['coarse', 'granular'])
ax2.spines.bottom.set_bounds(2, n_classes)