GoKawiilVirTues architecture
Multiplexed imaging data pose modality-specific challenges to the development of scalable machine learning algorithms. The images represent high-dimensional samples, characterized by a large number of measured channels and high spatial resolutions. Further, the total number as well as the combination of channels varies between datasets as studies use different marker panels. These characteristics of the data modality hinder the simple off-the-shelf application of established vision architectures. Both convolutional neural networks and standard ViTs require a constant number of input channels, typically red, green and blue (RGB), with a fixed semantic meaning. Moreover, in contrast to RGB channels, which combine to produce colours, multiplex channels convey distinct biological meanings and exhibit complex inter-relationships. To address the unique challenges posed by multiplex imaging data, we propose VirTues—an encoder–decoder model based on the ViT architecture. VirTues is designed for the efficient processing of highly multiplexed image data accommodating varying numbers and combinations of measured markers. Furthermore, VirTues incorporates the attribution of distinct biological meaning to each measured marker. VirTues operates on tokenized image crops of size d c × d c = 128 × 128. Restricting VirTues’ input to such crops increases the number and diversity of pretraining samples while decreasing the dimensionality per sample.
Tokenization
To preserve the biologically distinct meaning of each channel and allow for a flexible number of channels per image, we used a multi-channel tokenization procedure31,33. Each channel is divided spatially into patches of size \({d}_{{\rm{p}}}\times {d}_{{\rm{p}}}=8\times 8\), as this captures approximately one cell per patch. Flattening each patch results in a three-dimensional grid of image tokens \({\bf{x}}\in {{\mathbb{R}}}^{M\times H\times W\times {d}_{{\rm{p}}}^{2}}\), where \(M\) is the number of measured channels and \(H=W={d}_{{\rm{c}}}/{d}_{{\rm{p}}}\) the grid height and width. For all \(M\) markers measured by the channels of \({\bf{x}}\), we retrieve from a pre-computed lookup table the corresponding protein embeddings \(\pi \in {{\mathbb{R}}}^{M\times {d}_{{\rm{PLM}}}}\) given by the PLM (ESM-2 (ref. 18) with \({d}_{{\rm{PLM}}}=640\)). We refer to these embeddings as marker tokens. For each channel \(m\) and each grid position \((i,j)\), we project the image token \({{\bf{x}}}_{{mij}}\) and the corresponding marker token \({\pi }_{m}\) to the same dimension \({d}_{{\rm{model}}}\) using learnable linear projections, to get \({{\bf{x}}}_{{mij}}^{{\prime} }\in {{\mathbb{R}}}^{{d}_{{\rm{model}}}}\) and \({\pi }_{m}^{{\prime} }\in {{\mathbb{R}}}^{{d}_{{\rm{model}}}}\), respectively. The image and marker tokens are fused through summation, resulting in the image tokens \(\widetilde{{\bf{x}}}\in {{\mathbb{R}}}^{M\times H\times W\times {d}_{{\rm{model}}}}\), where \({\widetilde{{\bf{x}}}}_{{mij}}={{\bf{x}}}_{{mij}}^{{\prime} }+{\pi }_{m}^{{\prime} }\). The fusion of the marker token with the image tokens serves two main purposes: (1) enabling VirTues to differentiate the channel origins of input tokens, and (2) introducing a biologically informed prior, reflecting sequence-level protein relationships, which cannot be added through other marker tokenization schemes (such as one-hot or learnable marker embeddings). We note that this is the first of many building blocks enabling VirTues to generalize across unseen markers. Further, to allow VirTues to capture an aggregated representation for each patch, we introduce an additional layer of learnable patch summary tokens \({\bf{c}}\in {{\mathbb{R}}}^{H\times W\times {d}_{{\rm{model}}}}\), one for each spatial position. Each patch summary token \({{\bf{c}}}_{{ij}}\in {{\mathbb{R}}}^{{d}_{{\rm{model}}}}\) is initialized using the same weights.
Masking
During training, a portion of the image tokens \(\{{\widetilde{{\bf{x}}}}_{{mij}}\}\) is masked by replacing them with a special masking token \({\rm{\square }}\in {{\mathbb{R}}}^{{d}_{{\rm{model}}}}\) initialized with learnable weights. Masking is applied channel-wise by sampling a masking ratio \({r}_{{\rm{masking}}}\) between 60% and 100% and uniformly selecting the corresponding \(\lceil {r}_{\mathrm{masking}}HW\rceil \) tokens to mask within the channel. We denote the resulting three-dimensional binary mask by \({\bf{M}}\in \{0,1{\}}^{M\times H\times W}\), where the value \(1\) marks masking. Masked tokens remain linked to their specific markers, which is indicated by adding the marker tokens to the masked tokens.
VirTues encoder
The set of all non-masked image tokens \(\{{\widetilde{{\bf{x}}}}_{{mij}}| {{\bf{M}}}_{{mij}}=0\}\) and the set of patch summary tokens \(\{{{\bf{c}}}_{{ij}}\}\) is passed as an input to the VirTues encoder. This encoder is constructed by modifying the vision transformer’s architecture14, to adapt it to work with varying input markers efficiently, and capture marker correlations and spatial patterns separately. In contrast to standard ViTs, which use full multi-head self-attention where all tokens attend pairwise to each other, we use two specialized sparse multi-head self-attention mechanisms—marker attention and spatial attention, akin to space and time attention used in video transformers40. In marker attention, only tokens that are placed at the same spatial grid position attend to each other, thereby capturing inter-marker dependencies and correlations. We denote the set of input tokens to the \({\ell }\) th transformer block as \(\{{t}_{mij}^{{\ell }}\}\), where the token \({t}_{mij}^{{\ell }}\) is associated to the \(m\)-th channel and position \((i,j)\). In this notation, we treat the layer of patch summary tokens simply as a further channel. Then, a marker attention transformer block computes
$${\rm{\forall }}{i}^{\ast },{j}^{\ast }:\{{t}_{mij}^{{\ell }+1}\,|i={i}^{\ast },j={j}^{\ast }\}={\rm{M}}{\rm{H}}{\rm{S}}{\rm{A}}(\{{t}_{mij}^{{\ell }}\,|i={i}^{\ast },j={j}^{\ast }\}).$$
where MHSA denotes a transformer block with standard multi-head self-attention. By contrast, in spatial attention, only tokens belonging to the same channel attend to each other hence capturing spatial patterns across tissue. Following the notation for marker attention, a spatial attention transformer block computes
... continue reading