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

TL;DR

This paper introduces a tensor equivariance framework for AI-assisted MU-MIMO systems, enabling efficient and adaptable precoding and scheduling with near-optimal performance.

Contribution

It proposes a novel tensor equivariance framework and TENN modules for designing AI algorithms in MU-MIMO, improving efficiency and generalization.

Findings

Achieves near-optimal precoding and scheduling performance.

Reduces computational complexity significantly.

Demonstrates strong generalization across varying input sizes.

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.