socam: implement class balance weighting of loss

[sjfleming]

Add class balancing weights to the cross_entropy loss. This is built into pytorch's cross_entropy already, so this is mostly plumbing to accept a class_counts input which is a pandas Series whose index is cell type IDs and whose values are cell counts across the whole training dataset. If the class_counts input is None (default), then there is no class balancing used (current behavior).

This is a new input file the user would need to create from the training data manually.

[sjfleming]

@magrenimish does this look right? Would you be able to use this in your manuscript review experiments? We are running into the same thing in CAS... we probably need to try it.

[sjfleming]

@fedorgrab how would we actually create this file? Is this already part of the extract metadata, or would we have to do a specific TileDB query to get this?

[magrenimish]

when the cl_names_subset is expanded to include ancestors (that were not part of the training labels) during training, it looks like class_counts expects a count for these expanded classes too?
Also, should the count change based on the number of descendants a cell type has (should we be propagating the count too)?

[sjfleming]

Propagating the count is a very interesting concept

[fedorgrab]

@fedorgrab how would we actually create this file? Is this already part of the extract metadata, or would we have to do a specific TileDB query to get this?

Nexus currently doesn’t produce per-class cell counts. The easiest way would be, yes, I suppose, to count this via TileDB using the same filters that were used to create the extract. The extract metadata file does include those filters.

[sjfleming]

Hmmmm, I had a thought today and ran it past gemini. The core of the thought is this:
Even though we're adding a class balancing weight to balance cell types in the loss function, we are missing one idea. The idea that maybe each dataset should have kind of an equal vote about the cell types that it contains. If dataset A has 100x as many T cells as dataset B, should dataset A have 100x as much say about what a T cell looks like?

Gemini:

By applying a standard class weight to your cross-entropy loss, you solved the problem of rare cell types being ignored. But you did not solve the problem of small datasets being ignored within those cell types.

Here is exactly what is happening under the hood and how to fix it.

The Mechanics of the "Dataset Vote"
Let's assume your inverse-frequency class weighting assigns a high weight multiplier to the "Dendritic Cell" (DC) class because it is relatively rare across your 60M cells.

When your model calculates the loss for a batch, it applies that multiplier to every DC it sees. However:

If Dataset A has 9,000 DCs and Dataset B has 100 DCs, Dataset A still provides 90x more examples to the model.

During gradient descent, the model is updating its weights based on the sum (or average) of the loss.

98.9% of the gradients pushing the model to learn "what a DC looks like" are coming from Dataset A.

The consequence: Your classifier isn't actually learning a universal representation of a Dendritic Cell. It is learning Dataset A's specific batch-effected version of a Dendritic Cell, alongside Dataset A's specific annotation biases. Dataset B's 100 DCs are treated as statistical noise and functionally ignored.

[sjfleming]

For a given cell $i$ belonging to class $c$ from dataset $d$, you can set its sample weight in the loss function to: $W_{i} = 1 / N_{d,c}$ Where $N_{d,c}$ is the total number of cells of class $c$ in dataset $d$.

The Danger (and how to mitigate it):
If Dataset C has exactly one severely mis-annotated DC, this formula will assign an astronomical weight to that single garbage cell, forcing your model to aggressively learn it. To prevent this, you should apply a smoothing factor or a cap/floor to your weights. For example, calculate the inverse frequency, but clip the maximum possible weight so that ultra-tiny partitions cannot hijack the model's gradients. Alternatively, drop partitions with fewer than $n$ cells before training begins, ensuring the model only learns from datasets that have a statistically meaningful sample of that specific cell type.

It's worth considering implementing two things (optional and mutually exclusive):

• cell_type class balancing in the loss function
• (dataset_id, cell_type) balancing in the loss function

The latter is a bit harder to implement, but might be interesting.

[sjfleming]

The class_counts thing is broken as written. Internal ancestor nodes might have zero cell counts in the training data (some surely do because we filter them out). TileDB will firstly not give these counts, but secondly, even if we manually set them to zero, the weights computed would be div-by-zero. Fix this.

EDIT: actually... if we manually set them to zero, they will never be accessed in the code since no cell is actually annotated that way. So it will not lead to a div-by-zero. I think the thing to do is accept a subset of the list from the user (just observed cells) and assume missing classes have zero counts. Allow zeros and test this is okay.

[sjfleming]

TODO

• add the code that optionally propagates class counts
• test to ensure zero count internal nodes cause no problems

[sjfleming]

Alright @magrenimish and @fedorgrab this is ready for another round of review. Since your last look, I've added optional propagation of class counts.

The issue about weighting by (dataset_id, celltype) rather than by celltype has been put off for later #408