Distributed Cross-Channel Hierarchical Aggregation for Foundation Models

2025-06-26

AI Summary

To address computational and memory bottlenecks in cross-channel tokenization and aggregation for high-dimensional multimodal images (e.g., hyperspectral, meteorological), this paper introduces the first hierarchical cross-channel aggregation mechanism—compatible with arbitrary model-parallel strategies and Vision Transformer architectures—to jointly optimize channel dimensionality and computation graphs. We further integrate tensor parallelism, model sharding, and hierarchical channel clustering to design a lightweight cross-channel attention module. Evaluated on the Frontier supercomputer (1,024 AMD GPUs), our approach reduces memory footprint by 75% and achieves over 2× sustained throughput improvement, significantly accelerating hyperspectral analysis and weather forecasting tasks. Our core contribution is the first systematic solution to the distributed scalability challenge of large-scale cross-channel image token aggregation, enabling efficient, scalable multimodal vision modeling.

Technology Category

• Computer Vision: Multi-modal Vision
• Machine Learning: Multimodal Learning
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Web Mining and Content Analysis: Mining multimedia, multimodal, multilingual, cross-lingual Web data
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Vision-based scientific foundation models hold significant promise for advancing scientific discovery and innovation. This potential stems from their ability to aggregate images from diverse sources such as varying physical groundings or data acquisition systems and to learn spatio-temporal correlations using transformer architectures. However, tokenizing and aggregating images can be compute-intensive, a challenge not fully addressed by current distributed methods. In this work, we introduce the Distributed Cross-Channel Hierarchical Aggregation (D-CHAG) approach designed for datasets with a large number of channels across image modalities. Our method is compatible with any model-parallel strategy and any type of vision transformer architecture, significantly improving computational efficiency. We evaluated D-CHAG on hyperspectral imaging and weather forecasting tasks. When integrated with tensor parallelism and model sharding, our approach achieved up to a 75% reduction in memory usage and more than doubled sustained throughput on up to 1,024 AMD GPUs on the Frontier Supercomputer.

Problem

• Efficiently aggregating multi-channel images for foundation models
• Reducing compute intensity in tokenizing diverse image sources
• Improving scalability of vision transformers on large datasets

Innovation

• Distributed Cross-Channel Hierarchical Aggregation method
• Compatible with model-parallel strategies and transformers
• Reduces memory usage and increases throughput