Construct Gene Regulatory Network

This module utilizes layers/heads self attention of genes to determine the importance of gene relationships to each other. This can be utilized to create a directed graph for genes that “pay attention to each other” in pretrained model

from stFormer.network_dynamics.gene_regulatory_network import GeneRegulatoryNetwork

  1. Model & Tokenizer: Load pre-trained BERT with output_attention=True to extract per layer/per-head attention scores

  2. pad all sequences out to the maximum length in the filtered set (so they form a rectangular tensor

grn = GeneRegulatoryNetwork(
    model_dir = 'output/spot/models/250422_102707_stFormer_L6_E3/final',
    dataset_path = 'output/spot/visium_spot.dataset',
    model_type = 'Pretrained',
    metadata_column = 'subtype',
    metadata_value = 'TNBC',
    device='cuda',
    batch_size=24,
    nproc = 12
)

  1. Compute Attention:

    • Stack & average over heads and layers → a (batch, seq_len, seq_len) tensor.

    • Sum these per-batch attention matrices into one running attn_sum

    • Finalize by dividing attn_sum by total examples → your final attention_matrix (NumPy array)

grn.compute_attention()

  1. Build Graph:

    • Build a “counts” matrix C of shape (T, seq_len), where C[i,p] = how many times token i appears at position

    • Compute N = C @ attention_matrix @ C.T (a single high-performance mat-mul) and normalize by the outer product of occurrence counts to get the average attention from token i to token j.

    • Edge Selection: cutoff, percentile, top_k

    • add a directed edge for every selected token pair with aggregated weight

#Can take a long time, computing average i,j attention across all samples for all token pairs
grn.build_graph(
    percentile = 0.99999, #filter node-edges by attention weight (value above percentile)
    min_cooccurrence=100 #filter node-edges by number of samples expressing pair
)

write to file

grn.plot_network('output/spot/gene_network.png')

  1. Transcription Factor Analysis: Future Updates