# Deep Joint Source-Channel Coding for Adaptive Image Transmission over MIMO
Channels

Haotian Wu, , Yulin Shao, , Chenghong Bian,
Krystian Mikolajczyk, , Deniz Gündüz The authors are with the Department of
Electrical and Electronic Engineering, Imperial College London, London SW7
2AZ, U.K. (e-mails: {haotian.wu17, y.shao, c.bian22, k.mikolajczyk,
d.gunduz}@imperial.ac.uk). This paper was presented in part at IEEE
International Conference on Communications 2023.

###### Abstract

This paper introduces a vision transformer (ViT)-based deep joint source and
channel coding (DeepJSCC) scheme for wireless image transmission over
multiple-input multiple-output (MIMO) channels, denoted as DeepJSCC-MIMO. We
consider DeepJSCC-MIMO for adaptive image transmission in both open-loop and
closed-loop MIMO systems. The novel DeepJSCC-MIMO architecture surpasses the
classical separation-based benchmarks with robustness to channel estimation
errors and showcases remarkable flexibility in adapting to diverse channel
conditions and antenna numbers without requiring retraining. Specifically, by
harnessing the self-attention mechanism of ViT, DeepJSCC-MIMO intelligently
learns feature mapping and power allocation strategies tailored to the unique
characteristics of the source image and prevailing channel conditions.
Extensive numerical experiments validate the significant improvements in
transmission quality achieved by DeepJSCC-MIMO for both open-loop and closed-
loop MIMO systems across a wide range of scenarios. Moreover, DeepJSCC-MIMO
exhibits robustness to varying channel conditions, channel estimation errors,
and different antenna numbers, making it an appealing solution for emerging
semantic communication systems.

###### Index Terms:

Joint source-channel coding, semantic communication, vision transformer, MIMO,
image transmission.

##  I Introduction

The exponential growth of wireless multimedia applications, encompassing
augmented reality, virtual reality, real-time streaming, and edge
intelligence, has significantly heightened the demand for efficient wireless
transmission of image/video signals, particularly under strict delay
constraints[1, 2, 3]. Consequently, this has sparked a notable surge in
research focused on designing optimized image communication systems tailored
for wireless channels [4, 5, 6, 7, 8, 9, 10, 7, 11, 12, 13, 14].

The traditional solution from Shannon’s separation theorem is to independently
design source and channel coding, which is optimal in the asymptotic limit of
infinite block length for ergodic source and channel distributions[15].
However, the separation-based approach is known to be sub-optimal in the
practical finite block length regime, which becomes particularly limiting in
applications that impose strict latency constraints[16, 17]. Despite the known
theoretical benefits, designing practical joint source-channel coding (JSCC)
schemes has been an ongoing challenge for many decades. Significant progress
has been made in designing JSCC schemes over recent years thanks to the
introduction of deep neural networks (DNNs)[4, 5, 6, 7]. The first deep
learning-based JSCC (DeepJSCC) scheme for wireless image transmission was
presented in [4], and it was shown to outperform the concatenation of state-
of-the-art image compression algorithm better portable graphics (BPG) with
low-density parity-check (LDPC) codes. It was later extended to transmission
with adaptive channel bandwidth in[8] and [9] and the transmission over multi-
path fading channels in [5] and [18]. All of these prior works demonstrate the
advantage of the practical DeepJSCC scheme in low block length. However, these
existing DeepJSCC schemes solely consider single-antenna transmitters and
receivers, disregarding the potential of employing multiple-input multiple-
output (MIMO) systems to improve capacity and reliability. While there is a
growing literature successfully employing DNNs for various MIMO-related tasks,
such as detection and channel estimation [19, 20, 21], with the exception of
our initial results[22], no previous work has so far applied DeepJSCC to the
more challenging MIMO channels. JSCC over MIMO channels is initially studied
in [23] from a theoretical perspective. It is challenging to design a
practical JSCC scheme for MIMO channels, where the model needs to retrieve
coupled signals from different antennas experiencing different channel gains.

The advantages of DeepJSCC and end-to-end MIMO schemes in practical
transmission scenarios have prompted the investigation of JSCC over MIMO
systems by integrating deep learning (DL) technologies. Despite this, the
number of studies focusing on end-to-end MIMO communication schemes based on
DNNs remains limited. The first autoencoder (AE) based end-to-end MIMO
communication method is introduced in[24]. In the subsequent work [25], the
authors systematically establish the symbol error rate benchmarks for MIMO
channels by evaluating several AE-based models with channel state information
(CSI). A singular-value decomposition (SVD) based AE is proposed in [26] to
achieve the state-of-the-art bit error rate. These works indicate the
potential of end-to-end MIMO communication schemes to enhance transmission
quality when coupled with advancing DL technologies. However, it is essential
to note that these MIMO schemes solely consider the transmission of bits at a
fixed signal-to-noise ratio (SNR) value, ignoring the source signal’s
semantics and the model’s adaptability to varying channel conditions. Two
recent studies that are more relevant to this work are [27] and [28]. The
authors in [27] theoretically analyze the excess distortion exponent for
semantic-aware MIMO systems in the context of JSCC. The authors in [28]
propose a semantically coded transmission scheme utilizing MIMO, focusing
solely on the scenarios, where the CSI is available at the receiver.

One critical shortcoming of the aforementioned works is the lack of
consideration for essential factors of the proposed DeepJSCC model, such as
the generalizability, the ability to adapt to diverse channel conditions, and
the robustness to channel estimation errors. Addressing these aspects is
pivotal in ensuring the efficacy and practicality of the proposed solutions.
As for channel adaptability, [6] and [18] incorporates the attention mechanism
to adapt the DeepJSCC model to different channel SNRs. Regarding the
robustness against the channel estimation errors, [29] revealed that even with
imperfect CSI and minimal noise, the performance of the MIMO system can still
be enhanced. However, significant noise accompanying imperfect CSI can
significantly degrade the system’s performance. [30] considered the channel
estimation errors and explored DL methods for joint MIMO channel estimation
and signal detection, which demonstrated the potential of DL approaches in
improving the model robustness against channel estimation errors. To the best
of our knowledge, prior research has not delved into the specific exploration
of adaptive JSCC-MIMO systems designed for wireless image transmission,
encompassing both scenarios with open-loop MIMO with only receiver-side
channel state information (CSIR) and closed-loop MIMO with both transmitter
and receiver side channel state information (CSIT). Such a system would be
capable of leveraging the semantics of the source signal and channel
conditions simultaneously, and efficiently managing both fluctuating channel
conditions and potential channel estimation errors, all within a cohesive
framework.

In light of this research gap, our paper introduces an innovative DeepJSCC
scheme meticulously tailored for MIMO image transmission. Our approach
introduces a unified vision transformer (ViT)-based DeepJSCC scheme, named
DeepJSCC-MIMO, which accommodates both scenarios with CSIR and CSIT. Inspired
by the success of the attention mechanism in the development of flexible
communication schemes[18, 6, 31, 32], we leverage the self-attention mechanism
inherent in the ViT for adaptive wireless image transmission. Specifically, we
represent the channel conditions with a channel heatmap and adapt the JSCC
encoding and decoding parameters according to this heatmap. Our approach can
learn global attention between the source image and the channel conditions
across all the intermediate layers of the DeepJSCC encoder and decoder.
Intuitively, we expect this design to simultaneously learn channel symbol
mapping and power allocation, considering different channel conditions,
channel estimation errors, and different antenna number settings. Therefore,
unlike most existing literature on DNN-aided MIMO communications, we eliminate
the need for training separate DNN models for each different transceiver pair
and different channel conditions. The proposed DeepJSCC-MIMO is a practical
and unified scheme with robustness to channel estimation errors.

Our main contributions can be summarized as follows:

  * •

This paper presents the first DeepJSCC scheme over the MIMO system, DeepJSCC-
MIMO, for adaptive image transmission. The key innovation lies in the
utilization of a ViT-based model, which adeptly exploits both the semantic
features of the image and the CSI (subject to availability) through a self-
attention mechanism.

  * •

DeepJSCC-MIMO is a comprehensive and versatile solution that can be seamlessly
applied to both open-loop and closed-loop MIMO systems. Extensive numerical
evaluations validate the superiority of the proposed model, showcasing
substantial enhancements in transmission quality across a wide range of
channel conditions and bandwidth ratios when compared to traditional separate
source and channel coding schemes. Furthermore, the DeepJSCC-MIMO scheme
exhibits remarkable resilience to channel estimation errors.

  * •

Importantly, DeepJSCC-MIMO offers remarkable flexibility and adaptability, as
it can effectively accommodate varying channel conditions and configurations
of antenna numbers without the need for retraining. By leveraging the channel
heatmap and employing a self-attention mechanism, DeepJSCC-MIMO intelligently
captures and adapts to the characteristics of the channel conditions and
antenna configurations.

##  II System Model

We consider an $M\times M$ MIMO communication system, where an $M$-antenna
transmitter aims to deliver an image $\bm{S}\in\mathbb{R}^{h\times w\times 3}$
to an $M$-antenna receiver ($h$ and $w$ denote the height and width of the
image, while $3$ refers to the color channels R, G, and B). The transmitter
encodes the image into a vector of channel symbols
$\bm{X}\in\mathbb{C}^{M\times k}$, where $k$ is the number of channel uses.
Following the previous standard definition [4], we denote the bandwidth ratio
(i.e., channel usage to the number of source symbols ratio) by $R\triangleq
k/n$, where $n=3hw$ is the number of source symbols. Intuitively, $R$ reflects
the number of available channel symbols for each source symbol. The
transmitted signal $\bm{X}$ is subject to a power constraint $P_{s}$ as:

$\frac{1}{Mk}\|\bm{X}\|_{F}^{2}\leq P_{s},$ (1)

where $\|\cdot\|_{\text{F}}$ denotes the Frobenius norm, and we set $P_{s}=1$
without loss of generality.

The channel model can be written as:

$\bm{{Y}}=\bm{H}\bm{X}+\bm{W},$ (2)

where $\bm{X}\in\mathbb{C}^{M\times k}$ and $\bm{Y}\in\mathbb{C}^{M\times k}$
denote the channel input and output matrices, respectively, while
$\bm{W}\in\mathbb{C}^{M\times k}$ is the Additive White Gaussian Noise (AWGN)
term that follows independent and identically distributed (i.i.d.) complex
Gaussian distribution with zero mean and variance $\sigma_{w}^{2}$, i.e.,
$\bm{W}[i,j]\sim\mathcal{CN}(0,\sigma_{w}^{2})$. The entries of the channel
gain matrix $\bm{H}\in\mathbb{C}^{M\times M}$ follow i.i.d. complex Gaussian
distribution with zero mean and variance $\sigma_{h}^{2}$, i.e.,
$\bm{H}[i,j]\sim\mathcal{CN}(0,\sigma_{h}^{2})$. We consider a slow block-
fading channel model, in which the channel matrix $\bm{H}$ remains constant
for $k$ channel uses, corresponding to the transmission of one image, and
takes an independent realization in the next block.

Given the channel output $\bm{Y}$, the receiver reconstructs the source image
as $\bm{\hat{S}}\in\mathbb{R}^{h\times w\times 3}$. We use the peak signal-to-
noise ratio (PSNR) as the distortion metric:

$\text{PSNR}=10\log_{10}\frac{\|\bm{S}\|_{\infty}^{2}}{\text{MSE}(\bm{S},\bm{\hat{S}})}~{}(\text{dB}),$
(3)

where $\|\cdot\|_{\infty}$ is the infinity norm and
$\text{MSE}(\bm{S},\bm{\hat{S}})\triangleq\frac{1}{3hw}\|\bm{S}-\bm{\hat{S}}\|^{2}_{2}$
is the mean squared error (MSE) between $\bm{S}$ and $\bm{\hat{S}}$.

As shown in Fig. 1, there can be two approaches to solve this problem, i.e.,
the traditional separate source and channel coding scheme and the JSCC scheme.
This paper considers both the CSIR (open-loop MIMO) and CSIT (closed-loop
MIMO) scenarios for each transmission scheme.

###  II-A Open-loop MIMO with CSIR

In an open-loop MIMO system, CSI is only available at the receiver, enabling
it to equalize the transmitted signals and subsequently decode the image.

####  II-A1 Separate source and channel coding scheme

In an open-loop MIMO system with a separate source and channel coding
transmission scheme, the transmitter follows a sequential process, including
the source coding, channel coding, and the modulation to generate the channel
symbols $\bm{X}$, which will go through the MIMO channel, as in Eqn. (2). We
note that, due to the lack of CSI, the transmitter has to choose the
compression and channel coding rates and the constellation size independently
of the channel realization.

Given the received signal $\bm{Y}$ and the CSI, the receiver decodes the
source signal as $\bm{\hat{S}}$ by sequentially performing the MIMO detection,
demodulation, channel decoding, and source decoding operations.

####  II-A2 DeepJSCC scheme

Different from the traditional separation-based scheme, DL-based JSCC scheme
uses a joint source-channel encoder at the transmitter, denoted as
$f_{\bm{\theta}}$, to directly map the source signal $\bm{S}$ into the channel
input matrix $\bm{X}$ for MIMO channel transmission. This process can be
represented as:

$\bm{X}=f_{\bm{\theta}}(\bm{S}).$ (4)

In the receiver, a MIMO equalization algorithm, such as zero-forcing (ZF) or
minimum mean square error (MMSE) algorithm[33], is firstly employed to exploit
the CSI to decouple the entangled signal $\bm{Y}$ as $\bm{X^{\prime}}$.
Subsequently, a JSCC decoder, referred to as $f_{\bm{\phi}}$, is employed to
reconstruct the source signal based on the available CSI represented by
($\bm{H}$, $\sigma_{w}^{2}$), along with $\bm{X^{\prime}}$. This can be
expressed as:

$\bm{\hat{S}}=f_{\bm{\phi}}(\bm{X^{\prime}},\bm{H},\sigma_{w}^{2}).$ (5)

It is important to note that an alternative approach is to directly learn the
decoding process using the CSI without explicitly performing the channel
equalization operation. This alternative approach involves training the model
to minimize the reconstruction distortion in an end-to-end manner. However, in
order to achieve better transmission quality, we adopt a model-driven approach
that explicitly performs channel equalization prior to the decoding process.

###  II-B Closed-loop MIMO with CSIT

Within a closed-loop MIMO system, the CSI is accessible to both the
transmitter and receiver, which allows them to apply pre-coding and power
allocation at the transmitter, and MIMO equalization at the receiver, thereby
improving the image transmission quality.

Figure 1: Block diagram of the MIMO image transmission system: (a)
conventional separate source-channel coding scheme and (b) DeepJSCC scheme,
where the gray blocks with dashed lines are the additional operations for the
closed-loop MIMO system.

####  II-B1 Separate source and channel coding scheme

The transmitter sequentially performs source coding, channel coding, and
modulation to generate the channel input matrix $\bm{X}$, the elements of
which are constellations with average power normalized to 1. Additional
operations at the transmitter with CSIT are precoding and power allocation,
which can boost the communication rate. Specifically, given the CSI, we first
decompose the channel matrix by singular-value decomposition (SVD), yielding
$\bm{{H}}=\bm{U\Sigma V^{H}}$, where $\bm{U}\in\mathbb{C}^{M\times M}$ and
$\bm{V}\in\mathbb{C}^{M\times M}$ are unitary matrices, and $\bm{\Sigma}$ is a
diagonal matrix whose singular values are in descending order. We denote
$\bm{\Sigma}$ by $\text{diag}(s_{1},s_{2},\dots,s_{M})$, where $s_{1}\geq
s_{2}\geq\dots\geq s_{M}$.

Let us denote the power allocation matrix by $\bm{\Lambda}$, where
$\bm{\Lambda}$ is diagonal, with its diagonal elements being the power
allocation weights for the signal streams of separate antennas. $\bm{\Lambda}$
can be derived using the standard water-filling algorithm. With power
allocation and SVD precoding (precoding $\bm{X}$ into $\bm{VX}$), Eqn. (2) can
be rewritten as:

$\bm{{Y}}=\bm{H}\bm{V}\bm{\Lambda X}+\bm{W}=\bm{U\Sigma}\bm{\Lambda
X}+\bm{W}.$ (6)

Multiplying both sides of Eqn. (6) by $\bm{U}^{H}$ (MIMO detection) gives us:

$\bm{X^{\prime}}=\bm{\Sigma}\bm{\Lambda X}+\bm{U}^{H}\bm{W}.$ (7)

As can be seen, SVD-based precoding converts the MIMO channel into a set of
parallel subchannels with different SNRs. In particular, the SNR of the $i$-th
subchannel is determined by the $i$-th singular value of $s_{i}$ and the
$i$-th power allocation coefficient of $\bm{\Lambda}$.

Given $\bm{X^{\prime}}\in\mathbb{C}^{M\times k}$, the receiver performs
demodulation, channel decoding, and source decoding sequentially to
reconstruct the source image as $\bm{\hat{S}}$.

####  II-B2 DeepJSCC scheme

For the DeepJSCC scheme, we exploit DL technologies to parameterize the
encoder and decoder functions, which are trained jointly on an image dataset
and the channel model as in Eqn. (2). Let us denote the DeepJSCC encoder and
decoder by $f_{\bm{\theta}}$ and $f_{\bm{\phi}}$, respectively, where
$\bm{\theta}$ and $\bm{\phi}$ denote the network parameters. We have

$\bm{X}=f_{\bm{\theta}}(\bm{S},\bm{H},\sigma_{w}^{2}).$ (8)

Unlike the separate source-channel coding scheme, the transmitter does away
with power allocation. Instead, we leverage the DeepJSCC encoder to perform
feature extraction and channel symbol mapping, and power allocation, all at
once. Intuitively, the DNN is expected to transmit critical features over
subchannels with higher SNRs, thereby improving the transmission performance.

We note here that one option is to train the encoder/decoder networks
directly, hoping they will learn to exploit the spatial degrees-of-freedom
MIMO channel provides. We will instead follow the model-driven approach, where
we will exploit the SVD as is done above for the channel coding scheme and
convert the MIMO channel into subchannels. The received signal can then be
written as

$\bm{{Y}}=\bm{HVX+W}=\bm{U\Sigma}\bm{X}+\bm{W}.$ (9)

To simplify the training, we apply MIMO equalization by left multiplying both
sides of Eqn. (9) by $\bm{\Sigma}^{\dagger}\bm{U}^{H}$ to obtain

$\bm{X^{\prime}}=\bm{\Sigma}^{\dagger}\bm{U}^{H}\bm{Y}=\bm{X}+\bm{W^{\prime}},$
(10)

where
$\bm{W}^{\prime}\triangleq\bm{\Sigma^{\dagger}U^{H}W}\in\mathbb{C}^{M\times
k}$ is the equivalent noise term, $\bm{\Sigma^{\dagger}}$ is the Moore-Penrose
inverse of matrix $\bm{\Sigma}$, and $\bm{U}^{H}$ is the conjugate transpose
of matrix $\bm{U}$.

Finally, we feed both $\bm{X^{\prime}}$ and the CSI into the DeepJSCC decoder
to recover the image as

$\bm{\hat{S}}=f_{\bm{\phi}}(\bm{X^{\prime}},\bm{H},\sigma_{w}^{2}).$ (11)

##  III Proposed method

This section presents a novel DeepJSCC architecture called DeepJSCC-MIMO,
which aims to enable efficient image transmission in MIMO systems. The
DeepJSCC-MIMO scheme is mainly based on the vision transformer and
incorporates self-attention mechanisms to exploit the CSI. Notably, our
DeepJSCC-MIMO design exhibits versatility and can be employed in both open-
loop and closed-loop MIMO systems.

The pipeline of our DeepJSCC-MIMO scheme is illustrated in Fig. 2, and
presented in Algorithm 1, where DeepJSCC-MIMO utilizes a pair of ViT-based
encoder $f_{\bm{\theta}}$ and decoder $f_{\bm{\phi}}$, respectively. Further
elaborations on the inner structures of these ViTs are provided in Fig. 3. In
the following, we detail the pipeline of DeepJSCC-MIMO in five main steps:
image-to-sequence transformation, channel heatmap construction, ViT encoding,
ViT decoding, and the loss function.

Figure 2: The pipeline of the DeepJSCC-MIMO scheme, where the source image
$\bm{S}$ is encoded by a ViT-encoder and reconstructed by a ViT-decoder as
$\bm{\hat{S}}$. Precoding operation in dashed line will be performed if the
CSI is available at the transmitter, and the CSI ($\bm{H},\sigma_{w}^{2}$) is
fed to the encoder, if available, and decoder in the form of a “heatmap”
$\bm{M}$ to facilitate the JSCC encoding/decoding process. Figure 3: The
architecture of the ViT-based encoder and decoder, where both encoder and
decoder comprise linear projection layers, a positional embedding layer, and
multiple transformer layers.

###  III-A Image-to-sequence transformation

To construct the input of DeepJSCC-MIMO, we first convert the three-
dimensional input image $\bm{S}$ into a sequence of vectors, denoted by
$\bm{S_{s}}=\textit{Seq}(\bm{S})$. Specifically, given a source image
$\bm{S}\in\mathbb{R}^{h\times w\times 3}$, we divide $\bm{S}$ into a grid of
$p\times p$ patches, and reshape each patch into a vector of dimension
$\mathbb{R}^{\frac{3hw}{p^{2}}}$. In this way, $\bm{S}$ is converted to
$\bm{S_{s}}\in\mathbb{R}^{l\times c}$, where $l={p^{2}}$ is the sequence
length and $c\triangleq\frac{3hw}{p^{2}}$ is the dimension of each vector in
the sequence.

###  III-B Channel heatmap construction

To enable efficient training, we construct a channel heatmap from CSI,
indicating the effective noise variance faced by each channel symbol generated
by ViT. Let us define $\bm{P_{n}}\in\mathbb{R}^{M\times k}$ as the average
power of the additive noise term $\bm{W^{\prime}}$. Specifically, for
different scenarios, $\bm{{P}_{n}}$ is defined as:

$\bm{{P}_{n}}\triangleq\begin{cases}\sigma_{w}^{2}(\bm{H_{w}}\odot\bm{\bar{H}_{w}})\bm{J}_{M\times
k},&\text{MIMO with CSIR}\\\
{\sigma}_{w}^{2}\bm{\Sigma}^{\dagger}\bm{U}^{H}\bm{J}_{M\times k},&\text{MIMO
with CSIT},\\\ \end{cases}$ (12)

where $\odot$ denotes the Hadamard product and $\bm{\bar{H}_{w}}$ is the
element-wise conjugate of matrix $\bm{{H}_{w}}$, which is obtained from the
zero-forcing operation as: $\bm{H_{w}}\triangleq\bm{(H^{H}H)^{-1}H^{H}}$, and
$\bm{J}$ is the matrix of ones.

Then the heatmap $\bm{M}\in\mathcal{R}^{l\times\frac{2Mk}{l}}$ is constructed
as follows:

$\bm{{M}}\triangleq\texttt{reshape}\left(\texttt{concat}\left(\frac{1}{2}\bm{P_{n}},\frac{1}{2}\bm{P_{n}}\right)\right),$
(13)

where $\texttt{concat}(\cdot)$ and $\texttt{reshape}(\cdot)$ denote the
concatenation and reshape operations, respectively. We concatenate two
matrices $\frac{1}{2}\bm{P_{n}}$ to get the shape of $\mathbb{R}^{M\times
2k}$, and reshape it into $\bm{M}\in\mathcal{R}^{l\times\frac{2Mk}{l}}$, which
represents the equivalent noise term faced by each real encoder output
element.

We expect that feeding the CSI in the form of a heatmap will simplify the
training process as the model only needs to focus on the ‘additive’ noise
power faced by each channel symbol. As we will show later, this design enables
our model to be used under various antenna numbers without retraining.

###  III-C ViT encoding

For the two scenarios, the input sequences to our ViT encoder
$\bm{{S}_{in}}\in\mathbb{R}^{l\times c_{in}}$ are given by:

$\bm{{S}_{in}}=\begin{cases}\bm{S_{s}},&\text{MIMO with CSIR}\\\
\texttt{concat}(\bm{S_{s}},\bm{M}),&\text{MIMO with CSIT},\\\ \end{cases}$
(14)

where the dimensions of the vectors in these sequences are $c_{in}=c$ and
$c_{in}=c+\frac{2Mk}{l}$ for the MIMO with CSIR and MIMO with CSIT scenarios,
separately.

The architecture of our encoder is shown in Fig. 3, which mainly consists of
linear projection layers, a positional embedding layer, and several
transformer layers.

Linear projection and positional embedding: $\bm{S_{in}}$ firstly goes through
a linear projection operation with parameters
$\bm{W_{0}}\in\mathbb{R}^{c_{in}\times d}$ and a positional embedding
operation $P_{e}(\cdot)$ to get the initial input
$\bm{F_{0}}\in\mathbb{R}^{l\times d}$ for the following transformer layers:

$\bm{F_{0}}=\bm{S_{in}}\bm{W_{0}}+P_{e}(\bm{p}),$ (15)

where $d$ is the output dimension of the hidden projection layer. Positional
embedding $P_{e}(\cdot)$ serves the purpose of representing the spatial
arrangement of sequences for better performance, where two typical methods are
dense layer-based position embedding (DPE) [34] and conditional position
embedding (CPE) [35]. In specific, $P_{e}(\cdot)$ is implemented with a DPE
that employs a dense layer to embed the index vector $\bm{p}$ of each patch
into a $d$ dimensional vector.

Transformer Layer: As shown in Fig. 3, the intermediate feature map
$\bm{F_{i}}$ is generated by the $i\textit{-}th$ transformer layer by a multi-
head self-attention (MSA) block and a multi-layer perceptron (MLP) layer as:

$\bm{F_{i}}=MSA(\bm{F_{i-1}})+MLP(MSA(\bm{F_{i-1}})),$ (16)

where $\bm{F_{i}}\in\mathbb{R}^{l\times d}$ is the output sequence of the
$i$-th transformer layer, GeLU activation and layer normalization operations
are applied before each MSA and MLP block.

Each MSA block consists of $N_{s}$ self-attention (SA) modules with a residual
skip, which can be formulated as:

$MSA(\bm{F_{i}})=\bm{F_{i}}+[SA_{1}(\bm{F_{i}}),\cdots,SA_{N_{s}}(\bm{F_{i}})]\bm{W_{i}},$
(17)

where the output of all SA modules $SA(\bm{F_{i}})\in\mathbb{R}^{l\times
d_{s}}$ are concatenated for a linear projection
$\bm{W_{i}}\in\mathbb{R}^{d_{s}N_{s}\times d}$, $d_{s}=d/N_{s}$ is the output
dimension of each SA operation.

For each SA module, the operations are formulated as:

$SA(\bm{F_{l-1}})=softmax\left(\frac{\bm{qk}^{T}}{\sqrt{d}}\right)\bm{v},$
(18)

where the $\bm{q},\bm{k},\bm{v}\in\mathbb{R}^{l\times d_{s}}$ are the query,
key, and value vectors generated through three linear projection layers
$\bm{W_{q},W_{k},W_{v}}\in\mathbb{R}^{d\times d_{s}}$ as:

$\bm{q}=\bm{F_{l-1}W_{q}},~{}~{}\bm{k}=\bm{F_{l-1}W_{k}},~{}~{}\bm{v}=\bm{F_{l-1}W_{v}}.$
(19)

Linear projection and power normalization: After $L_{t}$ transformer layers,
we apply a linear projection $\bm{W_{c}}\in\mathbb{R}^{d\times\frac{2Mk}{l}}$
to map the output of the transformer layers $\bm{F_{L_{t}}}$ into the channel
symbols as:

$\bm{Z_{e}}=\bm{F_{L_{t}}}\bm{W_{c}},$ (20)

where $\bm{Z_{e}}\in\mathbb{R}^{l\times\frac{2Mk}{l}}$ is then reshaped and
normalized to satisfy the power constraints to form the complex channel input
symbols $\bm{X}\in\mathbb{C}^{M\times k}$.

###  III-D ViT decoding

Channel equalization: To simplify the decoding process, we perform an
equalization operation on the channel output $\bm{Y}$ to decouple the
transmitted data of each antenna as $\bm{X^{\prime}}\in\mathbb{C}^{M\times k}$
before feeding the signal into the ViT decoder. We apply different
equalization methods for the MIMO systems with CSIR and with CSIT:

$\bm{{X}^{\prime}}=\begin{cases}\bm{H_{w}}\bm{Y}+f_{res}(\bm{Y},\bm{H}),&\text{MIMO
with CSIR}\\\ \bm{\Sigma^{\dagger}U^{H}Y},&\text{MIMO with CSIT}.\\\
\end{cases}$ (21)

Figure 4: Structure of DL-aided channel equalization method, where the input
at the $i$-th channel use is the concatenation of all elements within $\bm{H}$
and the channel output

To perform MIMO equalization in an open-loop MIMO system with CSIR, we
designed a DL-aided equalization method (as detailed in Fig. 4) to retrieve
the transmitted signal $\bm{X^{\prime}}$ from $\bm{Y}$. Specifically, we add a
residual block, denoted as $f_{res}(\cdot)$, to learn the compensation after
the zero-forcing (ZF) channel equalization operation based on the received
signals and CSI, given as:

$\displaystyle\bm{{X^{\prime}}}$
$\displaystyle=\bm{H_{w}Y}+f_{res}(\bm{Y},\bm{H})$ (22)
$\displaystyle=\bm{X}+\bm{H_{w}W}+f_{res}(\bm{Y},\bm{H}),$ (23)

where
$\bm{{W^{\prime}}}\triangleq\bm{H_{w}W}\in\mathbb{C}^{M\times\frac{k}{M}}$ is
the equivalent additive noise term from the channel. The residual operation
$f_{res}(\cdot)$ encompasses two linear layers, parameterized with
$\bm{W_{r1}}\in\mathbb{R}^{(2M^{2}+2M)\times 128}$ and
$\bm{W_{r2}}\in\mathbb{R}^{128\times 2M}$, and a PReLU activation function. We
expect this additional residual block to learn to compensate for the ZF
equalization for a better transmission quality when faced with different
channel conditions and channel estimation errors, which will be detailed
later.

Note that we can also apply other MIMO equalization methods. However, we have
observed that a simple ZF estimation method, together with a learned
compensation, is sufficiently effective in the DeepJSCC-MIMO scheme. Moreover,
the ZF method converts the MIMO channels into equivalent parallel sub-
channels, which suits our approach of applying the self-attention mechanism
over different parallel sub-channels and source signals.

To perform MIMO equalization in a closed-loop MIMO system with CSIT, we
perform SVD decomposition and precoding as in Eqn. (9). At the receiver, we
perform equalization as in Eqn. (10) to get:

$\displaystyle\bm{{X^{\prime}}}$ $\displaystyle=\bm{\Sigma^{\dagger}U^{H}Y}$
(24) $\displaystyle=\bm{X}+\bm{\Sigma^{\dagger}U^{H}W},$ (25)

which converts the MIMO channel model into a set of parallel subchannels.

Given the equalized signal $\bm{X^{\prime}}$, which is reshaped into
$\bm{X^{{}^{\prime}}_{s}}\in\mathbb{R}^{l\times\frac{2Mk}{l}}$, and the noise
heatmap $\bm{M}$, a ViT-based decoder $f_{\bm{\phi}}$ is designed to recover
the source image as
$\bm{\hat{S}}=f_{\bm{\phi}}(\bm{X^{{}^{\prime}}_{s}},\bm{M})$. Our ViT-based
decoder consists of a Siamese layer, positional embedding layer, transformer
layer, and linear projection layer, elaborated upon below:

Algorithm 1 Algorithm of the DeepJSCC-MIMO scheme

For the transmitter:
Input: $\bm{S}\in\mathbb{R}^{h\times w\times 3}$,
$\bm{M}\in\mathbb{R}^{l\times\frac{2Mk}{l}}$ (MIMO with CSIT)
Output: $\bm{X}\in\mathbb{C}^{M\times k}$

1:$\bm{S_{s}}={Seq}(\bm{S})\in\mathbb{R}^{l\times c}$ $\triangleright$ Image
Sequentialization

2:if MIMO with CSIR then

3: $\bm{{S_{in}}}=\bm{S_{s}}$

4: $\bm{X}=f_{\theta}(\bm{S_{in}})\in\mathbb{C}^{M\times k}$ $\triangleright$
ViT encoding

5:else if MIMO with CSIT then

6: $\bm{{S_{in}}}={concat}(\bm{S_{s}},\bm{M})$

7: $\bm{X}=f_{\theta}(\bm{S_{in}})\in\mathbb{C}^{M\times k}$ $\triangleright$
ViT encoding

8: $\bm{{H}}=\bm{U\Sigma V^{H}}$ $\triangleright$ SVD decomposition

9: $\bm{X}=\bm{VX}$ $\triangleright$ Precoding

10:end if

Channel model:  $\bm{Y}=\bm{HX+W}$
For the receiver:
Input: $\bm{Y}\in\mathbb{C}^{M\times{k}}$, $\bm{H=U\Sigma V^{H}}$, $\bm{M}$
Output: $\bm{\hat{S}}\in\mathbb{R}^{h\times w\times 3}$

1:if MIMO with CSIR then

2: $\bm{{X^{\prime}}}=\bm{H_{w}Y}+f_{res}(\bm{Y,H})$ $\triangleright$ DL-aided
equalization

3:else if MIMO with CSIT then

4: $\bm{X^{\prime}}=\bm{\Sigma^{\dagger}U^{H}Y}$ $\triangleright$ SVD-based
equalization

5:end if

6:$\bm{X^{\prime}}\in\mathbb{C}^{M\times
k}\rightarrow\bm{X^{\prime}_{s}}\in\mathbb{R}^{l\times\frac{2MK}{l}}$
$\triangleright$ Reshape

7:$\bm{{\hat{S}}}=f_{\bm{\phi}}({concat}(\bm{X^{\prime}_{s}},\bm{M}),\bm{M})$
$\triangleright$ ViT decoding

Siamese layer and positional embedding: We design a weight-shared Siamese
layer, denoted as $\text{Siam}(\cdot)$, consisting of several linear
projection layers and GeLU activation functions. To form the input of the
Siamese layer $\bm{S_{d}}$, we concatenate $\bm{X^{{}^{\prime}}_{s}}$ and
$\bm{M}$ as:

$\bm{S_{d}}=\texttt{concat}(\bm{X^{{}^{\prime}}_{s}},\bm{M})\in\mathbb{R}^{l\times\frac{4Mk}{l}}.$
(26)

$\bm{S_{d}}$ multiplied by $1$ and $-1$ are fed into several linear projection
layers and GeLU functions, as illustrated in Fig. 3. In doing so, our networks
are tasked with handling the positive and the negative noise realizations
through two parallel branches, and subsequently, the resultant features are
aggregated by a linear layer to obtain the final output. We expect that these
GeLU functions and linear projection layers can learn to truncate excessive
noise realizations to bootstrap the performance, where similar designs can be
found in [31, 32]. To introduce the positional information in the decoding
process, we use the same positional embedding layer as in the Eqn. (15) to get
the output $\bm{D_{0}}\in\mathbb{R}^{l\times d}$:

$\bm{D_{0}}=\text{Siam}(\bm{S_{d}})+P_{e}(\bm{p}).$ (27)

Transformer layer: After the Siamese layer and positional embedding,
$\bm{D_{0}}$ is then passed through $L_{t}$ transformer layers, where

$\bm{D_{l}}=MSA(\bm{D_{l-1}})+MLP(MSA(\bm{D_{l-1}})),$ (28)

where $\bm{D_{i}}\in\mathbb{R}^{l\times d}$ is the output of the $i$-th
transformer layer at the decoder, and the $MSA$ and $MLP$ blocks share the
same structure as those in Eqn. (16).

Linear Projection: Given the output of the $L_{t}$-th transformer layer
$\bm{D_{L_{t}}}$, we apply a linear projection
$\bm{W_{out}}\in\mathbb{R}^{d\times c}$, and then reshape the output into a
matrix of size $\mathbb{R}^{h\times w\times 3}$ to reconstruct the input image
as:

$\bm{\hat{S}}=\texttt{reshape}(\bm{D_{L_{t}}}\bm{W_{out}}).$ (29)

(a) $R=1/24$

(b) $R=1/12$

Figure 5: Performance comparisons between the proposed DeepJSCC-MIMO and BPG-
Capacity over different SNR values and bandwidth ratios for the MIMO systems
with CSIR.

###  III-E Loss function

The encoder and decoder are optimized jointly to maximize the PSNR between
$\bm{S}$ and $\bm{\hat{S}}$. In specific, we adopt the mean square error as
the loss function:

$\mathcal{L}(\bm{\theta},\bm{\phi})=\text{MSE}(\bm{S},\bm{\hat{S}}),$ (30)

where we train the model to search for the optimal parameter pair
$(\bm{\theta}^{*},\bm{\phi}^{*})$ with a minimal
$\mathcal{L}(\bm{\theta},\bm{\phi})$ as:
$(\bm{\theta}^{*},\bm{\phi}^{*})=\mathop{\arg\min}_{\bm{\theta},\bm{\phi}}\mathbb{E}\big{[}\mathcal{L}(\bm{\theta},\bm{\phi})\big{]}$,
where the expectation is taken over the image and channel datasets.

##  IV Training and evaluation

This section conducts a set of numerical experiments to evaluate the
performance of our DeepJSCC-MIMO in various bandwidth and SNR scenarios.
Unless stated otherwise, we consider a $2\times 2$ MIMO system and the CIFAR10
dataset[36], which has $50000$ images with dimensions $3\times 32\times 32$
(color, height, width) as the training dataset and $10000$ images as the test
dataset. We consider both the practical separation-based codes and the
theoretical bounds (assuming capacity-achieving channel codes) as benchmarks.

###  IV-A Experimental setup

For the practical separation-based coding scheme, we use the image compression
algorithm BPG as the source coding method and low-density-parity check (LDPC)
codes for channel coding, denoted as the BPG-LDPC scheme. We consider the BPG
as the compression method due to its status as one of the most widely employed
benchmark codecs. Notably, it demonstrates competitive performance within our
evaluation dataset, particularly at low bit-per-pixel (bpp) values, when
compared to other DL-based compression methods that entail higher
complexities. For channel coding, we consider LDPC codes at rates
($1/2,2/3,3/4,5/6$) and various constellation sizes. In particular, we adopt
WiFi (IEEE 802.11n) LDPC code construction, featuring block lengths of 648,
1296, and 1944 bits, along with constellations of QPSK, 4QAM, 16QAM, and
64-QAM.

For the theoretical bound, we assume capacity-achieving channel codes for
transmission, denoted as the BPG-Capacity scheme. For the MIMO system with
CSIT, we apply the water-filling algorithm for power allocation among the
parallel channels formed after Eqn. 7.

Figure 6: Visual comparisons of images from the Kodak dataset transmitted over
the open-loop MIMO channel at SNR $5$dB and $12$dB with $R=1/24$. The first
and second columns are the original image and the original patches of the
bounding boxes.

Our model was implemented in Pytorch with two GTX 3090Ti GPUs. We use a
learning rate of $5e^{-5}$ and a batch size of $128$ with an Adam optimizer.
Models were trained until the performance on a validation set stopped
improving. Considering the model complexity and the performance, we set $p=8$,
$l=64$, $c=48$ for image vectorization, and $L_{t}=8$, $d=256$, and $N_{s}=8$
for each transformer layer of the ViT.

Each element of $\bm{H}$ is sampled from a complex Gaussian process as
$H[i,j]\sim\mathcal{CN}(0,1)$, where we set $\sigma_{h}^{2}=1$. To measure the
channel quality, we define the SNR value $\mu$ as:

$\mu\triangleq
10\log_{10}\frac{\mathbb{E}_{\bm{H,X}}[\bm{\|HX\|}^{2}_{F}]}{\mathbb{E}_{\bm{W}}[\bm{\|W\|}^{2}_{F}]}~{}(\text{dB})=10\log_{10}\frac{M}{\sigma_{w}^{2}}.$
(31)

###  IV-B Open-Loop MIMO system with CSIR

We first evaluate our DeepJSCC-MIMO scheme for the open-loop MIMO system with
CSIR. Specifically, we compare our DeepJSCC-MIMO, trained at specific SNRs,
with the BPG-Capacity and BPG-LDPC schemes in different SNR values and
bandwidth ratios.

For the BPG-Capacity benchmark here, we assume the transmitter knows the
capacity of each channel realization and can select the best compression rate
without power allocation as an upper bound. As to the BPG-LDPC scheme, we
apply the sphere decoding algorithm[37] to decode the signals

####  IV-B1 General performance

Fig. 5a and Fig. 5b show the performance of the DeepJSCC-MIMO, BPG-Capacity,
and BPG-LDPC schemes over the SNR values from $0$dB to $20$dB when the
bandwidth ratio is $R=1/24$ and $R=1/12$. We can observe that our DeepJSCC-
MIMO can generally outperform the BPG-Capacity and BPG-LDPC schemes in all SNR
values. We emphasize that the BPG-Capacity performance shown in these figures
is not achievable by practical channel coding schemes. Thus, it serves as an
upper bound on any separation-based scheme that employs BPG for compression of
the input image. We observe that DeepJSCC-MIMO provides significant
improvements (at least $1.3$dB and $5.05$dB) for $R=1/24$, compared with the
BPG-Capacity scheme and the BPG-LDPC scheme, respectively. For $R=1/12$, our
DeepJSCC-MIMO can still outperform the BPG-Capacity scheme (at least $0.93$dB)
and the BPG-LDPC scheme (at least $5.93$dB) in all SNRs. We observe that the
gap between DeepJSCC-MIMO and BPG capacity is larger in the lower bandwidth
ratio case.

TABLE I: Ablation study over DL-aided MIMO equalization methods for DeepJSCC-MIMO with $R=1/12$. Equalization methods | 1 dB | 5 dB | 10 dB | 15 dB | 19 dB
---|---|---|---|---|---
Pure ZF | 24.47 | 26.98 | 29.41 | 31.43 | 32.73
Pure MMSE | 24.61 | 27.04 | 29.48 | 31.50 | 32.74
DL-aided ZF | 24.95 | 27.56 | 30.03 | 31.81 | 32.78

####  IV-B2 SNR-adaptability

Note that each point on the DeepJSCC-MIMO curve is obtained by training a
separate encoder-decoder pair to be used exclusively for that training SNR. We
also consider a random SNR training strategy to evaluate the SNR adaptability
of the DeepJSCC-MIMO scheme, denoted as DeepJSCC-MIMO-universal. We train the
model with random SNR values uniformly sampled from $[0,22]$ dB and test the
well-trained model at different SNRs. The comparisons in Fig. 5 show that
there is a slight performance degradation compared to training a separate
DeepJSCC-MIMO encoder-decoder pair for each channel SNR (up to $0.99$dB for
$R=1/24$ and $2.02$dB for $R=1/12$); however, the DeepJSCC-MIMO-universal
brings significant advantage in terms of training complexity and storage
requirements. We can conclude that the DeepJSCC-MIMO-universal scheme learns
to adapt to different SNRs from the random training strategy at the expense of
a slight loss in PSNR, which tends to increase with SNR.

Figure 7: Performance comparison of DeepJSCC-MIMO and BPG-LDPC on the Kodak dataset for different bandwidth ratios at SNR=$5$dB. Ground Truth | Original Patch | $\sigma_{e}^{2}=0$ | $\sigma_{e}^{2}=0.1$ | $\sigma_{e}^{2}=0.2$ | $\sigma_{e}^{2}=1$
---|---|---|---|---|---
|  |  |  |  |
$512\times 768$ | PSNR (dB) | $29.44$ | $28.22$ | $27.78$ | $26.84$
|  |  |  |  |
$512\times 768$ | PSNR (dB) | $26.46$ | $25.44$ | $25.14$ | $24.35$
Figure 8: Visual comparisons of images reconstructed by DeepJSCC-MIMO-
universal in the presence of channel estimation errors, where the model is
trained on the ImageNet dataset and validated on the Kodak dataset at
SNR=$12$dB and $R=1/24$. The first and second columns are the source image and
the image patch from the red bounding box of the original image. The third to
sixth columns are the reconstructions of the same patches by DeepJSCC-MIMO
universal with different channel estimation errors, respectively. The average
PSNR metrics of the entire image are provided at the bottom of each visualized
patch.

####  IV-B3 Ablation experiments over the equalization methods

To evaluate the effectiveness of the DL-aided equalization method with the
residual block, we repeat the experiments in Fig. 5b with an MMSE receiver.
The experimental results in Table I show that the residual block improves the
performance (up to $0.62$dB). More interestingly, we observe that the ZF
combined with the residual block achieves the best performance, which shows
that the residual block helps to compensate for the limitation of ZF and
outperforms MMSE.

####  IV-B4 Different bandwidth ratios and datasets

To further evaluate the effect of the bandwidth ratio and data
generalizability for the DeepJSCC-MIMO, we evaluate the DeepJSCC-MIMO and BPG-
LDPC schemes with various bandwidth ratios and SNRs over the Kodak dataset. We
train the models with randomly cropped $128\times 128$ patches from the
ImageNet dataset, and evaluate the trained models on the Kodak dataset.

In Fig. 6, we present examples of images transmitted by DeepJSCC-MIMO scheme
and BPG-LDPC at $5$dB and $12$dB and bandwidth ratio $R=1/24$, where SNR
values and PSNR of the entire image are provided at the bottom of each
visualized patch. It can be observed that reconstructed by DeepJSCC-MIMO are
qualitatively better with more detailed high-frequency features. The
comparisons here can clearly show the advantages of the DeepJSCC-MIMO in high-
resolution datasets. We also conclude that training on a sufficiently large
dataset (ImageNet) can allow our DeepJSCC-MIMO scheme to perform well on a
never-seen dataset (Kodak) for a wide range of bandwidth ratios.

We also evaluate the model at various bandwidth ratios when SNR=$5$dB, as
shown in Fig. 7. We want to emphasize that the sphere decoding algorithm used
in BPG-LDPC is a much better MIMO detection method, compared with ZF employed
by DeepJSCC-MIMO. As shown in Fig. 7, DeepJSCC-MIMO outperforms the BPG-LDPC
in all bandwidth ratios, where the gap is even larger at low bandwidth ratios
(at least $2.95$ dB), which shows that DeepJSCC-MIMO outperforms BPG-LDPC
despite using a weaker detection algorithm.

Figure 9: Comparisons of DeepJSCC-MIMO-universal and BPG-LDPC schemes with
channel estimation errors in Kodak dataset

(a) $R=1/24$

(b) $R=1/12$

(c) $R=1/6$

Figure 10: Performance of the proposed DeepJSCC-MIMO model compared with the
BPG-Capacity benchmark at different channel SNR and bandwidth ratio scenarios
for the closed-loop MIMO.

####  IV-B5 Robustness to channel estimation errors

We take the same approach as in [29] to evaluate the robustness of MIMO
systems to channel estimation errors, where $\bm{H}$ is imperfectly estimated
as $\bm{\hat{H}}$ at the receiver. Similar to the work in[38, 39, 29], let
$\bm{E}_{n}\triangleq\bm{H}-\bm{\hat{H}}$, where the entries of $\bm{E}_{n}$
are zero-mean complex Gaussian with variance $\sigma^{2}_{e}$. Here
$\sigma^{2}_{e}$ is the ‘noise variance’ term to capture the channel
estimation quality, which can be appropriately selected depending on the
channel conditions and channel estimation method.

We evaluate the DeepJSCC-MIMO-universal and BPG-LDPC schemes with different
$\sigma^{2}_{e}\in\\{0,0.2,1\\}$ values over
$\text{SNR}_{\text{test}}\in[1,12]$dB on the Kodak dataset with $R=1/24$. For
the DeepJSCC-MIMO-universal scheme, we train and test the model both with
imperfect $\bm{\hat{H}}$. The experimental results are shown in Fig. 9. As
expected, for both schemes, the performance degrades as the channel estimation
error increases. In specific, for $\sigma^{2}_{e}=0.2$, the performance loss
of our DeepJSCC-MIMO scheme is up to $0.3$dB, compared with the performance
loss up to $5.9$dB from the BPG-LDPC scheme. For $\sigma^{2}_{e}=1$, the BPG-
LDPC scheme starts to fail to work. However, our DeepJSCC-MIMO scheme can
still work with a performance loss of less than $0.85$dB. To visualize the
performance of DeepJSCC-MIMO-universal in the presence of channel estimation
errors, we present the reconstructed images with different channel estimation
errors in Fig. 8. We can observe that the increasing channel estimation error
is reflected in the gradual loss of high-frequency details during
transmission. We want to emphasize that even with a significant channel
estimation error ($\sigma_{e}^{2}=1$), DeepJSCC-MIMO continues to function,
with only limited performance loss, compared with the significant performance
loss with the traditional separation benchmark. We can conclude that the
proposed scheme is more robust to channel estimation errors, as it can learn
to compensate for it during transmission.

###  IV-C Closed-Loop MIMO system with CSIT

In this section, we consider the closed-loop MIMO system with CSI at the
transmitter as well as at the receiver. We evaluate the DeepJSCC-MIMO trained
at specific SNRs with $\textit{SNR}_{\textit{test}}\in[0,20]$dB in three
bandwidth ratios $R=1/24$, $R=1/12$ and $R=1/6$.

####  IV-C1 General performance

As before, we first evaluate the DeepJSCC-MIMO model under the setup where the
training SNR, denoted by $\textit{SNR}_{\textit{train}}$ matches the test SNR,
$\textit{SNR}_{\textit{test}}$. In particular, we set
$\textit{SNR}_{\textit{train}}=\textit{SNR}_{\textit{test}}\in\\{1,5,10,15,19\\}$dB
and the bandwidth ratio $R\in\\{1/24,1/12,1/6\\}$.

The comparison between DeepJSCC-MIMO and the BPG-Capacity benchmark is shown
in Fig. 10. As can be seen, DeepJSCC-MIMO outperforms the separation-based
benchmark in all SNR and bandwidth-ratio scenarios. Specifically, from Fig.
10a and 10b, we can see that DeepJSCC-MIMO significantly outperforms the
benchmark by at least $1.98$dB and $1.78$dB when $R=1/24$ and $R=1/12$,
respectively. When $R=1/6$, improvements of up to $3.5$dB can be observed.
These significant improvements demonstrate the superiority of the DeepJSCC-
MIMO scheme and its capacity in extracting and mapping image features to the
available channels in an adaptive fashion.

Figure 11: Performance of DeepJSCC-MIMO and BPG-LDPC scheme with respect to
the bandwidth ratio at SNR=$10$dB. Figure 12: Visual comparison of different
schemes under various SNRs on the Kodak dataset with $R=1/24$. The first and
second columns are the original image, the original patch of the red bounding
box. The third, fourth, and fifth columns are the performance of the BPG-
Capacity, DeepJSCC-MIMO-universal, and DeepJSCC-MIMO schemes, respectively.

####  IV-C2 Different bandwidth ratios

To further evaluate the impact of the bandwidth ratio $R$ on the system
performance, we compare the DeepJSCC-MIMO and BPG-Capacity schemes over a wide
range of bandwidth ratios in Fig. 11, where the test SNR is set to $10$dB. We
observe that DeepJSCC-MIMO can outperform the BPG-Capacity benchmark at all
bandwidth ratios for $10$dB, with a gain of up to $2.2$dB in PSNR,
respectively. The gap between the two is more significant for low $R$ values.
We would like to emphasize that the BPG-Capacity performance shown in these
figures serves as an upper bound for separate source and channel coding
schemes employing BPG compression, and the real gap can be even larger,
especially in the short block length regime, i.e., low $R$. Thus, the results
here illustrate the clear superiority of DeepJSCC-MIMO for the CIFAR dataset
compared to separation-based alternatives.

####  IV-C3 SNR-adaptability

We also consider a random SNR training strategy to evaluate the channel
adaptability of the DeepJSCC-MIMO scheme. We train the DeepJSCC-MIMO model
with random SNR values uniformly sampled from $[0,22]$ dB and test the well-
trained model at different test SNRs, denoted by DeepJSCC-MIMO-universal. The
comparisons in Fig. 10 show that there is a slight performance degradation
compared to training a separate DeepJSCC-MIMO encoder-decoder pair for each
channel SNR (up to $0.65$dB, $0.83$dB, and $2.1$dB for $R=1/24$, $R=1/12$ and
$R=1/6$); however, the DeepJSCC-MIMO-universal brings significant advantage in
terms of training complexity and storage requirements.

We can observe more performance loss in the high SNR regime. However,
DeepJSCC-MIMO-universal still significantly outperforms the BPG-Capacity
benchmark. We conclude that the DeepJSCC-MIMO-universal scheme can adapt to
SNR variations with a slight loss in PSNR, especially in the high-SNR regime.

####  IV-C4 SVD ablation study

To evaluate the effectiveness of the SVD strategy in our scheme, we train the
models of DeepJSCC-MIMO-universal without SVD decomposition-based precoding in
different bandwidth ratios, denoted by DeepJSCC-MIMO-universal w/o SVD.
Specifically, we directly feed the channel heatmap $\bm{M}$ into both
transceivers and the performance is shown in Fig. 10. Compared with the
DeepJSCC-MIMO-universal, we can observe that the performance of DeepJSCC-MIMO-
universal w/o SVD has a gap of $0.97$dB, $1.6$dB and $2.9$dB for $R=1/24$,
$R=1/12$, and $R=1/6$, respectively. Although the networks can still learn to
communicate over the MIMO channel, the SVD-based model-driven strategy can
significantly simplify the training process and improve the performance, which
illustrates the importance of exploiting domain knowledge in the design of
data-driven communication technologies.

Figure 13: Performance of DeepJSCC-MIMO and BPG-LDPC scheme in the CelebA
dataset when $R=1/48$.

####  IV-C5 High resolution dataset

To evaluate the model generalizability, we validate the DeepJSCC-MIMO on Kodak
and CelebA datasets, where we train the DeepJSCC-MIMO and DeepJSCC-MIMO-
universal schemes with randomly cropped $128\times 128$ patches from the
ImageNet dataset, and evaluate the well-trained models on the Kodak and CelebA
[40] datasets. The detailed visualization and the PSNR performance of the
transmission, compared with BPG-Capacity, are shown in Fig. 12 and Fig. 13,
respectively.

Similarly to the CIFAR dataset, from Fig. 13, we can observe that our
DeepJSCC-MIMO can outperform BPG-Capacity at all test SNRs and generally
provide higher gains at lower SNRs (up to $3.03$dB) in the CelebA dataset.
There is a slight performance degradation of the DeepJSCC-MIMO-universal
scheme (up to $0.61$dB) compared with DeepJSCC-MIMO trained at specific SNRs;
however, the DeepJSCC-MIMO-universal can still outperform the BPG-Capacity
scheme at all SNRs.

We also visualize the results of the DeepJSCC-MIMO scheme on the Kadak dataset
and compare them with those of the BPG-Capacity scheme in Fig. 12. We can
observe that DeepJSCC-MIMO performs better in the low SNR regime ($1$ and $5$
dB), with more detailed high-frequency features. Comparing the performance of
DeepJSCC-MIMO-universal and DeepJSCC-MIMO schemes, we can observe that the
DeepJSCC-MIMO-universal can achieve comparable performance at the expense of a
slight drop in distortion.

Considering that the BPG-Capacity performance shown in these figures is an
upper bound on the performance of any separation-based scheme that employs the
BPG codec, the comparison here clearly shows the benefits of the DeepJSCC-MIMO
and DeepJSCC-MIMO-universal schemes in high-resolution datasets. We can also
conclude that a sufficiently large dataset (ImageNet) can allow our DeepJSCC-
MIMO scheme to perform well on other never-seen datasets (Kodak/CelebA) for a
wide range of channel conditions (SNRs).

Figure 14: Visualization of the channel heatmap and power allocation across
the channel symbols, where antenna number is set as $M=32$ and $R=1/6$

####  IV-C6 Visualization of the channel heatmap and power allocation

We visualize the channel heatmap and the power allocation of the channel input
matrix for a closed-loop MIMO system in Fig. 14. Within our channel heatmap,
we visualize the noise power experienced by each channel symbols, with the
color red indicating higher noise power, and hence, a poorer channel
condition. In the power allocation map, the color red denotes a higher
allocation of power to the respective channel symbols. We can observe that for
each SNR value, DeepJSCC-MIMO generally assigns more power to antennas with
better channel conditions. Specifically, in the low SNR regime, DeepJSCC-MIMO
concentrates most of the power on a few select antennas with exceptional
quality. In certain extreme cases, DeepJSCC-MIMO refrains from allocating
power to the sub-channels characterized by the worst channel conditions. As
the SNR increases, we note that DeepJSCC-MIMO demonstrates a tendency to
distribute power more evenly across a greater number of antennas, instead of
concentrating on a limited set of antennas with the worst channel conditions.

####  IV-C7 Adaptability to the antenna numbers

In all the above experiments, we have the same number of $M=2$ antennas at the
transmitter and receiver. In this subsection, we investigate the
generalizability of the DeepJSCC-MIMO architecture to different antenna
numbers. In particular, we repeat the experiment of DeepJSCC-MIMO in Fig. 10b
for different $M$ values. The performance of our DeepJSCC-MIMO models with
different antenna numbers $M\in\\{2,3,4\\}$ and channel SNRs
$SNR_{test}\in\\{1,5,10\\}$dB is shown in Fig. 15. We can observe that
DeepJSCC-MIMO can generally achieve better performance with the increase in
the number of antennas, thanks to its more flexible power allocation.

Figure 15: Performance of DeepJSCC-MIMO with different antennas number.

We also consider a random antenna number training strategy with a maximal
antenna number setting $M_{max}$, denoted as adaptive-M. We train the
DeepJSCC-MIMO model with the $M$ uniformly sampled from
$\\{2,\ldots,M_{max}\\}$ and test the well-trained model with different $M$
values. Specifically, our encoder generates the channel symbols
$\bm{X}\in\mathbb{C}^{M_{max}\times k}$, and only transmits first
$\mathbb{C}^{M_{i}\times k}$ symbols for the $M=M_{i}$ antenna scenario, while
the receiver pads the received signal into $\mathbb{C}^{M_{max}\times k}$ with
zeros for decoding. From experiments with $M_{max}=4$, we observe in Fig. 15
that the adaptive-M strategy can achieve similar performance compared with the
model trained for specific $M$ values: the performance loss is at most
$0.6$dB. We can conclude that DeepJSCC-MIMO is highly flexible and adaptive
that can be applied to different number of antennas without retraining.

##  V Conclusion

We introduced the first DeepJSCC-enabled MIMO communication system designed
for wireless image transmission. The proposed DeepJSCC-MIMO model utilizes a
vision transformer-based architecture to exploit contextual semantic features
and channel conditions through a self-attention mechanism. The results
demonstrate significant improvements in transmission quality across a wide
range of SNRs and bandwidth scenarios, with enhanced robustness to channel
estimation errors. Moreover, DeepJSCC-MIMO is a unified and adaptable design
capable of accommodating varying channel conditions and different antenna
numbers without the need for retraining.

Moving forward, there are several potential directions for future exploration
based on the findings of this paper:

  * •

JSCC with space-time coding. To harness the full potential of the spatial and
temporal dimensions in MIMO, the integration of DeepJSCC with space-time codes
presents a promising avenue for the open-loop scenario[41]. This synergy can
allow for even greater flexibility and performance enhancements, especially in
the low-SNR regime.

  * •

Variable length JSCC. Building upon the adaptability of ViT-based models to
varying antenna numbers, future investigations can extend this concept to
accommodate variable code lengths or even variable number of antennas, where
the channel resources are judiciously exploited depending on the channel state
as well as the input signal.

## References

  * [1] M. Jankowski, D. Gündüz, and K. Mikolajczyk, “Wireless image retrieval at the edge,” _IEEE Journal on Selected Areas in Communications_ , vol. 39, no. 1, pp. 89–100, 2020.
  * [2] J. Shao, Y. Mao, and J. Zhang, “Learning task-oriented communication for edge inference: An information bottleneck approach,” _IEEE Journal on Selected Areas in Communications_ , vol. 40, no. 1, pp. 197–211, 2021.
  * [3] W. Shi, S. Zhou, Z. Niu, M. Jiang, and L. Geng, “Multiuser co-inference with batch processing capable edge server,” _IEEE Transactions on Wireless Communications_ , vol. 22, no. 1, pp. 286–300, 2023.
  * [4] E. Bourtsoulatze, D. B. Kurka, and D. Gündüz, “Deep joint source-channel coding for wireless image transmission,” _IEEE Transactions on Cognitive Communications and Networking_ , vol. 5, no. 3, pp. 567–579, 2019.
  * [5] M. Yang, C. Bian, and H.-S. Kim, “OFDM-guided deep joint source channel coding for wireless multipath fading channels,” _IEEE Transactions on Cognitive Communications and Networking_ , vol. 8, no. 2, pp. 584–599, 2022.
  * [6] J. Xu, B. Ai, W. Chen, A. Yang, P. Sun, and M. Rodrigues, “Wireless image transmission using deep source channel coding with attention modules,” _IEEE Transactions on Circuits and Systems for Video Technology_ , 2021.
  * [7] J. Dai, S. Wang, K. Tan, Z. Si, X. Qin, K. Niu, and P. Zhang, “Nonlinear transform source-channel coding for semantic communications,” _IEEE Journal on Selected Areas in Communications_ , vol. 40, no. 8, pp. 2300–2316, 2022\.
  * [8] D. B. Kurka and D. Gündüz, “Bandwidth-agile image transmission with deep joint source-channel coding,” _IEEE Transactions on Wireless Communications_ , vol. 20, no. 12, pp. 8081–8095, 2021.
  * [9] M. Yang and H.-S. Kim, “Deep joint source-channel coding for wireless image transmission with adaptive rate control,” in _ICASSP 2022 - 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2022, pp. 5193–5197.
  * [10] Q. Hu, G. Zhang, Z. Qin, Y. Cai, G. Yu, and G. Y. Li, “Robust semantic communications with masked VQ-VAE enabled codebook,” _IEEE Transactions on Wireless Communications_ , 2023.
  * [11] H. Yoo, L. Dai, S. Kim, and C.-B. Chae, “On the role of vit and cnn in semantic communications: Analysis and prototype validation,” _IEEE Access_ , pp. 1–1, 2023.
  * [12] Y. Shao and D. Gunduz, “Semantic communications with discrete-time analog transmission: A PAPR perspective,” _IEEE Wireless Communications Letters_ , vol. 12, no. 3, pp. 510–514, 2023.
  * [13] E. Erdemir, T.-Y. Tung, P. L. Dragotti, and D. Gündüz, “Generative joint source-channel coding for semantic image transmission,” _IEEE Journal on Selected Areas in Communications_ , pp. 1–1, 2023.
  * [14] H. Wu, Y. Shao, E. Ozfatura, K. Mikolajczyk, and D. Gündüz, “Transformer-aided wireless image transmission with channel feedback,” _arXiv preprint arXiv:2306.09101_ , 2023.
  * [15] C. E. Shannon, “A mathematical theory of communication,” _The Bell system technical journal_ , vol. 27, no. 3, pp. 379–423, 1948.
  * [16] H. Xie, Z. Qin, X. Tao, and K. B. Letaief, “Task-oriented multi-user semantic communications,” _IEEE Journal on Selected Areas in Communications_ , vol. 40, no. 9, pp. 2584–2597, 2022.
  * [17] W. F. Lo, N. Mital, H. Wu, and D. Gündüz, “Collaborative semantic communication at the edge,” _IEEE Wireless Communications Letters_ , 2023\.
  * [18] H. Wu, Y. Shao, K. Mikolajczyk, and D. Gündüz, “Channel-adaptive wireless image transmission with OFDM,” _IEEE Wireless Communications Letters_ , vol. 11, no. 11, pp. 2400–2404, 2022.
  * [19] N. Samuel, T. Diskin, and A. Wiesel, “Learning to detect,” _IEEE Transactions on Signal Processing_ , vol. 67, no. 10, pp. 2554–2564, 2019.
  * [20] H. He, C.-K. Wen, S. Jin, and G. Y. Li, “Deep learning-based channel estimation for beamspace mmwave massive MIMO systems,” _IEEE Wireless Communications Letters_ , vol. 7, no. 5, pp. 852–855, 2018.
  * [21] M. B. Mashhadi and D. Gündüz, “Pruning the pilots: Deep learning-based pilot design and channel estimation for MIMO-OFDM systems,” _IEEE Transactions on Wireless Communications_ , vol. 20, no. 10, pp. 6315–6328, 2021.
  * [22] H. Wu, Y. Shao, C. Bian, K. Mikolajczyk, and D. Gündüz, “Vision transformer for adaptive image transmission over MIMO channels,” in _2023 IEEE International Conference on Communications (ICC)_ , 2023.
  * [23] D. Gunduz and E. Erkip, “Joint source–channel codes for MIMO block-fading channels,” _IEEE Transactions on Information Theory_ , vol. 54, no. 1, pp. 116–134, 2008.
  * [24] T. J. O’Shea, T. Erpek, and T. C. Clancy, “Deep learning based MIMO communications,” _arXiv preprint arXiv:1707.07980_ , 2017.
  * [25] J. Song, C. Häger, J. Schröder, T. J. O’Shea, E. Agrell, and H. Wymeersch, “Benchmarking and interpreting end-to-end learning of MIMO and multi-user communication,” _IEEE Transactions on Wireless Communications_ , vol. 21, no. 9, pp. 7287–7298, 2022.
  * [26] X. Zhang, M. Vaezi, and T. J. O’Shea, “SVD-embedded deep autoencoder for MIMO communications,” in _ICC 2022-IEEE International Conference on Communications_. IEEE, 2022, pp. 5190–5195.
  * [27] Y. Shi, S. Shao, Y. Wu, W. Zhang, X.-G. Xia, and C. Xiao, “Excess distortion exponent analysis for semantic-aware MIMO communication systems,” _arXiv preprint arXiv:2301.04357_ , 2023.
  * [28] S. Yao, S. Wang, J. Dai, K. Niu, and P. Zhang, “Versatile semantic coded transmission over MIMO fading channels,” _arXiv preprint arXiv:2210.16741_ , 2022.
  * [29] T. Yoo and A. Goldsmith, “Capacity and power allocation for fading MIMO channels with channel estimation error,” _IEEE Transactions on Information Theory_ , vol. 52, no. 5, pp. 2203–2214, 2006.
  * [30] H. He, C.-K. Wen, S. Jin, and G. Y. Li, “Model-driven deep learning for MIMO detection,” _IEEE Transactions on Signal Processing_ , vol. 68, pp. 1702–1715, 2020.
  * [31] Y. Shao, E. Ozfatura, A. Perotti, B. Popovic, and D. Gündüz, “Attentioncode: Ultra-reliable feedback codes for short-packet communications,” _IEEE Transactions on Communications_ , pp. 1–1, 2023.
  * [32] E. Ozfatura, Y. Shao, A. G. Perotti, B. M. Popović, and D. Gündüz, “All you need is feedback: Communication with block attention feedback codes,” _IEEE Journal on Selected Areas in Information Theory_ , vol. 3, no. 3, pp. 587–602, 2022.
  * [33] S. Yang and L. Hanzo, “Fifty years of MIMO detection: The road to large-scale MIMOs,” _IEEE Communications Surveys & Tutorials_, vol. 17, no. 4, pp. 1941–1988, 2015.
  * [34] S. Zheng, J. Lu, H. Zhao, X. Zhu, Z. Luo, Y. Wang, Y. Fu, J. Feng, T. Xiang, P. H. Torr _et al._ , “Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers,” in _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2021, pp. 6881–6890.
  * [35] X. Chu, Z. Tian, Y. Wang, B. Zhang, H. Ren, X. Wei, H. Xia, and C. Shen, “Twins: Revisiting the design of spatial attention in vision transformers,” _Advances in Neural Information Processing Systems_ , vol. 34, pp. 9355–9366, 2021.
  * [36] A. Krizhevsky, G. Hinton _et al._ , “Learning multiple layers of features from tiny images,” University of Toronto, Tech. Rep., 2009.
  * [37] Z. Guo and P. Nilsson, “Algorithm and implementation of the k-best sphere decoding for MIMO detection,” _IEEE Journal on selected areas in communications_ , vol. 24, no. 3, pp. 491–503, 2006.
  * [38] B. Hassibi and B. M. Hochwald, “How much training is needed in multiple-antenna wireless links?” _IEEE Transactions on Information Theory_ , vol. 49, no. 4, pp. 951–963, 2003.
  * [39] R. Etkin and D. Tse, “Degrees of freedom in some underspread MIMO fading channels,” _IEEE Transactions on Information Theory_ , vol. 52, no. 4, pp. 1576–1608, 2006.
  * [40] Z. Liu, P. Luo, X. Wang, and X. Tang, “Large-scale celebfaces attributes (celeba) dataset,” _Retrieved August_ , vol. 15, no. 2018, p. 11, 2018.
  * [41] C. Bian, Y. Shao, H. Wu, and D. Gunduz, “Space-time design for deep joint source channel coding of images over MIMO channels,” in _2023 IEEE International Workshop on Signal Processing Advances in Wireless Communications (SPAWC)_ , 2023.