Cell Discovery Catalog
Background: Variational Inference
Cell Type Labeling
Model Training
As mentioned above, once we have specifications for our likelihood and variational distribution, training consists of gradient descent on the ELBO. We train and evaluate this model on two datasets from cells coming from two different tissues types: blood and bone marrow. In the peripheral blood mononuclear cell (pbmc) dataset, there are only 200 labeled cells of 4 different cell types in a dataset of 20k samples with 20k+ genes. The marrow dataset has about 8000 samples with 300 genes and 12 uniques cell types.
There are many hyper-parameters we could choose to tune, but variational auto-encoders are known to be fairly insensitive to hyper-parameters in the unsupervised context, where we are only concerned with maximizing the log probability of the data. However, we can perform sweeps using weights and biases that help us find good settings for those hyper-parameters we might want to put more care into such as:
- The dimensionality of the latent space for z1z1 and z2z2
- α\alpha-the weighting of the classification loss in the objective
- batch size
- learning rate decay
In the semisupervised setting, we are optimizing for the ELBO as well as a classification loss. The supervised component is concerned with predicting the class yy of a sample based on its latent representation z2z_2. If z2z_2 is just normally distributed as in the mean-field case of the variational distribution, then it may not be an optimal representation such that the different classes are well separated in the latent space. A normalizing flow will allow z2z_2 to have arbitrary complexity in its representation so that the classifier network of the encoder q(y∣z2)q(y|z_2) can discriminate between the classes easier. Indeed, we see better classification performance when a normalizing flow is employed. A normalizing flow is sometimes overkill, but it can be useful to see quickly whether using one can offer any benefits using Weights and Biases.
Datasets
Model Evaluation
Model Interpretability
We can also inspect the latent variable z2z_2 to see how the model separates the classes for the labeled data and if the model can group unlabeled points appropriately. We can make this comparison by leveraging WandB's 2D Projection Plot to project z2z_2 down into two dimensions and then compare the labels to the model's predicted probability.