Hi everyone,
I have recently started working with neural nets and with pytorch, and I am trying to implement a Gumbel softmax VAE (based on the code here 11) to solve the following task:

Encode a one-hot array with length 10. Latent space has dimension 10, too.
Send a one-hot vector with length 10 to the decoder.
Decode

I would have expected that it is a simple task for the network to learn to reconstruct the input perfectly, by simply copying the input vector. There is no bottleneck because all the hidden layers are also as large as the input. And yet, almost always the network ends up in some local minimum, where multiple input vectors receive the same encoding.
Here is the code for the Gumbel part:
import torch
import torch.nn.functional as F
from torch import nn, optim

import numpy as np

def sample_gumbel(shape, eps=1e-20):
    """
    Parameters
    ----------
    shape: tuple of ints
        The dimension of the samples to take. 
    Returns
    -------
    Float tensor
        Each element in the tensor is a sample from a Gumbel(0,1) distribution.
    """
    # sample from a Gumbel(0,1) distribution can be obtained by sampling U from a Uniform(0,1)
    # and then doing -log(-log(U))
    U = torch.rand(shape)
    return -torch.log(-torch.log(U + eps) + eps)


def gumbel_softmax_sample(logits, temperature):
    """
    Takes a sample from a Gumbel distribution using the reparameterization trick.
    Uses equation 2 in https://arxiv.org/pdf/1611.01144.pdf
    
    Parameters
    ----------
    logits: array
        Class logprobabilities
    temperature: float
        Parameter of softmax
    Returns
    -------
    Float tensor
        An array containing a softmax version of samples from the Gumbel distribution.
    """
    y = logits + sample_gumbel(logits.size())
    # F.softmax applies a softmax function to a vector, which approximates an argmax function + one-hot
    return F.softmax(y / temperature, dim=-1)


def gumbel_softmax(logits, temperature):
    """
    Parameters
    ----------
    logits: array
        Shape (# instance, # classes). Contains the log-probabilities of each class for each instance
    temperature: float
        See gumbel_softmax_sample
    Returns
    -------
    tensor
        Every row is a one-hot vector.
    """
    
    y = gumbel_softmax_sample(logits, temperature)
    
    # find the argmax of each gumbel sample, for each row of logits
    ind = y.argmax(dim=-1)
    
    # create an array where each row is a one-hot vector for the respective index in ind
    y_hard = torch.zeros_like(y)
    y_hard.scatter_(1, ind.view(-1, 1), 1)
    
    # detach the samples from the gradient calculation. 
    # Not sure why it was written this way!
    y_hard = (y_hard - y).detach() + y
    
    return y_hard

And here is the model itself:
class VAE_gumbel(nn.Module):

    def __init__(self, temp, n_objects, n_signals):
        
        self.n_objects, self.n_signals = n_objects, n_signals
        
        super(VAE_gumbel, self).__init__()

        # encoder layers
        self.fc1 = nn.Linear(n_objects, 10)
        self.fc2 = nn.Linear(10, 10)
        self.fc3 = nn.Linear(10, n_signals)
    
        # decoder layers
        self.fc4 = nn.Linear(n_signals, 10)
        self.fc5 = nn.Linear(10, 10)
        self.fc6 = nn.Linear(10, n_objects)
        
        # non linear transformations
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def encode(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        return self.relu(self.fc3(x))

    def decode(self, x):
        x = self.relu(self.fc4(x))
        x = self.relu(self.fc5(x))
        # outputs a vector of probabilities
        return F.softmax(self.fc6(x), dim=-1)
    
    def probs_signals(self, x):
        return F.softmax(self.encode(x), dim=-1)
    
    @property
    def language(self):
        input_onehots = torch.eye(self.n_signals)
        probs_signals = self.probs_signals(input_onehots).detach().numpy()
        return probs_signals

    def forward(self, x, temp):
        
        # finds the vector of parameters to transform into logprobabilities
        q = self.encode(x)
        
        # these should be logprobabilities, because the gumbel_softmax distribution expects logprobabilities
        q_logprobs = F.log_softmax(q, dim=-1)
        z = gumbel_softmax(q_logprobs, temp)
        return self.decode(z), F.softmax(q, dim=-1)

Finally, the loss function and the function to run:
def loss_function(recon_x, x, qy):
    """
    The variational prior is a categorical distribution with uniform probabilities.
    
    Reconstruction + KL divergence losses summed over all elements and batch
    """
    BCE = F.binary_cross_entropy(recon_x, x, size_average=False)

    log_qy = torch.log(qy+1e-20)
    g = torch.log(torch.Tensor([1.0/x.shape[1]]))
    KLD = torch.sum(qy * (log_qy - g), dim=-1).mean()

    return BCE + KLD

def run(temp, epochs, n_objects, training_data_batches, n_signals, anneal_rate, temp_min, learning_rate=1e-3):
    
    model = VAE_gumbel(temp, n_objects, n_signals)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    for epoch in range(epochs):        
        
        total_loss = 0
        for batch_idx, data in enumerate(training_data_batches):
            optimizer.zero_grad()
            recon_batch, qy = model(data, temp)
            loss = loss_function(recon_batch, data, qy)
            
            loss.backward()
            optimizer.step()
            
            total_loss += loss
            # every 100 steps, diminish the temperature up until it's temp_min
            if batch_idx % 100 == 0:
                temp = np.maximum(temp*np.exp(-anneal_rate*batch_idx), temp_min)
                
        print(total_loss)
        print(temp)

    return model

I have ran it e.g. with the following parameters:
n_objects, n_signals = 10, 10

n_batches, size_batch, epochs = 500, 20, 20

temp = 1.
temp_min = 0.1
anneal_rate = 1e-4
learning_rate = 1e-2

But I tested various parameter combinations and it does not seem to make much of a difference. E.g. when running the following:
input_onehots = torch.eye(n_signals)
probs_signals = trained_model.language
plt.pcolormesh(probs_signals)
plt.xlabel("signal")
plt.ylabel("state")
plt.colorbar()
plt.show()

I get this type of language:

Clearly, my network is getting stuck despite the task having a simple solution. So, my question is: what is causing this problem? And related to this: how do I improve performance? is this type of problem discussed in some paper? Thank you for the help!