Fine-tuning Borzoi with Single-Cell Epigenetics

A practical dive into the cutting edge of biological foundation models

Introduction

Prior articles focused on current-generation foundation models in epigenetics: sequence-to-function models that learn tissue- and cell-specific epigenetic and expression states and their relationship to underlying DNA sequences, and function-to-function (”smoothing”) methylation models that predict methylation states at one locus based on a context window of DNA sequence and DNA methylation states.

This new blog focuses on the practicality of fine-tuning foundation models using new data. As our single-cell CUT&Tag and PairedTag (scCUT&Tag + RNA) assays produce cellular maps of histone modifications, transcription factor binding, and chromatin remodeler occupancy (with or without corresponding RNA expression), fine-tuning Borzoi (a sequence-to-function model) on Paired-Tag maps constitutes a natural use case.

We used a small dataset of human peripheral blood mononuclear cells (PBMCs; publicly available via our data request form) to fine-tune new cell-type-specific output heads for Borzoi, demonstrating the ability to predict functional consequences of mutations in hematological cell epigenetics. This practical article covers instance selection and setup, model loading and extension, data formatting, training, and prediction.

Background

Sequence-to-epigenetics foundation models learn relationships between DNA sequence and cell-specific or tissue-specific epigenetic states. For the current generation of models, cell specificity is not directly encoded in the latent sequence representation (or in the representation-determinative model weights; "encoder"), but instead in the output-specific components of the model ("decoder"). In existing foundation models, specificity stems from the output head weights, and each output head sees the same input representation. An accurate model indicates that the model has learned a set of sequence-derived features "useful" for predicting epigenetic states. The correlation of latent features with established biological genomic features and functional states is often used as evidence that the model captures known biology.

Foundation model attention weights can be found that correlate to functional genomic features. (Top) scDNAm-GPT attention weights at certain layers coincide with H3K4me3 tracks within the same cell type. (Bottom) Attention weights from Evo-2 partitions the genome into biologically relevant sections.

The design of such models makes extending them to new datasets straightforward: adding a new output head corresponds to a new cell-specific predictive model. If the pre-training successfully learned "good" features for sequence-to-function maps, only the new weights require updating, preserving the full model's performance on all other outputs.

Fine-Tuning Borzoi

AlphaGenome and Borzoi, as sequence-to-function models, use multi-head outputs for tissue- and modality-specific predictions. The same strategy can be applied to fine-tuning these models, but because AlphaGenome is a fully closed model (as neither the model itself nor the associated weights can downloaded) no one outside of Alphabet can fine tune it. Borzoi has two implementations: the original Borzoi and a faster version (Flashzoi), which applies Flash-attention to speed up training and inference.

According to the Borzoi article (”Using TensorFlow (v.2.11), backpropagation of this model on a 524 kb sequence maxes out the 40 GB of RAM of a standard NVIDIA A100 GPU.”), the researchers chose the specific architecture so that the model and minibatch could fit on the A100 GPU, which possesses 40 GB of GPU memory. The closest and cheapest instance to apply (at the time of writing) is the AWS g6e.xlarge, which employs an L40S GPU with 48 GB of memory ($1.861/hr). The H100 (successor to the A100) is approximately five times more expensive; however, the trick with the g6e.xlarge instance is to stay within the CPU memory limit, which is “only” 32 GB.

The Deep Learning (single-CUDA) Amazon AMI, which already includes the NVIDIA drivers and CUDA, represents the most straightforward base to build from, with the next step to install Torch, Flash-attention, and Flashzoi via the following code:

nvidia-smi --query-gpu=name,driver_version --format=csv,noheader
# NVIDIA L40S, 580.105.08

wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh
bash Miniforge3-Linux-x86_64.sh -b -p ~/miniforge3
~/miniforge3/bin/conda init bash && source ~/.bashrc

conda create -n torch python=3.12 -y
conda activate torch
pip install --upgrade pip
pip install --index-url https://download.pytorch.org/whl/cu126 torch==2.6.0 torchvision torchaudio
conda install -c conda-forge jupyter notebook -y

# transformers<4.51.0,>=4.34.1
pip install borzoi-pytorch
pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.2/flash_attn-2.8.2+cu12torch2.6cxx11abiTRUE-cp312-cp312-linux_x86_64.whl --no-deps --force-reinstall
pip install scipy

wget https://github.com/johahi/borzoi-pytorch/raw/refs/heads/main/wt_seq.npy

aws s3 cp s3://epigenome-data/flashzoi-train/train_bundle.tar.gz train_bundle.tar.gz
tar xvzf train_bundle.tar.gz
wget https://download.gcc.rug.nl/downloads/eqtlgen/cis-eqtl/2019-12-11-cis-eQTLsFDR0.05-ProbeLevel-CohortInfoRemoved-BonferroniAdded.txt.gz


To verify that the weights can be downloaded and that the model works, the following script downloads a Borzoi split (replicate-0) and employs the test data in the Flashzoi repo to produce test predictions. This process should run with deprecation warnings but no errors.

import torch
torch.backends.cudnn.allow_tf32 = False
torch.backends.cuda.matmul.allow_tf32 = False
from tqdm import tqdm
import numpy as np
from borzoi_pytorch import Borzoi

device = torch.device("cuda")
borzoi = Borzoi.from_pretrained('johahi/borzoi-replicate-0').to(device)
genome_1hot = np.load('./wt_seq.npy')
print(genome_1hot.shape)
genome_1hot = torch.as_tensor(genome_1hot).to(device).permute(1,0)
print(genome_1hot.shape)

tracks = borzoi(genome_1hot[None,...])
print(tracks.shape)
print(tracks.cpu().detach().numpy()[0,:,:][chip_idx,:].max(axis=0))
borzoi_type = 'HUMAN' if hasattr(borzoi, 'human_head') else 'MOUSE' # it can have both
print(tracks.cpu().shape)
borzoi.tracks_df.description.apply(lambda t: t.split(':')[0]).value_counts()


Fine-tuning Data

The data used for fine-tuning Flashzoi derived from eight Paired-Tag (scCUT&Tag+RNA) experiments performed on PBMCs from two healthy female donors, targeting the H3K27ac, H3K27me3, H3K4me1, and H3K4me3 histone modifications (data available upon request via our data request form).

Data processing applied standard CellRanger filters, merged single-nucleus (sn)RNA data from all eight experiments, identified highly variable genes using BigSur, applied PCA+Harmony for lane correction, clustered using Leiden clustering, and annotated major cell classes using established markers (e.g., CD3, CD4, CD8, and MS4A1). We defined six major cell types (excluding Mucosal-associated invariant T (MAIT) cells): T, B, dendritic, MAIT, natural killer, and monocytes), resulting in 30 fine-tuning tracks (6 x 4 = 24 histone marks; 6 x 1 = 6 RNA tracks).

Borzoi uses bins of 32 bp for quantification; after splitting the fragment files into their corresponding cell types, bedtools computed bin coverage, and the bins were normalized to fragments-per-million (“FPM-normalized”).

Formatting

We also required the curation and formatting of input sequences; the Borzoi article describes an input of 524,288 bp of one-hot-encoded sequence, corresponding to 16,384 32-bp bins; however, the alignment of these bins on the genome remains unknown. For instance, is the first bin used for training on the first non-N base on chr1, or is it used further along?

The Borzoi training data (see: https://console.cloud.google.com/storage/browser/borzoi-paper?project=borzoi-paper-project) provides a clue: data_hg38_contigs.bed contains entries such as chr1:1644384-2702781, which, at 1,058,396 bp, has exactly 32,768 32-bp bins. We can therefore guess that bins align with training contigs:

bedtools makewindows \
-b data_hg38_contigs.sorted.bed \
-w 32 | bgzip -c > tmp.bed.gz

bedtools intersect -a tmp.bed.gz \
-b data_hg38_contigs.sorted.bed \
-wa -wb -loj | \
awk '{print $1"\t"$2"\t"$3"\t"$4":"$5"-"$6"\t."}' | \
bgzip -c > flashzoi_bins.bed.gz

bedtools getfasta \
-fi $GDIR/GRCh38.primary_assembly.genome.fa \
-bed data_hg38_contigs.sorted.bed -name -tab | \
python onehot_enc.py

bedtools intersect \
-a flashzoi_bins.bed.gz \
-b raw/DC.H3K4me1.tsv.gz -c | \
bgzip -c > coverage/DC.H3K4me1.bed.gz


To handle the Borzoi training data we will restrict to individual contigs to help limit CPU memory usage (which must be below 32 GB). As the output coverage files align (same bins), we can concatenate them by column, subset the rows to only those bins that appear within each contig, and save all 30 tracks as contigs/{contig}.bin.npz. Following the Borzoi paper, we can transform RNA bins onto the squashed scale: x^(3/4) if x^(3/4) < 384, otherwise 384 + sqrt(x^(3/4)-384). Copying the resulting data to S3 then allows downloading onto the GPU instance.

Fine-tuning Setup

The fine-tuning mechanism remains conceptually simple: first, wrap the existing Flashzoi model as a "backbone," delete the human and mouse heads, and add a new convolutional output head.

class BorzoiFineTune(nn.Module):
def __init__(self, backbone, new_out_channels: int):
super().__init__()
self.backbone = backbone # this is a fully loaded Borzoi instance
self.eps = 1e-8 # clamp value

# remove original heads if present (prevents accidental use + removes params)
if hasattr(self.backbone, "human_head"): delattr(self.backbone, "human_head")
if hasattr(self.backbone, "mouse_head"): delattr(self.backbone, "mouse_head")
self._new_out_ch = new_out_channels
self.new_outputs = nn.Conv1d(in_channels=1920,
out_channels=new_out_channels, kernel_size=1)


def forward(self, x, return_embeddings=False):
x = self.backbone.get_embs_after_crop(x)
embs = self.backbone.final_joined_convs(x)
y = F.softplus(self.new_outputs(embs))

if return_embeddings:
return y, embs
return y

def loss(self, preds, targets, mask=None):
# poisson-multionomial loss
# "we cropped from each side to focus the loss computation on the center 196,608 bp"
# -> 6144 bins / 16384 bins
# -> -5120 bins / each side
# -> preds already has the cropping layer
targets = targets[..., 5120:-5120]
lam = preds.clamp_min(self.eps)
poisson_nll = lam - targets * torch.log(lam)
lam_sum = lam.sum(dim=-1, keepdim=True).clamp_min(self.eps)
y_sum = targets.sum(dim=-1, keepdim=True).clamp_min(self.eps)
p_hat = lam / lam_sum
p_true = targets / y_sum
multinom_ce = -(p_true * torch.log(p_hat.clamp_min(self.eps))).sum(dim=-1, keepdim=True)
poisson_term = poisson_nll.mean(dim=-1, keepdim=True)
loss = poisson_term + 0.2 * multinom_ce

return loss.sum()

Freezing the backbone ensures that only the head undergoes tuning; the weights of the pre-trained Flashzoi model remain unmodified, and only the parameters of the new PBMC tracks become updated.

def freeze_trunk(self):
for p in self.backbone.parameters():
p.requires_grad = False
for p in self.new_outputs.parameters():
p.requires_grad = True

def unfreeze(self):
for p in self.backbone.parameters():
p.requires_grad = True

This approach also requires a straightforward means of instantiating the model from HuggingFace weights, which calls into the Flashzoi implementation to set up the backbone:

@classmethod
def from_pretrained(cls, repo_id: str, new_out_channels: int, **kwargs):
backbone = Borzoi.from_pretrained(repo_id, **kwargs) # already works for you
model = cls(backbone=backbone, new_out_channels=new_out_channels)
return model

For the training, both the 1-hot encoded reference and the 32bp-binned histone (and RNA) tracks for each cell type are saved as sparse matrices (as 0 values are typical). A simple iterator over the bins and contigs is given by:

def read_contig_dat(contig):
print(f' Loading contig {contig} ...')
refdata = sp.sparse.load_npz(f'contigs/{contig}.1hot.npz')
bindata = sp.sparse.load_npz(f'contigs/{contig}.bin.npz')
return bindata, refdata

def get_data(bin_offset, bin_dat, ref_dat):
# window is 524288bp = 16384 bins
bin_start = bin_offset - 8192
bin_end = bin_offset + 8192
track_dat = bin_dat[bin_start:bin_end,:]
ref_start = 32 * bin_start
ref_end = 32 * bin_end
onehot_dat = ref_dat[ref_start:ref_end,:]

return onehot_dat, track_dat


def iter_training(contig_list=None, thin=None):
state = np.random.RandomState(24)
if contig_list is None:
contig_list = [x.split('.1hot')[0] for x in os.listdir('contigs') if '1hot' in x]
state.shuffle(contig_list)
# read in all the bin data
for contig in contig_list:
bin_data, ref_data = read_contig_dat(contig)
offsets = np.arange(8192, bin_data.shape[0]-8192)
if thin is not None:
offsets = offsets[state.random(offsets.shape[0]) <= thin]
for bin_offset in offsets:
in_ref, out_tracks = get_data(bin_offset, bin_data, ref_data)
yield in_ref, out_tracks

def batch_data(iter_training, contigs=None, k=16, thin=1e-2):
Xl, Yl = list(), list()
for X, y in iter_training(contigs, thin=thin):
Xl.append(torch.as_tensor(X.T.todense(), dtype=torch.float32))
Yl.append(torch.as_tensor(y.T.todense(), dtype=torch.float32))

if len(Xl) == k:
Xb = torch.stack(Xl, dim=0)
Xb.pin_memory()
Xl.clear()
Yb = torch.stack(Yl, dim=0)
Yb.pin_memory()
Yl.clear()
yield Xb, Yb
if Xl:
Xb = torch.stack(Xl, dim=0)
Xb.pin_memory()
Xl.clear()
Yb = torch.stack(Yl, dim=0)
Yb.pin_memory()
Yl.clear()
yield Xb, Yb


"Contigs" and "thin" limit the training dataset; on the L40S, a full epoch of training takes around 140 hours (an excessive amount for this exercise). Compared to the ~50M parameters in the Borzoi model, the output weights (1920 x 30 x 1 = 57,600) suggest that the training should converge more rapidly. A note on the data shape: we employ "genomic" convention (position x channel), and, as the TensorFlow convention is (batch, channel, position), we need to transpose:

from train_util import *
from eval_util import *
from contextlib import nullcontext

itr = iter_training()
X, y = next(itr)
X.shape, y.shape # (524288,4) | (16384,30)
# torch wants (batch, channel, feature)
(y.T.todense())[None,:,:].shape
$
Loading contig chr1:23200698-24788581 ...
(1, 30, 16384)

Prior to training, the output head weights remain at their default values (0 or random), so the outputs should bear little resemblance to the actual data. We can check this via:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = BorzoiFineTune.from_pretrained('johahi/flashzoi-replicate-0', new_out_channels=30)
model.freeze_trunk()

loc = ('chr12', 69747834, 'C', 'T')
model.to(device)
dat = get_raw_pred_for_locus(loc, model, device)
plot_locus(dat['target'], loc, 'H3K4me3')
plot_locus(dat['predicted'], loc, 'H3K4me3', True)

Prior to fine-tuning, predicted H3K4me3 occupancy (right) across all cell types is uncorrelated with the true occupancy (left).

Interestingly, while we confirmed that the predictions lie far from the true data, we also observed potentially interesting "goings-on" at a nearby transcription start site (TSS; H3K4me3 peak), indicating that a substantial portion of the variance in the 1920 channels localizes to the TSS region.

Fine-tuning

The actual fine-tuning of our BorzoiFineTune model employs the standard approach: the AdamW optimizer with weight decay. As we significantly limit the time spent training, we radically thinned the dataset to a small fraction of bins and contigs to obtain random "representative" samples of the genome (rather than training on all bins from a single small contig).

import time
from tqdm.auto import tqdm

optimizer = torch.optim.AdamW(
filter(lambda p: p.requires_grad, model.parameters()),
lr=1e-4 if 'optimizer' not in dir() else 0.5 * optimizer.param_groups[0]['lr'],
weight_decay=1e-8,
)
model.train()
model.to(device)
optimizer.zero_grad(set_to_none=True)

batch_size=14
thin_frac=0.0025
n_contigs = 25
scaler = torch.cuda.amp.GradScaler(enabled=True)
if 'epochs' in dir(): # in case we want to add more
old_e, epochs = epochs, epochs + 5
else:
old_e, epochs = 0, 10
manual_log = list()
global_step = 0
for epoch in range(old_e, epochs):
contigs, train_size = get_contigs(n_contigs)
pbar = tqdm(batch_data(iter_training, contigs=contigs, k=batch_size, thin=thin_frac),
desc=f"epoch {epoch+1}/{epochs}", total=int(train_size*thin_frac/batch_size))
running_loss = 0.0
running_count = 0
t0 = time.time()
for t, (X, Y) in enumerate(pbar):
optimizer.zero_grad(set_to_none=True)
X = X.to(device)
Y = Y.to(device)
with torch.cuda.amp.autocast(enabled=True):
preds = model(X)
loss = model.loss(preds, Y)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

running_loss += loss.item()
running_count += 1
global_step += 1
it_per_s = running_count / max(time.time() - t0, 1e-6)
postfix = {
"loss": f'{loss.item():.4g}',
"loss (avg)": f"{(running_loss/running_count):.4g}",
"bs": batch_size,
"it/s": f"{it_per_s:.2f}",
}
if torch.cuda.is_available() and str(device).startswith("cuda"):
postfix["memGB"] = f"{torch.cuda.max_memory_allocated()/1e9:.2f}"

if (t < 200) or (global_step % 150 == 0):
manual_log.append({
'epoch': epoch+1,
'step': t+1,
'loss': loss.item(),
'running_loss': running_loss,
'global_step': global_step
})
pbar.set_postfix(postfix)


Despite the need for a 40+ GB GPU, the GPU should hover around 28 GB during training as: i) we eliminated ~1000 output heads from the baseline model, and ii) the gradients only need to be computed/stored for the output layer, which significantly reduces memory usage. Note that omitting backbone freezing limits the batch size to around four data points due to additional GPU gradients.

Within the first few epochs, the loss converges to around 875 and stays fixed, even after shrinking the learning rate a few times via re-running the cell:

By epoch 20, even with heavy down-sampling, the model has converged and is no longer improving.

Moreover, the model has indeed converged to mimic the scale and shape of the input FPM tracks:

After a short amount of training, the scale and location of occupancies closely match the input data.

Of note, the dendritic cell (DC) head provided evidence for a “noisy” DC track (due to its much smaller number of nuclei – roughly 300); therefore, the output head appears to apply significant “shrinkage” to the raw data. However, although training employed 35 “epochs”, each epoch used <2.5% of all contigs and 0.25% of all bins; therefore, there exists an approximately 50% chance that any particular visualization of 500 bins was used.

The model does learn to reconstruct cell-specific/cell-enriched loci; the following position depicts the model recapitulating both T-cell-enriched H3K4me1 and T-cell-enriched RNA expression:

RNA predictions re-capitulate T-cell-upregulated gene expression profiles.

However, in other cases, the model overrides true cell-specificity and hallucinates baseline expression:

Despite higher expression in DCs, the fine-tuned model predicts similar levels of expression across all cell types at this locus -- a failure mode likely due to the small DC population and thus higher "noise" and "shrinkage" applied.

Downstream Application: Quantitative Trait Locus Impacts

As Borzoi is a sequence-to-function model, we can interrogate quantitative trait loci (QTLs) for their predicted impacts on expression or epigenetic states. This constitutes a straightforward: given a point mutation (chr, pos, ref, alt):

1. Identify the corresponding contig and bin

2. Obtain the 8192 bins prior to and 8191 subsequent bins

3. Obtain the corresponding reference sequence a. (ref_start, ref_end) = (bin start – 819232: bin_start + 819232)

4. Mutate the base at (pos - ref_start) by setting the ref base to 0 and the alt base to 1

5. Call model.predict

One downloaded file contains summary statistics for blood expression (e)QTLs from a large consortium eQTLGen. For example, we chose rs2073529 - a mutation in the intron of BTN3A2, where the alternate allele "C" lowers gene expression. Importantly, the mutation locates to an H3K4me1 peak in the intron. The fine-tuned model predicts not only lower expression due to the mutation, but also a highly focal loss of H3K4me1 occupancy at the mutation site, strongly suggesting that the eQTL modulates expression by lowering enhancer activity at its locus.

Mutating rs2073529 T to C in silico results in a large predicted change in H3K4me1 occupancy - specifically lower - suggesting the eQTL is a functional hQTL associated with an enhancer reduction of function.

Conclusions

Output-head fine-tuning suffers from limitations compared to fine-tuning the entire model; for example, Borzoi training used CAGE and poly (A)- capture bulk tissue RNA. The snRNA data from Paired-Tag likely contain significantly more intronic counts from nascent RNA than the model used for training; therefore, the relative attention given to splicing and genic features may mismatch and could improve with whole-model fine-tuning. Similarly, the ENCODE ChIP data used represented fold-change-over-control (input) tracks, which, although correlated with CUT&Tag, may have prompted the model to incorporate sequence features specific to the ChIP fragment size and sonication biases.

Overall, their very structure clearly limits these models; they can only demonstrate the relationship between sequence changes and molecular outcomes, but not direct relations between epigenetic states themselves, which would require re-architecting the model to look more akin to the methylation models, incorporating both sequence and epigenetic tracks as inputs and outputs.

This all provides evidence that the distal effects of gene knockouts lie beyond the model’s capacity, as for epigenetic silencing effects. The field of cellular foundation models remains in its infancy; however, single-cell epigenetic assays such as scCUT&Tag and PairedTag (scCUT&Tag+RNA) remain critical to their future development.