Aligning RNA Expression Between CCLE Cancer Cell Lines and TCGA Tumors: a Statistical (Celligner) vs. a Deep Learning (Autoencoder) Approach

Summary

Screening compounds against cancer cell line panel for anti-proliferation is a powerful approach to discover drug sensitivity biomarkers. In vitro drug response data in the form of AUC or IC50, when paired with RNA expression data, can be used to build machine learning models for sensitivity prediction. Such a predictive model can be applied to tumor transcriptome to predict patient response. During my time in the industry, I’ve seen two challenges in this field: (1) the limited number of cell lines available for screening due to cost constraints, which limits statistical power for biomarker discovery, and (2) the lack of translatability of results between cell line and tumor samples, due to the poor alignment between cell line and tumor RNA expression data. This poor alignment is likely caused by biological (e.g. growth conditions, tumor cell purity & immune infiltration) and/or technical (e.g. sample preparation/NGS batch effect) confounders. In a two-part series, I aim to tackle these problems. In this first part, I use two popular approaches to align RNA expression between CCLE cancer cell lines and TCGA tumors: Celligner (Warren et al. 2021), and an autoencoder (AE)-based approach (similar to CODE-AE-ADV; He et al. 2022). Both models show excellent alignment performance, as shown in the t-SNE and UMAP plots below, with the AE model showing slightly better alignment quality as measured by k-NN domain mixing score. In a future part II, I will use aligned RNA expression (from Celligner) or latent representation (from AE) to build drug response prediction models, and evaluate their performance on held-out CCLE lines and TCGA tumors. I will also determine what is an “adequate” sample number for a cell panel screen to achieve good predictive performance. This decision is important for drug discovery, as it can help reduce the cost of screening by determining the minimum number of cell lines needed to achieve good predictive performance.

I am a US-based independent consultant specialized in NGS, omics, statistical & ML analysis. Contact me via LinkedIn if you think you have a project that I can help with! https://www.linkedin.com/in/jay-cy-chung/

Results

Figure 1. t-SNE (left) or UMAP (right) representation of raw RNA expression between CCLE (n=1,483; orange dots) vs. TCGA (n=11,274; blue dots). Note that CCLE and TCGA clusters are well separated and thus not aligned.

Figure 2. t-SNE or UMAP representation of Celligner aligned expression (18,627 aligned genes). Note the better mixing between domains across cell clusters (cancer lineage clusters).

Figure 3. t-SNE or UMAP representation of AE aligned embeddings (128 dimensions). AE also shows good mixing but with less granular clustering structure likely due to the compression of full gene expression into 128 embeddings. Whether this affects drug sensitivity prediction will be tested in my next post.

Figure 4. Gene expression alignment quality measured by k-NN domain mixing score (on average, how many of each sample’s nearest neighbors are from another data domain?). Both models show good mixing quality, with AE slightly outperforming Celligner. Dashed lines indicate mixing score under perfect mixing condition.

Methods

Data processing

CCLE (24Q2) and TCGA (pan-cancer atlas) TPM expression data were downloaded from DepMap portal https://depmap.org/portal/download/ and GDC https://gdc.cancer.gov/about-data/publications/pancanatlas/, respectively. CCLE and TCGA expression data were merged by gene symbol, and only genes with non-zero expression in both datasets were kept. The final dataset contains 18,627 genes across 1,483 CCLE cell lines and 11,274 TCGA tumors, which is used as the input for Celligner alignment. For AE alignment, the genes were further selected by variance, and the top 4,000 most variable genes were used as input for AE training.

Celligner

Celligner is now available as a Python package on PyPI https://pypi.org/project/celligner/. Under the Colab environment, it used 20-30 GB RAM and took 15 mins to run for the full CCLE-TCGA alignment.

from celligner import Celligner

my_celligner = Celligner(low_mem=True)
my_celligner.fit(ccle_log)
my_celligner.transform(tcga_log)

combined = my_celligner.combined_output
celligner_ccle_df = combined.loc[combined.index.intersection(ccle_log.index)]
celligner_tcga_df = combined.loc[combined.index.intersection(tcga_log.index)]
print("Celligner aligned -- CCLE:", celligner_ccle_df.shape, " TCGA:", celligner_tcga_df.shape)

Output:

Doing PCA..
Computing neighbors..
Clustering..
Running differential expression on 48 clusters..
Running limmapy..
Doing PCA..
Computing neighbors..
Clustering..
Running differential expression on 68 clusters..
Running limmapy..
Running cPCA..
Regressing top cPCs out of reference dataset..
Regressing top cPCs out of target dataset..
Doing the MNN analysis using Marioni et al. method..
  Looking for MNNs...
  Found 10258 mutual nearest neighbors.
Done
Celligner aligned -- CCLE: (1483, 18627)  TCGA: (11274, 18627)

Autoencoder (AE)

The AE model architecture is inspired by CODE-AE-ADV (He et al. 2022). This model is trained from scratch and implemented in PyTorch. The model takes in CCLE and TCGA expression data, and learns a latent representation that aligns the two domains. The training process involves minimizing reconstruction loss, orthogonality loss and domain adversarial loss to ensure that the latent space captures shared features between the two datasets.

Figure 5. Architecture of the AE model.

At the heart of the domain adversarial loss is a gradient reversal layer (GRL) and a domain discriminator. The gradient reversal layer is implemented as a custom autograd function in PyTorch, which reverses the gradient during backpropagation. The domain discriminator is a simple feedforward neural network that predicts the domain (CCLE or TCGA) of the latent representation. Due to the GRL mechanism, the shared encoder is trained to make the discriminator’s job impossible, thus aligning the two domains’ latent representations.

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

class GradReverse(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, lambd):
        ctx.lambd = lambd # store lambd
        return x.view_as(x) # return x as is, and tell pytorch this is an output
    @staticmethod
    def backward(ctx, grad_output):
        return -ctx.lambd * grad_output, None # Return lambd as None

def grad_reverse(x, lambd=1.0):
    return GradReverse.apply(x, lambd)
  
class DomainDiscriminator(nn.Module):
    def __init__(self, in_dim, hidden=64):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(in_dim, hidden), nn.ReLU(), nn.Linear(hidden, 1))
    def forward(self, z):
        return self.net(z).squeeze(-1)

A custom orthogonality loss is implemented to encourage the shared and private latent representations to be orthogonal, which helps in disentangling the shared and private features between CCLE and TCGA datasets. The orthogonality loss is computed as the squared Frobenius norm of the dot product between the normalized shared and private representations, averaged over the batch size. Note that the normalization step is crucial here as it ensures that the encoder does not take a shortcut to push the shared and private representations to zero to minimize the loss.

def orthogonality_loss(z_shared, z_private):
    zs = nn.functional.normalize(z_shared, p=2, dim=1) # sample-wise L2 norm
    zp = nn.functional.normalize(z_private, p=2, dim=1)
    return (zs.t() @ zp).pow(2).sum() / zs.size(0)

The MLP and the AE model are as below. It is very important to use nn.LayerNorm instead of nn.BatchNorm1d in the MLP, as the latter will weaken the batch effect between CCLE and TCGA only during training (as x_ccle and x_tcga each gets normalized independently), and it will come back during eval time, when BatchNorm would use the blended running statistics - making the “batch effect” reappear and eliminating any alignment quality from the encoder.

class MLP(nn.Module):
    def __init__(self, in_dim, hidden_dims, out_dim, dropout=0.1):
        super().__init__()
        layers, d = [], in_dim
        for h in hidden_dims:
            layers += [nn.Linear(d, h), nn.LayerNorm(h), nn.ReLU(), nn.Dropout(dropout)]
            d = h
        layers.append(nn.Linear(d, out_dim))
        self.net = nn.Sequential(*layers)
    def forward(self, x):
        return self.net(x)
      
class CodeAE(nn.Module):
    def __init__(self, in_dim, shared_dim=128, private_dim=64, hidden=(512, 256), dropout=0.1):
        super().__init__()
        self.shared_encoder        = MLP(in_dim, hidden, shared_dim, dropout)
        self.private_encoder_ccle  = MLP(in_dim, hidden, private_dim, dropout)
        self.private_encoder_tcga  = MLP(in_dim, hidden, private_dim, dropout)
        self.decoder_ccle          = MLP(shared_dim + private_dim, hidden[::-1], in_dim, dropout)
        self.decoder_tcga          = MLP(shared_dim + private_dim, hidden[::-1], in_dim, dropout)
        self.discriminator         = DomainDiscriminator(shared_dim)

    def encode_shared(self, x):
        return self.shared_encoder(x)

    def forward(self, x_ccle, x_tcga, grl_lambda=1.0):
        z_s_ccle = self.shared_encoder(x_ccle)
        z_s_tcga = self.shared_encoder(x_tcga)
        z_p_ccle = self.private_encoder_ccle(x_ccle)
        z_p_tcga = self.private_encoder_tcga(x_tcga)

        recon_ccle = self.decoder_ccle(torch.cat([z_s_ccle, z_p_ccle], dim=1))
        recon_tcga = self.decoder_tcga(torch.cat([z_s_tcga, z_p_tcga], dim=1))

        domain_logits_ccle = self.discriminator(grad_reverse(z_s_ccle, grl_lambda))
        domain_logits_tcga = self.discriminator(grad_reverse(z_s_tcga, grl_lambda))

        return dict(z_s_ccle=z_s_ccle, z_s_tcga=z_s_tcga, z_p_ccle=z_p_ccle, z_p_tcga=z_p_tcga,
                    recon_ccle=recon_ccle, recon_tcga=recon_tcga,
                    domain_logits_ccle=domain_logits_ccle, domain_logits_tcga=domain_logits_tcga)

The training steps are as below. Note that the total loss = reconstruction loss + orthogonality loss + domain adversarial loss. The GRL lambda is ramped up during the first 10% of training epochs, and then kept constant at 1.0. This allows the discriminator to get better before the shared encoder is forced to fool it.

def train_code_ae(ccle_X, tcga_X, ccle_X_val=None, tcga_X_val=None,
                   n_epochs=200, batch_size=128, lr=1e-3,
                   shared_dim=128, private_dim=64, hidden=(512, 256),
                   w_recon=1.0, w_adv=0.1, w_orth=0.1, grl_max=1.0, grl_ramp_frac=0.1,
                   patience=15, min_delta=1e-4, seed=SEED):
                     
    torch.manual_seed(seed)
    model = CodeAE(ccle_X.shape[1], shared_dim, private_dim, hidden).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)
    bce, mse = nn.BCEWithLogitsLoss(), nn.MSELoss()

    ccle_t = torch.tensor(ccle_X, dtype=torch.float32)
    tcga_t = torch.tensor(tcga_X, dtype=torch.float32)
    n_ccle, n_tcga = ccle_t.shape[0], tcga_t.shape[0]

    has_val = ccle_X_val is not None and tcga_X_val is not None
    if has_val:
        ccle_val_t = torch.tensor(ccle_X_val, dtype=torch.float32).to(device)
        tcga_val_t = torch.tensor(tcga_X_val, dtype=torch.float32).to(device)

    history = {"train_loss": [], "val_recon_loss": [], "z_shared_var": [],
               "loss_recon": [], "loss_adv": [], "loss_orth": [], "disc_acc": []}
    best_val_loss = float("inf")
    best_epoch = None
    best_state = None
    epochs_no_improve = 0

    for epoch in range(n_epochs):
        model.train()
        perm_ccle = torch.randperm(n_ccle)
        perm_tcga = torch.randperm(n_tcga)
        n_batches = max(n_ccle, n_tcga) // batch_size
        grl_lambda = grl_max * min(1.0, epoch / max(1, int(n_epochs * grl_ramp_frac)))

        epoch_loss = epoch_recon = epoch_adv = epoch_orth = epoch_disc_correct = epoch_disc_n = 0.0
        last_train_out = None
        for b in range(n_batches):
            idx_c = perm_ccle[(b * batch_size) % n_ccle: (b * batch_size) % n_ccle + batch_size]
            idx_t = perm_tcga[(b * batch_size) % n_tcga: (b * batch_size) % n_tcga + batch_size]
            if len(idx_c) < 2 or len(idx_t) < 2:
                continue
            xc, xt = ccle_t[idx_c].to(device), tcga_t[idx_t].to(device)

            out = model(xc, xt, grl_lambda=grl_lambda)
            loss_recon = mse(out["recon_ccle"], xc) + mse(out["recon_tcga"], xt)
            dom_c = torch.zeros(xc.size(0), device=device)
            dom_t = torch.ones(xt.size(0), device=device)
            loss_adv = bce(out["domain_logits_ccle"], dom_c) + bce(out["domain_logits_tcga"], dom_t)
            loss_orth = (orthogonality_loss(out["z_s_ccle"], out["z_p_ccle"]) +
                         orthogonality_loss(out["z_s_tcga"], out["z_p_tcga"]))

            loss = w_recon * loss_recon + w_adv * loss_adv + w_orth * loss_orth
            opt.zero_grad(); loss.backward(); opt.step()

            epoch_loss += loss.item()
            epoch_recon += loss_recon.item()
            epoch_adv += loss_adv.item()
            epoch_orth += loss_orth.item()
            with torch.no_grad():
                correct = ((out["domain_logits_ccle"] < 0).float().sum() +
                           (out["domain_logits_tcga"] > 0).float().sum())
                epoch_disc_correct += correct.item()
                epoch_disc_n += xc.size(0) + xt.size(0)
            last_train_out = out

        n_b = max(1, n_batches)
        train_loss = epoch_loss / n_b
        history["train_loss"].append(train_loss)
        history["loss_recon"].append(epoch_recon / n_b)
        history["loss_adv"].append(epoch_adv / n_b)
        history["loss_orth"].append(epoch_orth / n_b)
        history["disc_acc"].append(epoch_disc_correct / max(1, epoch_disc_n))

        model.eval()
        with torch.no_grad():
            if has_val:
                val_out = model(ccle_val_t, tcga_val_t, grl_lambda=grl_lambda)
                val_recon = (mse(val_out["recon_ccle"], ccle_val_t) + mse(val_out["recon_tcga"], tcga_val_t)).item()
                z_shared_for_var = torch.cat([val_out["z_s_ccle"], val_out["z_s_tcga"]], dim=0)
            else:
                val_recon = None
                z_shared_for_var = torch.cat([last_train_out["z_s_ccle"], last_train_out["z_s_tcga"]], dim=0)
            z_var = z_shared_for_var.var(dim=0).mean().item()
            history["val_recon_loss"].append(val_recon)
            history["z_shared_var"].append(z_var)

        if epoch % 10 == 0 or epoch == n_epochs - 1:
            val_str = f"  val_recon {val_recon:.4f}" if has_val else ""
            print(f"epoch {epoch:4d}  recon {history['loss_recon'][-1]:.4f}  "
                  f"adv {history['loss_adv'][-1]:.4f}  disc_acc {history['disc_acc'][-1]:.3f}"
                  f"{val_str}  z_shared_var {z_var:.4f}  grl_lambda {grl_lambda:.2f}")

        if has_val:
            if val_recon < best_val_loss - min_delta:
                best_val_loss, best_epoch = val_recon, epoch
                best_state = {k: v.detach().clone() for k, v in model.state_dict().items()}
                epochs_no_improve = 0
            else:
                epochs_no_improve += 1
                if epochs_no_improve >= patience:
                    print(f"Early stopping at epoch {epoch} -- restoring weights from epoch {best_epoch} (val_recon={best_val_loss:.4f})")
                    break

    if has_val and best_state is not None:
        model.load_state_dict(best_state)
        history["best_epoch"] = best_epoch
        history["best_val_recon_loss"] = best_val_loss
        history["stopped_epoch"] = epoch

    return model, history

To ensure that the model does not overfit, I split the data into training and validation sets, and used early stopping based on the validation reconstruction loss. The final embeddings for the full dataset were obtained by passing the entire CCLE and TCGA datasets through the trained model. I also monitored the variance of the shared latent representation during training, as a low variance may indicate that the model is not learning meaningful features.

model, history = train_code_ae(
    ccle_z_train.values, tcga_z_train.values,
    ccle_X_val=ccle_z_val.values, tcga_X_val=tcga_z_val.values,
    n_epochs=200, patience=15,
)

Figure 6. Training history of AE model. The training full loss (blue) and validation reconstruction-only loss (orange) are shown on the left. The dashed line marks the epoch with the best validation reconstruction loss (the epoch whose weights were restored). On the right is the shared embedding variance throughout training. The fact that it is well above zero shows that the model did not take the shortcut of shrinking z_shared to zero.

The transcriptome data from CCLE and TCGA were compressed into 128-dimensional embeddings by the encoder:

model.eval()
with torch.no_grad():
    z_ccle = model.encode_shared(torch.tensor(ccle_z.values, dtype=torch.float32).to(device)).cpu().numpy()
    z_tcga = model.encode_shared(torch.tensor(tcga_z.values, dtype=torch.float32).to(device)).cpu().numpy()

codeae_ccle_df = pd.DataFrame(z_ccle, index=ccle_log.index, columns=[f"CODEAE_{i}" for i in range(z_ccle.shape[1])])
codeae_tcga_df = pd.DataFrame(z_tcga, index=tcga_log.index, columns=[f"CODEAE_{i}" for i in range(z_tcga.shape[1])])
print(codeae_ccle_df.shape, codeae_tcga_df.shape)

Output:

(1483, 128) (11274, 128)

Disclosures

The code was written with the aid of Claude Sonnet 5, and the computation was performed on Google Colab Pro with a Tesla T4 GPU with high RAM. The accuracy of the code was verified by the author.

References

Ganin, Y., Ustinova, E., Ajakan, H., Germain, P., Larochelle, H., Laviolette, F., Marchand, M., & Lempitsky, V. (2016). Domain-adversarial training of neural networks. Journal of Machine Learning Research, 17(59), 1–35.

Ghandi, M., Huang, F. W., Jané-Valbuena, J., Kryukov, G. V., Lo, C. C., McDonald, E. R., et al. (2019). Next-generation characterization of the Cancer Cell Line Encyclopedia. Nature, 569(7757), 503–508. https://doi.org/10.1038/s41586-019-1186-3

He, D., Liu, Q., Wu, Y., & Xie, L. (2022). A context-aware deconfounding autoencoder for robust prediction of personalized clinical drug response from cell-line compound screening. Nature Machine Intelligence, 4(10), 879–892. https://doi.org/10.1038/s42256-022-00541-0

Hoadley, K. A., Yau, C., Hinoue, T., Wolf, D. M., Lazar, A. J., Drill, E., et al. (2018). Cell-of-origin patterns dominate the molecular classification of 10,000 tumors from 33 types of cancer. Cell, 173(2), 291–304. https://doi.org/10.1016/j.cell.2018.03.022

Warren, A., Chen, Y., Jones, A., Shibue, T., Hahn, W. C., Boehm, J. S., Vazquez, F., Tsherniak, A., & McFarland, J. M. (2021). Global computational alignment of tumor and cell line transcriptional profiles. Nature Communications, 12, 22. https://doi.org/10.1038/s41467-020-20294-x

van der Maaten, L., & Hinton, G. (2008). Visualizing data using t-SNE. Journal of Machine Learning Research, 9, 2579–2605.

McInnes, L., Healy, J., & Melville, J. (2018). UMAP: Uniform manifold approximation and projection for dimension reduction. arXiv:1802.03426.