Structure-Guided Transformer Clustering for High-Dimensional Data

Clustering high-dimensional data is challenging when informative signals are distributed across dependent features and cluster structure is nonlinear. We introduce Structure-Guided Transformer Clustering (SGTC), which estimates a feature conditional-dependence graph, partitions related variables into groups, guides the splitting of Transformer heads, and learns a prototype-based latent partition. Across simulations, SGTC achieved higher agreement with ground-truth clusters than an unstructured Transformer, K-means, spectral clustering, and UMAP-based clustering, particularly as dimensionality increased. Its advantage was strongest when the sample size was sufficient to estimate useful dependence structure and diminished in sample-limited settings. SGTC showed robust performance under random edge deletions and additions to the feature-dependence graphs. We further applied SGTC to miRNA expression profiles from 418 patients with lung squamous cell carcinoma in The Cancer Genome Atlas. Among the evaluated methods, the clusters identified by SGTC showed the strongest overall survival separation (log-rank p=0.0066). The pathway enrichment analysis of the top driven miRNAs suggested more plausible pathways based on SGTC than the classic transformer.