I am trying to implement a Deep Embedded Self-Organizing Map (DESOM) which is an Autoencoder together with a trainable SOM as a trainable layer which I implement using a Linear layer:

class SOM(nn.Module):
    def __init__(
        self, map_height = 10, map_width = 10,
        latent_dim = 50, p_norm = 2
    ):
        super().__init__()
        
        self.map_height = map_height
        self.map_width = map_width
        self.latent_dim = latent_dim
        self.p_norm = p_norm
        self.som_nodes = self.map_height * self.map_width

        # Uniform sampling for SOM (flattened) weights initialization-
        # self.som_wts = torch.distributions.uniform.Uniform(low = - 1 / np.sqrt(latent_space_dim), high = 1 / np.sqrt(latent_space_dim)).sample((m * n, latent_space_dim))

        # create the embedding dictionary-
        # self.embedding = nn.Embedding(self.som_nodes, self.latent_dim)
        # self.embedding.weight.data.uniform_(-np.sqrt(1 / self.latent_dim), np.sqrt(1 / self.latent_dim))

        # Create SOM using linear layer (without bias)-
        self.som_wts = nn.Linear(in_features = self.latent_dim, out_features = self.map_height * self.map_width, bias = False)
        self.som_wts.weight.data.uniform_(-np.sqrt(1 / self.latent_dim), np.sqrt(1 / self.latent_dim))


    def forward(self):
        pass

The entire model combining Autoencoder with SOM is:

class DESOM(nn.Module):
    def __init__(
        self, latent_dim = 50,
        capacity = 16, map_height = 10,
        map_width = 10, p_norm = 2,
    ):
        super().__init__()
        self.latent_dim = latent_dim
        self.capacity = capacity
        self.map_height = map_height
        self.map_width = map_width
        self.p_norm = p_norm

        # tot_train_iterations = num_epochs * len(train_loader)
        # self.decay_vals = list(scheduler(it = step, tot = tot_train_iterations) for step in range(1, tot_train_iterations + 6))
        # self.decay_vals = torch.tensor(np.asarray(decay_vals))

        self.encoder = Encoder(latent_dim = self.latent_dim, capacity = self.capacity)
        self.decoder = Decoder(latent_dim = self.latent_dim, capacity = self.capacity)
        self.som = SOM(map_height = self.map_height, map_width = self.map_width, p_norm = self.p_norm)


    def forward(self, x):
        z = self.encoder(x)
        x_recon = self.decoder(z)
        return z, x_recon

# Specify SOM hyper-params-
# m = SOM height-
som_height = 20

# n = SOM width-
som_width = 20

latent_space_dim = 50

# Initialize model-
# Initialize DESOM model-
model = DESOM(
    latent_dim = latent_space_dim, capacity = 16,
    map_height = som_height, map_width = som_width,
    p_norm = p_norm
)

model.som.som_wts.weight.shape
# torch.Size([400, 50])

# Randomly initialized weights-
model.som.som_wts.weight.min().item(), model.som.som_wts.weight.max().item()
# (-0.14141587913036346, 0.14140239357948303)

Within train_one_epoch() function, the primary code to train both the Autoencoder and SOM layer is:

# Get latent code and reconstruction-
z, x_recon = model(x)

optimizer.zero_grad()

# Autoencoder reconstruction loss-
recon_loss = F.mse_loss(input = x_recon, target = x)

# SOM  training code-
l2_dist_z_soms = torch.cdist(x1 = z, x2 = model.som.som_wts.weight, p = p_norm)
mindist, bmu_indices = torch.min(l2_dist_z_soms, -1)
bmu_locations = locations[bmu_indices]
squared_l2_norm_dists = torch.square(torch.cdist(x1 = bmu_locations, x2 = locations, p = p_norm))

# Compute sigma for current iteration/step-
global step
curr_sigma_val = sigma_0 * torch.exp(-step / lmbda_val)
step += 1

# Compute Gaussian topographic neighborhood-
topo_neighb = torch.exp(-squared_l2_norm_dists / ((2 * torch.square(curr_sigma_val)) + 1e-6))

# Compute topographic loss-
topo_loss = topo_neighb * squared_l2_norm_dists

# Sum along all SOM units and mean along batch-
topo_loss = topo_loss.sum(1).mean()

# Compute total loss-
total_loss = recon_loss + (gamma * topo_loss)
# gamma = 0.001

# Compute gradienst wrt computed loss-
total_loss.backward()
        
# Perform one step of gradient descent-
optimizer.step()

The entire code can be referred here. I omit other parts for brevity.

But after the training is done, when I see the trained weights for SOM layer-

model.som.som_wts.weight.min().item(), model.som.som_wts.weight.max().item()
# (-0.14141587913036346, 0.14140239357948303)

model.som.som_wts.weight.requires_grad
# True

The SOM as a linear layer does not appear to not being trained! What’s preventing it from being trained?