In an earlier post I demonstrated that when a cell has low UMI counts, an SCVI model will place them in the center of the representation space. In that post, I used a small subset of genes and looked at how representations in a 2-dimensional SCVI model depended on UMI depth.

In that post, the low-count behavior could be explained by the balance between the negative binomial reconstruction likelihood and the KL divergence with the prior, which pulls representations to the center. When counts were low the KL term dominated.

Using the small set of 141 genes, you could see which total UMI counts you needed to beat the KL term and move cells away from the prior in the center of representation space. As a follow-up, I wanted to understand what kind of UMI library sizes you need to beat the prior in actual transcriptome wide data that I’d use for analysis. This way I could see what lower limits of useful total UMI counts are, and maybe I could put a lower weight on the KL term in the SCVI model to make it work better for cells with lower counts. This would enable analysis with cheaper and lower quality data.

It turns out, the KL term is not responsible for the low-UMI clustering behavior! In fact, low-UMI cells don’t converge towards the center. Instead, the SCVI encoder defines a bias point f([0, ..., 0]), where a hypothetical cell with zero counts from all genes is placed. This bias point is implicitly learned by the SCVI encoder during training, and it can be as far away from the center of representations space as high-UMI cells.

It turns out the default KL term is far too weak to assert meaningful influence on the representations. To force convergence towards the center of representations space it needs to be upweighted about 100X. At that point, you will also start observing posterior collapse where the model fails to learn meaningful representations.

Lowering the KL weight on the other hand, does _not_ reduce convergence towards the bias point as total UMIs per cell decreases. So if we want to improve SCVI for cells with lower UMI counts, what can we do instead?

The way I investigated the behavior was by subsampling UMI counts of cells using binomial thinning. One idea was to include this thinning procedure during training, increasing the depth-diversity the model sees. This didn’t improve the the encoder: it still showed similar or worse performance on downstream tasks.

Instead of just showing the model more diverse training data, I added a joint embedding loss. By passing original counts y and thinned counts y* through the encoder to make representations z and z*, we can then use a loss that penalizes the model when it embeds z and z* far apart from each other. In particular I used a cross correlation loss to simultaneously encourage close representations of z and z* and avoid posterior collapse.

This strategy makes the model learn encoders that slows down the convergence towards the bias point as UMIs per cell decreases. This directly leads to improved performance in terms of preserving cell type identity or experimental conditions in the learned representations.

Looking at reconstruction performance, the inclusion of the joint embedding loss has minimal impact on the ability to normalize gene expression or simulate gene UMI counts.

With this additional loss, QC thresholds for data used in analysis can be lowered, and data from precious archival samples can be used more effectively.

I also investigated if this joint embedding loss is sufficient. Maybe we don’t need to perform reconstruction at all to learn useful biological cell representations?

It turns out including the reconstruction loss is crucial. Models trained with joint embedding loss alone fails completely on downstream tasks like classifying clusters or separating experimental conditions.

The full paper is available at https://www.biorxiv.org/content/10.64898/2026.02.11.705441v1. The models are implemented in the SCVI branch at https://github.com/vals/scVI/tree/scvi-joint-embedding, and code for analysis and figures are available at https://github.com/vals/scvi-joint-embedding-reproducibility.

From around the web

• Why we (usually) don’t have to worry about multiple comparisons

• Chance favors the (theoretically) prepared mind

• The Origins of Adjuvants

• The Misuses of the University

• Interferons in health and disease

To get intuition about latent representations and embeddings when modeling data, it is often useful to produce 2-dimensional visualizations using TSNE. These can be used to explore what aspects of data the models are learning.

Creating TSNE visualizations was initially very time consuming and only worked for small datasets. Over time, data structures, algorithmic innovations, and, approximations has increased the scale of data you can investigate with TSNE in shorter time.

Recently, RapidsAI made an extremely fast TSNE implementation on Nvidia GPUs. If you have access to a consumer level Nvidia GPU, hundreds of thousands of data points can be visualized in a few seconds.

If you do not have access to an Nvidia GPU, you can instead use openTSNE, a package with several highly optimized implementations of TSNE. These implementations are quite fast, but still an order of magnitude slower than the RapidsAI TSNE.

This makes it attractive to rent a GPU instance on a cloud vendor for interactive work. When exploring models and data, being able to quickly iterate between options and visualize results is very valuable. However, moving your work to a cloud instance means a lot of overhead. You also need to pay for running the GPU instance even though the GPU itself is only being used in occasional bursts.

Serverless deployment of GPU accelerated TSNE

As an alternative to renting a GPU instance, we can use a ‘serverless’ GPU provider just for the TSNE visualization that leverages a GPU.

The GPU provider Modal lets you define small applications consisting of single functions that can be executed using their API. When you call the function, a GPU instance starts up in a few seconds, runs the function on the input, sends back the result to the client, then shuts down the instance after a few minutes of activity. In this model you only pay for the few seconds that the function takes to run.

I created a small Modal service which performs TSNE using a serverless GPU with a local Python client library: https://github.com/vals/gpu-embedding-service. With this, running RapidsAI’s TSNE is just a matter of doing

from gpu_embedding import gpu_tsne
coords = gpu_tsne(X)

No local GPU necessary. The data in X is sent to the service, which starts up in ten seconds if it’s not started, runs TSNE, then sends back the result as an array to the local client.

Serverless TSNE give results four times faster than local

Even with the overhead of needing to start the service and transfer the data it is about four times faster than running the equivalent TSNE locally on my Mac using 12 cores on 10-dimensional input embeddings.

As input data grows larger (either in number of observations or dimensionality), the data transfer overhead will increase too, leading to diminishing returns. For many high-dimensional embeddings, this strategy won’t necessarily be useful compared to running openTSNE locally. However, in many cases the overhead is worth it.

For this benchmark I used 10-dimensional SCVI embeddings of the Asian Immune Diversity Atlas with variational batch encodings.

Do you ever need to run TSNE on a million observations? Not really. With a random sample of ~100,000 observations from the data you probably have enough points to notice outlier populations even when faceting to ~10 categories. More data points allow you to facet the data more, but there is also a point when it is hard to read and interpret too many facets.

Scripts for benchmarking and plotting are available on GitHub at https://github.com/vals/Blog/tree/master/260120-serverless-tsne

From around the web

• The spectrum of inflammatory responses

• We predict that the transcript abundance from the gene we are knocking down will (checks notes) decrease

• Life’s most important problems

About a year ago, scvi-tools added support for metal performance shaders (MPS), a compute accelerator on M-series Mac processors. If you’re analyzing single cell data on a modern Mac, installing scvi-tools with MPS support and using accelerator="mps" substantially speeds up SCVI model fitting without needing to move your work to a Linux server.

After optimizing a couple of training parameters (batch size: 512, learning rate: 0.004), useful models on typical datasets can be trained with MPS acceleration in about five minutes.

Training SCVI on CPU vs MPS

To test the improvement in training time we use a dataset by Chen et al. 2023, where the authors collected 90,852 cells from mouse spleens to study the immune response to surgically induced sepsis.

Since we are just interested in the training times, we can fit the model for just five epochs to quickly get timings, as well as getting data on initial training loss dynamics. (For analysis, I would recommend at least 25 epochs for a dataset with 100,000 cells).

Training five epochs with CPU takes 6m 55s, while using the MPS accelerator only takes 1m 40s. This alone makes training four times faster, but we can optimize it a bit further.

Increase batch size for faster MPS training

On Macs with M-series processors, MPS uses the same unified RAM as the CPU. This means the accelerator has a large amount of memory to work with (at least relative to consumer level GPUs). This means we can send more data at a time to the accelerator in batches. The default batch size in the SCVI trainer is 128.

When increasing the batch size to 2,048 we get the training time for five epochs down to 1m 17s.

Recover loss performance by scaling learning rate

Increasing the batch size comes at a cost. The model is only updated once per batch, and larger batches mean fewer updates. As a consequence, the optimizer explores the parameter space slower. This leads to worse models with the same number of epochs. We also observe this over the five epochs in the benchmark training models.

To adjust for this effect, we can scale the learning rate of the training optimizer relative to the batch size. This way the model will take larger steps when there is a smaller number of updates.

A common and effective learning rate scaling approach is just to proportionally increase it with the batch size. Since the default batch size of 128 and default learning rate 0.001 in SCVI generally work very well for training, we can set a batch size scaled learning rate as batch_size * 0.001 / 128.

Changing the learning rate does not affect the training time, but will change the training loss dynamics over epochs. The proportional learning rate scaling does not recover the loss performance for the largest batch sizes, but we can identify a balance of training time and performance.

Conclusion

Based on the experiments, we get good training dynamics with a batch size of 512 when scaling the learning rate to 0.004. With these settings, training the five epochs takes 1m 23s. For a dataset with ~100,000 cells you usually get a useful SCVI model at 20-25 epochs, which would take around five minutes to train.

Batch Size    LR      Time     Val Loss
128           0.001   100.0s   4911
256           0.002   89.0s    4880
512           0.004   83.1s    4856
1024          0.008   80.5s    4925
2048          0.016   77.7s    4968
4096          0.032   79.2s    5228

I tried using SCVI with MPS a bit less than a year ago, but back then I couldn’t get it to work. If I had to guess, I had probably messed up my PyTorch installation. This time I didn’t have any issues getting it working. You can install scvi-tools with MPS support using pip install -U scvi-tools[metal].

Over the last year I have been using the free tier of LightningAI to train SCVI models for hobby projects. I like it, and might use it if I want to test something on larger data. But it will be nice to try simpler experiments locally.

Scripts for these benchmarks are available on GitHub: https://github.com/vals/Blog/tree/master/260114-scvi-metal

From around the web

• Selective remodelling of the adipose niche in obesity and weight loss

• Assessing a Virtual Cell’s utility

• Contagion (2011)

In cell biology, we think of discrete populations of cells, that we say do or do not express some (ideally) specific marker.

When we say that immune cells express CD45, that means “more CD45 than we can see in other cells from the same context”. T cells are considered expressing CD3, which means “more CD3 than other immune cells”, and defining CD4+ T helper cells means “T cells that express more CD4 than other T cells”.

Mentally, we are always operating on a relative scale. We work at the very edge of what technologies can detect. To illustrate that something exists, it is in comparison to a negative control where we know it doesn’t exist.

In single cell genomics, a gene being highly expressed in a cell can have two different meaning: the gene is highly expressed in that cell relative to other cells, or the gene is highly expressed relative to other genes in that cell.

This relative relation between gene expression levels in individual cells is explicitly modeled in SCVI. The SCVI model takes observed UMI counts as input, encodes them to representation vectors z, then decodes them to relative abundance levels rho for each gene. A softmax operation at the end of the decoder ensures that the relative abundances sum to 1:

During training, the relative abundances, rho, are combined with the observed library size, l, to quantify whether the predicted relative abundance is compatible with the observed UMI counts of a gene g in a cell:

The predicted rho values can be treated as normalized expression levels. Normalized, since they aim to remove the observed variation due to differences in library size. They also tend to have higher resolution than observed UMI counts, summarizing informative coexpression relationships between all the observed input genes through the encoder and decoder, resulting in continuous rather than discrete values.

These normalized expression values are great for exploring single cell RNA-sequencing data. Many cell types and cell states are defined by having ‘high’ expression of some gene, then additionally ‘high’ or ‘low’ expression of some other gene.

As an example, we can use a recent dataset by Michki et al 2025. In the paper, the authors investigate the effects of e coli infection in juvenile and neonatal mouse lungs. They collected lungs at 48 hours and 96 hours after infecting the mice with e coli, as well as negative control samples that can be used for comparison.

After fitting and saving an SCVI model, we can use the model to generate normalized expression levels rho for any genes we are interested in. If we want to study T cells, we might produce normalized expression levels for the genes Ptprc and Cd3d. Ptprc is the gene encoding CD45, a transmembrane receptor which enables fundamental external-to-internal protein-protein interactions required for almost all immune cells. Its presence is a classical global marker for immune cells. The gene Cd3d encodes the delta chain of the CD3 protein complex, which together with the T cell receptor complex enables antigen recognition by T cells.

To explore the lung cells from the experiment, we can plot the bivariate distribution of Ptprc and Cd3d expression.

We can clearly see a population of cells with low Ptprc expression, and another population with high Ptprc expression. This is most likely a population of immune cells. We can also see, within the Ptprc high population, there is a subpopulation of cells with ‘high’ Cd3d expression. These are likely T cells.

The plot is faceted over experimental conditions to highlight the consistency of these ‘high’ and ‘low’ expressing populations are between samples. Nested within these highlighted experimental conditions are also neonatal and juvenile conditions, and biological replicates.

With the scRNA-seq data though, we can dive further into the Cd3d-high T cells. There are multiple subsets of T cells. This includes Th17 cells, an effector subset of CD4+ T helper cells which provides protection at barrier tissues (like in lungs) by secreting the cytokine IL-17, which in turn activates nearby cells to produce chemokines to recruit neutrophils.

Production of IL-17 is regulated by the transcription factor RORγt, which is encoded by the gene Rorc in mice. Expression of Rorc in T cells defines the ability to perform the IL-17 production function of Th17 cells. So we can look at the bivariate distribution of Cd3d expression and Rorc expression.

Within the Cd3d-high population of cells, we can see a number of cells with ‘high’ expression of Rorc, most likely indicating Th17 cells.

We can note from these examples that on the normalized relative abundance scale of gene expression, expression levels are never exactly 0. This is a fundamental property of the relative abundance scale the rho values are defined on that is used to model the observed counts.

The rho values can be interpreted as the probability of seeing a molecule from a given gene in a cell if you’d sample a new molecule. In this setting, a rho value of 0 is a very strong statement: it means it is impossible to see molecules from that gene. Instead, the values can be very low, making it extremely unlikely we’d see a molecule, yet still possible. A rho value of 1e-9 means that if we’d sample a billion molecules, we should only expect to see one molecule produced from the gene.

On the other hand, we can easily detect populations with ‘low’ or ‘high’ normalized expression in different genes. The subdivisions into subpopulations make sense on the relative scale defined by the full dataset. Even though we don’t have a clear threshold in the form of ‘if a gene has normalized expression value above X in a cell, it expresses the gene’.

What if we could have the model infer and report at which level a gene is definitely expressed?

With this, we could automatically set thresholds for when genes are actually expressed, and distinguish cells where we definitely should not expect to see a molecule from a gene, ever.

To learn what rho values an SCVI model considers as ‘0’, we can fit the model on a modified dataset where we artificially add a gene called ‘NULL_CONTROL’ that has observed 0 expression in all cells. Then we can decode the normalized expression rho for the ‘NULL_CONTROL’ in every cell.

For any real gene, we can compare it’s expression level in a cell to the expression level of ‘NULL_CONTROL’ in the same cell. If the gene expression is similar to the ‘NULL_CONTROL’ expression, the model can’t distinguish the expression level from 0 given the data.

What we see though, is that the predicted expression levels from the artificial all-zero ‘NULL_CONTROL’ are extremely low in all cells. In the plots, the dashed line is a unit line, and Rorc expression is higher than ‘NULL_CONTROL’ expression in all cells.

Unfortunately, this means that this strategy won’t work to automatically define cells truly expressing a gene. The way the model sees it, you are more likely to see molecules from Rorc than from ‘NULL_CONTROL’. This is probably due to some very small number UMI’s being counted in a way that doesn’t respect the cell populations. This can happen due to experimental measurement reasons or spurious read assignment reasons. In the end, the data is informing the model that regardless the cell, there is a small (yet nonzero) chance to observed Rorc molecules. In a way, it is positive that the model is respecting the uncertainty in the true data.

Yet, our lives would be easier if we had a clear threshold for defining, for example, CD4+ T cells, or other cell types with known markers. Empirically, I have found that a threshold of 1e-4 is a good first guess at separating ‘high’ from ‘low’ expressing cells (meaning, it often reproduces known cell type hierarchies). The scale of the rho values will depend on how many genes are in the dataset the model is fitted to, and generally the complexity of the full population of cells.

This was a pretty quick experiment of adding the artificial ‘NULL_CONTROL’ before fitting an SCVI model. I felt the post needed to be quite long though, because I never see people describing these gene distribution explorations of scRNA-seq data, and outside the context of this relative expression -based workflow it’s not so clear why it would be useful to obtain true 0 thresholds.

Code and notebook for this analysis are available on Github

From around the web

• Interpretable Inflammation Landscape of Circulating Immune cells

• Virtual Cell Perturbation Metrics Reloaded

Part of this is an agent skill to perform QC analysis for scRNA-seq data using scanpy. Agent skills were also introduced recently. These are folders with instructions about when to use the skill, what to do in the cases where the skill gets triggered, and helper tools in the form of CLI scripts to perform the tasks needed.

When I want to explore some idea in scRNA-seq analysis, I usually start with finding a new dataset to work with. I get tired of using the same dataset all the time, and different experimental designs are good for trying different ideas.

Once I get a new interesting dataset I usually try to figure out the design of the experiment based on how the original researchers structured the data. There are many pieces of information that often tell the story of what was done. The data might be in different files, where the file names are informative. The different barcodes can indicate which experimental conditions were performed first or last. Naming conventions in samples tell you what original idea was, what might have been tacked on later, and comparisons that make sense that the researchers didn’t report on for some reason. It is a puzzle to solve.

This can take some time, so I created a Claude agent skill to do this, available at https://github.com/vals/anndata-design-inspector.

The skill provides helper scripts that uses hdf5-tools to explore column names and categories of a provided .h5ad file straight from the command line without needing to load the full dataset. If the tools aren’t available, the skill will install them.

The goal of the skill is to describe the experimental design of the dataset. I wanted this to be concise and easy to compare between datasets. There has been some work on this, for example using Hasse diagrams. The currently available options are very flexible and complete, but unintuitive in the settings I usually encounter. To help with this, I created a domain specific language that defines a grammar for experimental designs. The goal of the agent then becomes to produce a string in this grammar that describes the experiment. The package with the grammar definition also includes code for parsing the grammar into an ASCII-art visualization of the structure of the experiment. This package is also installed by the skill to create visualizations.

┌─────────────────────── Design Structure ───────────────────────┐
│                                                                │
│ ProcessingBatch(6)═════════════════════════════════════════╗   │
│                                                            ║   │
│ Center(3)  ≈≈≈≈  Protocol(2)                               ║   │
│    ↓                ↓                                      ║   │
│ Patient([30 | 25 | 18])                                    ║   │
│    ↓                                                       ║   │
│ Sample(2) ═════════════════════════════════════════════════╝   │
│    ↓                                                           │
│ Cell(~5000)                                                    │
│    :                                                           │
│ CellType(42)                                                   │
│                                                                │
│   Confounded: Center ≈≈ Protocol                               │
│   Batch: ProcessingBatch ══ Sample                             │
└────────────────────────────────────────────────────────────────┘

In addition to puzzling out the experimental design, the agent using the skill infers the research context of the experiment and provides some comments about the biological theory behind the experiment.

This all gets reported in a summarized ‘experiment card’ markdown file following a standard structure. With these markdown files I can get an idea of what the experiment was at a glance. I put up a couple of examples of experiment cards, for the datasets GSE166504 and GSE290106, as gists.

The visualizations simplify all the subtleties of the design, but are design to highlight important aspects of the design such as crossing and nesting at different levels, which helps you think about how to handle the different types of variation and confounding present in the data.

I’ve used it on ~20 different datasets, and usually performs quite well. Sometimes it over-summarizes the experiment so you might miss some hierarchical structures, but it tends to give you a lot of information about what was going on in it.

From around the web

• Deep generative modeling of sample-level heterogeneity in single-cell genomics

• Hollywood Archaeology: The Super Mario Bros. Movie

• The Truth Is Out There, Part 4: The Downside of Belief

• We still can’t predict much of anything in biology

• What will it take for AI to change drug discovery?

Previously, we looked at how posterior predictive distributions change during model fitting, where we saw that they stabilized after relatively few epochs.

Here, we want to see how differential expression results change as an SCVI model fits to a dataset.

As an example, we can use a recent dataset from Li et al. 2024. The paper describes identification of a population of stem-like T cells that are enriched in inflamed areas of the colon in patients with ulcerative colitis. The authors isolated T cells from colons of healthy donors and ulcerative colitis patients. From patients with active ulcerative colitis, they collected T cells from non-inflamed tissues and inflamed tissues from 10 patients. This is a nice experimental design that allows comparison of gene expression levels between the inflamed and non-inflamed tissues.

The authors annotated 11 different T cell subsets, but applying a filter to retain only T cell populations with at least 50 cells in both inflamed and non-inflamed tissues leaves eight subsets.

We fit the SCVI model to the full dataset, 15,411 cells with measurements of 14,203 genes. For the differential expression analysis we limit the cells only to the eight T cells subsets with sufficient cell numbers and the paired inflamed/non-inflamed tissues. We fit the SCVI model for a total of 100 epochs.

We can see in the loss curves that training set loss is still decreasing, but the validation set loss appears to have converged after about 50 epochs.

After each epoch, we stop training and perform differential expression analysis per cell type between inflamed and non-inflamed tissue. The results are stored, then visualized as volcano plots for each epoch. To see how the results change, we convert them to an animation that can be viewed below:

Over the first 10 epochs, the results are changing dramatically. Following these epochs though, the results do not appear to change systematically. We do see that estimated fold changes and p-values vary between epochs. P-values change more than log2 fold changes between epochs; after convergence (epoch 50), log2 fold changes have a standard deviation of 0.023 on average between epochs, which is well below observable variation. In this case, no genes have p-values smaller than 0.05

The log2 fold change estimation in the SCVI model stabilized faster than I had expected.

Scripts for performing the analysis and creating the visualizations are available on Github.

From around the web

• A single-cell spatial chart of the airway wall reveals proinflammatory cellular ecosystems and their interactions in health and asthma

• A single-cell spatial chart of the airway wall reveals proinflammatory cellular ecosystems and their interactions in health and asthma

For a long time, I have believed that the lack of high resolution temporal data of gene expression has been a bottleneck in learning accurate regulatory networks of transcriptional regulation (this was the basis for my research program before going to the therapeutics industry). Since technologies for these measurement don’t exist, I have kept an eye out for other, comparable, datasets to learn how analysis methods for temporal causal inference works.

A while ago a got a Home Assistant server and connected multiple thermometers throughout my non-air conditioned apartment to it. Sensors, like thermometers, are recorded over time in Home Assistant, allowing you to look up history at high resolution for long periods of time.

This seemed like an interesting opportunity to test out Granger analysis, a classical method to identify causality from time series data.

Not having an HVAC system, I have a pretty good idea of how air flows between the rooms. This should determine heat exchange between them. Which rooms heat up first throughout the day and transport heat to the other rooms is not obvious though.

I downloaded the temperature sensor data from Home Assistant for the last few months, comprising 37,110 measurements from ten thermometers.

Longer term seasonal trends can be seen through correlated increases in average temperature across ambient air thermometers. When looking at the full dataset, it is hard to daily fluctuations in temperature.

I have split up the thermometers into three categories: environmental sensors, refrigeration, and NAS. Environmental sensors measures the temperature of the air in the rooms or closets. Refrigeration thermometers measures the temperatures in fridges, and should (ideally) not be affected by room temperatures, and can be thought of a negative controls for a causal analysis. NAS is the temperature of my NAS which is hidden away in an IT closet. The temperature of the NAS is determined by the load as it is being used.

The oscillatory temperature changes in the fridges is due to the defrost cycle. Modern fridges and freezers vary the temperature to prevent frost buildup. This is why labs need to buy special, pricier, -20 freezers which maintains temperature more consistently.

If you look closely at the environmental sensors, you might see that temperatures in some rooms peak before the temperature in other rooms. This property is what we’re hoping to use to learn directionality in heat transfer between rooms.

The classical method for this is Granger analysis. The idea is that you fit two models to your data: one simple autoregressive model,

and one bivariate autoregressive model,

where x are measurements from the potentially causal source. The null hypothesis, that x does not cause y, means that all beta coefficients are 0. Or, in other words, predictions when including x-values do not improve upon using only y-values.

The models can be compared using an F-statistic that compares the residual sum of squares for the two different models:

The simple autoregressive model is restricted since it uses a subset of the predictors available for the bivariate autoregressive model which also includes potentially causal measurements.

In the example above, the top panel uses three previous time points of office temperatures to predict temperature at a given time point, while the bottom panel can also use temperatures from three prior time points from the living room. The qualitative difference in predictions are hard to see by eye in the figure, but it has an F-statistic of 210, corresponding to an error reduction from an RMSE of 0.0970°C to an RMSE of 0.0671°C.

To build a causal network, we can pairwise fit these models to all possible combinations of temperature sensors, meaning 10 * 9 = 90 possible pairs.

There are some additional implementation details for the analysis.

To avoid spurious correlations from seasonality, models are fitted to first differences of the temperatures:

We also perform the analysis with multiple lag orders, between one and five hours at one-hour increments. That is, we fit one model using only one hour prior as predictor, another model with one and two hours prior as predictor, etc., up to five hours. The hourly intervals and limiting to five hours was chosen based on it being plausible for temperatures to change at these time scales.

After fitting all the models and getting all the F-statistics, we can use them to visualize the causal graph of how temperature flows through the apartment based on the sensor data.

This diagram only shows relations with F-statistics larger than 400, a very strict threshold corresponding to a p-value of about 10^-85. Almost all relations between sensors pass a p-value threshold of 0.05. Since we have a very large number of observations, the test is powered to detect extremely small potential causes in a null hypothesis framework, and it is probably more productive to think about strength of evidence rather than null hypothesis significance testing.

The strongest link is from the living room to the office, which is believable. On the other hand, the 3D printer closet is physically located inside the office, and we are finding a stronger causal link between the living room and the 3D printer closet than from the office to the 3D printer closet. This seems like a false positive.

I would have expected the outdoors temperature to be a strong predictor throughout the system, but it’s not unreasonable that it’s not directly causal: sunlight hitting the apartment is likely to heat up the rooms before the outside temperature rises.

The lack of sunlight as a predictor is a general issue with the analysis. When there is an unobserved underlying causal factor, Granger analysis tend to produce false positive results. It is pretty likely that sunlight directly affects three of the rooms acting as a hidden confounder.

The wine fridge is affected by a lot sources. This is somewhat believable, because wine fridges are usually of much lower quality than standard fridges and don’t maintain consistent temperature well. It is not very believable that the IT closet or the 3D printer closet temperatures directly affect the temperature of the wine fridge; they should at least act through the office or living temperatures.

This highlights another shortcoming of this strategy. We are simply extracting the pairwise causal results. If there is an intermediate node, we are not testing if mediation through that node better explains the system than direct connections. Mediation should be visible in the system though, by triplets of nodes A, B, C being connected as A → B, A→ C, B → C, if A causes C mediated through B.

In addition to missing sunlight, other potential interventional factors like opening windows or cooking or running the 3D printer are not captured here. Though, I don’t do those very often, so I don’t think they affect the data at the majority of the time points.

Takeaways

Compared to the kind of data that can be generated for gene expression, this data is extremely high resolution, has an enormous amount of observations, and very little observational noise. The system is also very small compared to transcriptional regulation. But even with this data, accurately identifying a causal network is challenging.

It leaves me a bit pessimistic about the potential to identify transcriptional regulatory networks even with measurement technologies that currently don’t exist.

Granger analysis is a classical method that has been around for over half a century. Shortcomings of Granger analysis are well-known. I was hoping to find modern alternatives, and in the last few years the statistics community has had a large focus causality identification. There are many new methods and strategies, but to my surprise they tend to not make use of temporal information.

Scripts and data for this post are available on Github.

From around the web

• Target Deconvolution Through Data Integration: Unifying Drug and Genetic Perturbations

• Abnormal mapping

• Side Story

• Shelved By Genre

• Remap radio

Differences in counts

When analyzing scRNA-seq data, we are ultimately comparing molecule counts from different sources (genes, cell types, experimental conditions).

There are many other situations where your observations are counts.

As we learn to interpret data, we get good intuition about variabilities and magnitudes for measured values. These intuitions are typically based on more continuous values, and unfortunately counts behave differently from these.

Specifically, comparing magnitudes of observed counts between two groups of observations seems like it’s straightforward. Unfortunately, with a small number of samples (less than around ten) with observed counts of small magnitudes (also when typical counts are less than around ten) these comparisons are very challenging. Standard strategies for estimating means and variances from the data will fail, leading to inaccurate effect size estimation and p-values that can drastically change in magnitude from small sampling variation in the data.

The typical tool to solve these issues is negative binomial regression. This is used to estimate parameters needed to compare the groups of samples by directly modeling the properties of discrete counts. The negative binomial distribution has two parameters: mean and over-dispersion. Means are used for effect sizes. Over-dispersion indicates how much unidentified variation there is between samples on top of the variation from the counting process, which is needed to get p-values for the effect size.

Negative binomial regression is not magic though. Estimating the over-dispersion is very difficult with few observations and small counts. Such a small amount of data contains very limited information. Still, it is the best we have.

There is another issue however: the standard procedure to estimate the parameters of the negative binomial regression model is computationally expensive. For an individual comparison, computation is negligible, finishing in a millisecond. However, the ability to generate data through, for example, genome-wide CRISPR screens, and similar technologies, means that we need to do a huge number of comparisons. In the CRISPR screen example, if we knock down 20,000 genes and read out expression for 20,000 genes, we need to perform 400,000,000 comparisons. Even if a comparison takes a millisecond, this will still take over a hundred hours.

Iterative and non-iterative statistics

The reason that negative binomial regression is slow is because it typically uses an iterative method to estimate the parameters. It maximizes that likelihood by gradually improving the parameters estimates using analytical gradients.

These iterative parameter estimation methods have been around for a century, but were labor intensive before the invention of computers. A large amount of statistical methods research aimed to estimate parameters without iterative updates, referred to as non-iterative statistics. These methods use simple operations in one step to estimate statistics. For example, the mean in a normal distribution is non-iterative statistic; you sum the values and divide by the number of observations. You can estimate the mean with iterative maximum likelihood optimization, but it’s unnecessary, and might even produce worse results from numerical errors.

The standard iterative method to estimate the parameters in negative binomial distribution is called ‘iteratively reweighted least squares’.

Another way to estimate parameters is through the method of moments. Moments are values that describe properties of a probability distribution, such as mean, variance and skewness. For a given parameterized probability distribution, the observed moments can be expressed as functions of the parameters. In special cases, you can then solve for the parameters of interest with respect to the observed moments.

Up until the beginning of the 1990’s, there was increasing interest in the method of moments. In the early era of data collection, performing a statistical analysis was expensive. Puzzling out method of moments estimators for a given problem was a worthwhile effort, because the computational resources needed to do statistical inference were limited. As computational hardware became more affordable, the need for these custom solutions went away.

In addition to the work of finding method of moments estimators, maximum likelihood is optimally asymptotically efficient. It requires the the least possible amount of observations to bound the error of a parameter estimate within some interval. It is usually considered that method of moment estimation requires a larger number of observations for accurate results.

As our ability to collect data is outpacing performance improvements in computational hardware, we might want to revisit the concept of non-iterative statistics, to quickly get answers to the questions we have of the data.

Pair-set transformer for comparative statistics

In a previous post I explored the idea of a pair-set transformer that can be pre-trained for the task of statistical estimation. The same architecture can be adapted to estimate the statistics needed for negative binomial regression and inference.

The general idea is to explore if modern machine learning frameworks such as transformers can be combined with the 1980’s concept of non-iterative statistics. Instead of deriving equations for particular special case problems, we can use a generic transformer-based neural network architecture combined with synthetic data generation to learn a network that predicts statistical parameters of interest.

While transformers are expensive to execute, relatively small ones can be more efficient than iterative parameter estimation. Training it will be expensive, but once it’s trained it can be re-used forever.

You simply show it the data, and it will predict what the statistical parameters of interest are.

The asymptotic optimality of maximum likelihood estimation means it would not be able to beat it in numeric performance, but may be a lot faster, and at least better than method of moments estimation.

I created a variation of the transformer architecture described in the previous post, designed for the task of estimating parameters needed to learn magnitude and significance in differences in counts. I trained the model and wrote up a paper with detailed methods and findings.

Method of moments actually performs really well

I compared the transformer-based method to the standard approach of estimating parameters through iterative maximum likelihood, as well as an analytical method of moments solution for the parameters. I did this because I figured the analytical solution would be extremely fast but have terrible performance.

It did not!

The method of moments solution had nearly identical estimation error as iterative maximum likelihood, much better P-value calibration than either the transformer method or the iterative method, and the most statistical power.

This was a big surprise for me.

For this analysis, I used practically realistic sample sizes: between two and ten replicates per condition. This is likely far below where the asymptotic optimality of iterative maximum likelihood starts being relevant.

While I found it interesting to work out the transformer architecture and training strategy, the more valuable part of this work was in the benchmarking. It identified that method of moments solution didn’t have the problems I expected it to have.

From around the web

• Sequences and consequences - https://royalsocietypublishing.org/doi/10.1098/rstb.2009.0221

• Optimizing murine sample sizes for RNA-seq studies revealed from large-scale comparative analysis - https://www.biorxiv.org/content/10.1101/2024.07.08.602525v1

• A recap of virtual cell releases circa June 2025 - https://ekernf01.github.io/virtual-cell-june-2025

• A bioinformatician, computer scientist, and geneticist lead bioinformatic tool development—which one is better? - https://academic.oup.com/bioinformaticsadvances/article/5/1/vbaf011/7989318

• Which Kind of Science Reform - https://elevanth.org/blog/2025/07/09/which-kind-of-science-reform/

• Using hierarchical modeling to get more stable rankings of gene expression - https://statmodeling.stat.columbia.edu/2025/07/30/using-hierarchical-modeling-to-get-more-stable-rankings-of-gene-expression/

• (1) Fitting hierarchical models in genetics, (2) A Stan model that runs faster with 400,000 latent parameters, (3) Super-scalable penalized maximum likelihood inference for biome problems, (4) “In the end, I basically gave up working on biology bec... - https://statmodeling.stat.columbia.edu/2025/07/31/bio/

• Overinterpreting underpowered multi-omics experiments - https://thecodon.substack.com/p/falling-in-the-trap-of-overinterpreting

• The pharma industry from Paul Janssen to today: why drugs got harder to develop and what we can do about it - https://atelfo.github.io/2023/12/23/biopharma-from-janssen-to-today.html

• The therapeutic potential of stem cells - https://royalsocietypublishing.org/doi/10.1098/rstb.2009.0149

• Animals as chemical factories - https://worksinprogress.co/issue/animals-as-chemical-factories/

• The immunology of asthma - https://www.nature.com/articles/s41590-025-02212-9

• What's going on with gene therapies? (Part one) - https://nehalslearnings.substack.com/p/whats-going-on-with-gene-therapies

• The Day Novartis Chose Discovery - https://www.alexkesin.com/p/the-day-novartis-chose-discovery

• A Visual Guide to Gene Delivery - https://press.asimov.com/articles/gene-delivery

• How pour-over coffee got good - https://worksinprogress.co/issue/how-pour-over-coffee-got-good/

• THE LOTTERY OF FASCINATIONS - https://slatestarcodex.com/2013/06/30/the-lottery-of-fascinations/

• The Truth Is Out There, Part 3: The Game of Belief - https://www.filfre.net/2024/10/the-truth-is-out-there-part-3-the-game-of-belief/

• Alien: Earth (2025) - https://www.themoviedb.org/tv/157239-alien-earth

• The immunology of asthma and chronic rhinosinusitis - https://www.nature.com/articles/s41577-025-01159-0

• Hanseatic League - https://en.wikipedia.org/wiki/Hanseatic_League

Some series on GEO have had anndata-based .h5ad files submitted. These are much more likely to have complete metadata and annotations included in them. For re-use, this is by far more valuable than the basic CellRanger output! In particular the cell type annotation process is very fraught and time consuming (you can usually infer experimental conditions from file names).

I was happily surprised that the number of GEO series with .h5ad files seem to be increasing over time!

From around the web

• The Truth Is Out There, Part 1: The Will to Believe

• The Truth is Out There, Part 2: The Power of Belief

• On Progress Studies

• Creating therapeutic abundance

• The state of research on virtual cell modeling

• Reading Neuromancer for the very first time in 2025

• The Elite Overproduction Hypothesis

• The elusive virtual cell

• IC50 is a deep rabbit hole

With the introduction of transformers, the machine learning community took another large step. Transformers have sets as domains. With transformers we can learn to evaluate functions that take a set as an input, and e.g., produces a number as an output. Many of the ‘functions’ we typically use are classically defined by multi-step algorithms. If you have enough data, you can learn to approximate these very complicated functions.

In many cases it’s possible to get arbitrarily large amounts of data through synthetic data generation.

As an example, we can think of performing a t-test. When you perform a t-test, you have two relatively small sets of numbers. From these two sets you calculate the t-statistic, then the degrees of freedom, and from that you can obtain a p-value. To simplify the problem, let’s consider the problem of calculating the t-statistic. Calculating the t-statistic isn’t a hard problem, but can serve as a useful example.

If we target the typical situation of having between two and ten observations per group, we can define a fairly simple transformer architecture that take these two sets as input.

This transformer architecture has the ability to learn how to interpret within-group variation in the group, and how it relates to between-group variation. To train the model, we can sample random values for X and Y, of varying sizes, calculate our ground truth t-statistic, then give these as training data.

The model is relatively small, but still requires a large amount of training. I used ~20 million synthetically generated examples to train the model in total for the task of calculating t-statistics. Training took about six and a half hours on my Mac Mini. But after training, the pre-trained model can be reused to predict t-statistics.

The parameters for the pre-trained t-test transformer model takes about ~5 MB of storage (or ~2 iPhone photos).

I put together a package with the pre-trained t-test transformer in a Github repo: https://github.com/vals/TTT. I don’t suggest replacing standard t-statistics calculations, but the package illustrates the model architecture and training scripts. It is exciting thinking about other of functions defined with sets as domains we can learn.

One of the slot machines is broken until you figure out how to fix it, which requires sacrificing a valuable item.

The difference between the standard slot machines and the locked slot machine is that the locked slot machine allows you to respin individual reels five times instead of three times.

There are a lot of satisfying puzzles in Blue Prince, as well as random draws that involve setting up synergistic benefits. I figured this would also be true for the slot machine that requires you to sacrifice the valuable item. But after playing it many times I couldn’t see a clear strategy for the slot machine that would beat it.

Because the difference was in the number of respins, I figured there must be some combinations of symbols on the reels that dramatically increase in expected values when you respin more times, giving higher expected yields. I thought maybe something with the probabilities of getting snake and net combinations.

People online have been collecting statistics from the slot machine estimated probabilities for the different symbols on the reels

• Blank: 28%

• Coin: 30%

• Coin Stack: 10%

• Snake: 10%

• Net: 4%

• x2: 9%

• Clover: 1%

• Crown: 8%

With the probabilities and the rules I could implement simulations of playing the game to test different strategies.

I thought I could use Q-learning to identify non-obvious strategies buried deep in the game design. After some iterations on different Q-learning strategies, it wasn’t able to find any particularly revolutionary strategies.

To get some points of comparison, I also created a JavaScript version of the game that I could multiple times to collect data on the success on strategies I had intuited from playing the game many many times.

I collected empirical data for 25 runs (until bankruptcy) while. These empirical data can be compared with having the Q-learning agent play the game for 25 runs, then plotting the total balance over the turns in playing the slot machine.

The Q-learning agent had learned some strategies, but it seems I took more risky bets that worked out. On average, however, the performance is about equal.

I wanted to see how this compared to a naive strategy that never uses the respins, and instead alway cashes out after the initial spin.

This performed much better than the Q-learning agent or my strategies!

Generally, the cost of performing a respin does not beat out the expected values for the majority of symbol combinations on the reels. There are a few combinations that do however. For example, if you have three crowns, respinning the remaining reel at a cost of 1 gold has an expected value of 0.08 * 100 = 8 gold. I created a manual strategy that cashes out unless there is a respin available with a positive expected value, and performs that instead.

The performance of the optimal expected value strategy wasn’t much better than the cash-out strategy.

So why would you sacrifice the value item in order to use a slot machine with more respins, if using respins always end up losing you gold?

There is an achievement in the game to collect the jackpot from the slot machine (getting four crowns). If you pull three crowns, and you have three respins, the chance of getting the jackpot is 1 - 0.92^3 = 22.13%, but with five respins it is 1 - 0.92^5 = 34.09%. That is a large different if you really want to get this achievement; you end up having to play many fewer games.

The locked slot machine ends up not being gains in gold but part of ‘completing’ the game.

Code and notebook for this post are available at https://github.com/vals/Blog/tree/master/250705-slot-machine-strategy

From around the web

• What Are Virtual Cells?

• Soviet Anime?

• Drug Discovery may be in the Cold War Era

These models allow you to ask questions of the data you have trained the models one. But what if you have a pre-trained model and want to apply it to new data? You will not be able to integrate out new, unseen, batches. In this post I am exploring a strategy to solve this problem.

As a reminder, the conditional variational autoencoder used for batch integration in SCVI can be written as

In this formulation, z_n are the learned representations of the single cells, while s_n is an indicator for the batch identity. This allows the neural network f() to decode gene expression differently depending on the interaction between single cell representations and batch identities.

Practically, the default in SCVI is to represent a batch k as a one-hot encoding vector. The vector will have the length corresponding to the total number of batches K, and the element k being 1 in the vector indicates batch k (and all other entries are 0). Thus “one-hot”.

Of course, the vector s_n having length K means you can’t really do anything if you want to use the model with a new K+1’st batch from some new data. You are limited to the batches you trained the original model on.

There is an experimental option in the SCVI model called batch_representation with the two options 'one-hot' and 'embedding'. The default 'one-hot' option implements the behavior described above. The option 'embedding' implements a new behavior that learns a low-dimensional embedding for each batch. Practically, these embeddings are lookup tables that takes a batch indices and return the low-dimensional vector representing the batch. In this setting, s_i[n] are then these continuous vectors, whose values are parameters that are learned during training.

The 'embedding' option has an interesting feature: in theory, ‘similar’ batches should get proximal embedding vectors. This could be used to answer questions about batches. For example, if batches are patient samples, and patients have different diagnoses, you could learn which diagnoses are similar to each other.

These embeddings are finite-size lookup tables, and they are learned during training. We are still in the position that we can’t add new batches. It would be possible to extend look-up table with the new batch and run some training to learn the parameters for this batch, but this is impractical.

We can use this batch representation strategy as inspiration for a next step.

Remember that inference in the SCVI model is done through an inference model that takes as input observed data,

Here y_n is an array of observed UMI counts for cell n.

What if we can make a version of the SCVI model where batch embeddings s_n are encoded using another inference model?

This would give us variationally encoded batches. The known data X need to contain rich enough information for an inference model to be learned to encode the data to the variational batch representations.

A good candidate for the data X is the pseudobulk of all cells in batch i.

Let’s try this out!

Results

As in the previous posts, we will use the Asian Immune Diversity Atlas (AIDA, Tian et al 2024). The dataset has 1.1 million blood cells, but particularly important here, it has samples from 503 donors, making it realistic to learn an inference model even after holding out 50 donors as a test set.

After implementing the model with variational batch encoding, I split out all data from 50 donors as a test set. The remaining dataset was used to train a model with the default 'one-hot' option, a model with the experimental 'embedding' option, and finally a model with our new 'variational' option for batch_representation.

All models used 10-dimensional cell representations, the default for SCVI models. Both the 'embedding' and 'variational' batch representation models used 5-dimensional representations for donors. This is the default for the experimental 'embedding' option. All models were trained for 20 epochs.

In each training run the training data is further randomly split into training and validation. To make performance comparable, we evaluate the ELBO (evidence lower bound) for the full non-test fraction of the data. This will not let us evaluate overfitting, but it it does give us a global picture of average performance.

The default 'one-hot' option performs the best in terms of fitting the model to the data. The new 'variational' option performs the worst.

To get a feel for the batch integration performance of the different options, we can quickly visually inspect the cell embeddings using TSNE.

The models were set up to learn cell representations that integrate out difference between donors. The ideal results in this case would show unmixed coloring by cell type (left column), and highly mixed coloring by donor ID (right column). The 'one-hot' and 'embedding' options seem to mix donors about equally well, but the new 'variational' option is not as clearly integrating out variation due to donor.

Variationally encoding unseen batches

The main aim of the new 'varitional' batch representation option is to be able to apply a trained model on new data which has new batches. The model uses a low-dimensional continuous representation of batches. Are these representations meaningful? When we tested the different alternatives above we held out all cells from 50 donors as a test set. Do ‘similar’ unseen donors get variationally encoded to ‘similar’ seen donors?

In the AIDA data, blood samples from donors were processed at three different institutes: the genome institute of Singapore (GIS), RIKEN (RIK), and Samsung genome institute (SGI). We can investigate if donors processed at the different institutes provide similar batch effects by performing TSNE on the 5-dimensional variational batch embeddings representing the donors. In particular, the design of the 'variational' batch representation option also lets us encode donors that were held out from training.

The inference network for the batch representations learns to encode donors from GIS differently from donors from RIK or SGI. In particular, we can see that the batch inference network has generalized, and has successfully encoded unseen donors so they are similar to seen donors from the same institute.

How do the encoded batch embeddings compare to embeddings that are learned through optimization using the 'embedding' batch representation option? We can similarly perform TSNE visualization on these learned embeddings, with the caveat that we can only get these representations for the training data with 453 donors.

Similar to how the 'embedding' option seems better at learning to represent the single cell data, it also appears more effective at learning different types of batch effects. We can see clustering into subpopulations even within institutes in the learned representations.

Conclusions

It is disappointing that this new model using variational encoder for batch representation has lower performance than the other options. It is still very encouraging that the ultimate goal, the ability to represent unseen batches, is working correctly!

It is very likely that 450 training points is insufficient to learn an inference model to encode the batches accurately. It is probably worth revisiting this strategy with a much larger set of batches. There are very few individual datasets with this many batches, but could potentially be useful if training on e.g., all cells in cellxgene and you want to be able to inject new data over time.

A fork and branch of scvi-tools that implements the 'variational' batch representation option is available on github at https://github.com/vals/scVI/tree/codex/add-pseudobulk-generator-and-integrate-with-vae. Scripts and notebooks for training and generating results are available on github at https://github.com/vals/Blog/tree/master/250621-variational-batch-encoding.

References

Tian, Chi, Yuntian Zhang, Yihan Tong, Kian Hong Kock, Donald Yuhui Sim, Fei Liu, Jiaqi Dong, et al. 2024. “Single-Cell RNA Sequencing of Peripheral Blood Links Cell-Type-Specific Regulation of Splicing to Autoimmune and Inflammatory Diseases.” Nature Genetics 56 (12): 2739–52.

From around the web

• The duplication crisis: the other replication crisis

• Figma Slides is a Beautiful Disaster

Initially, all fields were entered manually. Eventually, I realized that a lot of article information could be automatically populated from the article DOI. This is through a macro in Google Apps Script which queries crossref.org, receives an object containing the result data, pulls out the useful parts, formats them, and puts them in corresponding columns of the spreadsheet. This made it a lot easier to quickly add studies. This takes care of authors, journal, title, and publication date.

All other fields are still entered manually.

It is fascinating to see how the scale of studies are growing in size over the years. The genesis of tracking this trend was while writing a perspective article for Nature Protocols focused on technological development in the single cell genomics field, where we noted the exponential growth in scale for individual studies (Svensson, Vento-Tormo, and Teichmann 2018).

Over the years, I have learned some shortcuts for identifying the total cell numbers assayed for individual papers. Sometimes they are written in the abstract, often not. A quick search for the word ‘total’ in the paper can lead to sentences like ‘in total, we generated transcriptomes for 12,000 cells from liver and 20,000 cells from spleen’, then I can add them up to 22,000. If that doesn’t work, I tend to search for ‘cells’, which is obviously used everywhere in the paper, but the beginning of the ‘results’ section typically has numbers for collected cells. Sometimes the information is in figure legends. More often, it is just written as annotations in the figures themselves. For papers in more prestigious journals like Nature or Science, these kinds of details are more likely to be find in the supplemental material, which actually tend to a good thing, because those are usually available for free even if you don’t have access to the journal article itself. I also do some quick skimming of the paper, to see if I missed any way of getting these numbers that I didn’t think of. If I still can’t find it, I usually try to find the data accession ID, download the data, and try to see if I can easily load it and learn the numbers manually that way.

Recently, there has been an explosion of tools for data extraction from web sites using large language models. I wanted to try this strategy and see how far off they would be from my manually extracted cell numbers.

Evaluating Firecrawl for cell number extraction

The general pipeline is to have a scraper convert a query website into a local Markdown file, then pipe that file to an LLM together with a prompt to extract the information of interest. (It is healthy to keep in mind that a ‘paper’ is just a post on website by some guest authors, that happens to also have a more readable PDF version.)

There is a popular free tool called Crawl4AI, which has a handy CLI utility as well as a Python API. I played around with it for a bit, but while reading up on it I came across Firecrawl Extract.

In terms of functionality, Crawl4AI and Firecrawl Extract can do the same thing, with the key difference that you bring your own LLM API key or local LLM for a model of your choice to Crawl4AI, while this is built into extraction queries in Firecrawl Extract. For a quick test, I found Firecrawl Extract to be easier to get started with and getting results.

I sampled 97 scRNA-seq studies from the ‘Single Cell Studies Database’ where I had annotated the ‘Reported cells total’ column to use as an evaluation set.

Firecrawl Extract takes as input a URL, a prompt, and an output schema for structured outputs.

class ExtractSchema(BaseModel):
    total_number_cells: int
    failure_reason: str
    url: str

prompt = '''
The page is a paper where the authors performed single-cell RNA-sequencing.
Extract the total number of cells the authors report to have collected.
Store the number in the total_number_cells field.

If you cannot extract the number of cells, give a brief reason in the failure_reason field.

Report the URL that was called to extract in the url field.
'''

Every DOI is a valid URL that resolves correctly when prepended with

https://doi.org/

. Even so, I had some hiccups where every ~14 papers or so my extraction loop exited with ValueError: ('Extract job failed. Error: All provided URLs are invalid. Please check your input and try again.', 500). Simply rerunning the query would result in a successful extraction job, so I’m not sure what was happening there. It could be due to rate limiting, but I put a pretty generous sleep.time(10) between each paper.

Results

After extraction I could simply compare my manual annotations with the extractions from FireCrawl. In 38% of cases, Firecrawl Extract was not able to find a cell number and gave the reason that the number was not specified in the provided content. I suspect in most cases this is due to papers being behind paywalls. If I would rerun this experiment I would add further instructions to the prompt to report whether the URL has a paywall. Papers where Firecrawl Extract reported to not find cell numbers are included in the figure below as red dots.

The results had the exact same counts for 27% of the papers. Some of the 35% of papers with different cell counts can be explained by a few specific observations: a couple of the ‘Reported cells total’ values have very round numbers (e.g., 24,000) while Firecrawl Extract numbers or nearby but less round (e.g., 24,023). In these cases the manual curation is probably inaccurate. In some more cases the situation is reversed, for example, manual curation has 1,027,401 while Firecrawl reports 1,000,000. I think these are cases of Firecrawl picking up rounded numbers from the abstract (’In this study we profile a million cells …’ or so). There are a number of cases where Firecrawl reports 0 cells in the papers, but does not indicate that the information is missing. I’m not sure why this is happening.

Excluding cases where no cell numbers are reported, or numbers with 0 reported cells, the average fold change error in cell numbers was 1.4x.

The results are pretty binary. Either the extraction fails to find a cell number, or it is pretty likely to be near the actual number.

For now I think I will hold off on using this automation. I think the results are pretty useful, but the practical bottleneck is that if I was impromptu adding a paper to the spreadsheet, I would need to either use a CLI tool or go to the Firecrawl Extract dashboard and set up an extract job, and that is just too much friction.

I also like having confirmed manually curated numbers in the database for these sorts of evaluations. If I start using this automation I would want to add some kind of tag or additional column to indicate if the information was manually curated or automatically extracted.

Notebooks for this post are available on Github at https://github.com/vals/Blog/tree/master/250525-cell-number-extraction.

References

Svensson, Valentine, Eduardo da Veiga Beltrame, and Lior Pachter. 2020. “A Curated Database Reveals Trends in Single-Cell Transcriptomics.” Database: The Journal of Biological Databases and Curation 2020 (November). https://doi.org/10.1093/database/baaa073.

Svensson, Valentine, Roser Vento-Tormo, and Sarah A. Teichmann. 2018. “Exponential Scaling of Single-Cell RNA-Seq in the Past Decade.” Nature Protocols 13 (4): 599–604.

Once data is collected for training the models, the training itself needs to happen as well, requiring compute resources that are hard to come by.

Meanwhile, researchers in specific biological fields who analyze the data for specific questions might fit smaller models for their purposes. The authors of scvi-tools recently created infrastructure to share pre-trained SCVI models on Hugging Face (Ergen et al. 2024).

Can we reduce training work by making use of multiple models trained individually on smaller datasets?

Model merging

When generative image models became popular a few years ago, many different models were trained and fine-tuned to be able to generate images in specific domains. Some models were good at generating landscapes, others were good at generating drawings.

What if you wanted a model that was good at generating drawings of landscapes?

Ideally you would get the combined training sets, and fine-tune a stable diffusion model with them. But this is resource intensive. As an alternative, people found that taking the two fine-tuned models and ‘merging’ them can make a new model that leverages the strengths of both models (Wortsman et al. 2022).

In the generative image model and large language model community there are multiple platforms and packages for model merging. It is a popular way for hobbyists to customize models since it requires no training resources. One package, mergekit, designed for merging popular large language models, suggest in their documentation to start with linear merging and NuSLERP merging.

Linear merging simply means taking weighted averages of weights of two trained models,

NuSLERP is a more complicated merge method, which normalizes weights to lie on a unit sphere, then perform interpolation in spherical geometry,

It is hard to find proper evaluations of whether these merging strategies actually work. In particular, in the generative image model community users typically experiment iteratively until they get desirable images. If we want to apply these techniques to SCVI models, we can evaluate them using the reconstruction error.

Testing model merging on SCVI models

To evaluate whether we can combine strengths of two independent SCVI models, we want to model two very distinct pieces of biology with them, and see if merging the models will provide a clear benefit.

Brain cells, developed from the ectoderm, and liver cells, developed from the endoderm, perform extremely different functions. The main cell types, neurons and glia for the brain, hepatocytes for the liver, use distinct and specific transcriptional regulation to arrive at their final cell states over development.

Su et al. collected 82,168 cells from mouse liver to nonalcoholic fatty liver disease (NAFLD).

Hahn et al. collected 109,826 cells from mouse brain to study the effects of aging across the brain.

We can use these datasets together to learn how well model merging works by training models on them and evaluate how well a merged model works.

To ensure that both the liver model and the brain model learned the same amount of information from the data, the datasets were limited to the intersection of 21,576 measured genes, and both datasets were downsampled to 80,000 cells. For evaluation, 10,000 of those cells are held out as a test set from each of the datasets, leaving 70,000 cells each for training.

Both the liver model and the brain model were trained for 20 epochs on the 70,000 cells.

A even weighting of 0.5 was used when merging the models, both with linear merging and NuSLERP merging.

For a comparison, an SCVI model was trained on the combined dataset of 140,000 cells. The most fair comparison is to allot this combined model with the same fitting budget as the two individual models, meaning it was only allowed to train for 10 epochs. For some additional comparison, a fourth SCVI model was allowed to train for 20 epochs.

After training, performing linear merge, and NuSLERP merge, the two held out test sets with 10,000 cells from each of brain and liver were pushed through the six different models and used to calculate reconstruction error.

The results give us some interesting information.

Applying a model trained on one organ to another organ is a bad idea. The individual models are highly specific to domain they were trained on.

Model merging substantially degrades performance on in-domain data, but substantially improves out-of-domain performance. The resulting reconstruction errors appear similar to what you would get if you evaluated both models on the test data and averaged the results (somewhat better for the brain test data, 16,000 instead of 22,000).

You obtain a far better, more general, model by combining data then fitting a model, given equal fitting budget.

Even doubling the fitting budget, a model fitted on the combined data is marginally worse than a model fitted on the individual dataset. This was a surprise: the difference is very small, but it is still present. The fitting was not repeated, the difference might be explained by stochasticity.

Unfortunately, model merging does not appear to be a good alternative to collecting large amounts of data and fitting a model on all of it.

Notebooks with code for the analysis described here are available on Github at https://github.com/vals/Blog/tree/master/250520-model-merging

References

Ergen, Can, Valeh Valiollah Pour Amiri, Martin Kim, Aaron Streets, Adam Gayoso, and Nir Yosef. 2024. “Scvi-Hub: An Actionable Repository for Model-Driven Single Cell Analysis.” bioRxiv. https://doi.org/10.1101/2024.03.01.582887.

Wortsman, Mitchell, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S. Morcos, Hongseok Namkoong, et al. 2022. “Model Soups: Averaging Weights of Multiple Fine-Tuned Models Improves Accuracy without Increasing Inference Time.” arXiv [Cs.LG]. arXiv. http://arxiv.org/abs/2203.05482.

Su, Qi, Sun Y. Kim, Funmi Adewale, Ye Zhou, Christina Aldler, Min Ni, Yi Wei, et al. 2021. “Single-Cell RNA Transcriptome Landscape of Hepatocytes and Non-Parenchymal Cells in Healthy and NAFLD Mouse Liver.” iScience 24 (11): 103233.

Hahn, Oliver, Aulden G. Foltz, Micaiah Atkins, Blen Kedir, Patricia Moran-Losada, Ian H. Guldner, Christy Munson, et al. 2023. “Atlas of the Aging Mouse Brain Reveals White Matter as Vulnerable Foci.” Cell 186 (19): 4117-4133.e22.

From around the web

• Multiplexed perturbations enable massive scale ... but how big can we go?

• Breaking Away from Disney Animation

• How did places like Bell Labs know how to ask the right questions?

• Why ‘Hey Arnold!’ Sounded Like That

• What is a scientific field?

• What happened to pathology AI companies?

A few years ago, fast CUDA-based GPU implementations of these tools became available from cuML and PyMDE, and I got used to being able to make a visualization in just a couple of seconds. Having switched to a Mac desktop for hobby projects, the most frustrating loss has been the ability to make these quick exploratory visualizations on the fly.

The extremely fast TSNE implementation in cuML from rapids.ai uses the FFT-accelerated approximation of tSNE originally published and implemented as FI-tSNE. Unfortunately, this implementation is challenging to install. To make optimized TSNE implementations available for use and experimentations, several variations of TSNE were implemented in the more user friendly package openTSNE, including the version using FFT approximation.

The M-series processors in modern Macs have unified CPU/GPU architectures, where GPU cores can be used for computationally intensive tasks. Apple has released the MLX library to enable general computations to take advantage of the GPU cores on the Apple silicon processors.

I was wondering if the FFT-accelerated TSNE could be implemented in MLX and speed up TSNE visualizations. I was also curious about using Claude Code for translating implementations to new frameworks, so this seemed like a good opportunity.

Using the openTSNE package as context, I had Claude Code create a native MLX implementation of TSNE (MLXNative), and, following that, an MLX implementation of the FFT approximation (MLXFFT).

The MLXNative implementation worked fine, while the MLXFFT implementation ended up with a some edge cases leading to outliers and some uncontrolled gradients. It also seems MLX is lacking some FFT functionality compared to the FFT implementations used in cuML and FI-tSNE/openTSNE.

My primary interest was the runtime speed when making use of the M4 Pro GPU. If runtimes were promising, I figured it would be worth digging in to minor, potentially fixable, issues.

I investigated how the runtime to create TSNE embeddings depended on dataset size for the implementations, and compared that to the various implementations in openTSNE, using either one or twelve cores. I also compared runtimes with the default scikit-learn implementation of TSNE.

In addition, I benchmarked the recently announced ‘zero code change’ cuML acceleration of the scikit-learn TSNE function. This used the same data, but was run on an L40S node on lightning.ai instead of locally on the Mac.

Seeing the speed for the MLXNative implementation for small test sets was exciting and promising. Increasing the data set sizes indicated poor scaling.

Moving on to the MLXFFT implementation, the runtimes were very disappointing. I believe this is mostly due to lacking necessary low level functionality in the FFT library in MLX.

Both the MLXNative and MLXFFT had ~90% GPU utilization when running the benchmarks, so it does seem compute streams were correctly routed.

The only other implementation with as poor performance as MLXFFT was the openTSNE implementation of BH-TSNE when using 12 cores. The single core version actually performs a lot better.

So in the end, unfortunately, simply moving computation to the GPU cores of the M4 Pro doesn’t immediately provide automatic performance gains.

Both openTSNEFFT and opentTSNENNDescent with 12 cores are pretty usable (but still 3x slower than cuML-accelerated TSNE).

In practice, when I have datasets with 200k+ points I want to explore, I will probably use some cloud computing (or subsample to ~100k).

On a positive note, using Claude Code to create these implementations (including debugging, refactoring, benchmarking, optimization) was quite straightforward.

The MLX implementations of TSNE, along with benchmarking scripts, are available on this branch on Github: https://github.com/vals/openTSNE/tree/mlx-acceleration. A notebook for producing the results figure is available at https://github.com/vals/Blog/tree/master/250331-mlx-tsne.

From around the web

• Severed Edits

• “The Third University of Cambridge”: BBN and the Development of the ARPAnet

• Wax Tailor - Que Sera (Phonovisions Symphonic Orchestra) [YouTube]

I usually focus on this generative part of the model. It is what we use to interpret the model. Variational autoencoders (Kingma & Welling, 2013) are divided into two parts: the generative model written above, and an inference model that is used to learn posterior distributions of z from the data. In my mind, I think of it as the part the matters the most.

Technically, you don’t need an inference model; you could consider just the generative model and find posterior distributions for all the latent z variables, using the observed data and the likelihood with Bayesian posterior sampling methods. But this is computationally impossible.

Stepping back a bit, the posterior distributions of the z latent variables could, slightly more efficiently, be estimated by variational distributions: simplified parametric distributions with the goal of being as similar to the true posterior distributions as possible:

Variational inference is the strategy of approximating posterior distributions, which is an integration problem, by converting it to an optimization problem, where you are optimizing the parameters of the variational distributions.

If you have single-cell RNA-seq dataset with 100,000 cells, and you want to learn 10-dimensional representations of your cells, you would need to optimize 2 * 10 * 100,000 = 2,000,000 variational parameters, which ends up being a very difficult optimization problem.

One of the revolutionary insights in machine learning is that this optimization problem can be solved with neural networks (Kingma & Welling, 2013). Instead of optimizing the variational parameters, you optimize the weights of neural networks that outputs the variational parameters given the observed data

This is known as ‘autoencoding variational Bayes’ or ‘amortized inference’, and is the other half of the variational autoencoder. We are very good at training neural networks, so turning a problem into a neural network training task ends up being a good solution. These neural networks represent the inference model in the variational autoencoder model.

In the case of the conditional variational autoencoder, there are two valid options for implementing the inference models:

In option (1), the inference model is not aware of the batches to integrate. Only the generative model will be aware of the batches. The effect, though, is that the inference networks g will learn to map observed data to appropriate representations in an unsupervised manner. A strength with this choice is that you can infer z representations for new data without knowing if any particular batch category is equivalent.

In option (2), the inference model is explicitly informed of which batch the data comes from. The inference networks will use the interactions between the values in the observed data and the batches to learn potentially more efficiently to map observations to the appropriate representations. A strength with this choice is that you can obtain counterfactual representations. A weakness is that the model will not know how to embed data from a new batch without performing architecture surgery.

The default in the SCVI models from scvi-tools with batch integration is option (1), but option (2) can be enabled with the encode_covariates = True option.

For a demonstration, we can replicate the previous post, and compare representational embeddings between an SCVI model that does not integrate out the batches, and an SCVI model that integrates out batches, but in this case using the encode_covariates = True option.

This can be illustrated as in the previous post using the AIDA data (Tian et al. 2024).

The batch-integrating version of the SCVI model with batch-aware encoding leaves out donor-to-donor variation in the representations. The effect is the same when encoding batches in integration as with the default integration explored in the previous post.

Optional integration using a conditional inference model

In the previous post, we discussed how the choice of learning representations with or without contribution from known factors depends on the questions you aim to answer. In particular, we discussed how the MrVI model lets us work in both frameworks with a single model.

We can think about this problem when considering the SCVI model with encoded batches:

Imagine having the option to ‘switch off’ the s_n contribution to the encoders g and decoder f. Then this batch-integrated model would turn into the unintegrated model. Of course, the model will not know what to do if you simply remove the s_n input to the models.

We can give the model this ability by augmenting the training data. The original data can be expanded, so that each Y observation is used to train the model twice in an epoch: once with the observed batch category, and once with a new dummy batch category we can call 'unintegrated'.

Now, for the added data points where the batch category is 'unintegrated' the model and loss will be equivalent to an unintegrated model,

After training the model with this expanded training data, we can create two sets of cell embeddings by running the observed data through the encoder g_mu with different settings for the batch categories:

We can try this approach using the AIDA data, and see how these two different versions of the cell embeddings do or do not contain variation due to donor IDs.

These results were surprising to me! My expectation was that encoding batches -case would integrate out differences between donors, while using the fixed 'unintegrated' label would retain variation between donors. Instead we are seeing the opposite. When batches are provided to the encoder, it introduces variation between batches into the representations.

On the positive side, there is a substantial difference between the two versions of cell embeddings. I don’t understand why the model ends up having this behavior.

I have always viewed the inference models in VAE-based models as a clever solution to the optimization problem, focusing on the generative model. Through this experiment I have gotten an appreciation for how the encoders can be potentially be leveraged to solve problems. In particular I think it is interesting how we can expand the capabilities of a model by expanding and augmenting the training data.

Notebooks with analysis code are available on Github: https://github.com/vals/Blog/tree/master/250327-optional-integration

References

Kingma, Diederik P., and Max Welling. 2013. “Auto-Encoding Variational Bayes.” arXiv [Stat.ML]. arXiv. http://arxiv.org/abs/1312.6114v10.

Sohn, Kihyuk, Honglak Lee, and Xinchen Yan. 2015. “Learning Structured Output Representation Using Deep Conditional Generative Models.” Neural Information Processing Systems, December, 3483–91.

Tian, Chi, Yuntian Zhang, Yihan Tong, Kian Hong Kock, Donald Yuhui Sim, Fei Liu, Jiaqi Dong, et al. 2024. “Single-Cell RNA Sequencing of Peripheral Blood Links Cell-Type-Specific Regulation of Splicing to Autoimmune and Inflammatory Diseases.” Nature Genetics 56 (12): 2739–52.

From around the web

• The Process of 'NIMH’

• Who Cares About the Disney Method?

• Managing Lockheed’s Skunk Works

A common question is ‘which cells are similar to each other?’

In the SCVI modeling framework, vector representations of cells are learned so that similar vectors generate statistically similar transcriptional profiles. With some simplifications, the SCVI model is

The representations reflect the varying kinds of transcriptional profiles we can expect to see given the data. Groups of similar cells form dense regions in the representation space, which are often used to assign cell types or cell states that are not known when the data is collected.

On the other hand, information we do know about the cells is also often reflected in the representation space. If you collect cells from a healthy person, as well as from a person undergoing an immune challenge, many of the immune cells in the blood will have different transcriptional profiles, reflecting their state changes as part of the immune response. In cell representation space, B cells from the healthy person will be in a different area than B cells from a person with an infection.

In the SCVI modeling framework, if you have kept track of which cell comes from which person, you can learn cell representation that does not include the variation due to the difference between the persons. With some simplifications, the SCVI model with integration can be be written as

Should you do this?

This depends on the questions you want to find answers to using the cell representations.

You might want to learn which cells are B cells, and count how many there are from each person. This assumes there is a definition of B cells that is shared between the two people. At this point you are not interested in the different states of the B cells from the two people. If you learn representations that integrate out the differences between the blood donors, B cells will form densities in cell representation space that is consistent between the donors.

Or, you might want to define specifically a B cell state that is present in a particular infection. Not all cell types may change state due to an infection, and you might want to find which cell types react to an infection by explicitly asking the question of whether the cells from the two people are separated in representation space when looking at a specific cell type.

This leads to a situation where you have to use two different models depending on the questions you want to answer.

It would be better to be able to answer all our different questions using a single model.

As an example, we can use the recently published Asian Immune Diversity Atlas (AIDA, Tian et al 2024). The complete dataset has 1.1 million blood cells collected from 503 donors. If we sample 50 donors, we end up with ~100k cells with 33 different cell types annotated. We can fit one SCVI model without integration, and one model where the variation between donors is integrated out.

In the unintegrated version of the model, cell types split up by donor, reflecting variation between individuals, while in the integrated model we no longer see differences between individuals, but cell types are consistently defined.

Do both at once

The new MrVI model in scvi-tools solves this problem by learning two levels of variational representations of the cells (Boyeau et al 2024). With some simplifications, we can think of the MrVI model as

There are some additional important architecture differences from the SCVI model, but for the sake of the issue discussed here, this reflects the intuition behind the model. Importantly, there are two representation spaces for the cells: U and Z. The U-space reflects cell-cell variation with variation between samples integrated out. The interaction between sample-sample variation and cell-cell variation gets introduced when moving to the Z-space.

The result of this is that consistent definitions of cell types will be reflected in the U-space, while different states of the cell types will be reflected in the Z-space.

Now we can fit a single MrVI model to the AIDA data, and visualize the two different representations for the single cells.

Now with the MrVI model, the U-space have consistent cell types are consistently defined for all donors, while the Z space reflects variation between individuals for the different cell types.

References

Boyeau, Pierre, Justin Hong, Adam Gayoso, Michael I. Jordan, Elham Azizi, and Nir Yosef. 2022. “Deep Generative Modeling for Quantifying Sample-Level Heterogeneity in Single-Cell Omics.” bioRxiv. https://doi.org/10.1101/2022.10.04.510898.

Tian, Chi, Yuntian Zhang, Yihan Tong, Kian Hong Kock, Donald Yuhui Sim, Fei Liu, Jiaqi Dong, et al. 2024. “Single-Cell RNA Sequencing of Peripheral Blood Links Cell-Type-Specific Regulation of Splicing to Autoimmune and Inflammatory Diseases.” Nature Genetics 56 (12): 2739–52.

From around the web

• Deep Learning is Not So Mysterious or Different

• Half-Life

• Strange New Rules

• Discworld Rules

• RNA

Meanwhile, I have used Mac laptops and desktops for work for a long time. The ease of doing simple tasks, using third party software like the Adobe suite, or making use of external hardware like cameras and audio devices, is substantial compared to doing the same on Windows.

My main motivation to keep using Windows was to use GPU acceleration, which is incredible even with cheaper devices. For example, Lightroom uses GPU acceleration for image processing.

One of my hobbies is to play with machine learning models like SCVI, which is painfully slow without the ability to use a GPU.

After Apple moved to M series processors, a huge amount of the things I used the GPU for was incredibly fast on the M series CPUs. So I got rid of my Windows PC, and eventually got a Mac desktop.

While hardware acceleration for machine learning models on the M series processors is developing, it is still very slow and unreliable compared to GPU acceleration.

To keep playing with machine learning models, I looked around for cloud alternatives. I tried a few and found user testimonials of several. For now I have settled on using lightning.ai.

On lightning.ai you have four GPU options on the basic account: the L40S, L4, A10G, and T4 Nvidia GPUs. Since the main models I play with are related to scvi-tools, I wanted to compare how the performance of those available GPUs compare relative to typical prices (which vary daily, but some are typically cheaper). The SCVI-based models are quite small by modern standards, so higher specs such as GPU memory is not necessarily more useful for these models.

To benchmark the GPUs, I took a dataset with about 200k cells and trained a standard SCVI model for 10 epochs, noting the training times. I also ran the same training on my Mac desktop for comparison.

All the GPU options are substantially faster than the M4 Pro in the Mac. For my purposes, the L4 GPU, which ran the 10 epochs of training in 2:30, is the optimal choice. Only the L40S (which has 48 GB of memory compared to the 24 GB in the L4) was faster, at 2:16 for 10 epochs.

In this video I model the scaling strengths of the different poker hands in Balatro.

Subscribe now