Feature Importance of a Pytorch AutoEncoder

JAPardo: <p>I need to get from my Pytorch AutoEncoder the importance it gives to each input variable. I am working with a tabular data set, no images.</p>
<p>My AutoEncoder is as follows:</p>
<pre><code class="lang-auto">class AE(torch.nn.Module):
    def __init__(self, input_size, hidden_layer, latent_layer):
        super().__init__()

        self.encoder = torch.nn.Sequential(
            torch.nn.Linear(input_size, hidden_layer),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_layer, latent_layer)
        )

        self.decoder = torch.nn.Sequential(
            torch.nn.Linear(latent_layer, hidden_layer),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_layer, input_size)
        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded
</code></pre>
<p>To save unnecessary information, I simply call the following function to get my model:</p>
<pre><code class="lang-auto">average_loss, model, train_losses, test_losses = fullAE(batch_size=128, input_size=genes_tensor.shape[1],
                                 learning_rate=0.0001, weight_decay=0,
                                 epochs=50, verbose=False, dataset=genes_tensor, betas_value=(0.9, 0.999), train_dataset=genes_tensor_train, test_dataset=genes_tensor_test)
</code></pre>
<p>Where “model” is a trained instance of the previous AutoEncoder:</p>
<pre><code class="lang-auto">model = AE(input_size=input_size, hidden_layer=int(input_size * 0.75), latent_layer=int(input_size * 0.5)).to(device)
</code></pre>
<p>Well now I need to get the importance given by that model to each input variable in my original “genes_tensor” dataset, but I don’t know how. I have researched how to do it and found a way to do it with shap software:</p>
<pre><code class="lang-auto">e = shap.DeepExplainer(model, genes_tensor)

shap_values = e.shap_values(
    genes_tensor
)

shap.summary_plot(shap_values,genes_tensor,feature_names=features)
</code></pre>
<p>The problem with this implementation is the following: 1) I don’t know if what I am actually doing is correct. 2) It takes forever to finish, since the dataset contains 950 samples, I have tried to do it with only 1 sample and it takes long enough. The result using a single sample is as follows:</p>
<p>I have seen that there are other options to obtain the importance of the input variables like Captum, but Captum only allows to know the importance in Neural Networks with a single output neuron, in my case there are many.</p>
<p>The options for AEs or VAEs that I have seen on github do not work for me since they use concrete cases, and especially images always, for example:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/peterparity/PDE-VAE-pytorch">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/peterparity/PDE-VAE-pytorch" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/344;"><img src="https://opengraph.githubassets.com/1a1bfaa880417b9515f025022ad7adddc1874b3ddf9be4df2580d2ec894301c8/peterparity/PDE-VAE-pytorch" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/peterparity/PDE-VAE-pytorch" target="_blank" rel="noopener nofollow ugc">GitHub - peterparity/PDE-VAE-pytorch: PDE-VAE: Variational Autoencoder for...</a></h3>

  <p>PDE-VAE: Variational Autoencoder for Extracting Interpretable Physical Parameters from Spatiotemporal Systems using Unsupervised Learning - peterparity/PDE-VAE-pytorch</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/FengNiMa/VAE-TracIn-pytorch">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/FengNiMa/VAE-TracIn-pytorch" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/344;"><img src="https://opengraph.githubassets.com/626e977c1729c75c24a28a4c56d94e35f21052a44dc4a099677ae88b073ebf46/FengNiMa/VAE-TracIn-pytorch" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/FengNiMa/VAE-TracIn-pytorch" target="_blank" rel="noopener nofollow ugc">GitHub - FengNiMa/VAE-TracIn-pytorch: Official PyTorch implementation for...</a></h3>

  <p>Official PyTorch implementation for "Understanding Instance-based Interpretability of Variational Auto-Encoders." - FengNiMa/VAE-TracIn-pytorch</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Is my shap implementation correct, should I continue with shap even though it takes eons or should I switch to another method?</p>
<p>Thank you very much in advance.</p>
<p>Best regards.</p>