LGAI-EXAONE/

EXAONE-Path-MSI

Introduction

1. Patch-wise feature extraction – Each WSI is tiled into 256 × 256 px patches, which are embedded into 768-dimensional vectors using the frozen EXAONE Path encoder.
2. Slide-level aggregation – The patch embeddings are aggregated using a Vision Transformer, producing a unified slide-level representation that a lightweight classification head transforms into task-specific probabilities.

Key Improvements

• FlexAttention + torch.compile
  What changed: Replaced vanilla multi‑head self‑attention with IO‑aware FlexAttention kernels and enabled torch.compile to fuse the forward/backward graph at runtime. The new kernel layout dramatically improves both memory efficiency and training-and-inference throughput.
• Coordinate‑aware Relative Bias
  What changed: Added an ALiBi‑style distance bias that is computed from the (x, y) patch coordinates themselves, allowing the ViT aggregator to reason about spatial proximity.
• Scalable Mixed‑Omics Encoder (Token‑mixing Transformer)
  What changed: Each omics modality is first tokenised into a fixed‑length set. All modality‑specific tokens are concatenated into a single sequence and passed through a shared multi‑head self‑attention stack, enabling direct information exchange across modalities in one shot. The aggregated omics representation is subsequently fused with image tokens via cross‑attention. This release uses three modalities (RNA, CNV, DNA‑methylation), but the design is agnostic to modality count and scales linearly with token number.

Quick Start

Requirements

• NVIDIA GPU is required
• Minimum 40GB GPU memory recommended
• Tested on Ubuntu 22.04 with NVIDIA driver version 550.144.03

Installation

Quick Inference

Command‑line

Model Performance Comparison

License