Towards Unified AI Models for MU-MIMO Communications: A Tensor Equivariance Framework

Yafei Wang,Hongwei Hou,Xinping Yi,Wenjin Wang,Shi Jin

2024-06-13

Abstract:In this paper, we propose a unified framework based on equivariance for the design of artificial intelligence (AI)-assisted technologies in multi-user multiple-input-multiple-output (MU-MIMO) systems. We first provide definitions of multidimensional equivariance, high-order equivariance, and multidimensional invariance (referred to collectively as tensor equivariance). On this basis, by investigating the design of precoding and user scheduling, which are key techniques in MU-MIMO systems, we delve deeper into revealing tensor equivariance of the mappings from channel information to optimal precoding tensors, precoding auxiliary tensors, and scheduling indicators, respectively. To model mappings with tensor equivariance, we propose a series of plug-and-play tensor equivariant neural network (TENN) modules, where the computation involving intricate parameter sharing patterns is transformed into concise tensor operations. Building upon TENN modules, we propose the unified tensor equivariance framework that can be applicable to various communication tasks, based on which we easily accomplish the design of corresponding AI-assisted precoding and user scheduling schemes. Simulation results demonstrate that the constructed precoding and user scheduling methods achieve near-optimal performance while exhibiting significantly lower computational complexity and generalization to inputs with varying sizes across multiple dimensions. This validates the superiority of TENN modules and the unified framework.

Signal Processing

What problem does this paper attempt to address?

The main focus of this paper is on how to utilize artificial intelligence (AI) design assistance techniques in multi-user multi-input multi-output (MU-MIMO) communication systems. The paper presents a unified framework based on tensor equivariance to address the pre-coding and user scheduling problems in MU-MIMO systems. The authors define multi-equivariance, high-order equivariance, and multi-invariance and prove the mapping from channel information to optimal pre-coding tensor, pre-coding auxiliary tensor, and scheduling indicator possesses these equivariances.

The paper proposes a series of plug-and-play tensor equivariant neural network (TENN) modules to model equivariant mappings by transforming complex parameter sharing patterns into concise tensor operations. Based on these TENN modules, a unified equivariant framework is constructed, which is applicable to various communication tasks, and corresponding AI-assisted pre-coding and user scheduling schemes are designed. Simulation results demonstrate that the proposed approach achieves comparable performance with significantly reduced computational complexity and good generalization capability to varying input sizes.

The main contributions of this paper are as follows:
1. Introducing the concept of tensor equivariance (TE) for capturing multi-equivariance, high-order equivariance, and invariance.
2. Designing plug-and-play TENN modules for efficient utilization of TE.
3. Proposing a unified TE framework for easily designing neural networks suitable for various tasks, including pre-coding and user scheduling.

Through this framework, the paper addresses the challenges of high computational complexity and limited generalization capability faced by traditional methods in dealing with MU-MIMO systems, providing new insights for the AI-assisted design of future wireless communication systems.