Optimizing Large Railway Vision Models for Efficient In-Context Learning

Large railway vision models (LRVMs) have exhibited remarkable performance in tackling diverse railway-related vision-based tasks, attributed to their capacity for in-context learning (ICL). However, in terms of speed and memory usage, these models suffer from inefficient hardware utilization, particularly when training from scratch with suboptimal batch sizes. Consequently, LRVMs face challenges in achieving high throughput compared to state-of-the-art, task-specific railway models. To address this problem, we propose a novel framework with cross-block feature union (CBFU) designed to optimize the computation time and GPU memory allocations for LRVMs, in both the training and inference phases while preserving accuracy. Our key innovation lies in the synergistic integration of CBFU with cascaded group attention (CGA) and multi-level feature concatenation, which collectively address the inefficiencies of prior LRVMs—a challenge unresolved in existing work. Drawing inspiration from Efficient-ViT, we incorporate the CGA layer as part of the LRVM’s building blocks to alleviate computational bottlenecks of the multi-head self-attention (MHSA). Additionally, we implement a multi-level feature concatenation technique to enhance transferability across different railway-related tasks and reduce the computation cost. Compared to existing LRVMs, our model demonstrates significant improvement in computational efficiency, achieving a 19-fold reduction in FLOPs. Furthermore, the size of our model is about $9\times $ smaller and runs 4.1 and $27\times $ faster during training and inference, respectively. These advancements uniquely position our framework as the first to balance accuracy, speed, and memory efficiency for railway vision systems, enabling real-world deployment where prior LRVMs failed. Overall, the proposed advancements enable the usage of LRVMs in resource-limited railway scenarios.