hZACH-ViT: Curved Latent Geometry for Compact Vision Transformers in Low-Data Medical Imaging

2026-05-30

AI Summary

This work addresses the underexplored potential of non-Euclidean geometries in few-shot medical image classification, where existing compact Vision Transformers default to Euclidean latent spaces. Without altering the ZACH-ViT backbone, the authors replace only the final representation space and classification head to systematically evaluate, within a unified framework, the impact of hyperbolic (Poincaré and Klein) and spherical geometries. Experiments across seven MedMNIST datasets demonstrate that the optimal non-Euclidean configuration consistently outperforms the Euclidean baseline, yielding an average improvement of 0.021 in the primary metric and up to 0.055 in Macro-F1 on OCTMNIST. Low curvature settings (c = 0.1 or 0.2) prove generally effective, underscoring that both geometry type and curvature should be treated as data-dependent design choices.

Abstract

Compact Vision Transformers are attractive for medical imaging in low-data and resource-constrained settings, but most existing variants assume that Euclidean latent geometry is sufficient for organizing image representations. We introduce hZACH-ViT, a family of curved-geometry extensions of ZACH-ViT, a compact zero-token Vision Transformer that removes positional embeddings and the class token and relies on global average pooling over patch representations. To isolate the role of geometry, we preserve the verified ZACH-ViT backbone and modify only the final representation space and prototype-based classifier head, enabling a controlled comparison between Euclidean, hyperbolic, and spherical latent geometries. We evaluate Poincaré, Klein, and spherical hZACH-ViT heads on seven MedMNIST datasets under an identical few-shot protocol with 50 samples per class and five random seeds. The completed benchmark contains 770 training runs spanning seven datasets, three non-Euclidean geometries, seven curvature magnitudes, and a Euclidean baseline. Across all seven datasets, the best non-Euclidean hZACH-ViT configuration improves over Euclidean ZACH-ViT, with an average gain of +0.021 in the dataset-specific primary metric and the largest improvement on OCTMNIST (+0.055 MacroF1). Fixed low-curvature configurations retain positive gains on the majority of datasets, and low curvature values (c = 0.1 or 0.2) account for six of the seven dataset-level winners. Rather than identifying a universally optimal manifold, our results establish geometry and curvature as dataset-dependent model-selection variables, with fixed low-curvature analyses confirming that gains persist beyond exhaustive per-dataset tuning.

Problem

• Vision Transformers
• low-data medical imaging
• latent geometry
• non-Euclidean geometry
• compact models

Innovation

• curved latent geometry
• compact Vision Transformer
• non-Euclidean geometry
• low-data medical imaging
• prototype-based classification