# MambaJSCC: Deep Joint Source-Channel Coding with Visual State Space Model

Tong Wu∗, Zhiyong Chen∗, Meixia Tao∗, Xiaodong Xu†, Wenjun Zhang∗, and Ping
Zhang†
∗Cooperative Medianet Innovation Center, Shanghai Jiao Tong University,
Shanghai, China
${\dagger}$Beijing University of Posts and Telecommunications, Beijing, China
Email: {wu_tong, zhiyongchen, mxtao<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Lightweight and efficient deep joint source-channel coding (JSCC) is a key
technology for semantic communications. In this paper, we design a novel JSCC
scheme named MambaJSCC, which utilizes a visual state space model with channel
adaptation (VSSM-CA) block as its backbone for transmitting images over
wireless channels. The VSSM-CA block utilizes VSSM to integrate two-
dimensional images with the state space, enabling feature extraction and
encoding processes to operate with linear complexity. It also incorporates
channel state information (CSI) via a newly proposed CSI embedding method.
This method deploys a shared CSI encoding module within both the encoder and
decoder to encode and inject the CSI into each VSSM-CA block, improving the
adaptability of a single model to varying channel conditions. Experimental
results show that MambaJSCC not only outperforms Swin Transformer based JSCC
(SwinJSCC) but also significantly reduces parameter size, computational
overhead, and inference delay (ID). For example, with employing an equal
number of the VSSM-CA blocks and the Swin Transformer blocks, MambaJSCC
achieves a 0.48 dB gain in peak-signal-to-noise ratio (PSNR) over SwinJSCC
while requiring only 53.3% multiply-accumulate operations, 53.8% of the
parameters, and 44.9% of ID.

##  I Introduction

Deep learning-based joint source-channel coding (JSCC), which combines source
and channel coding into a single process using neural networks to optimize
end-to-end transmission over wireless channels, has recently gained
significant attention as a means to increase the efficiency of information
transmission. With the aid of deep learning, JSCC can effectively extract and
utilize the semantic features of data, which are critical for understanding
the meaning of information based on specific task requirements. Therefore,
JSCC is emerging as a key technology for enabling semantic communications [1,
2], a pivotal element in the development of sixth-generation (6G) wireless
networks[3, 4, 5].

Most of the recent research on JSCC focuses on finding suitable schemes to
maximize end-to-end performance. For text transmission, [6] proposes a JSCC
framework with a joint iterative decoder to enhance text transmission quality.
For video transmission, [7] develops a deep video semantic transmission system
based on JSCC to maximize overall performance. For image transmission,
DeepJSCC, proposed in [8], is based on convolutional neural networks (CNNs)
and outperforms traditional separation-based schemes, e.g., JPEG2000 and
capacity achieving channel code. To further enhance the performance of JSCC,
SwinJSCC is proposed in [9], which replaces CNNs with Swin Transformers[10].

Designing deep learning-based JSCC entails considerations beyond end-to-end
performance, and it also necessitates evaluating the complexity of artificial
intelligence (AI) algorithms and their adaptability to the dynamic changes of
wireless channels. The performance of DeepJSCC is considered unsatisfactory
due to the limited capacity of CNNs. In contrast, SwinJSCC achieves better
performance but incurs higher computational overhead and has a larger model
parameter size. To overcome performance degradation caused by dynamic channel
conditions, attention modules are introduced into the JSCC architecture in
[11, 9], enabling it to incorporate channel state information (CSI) for
adapting to changes in the channel signal-to-noise ratio (SNR). However, these
attention modules introduce a significant number of additional parameters and
computational overhead, greatly increasing the complexity of a single model.

To this end, this paper proposes a novel lightweight and efficient deep JSCC
scheme, named MambaJSCC, using the visual state space model (VSSM) [12] and
CSI embedding for image transmission over wireless channels. Recently, Mamba
is proposed in [13] based on state space models (SSMs), where the state space
is adjusted to a selective structure state space, enabling the models to focus
on relevant information in an input-dependent manner. Mamba achieves
remarkable performance in natural language processing tasks with linear
complexity. Motivated by Mamba, we design a VSSM with channel adaptation
(VSSM-CA) block as the backbone for MambaJSCC. Specifically, VSSM-CA
integrates 2D images with state space to capture their global information.
Meanwhile, a novel CSI embedding method is proposed for VSSM-CA, which injects
CSI into the VSSM-CA block in a low-complexity manner, allowing it to adapt to
changes in wireless channels. Leveraging the powerful VSSM-CA backbone, we
design MambaJSCC, which features a hierarchical structure with varying numbers
of patches of different sizes to effectively encode and decode images. We also
develop a 2D patch division module in MambaJSCC that reverses the process of
patch merging in the encoder to reconstruct the source image. Experimental
results demonstrate that the proposed MambaJSCC outperforms SwinJSCC, and it
has significantly lower parameter size, computational overhead, and inference
latency (ID) compared to SwinJSCC.

Figure 1: (a) The overall architecture of the proposed MambaJSCC. (b) The
structure of the VSSM-CA block.

##  II System Model of MambaJSCC

###  II-A State Space Model

State space models are a class of sequential neural network models inspired by
linear time-invariant (LTI) systems and are widely related to CNNs and
recurrent neural networks (RNNs). Generally, LTI systems map a one-dimension
function or sequence $x(t)\in\mathbb{R}$ to $y(t)\in\mathbb{R}$ through a
latent state space ${h}(t)\in\mathbb{R}^{N}$. Here, $N$ is the dimension of
the latent state space. The process can be formulated using linear ordinary
differential equations (ODEs) as follows:

$\displaystyle{h^{\prime}}(t)=\mathbf{A}h(t)+\mathbf{B}x(t),\leavevmode\nobreak\
\leavevmode\nobreak\ y(t)=\mathbf{C}h(t)+\mathbf{D}x(t),$ (1)

where $\mathbf{A}\in\mathbb{R}^{N\times N}$, $\mathbf{B}\in\mathbb{R}^{N\times
1}$, $\mathbf{C}\in\mathbb{R}^{1\times N}$, and
$\mathbf{D}\in\mathbb{R}^{1\times 1}$ are the state matrix, input matrix,
output matrix, and feed-forward matrix, respectively.

The OEDs represent continuous-time systems that require a discretization
process to be integrated into deep learning frameworks. A widely-adopted
discretization method is the zero-order hold (ZOH) rule, which utilizes a
timescale parameter $\mathbf{\Delta}\in\mathbb{R}^{N\times N}$ to discretize
the continuous parameters $\mathbf{A}$ and $\mathbf{B}$ into their discrete
parameters $\mathbf{\bar{A}}$ and $\mathbf{\bar{B}}$. The process can be
formulated as follows:

$\displaystyle\mathbf{\bar{A}}=\exp(\mathbf{\Delta}\mathbf{A}),\leavevmode\nobreak\
\leavevmode\nobreak\
\mathbf{\bar{B}}=(\mathbf{\Delta}\mathbf{A})^{-1}(\exp(\Delta\mathbf{A})-\mathbf{I})\mathbf{\Delta}\mathbf{B}.$
(2)

Then, (1) can be rewritten in discrete-time form as follows:

$\displaystyle
h_{t}=\mathbf{\bar{A}}h_{t-1}+\mathbf{\bar{B}}x_{t},\leavevmode\nobreak\
\leavevmode\nobreak\ y_{t}=\mathbf{C}h_{t}+\mathbf{D}x_{t}.$ (3)

###  II-B An Overview of the Proposed MambaJSCC

An overall architecture of MambaJSCC is illustrated in Fig. 1(a). In the
transmitter, the input image $\mathbf{s}\in\mathbb{R}^{3\times H\times W}$ is
encoded into the channel input signal $\mathbf{tr}\in\mathbb{C}^{c\times
h\times w}$ by the MambaJSCC encoder, where $H$ and $W$ represent the height
and width of the source image, respectively. This encoding process involves
one CSI encoding module, $l_{e}$ coding stages, and a convolution compression
module. The CSI encoding module encodes the CSI into a CSI vector
$\mathbf{u_{e}}\in\mathbb{R}^{m}$, where $m$ is the vector length. The vector
is then fed into all the VSSM-CA blocks to adapt to channel conditions.

In the first coding stage, the input image $\mathbf{s}$ is initially fed into
a 2D patch embedding module, which outputs $c_{1}^{e}$ patches, each of size
$h_{1}\times w_{1}$. Subsequently, $n_{1}^{e}$ VSSM-CA blocks, which
incorporate the CSI vector $\mathbf{u_{e}}$, are employed to process these
patches. These blocks adaptively extract and encode the features from the
patches into new patches based on the CSI, with the number and size of the
patches remaining unchanged. The output patches are then fed into Stage 2,
where the first module is the patch merging module. This module merges the
$c_{1}^{e}$ patches of size $h_{1}\times w_{1}$ into a greater number of
patches $c_{2}^{e}$, each of a smaller size $h_{2}\times w_{2}$, aiming to
capture more refined information. Following this, $n_{2}^{e}$ VSSM-CA blocks,
with the additional input of the CSI vector $\mathbf{u_{e}}$, are employed on
these patches. This stage is repeated ($l_{e}-1$) times in the MambaJSCC
encoder as Stage $k$, $k=2,3,...,l_{e}$. Each involves the patch merging
module outputting $c_{k}^{e}$ patches of size $h_{k}\times w_{k}$ and
$n_{k}^{e}$ VSSM-CA blocks processing these patches. At the end of the
encoder, the output patches are fed into the convolution compression module,
which uses a CNN to compress the patches to the designated number of channel
uses. After that, the compressed patches are normalized for power and
transmitted through the wireless channel.

In this paper, we consider an end-to-end transmission through the Additive
White Gaussian Noise (AWGN) channel or the Rayleigh fading channel. At the
receiver, the MambaJSCC decoder first applies an equalizer, such as the
minimum mean square error (MMSE) equalizer, to equalize the received signal
$\mathbf{r}$ in the case of the Rayleigh fading channel. It then decodes the
signal into $\mathbf{\hat{s}}\in\mathbb{R}^{3\times H\times W}$ for
reconstruction. The MambaJSCC decoder includes a CSI encoding module, $l_{d}$
decoding stages, and a convolution expanding module. Specifically, following
the equalizer, the convolution expanding module first expands the equalized
signal into $c^{d}_{l_{d}}$ 2D patches, each with size $h_{l_{d}}\times
w_{l_{d}}$, to match the output patches from Stage $l_{e}$ in the encoder.
These patches are then fed into $l_{d}$ concatenated stages, denoted as Stage
$k,k=l_{d},l_{d}-1,...,1$. For $k\geq 2$, each stage initially decodes the
patches using $n_{k}^{d}$ VSSM-CA blocks along with the CSI vector
$\mathbf{u_{d}}$, which is the output from the CSI encoding module in the
decoder. The output patches are then divided into fewer but larger patches
$c_{k-1}^{d}$ with size $h_{k-1}\times w_{k-1}$. This step is performed by a
2D patch division module that first applies layer normalization to the input
patches, followed by a full-connection (FC) layer that expands the number of
patches from $c^{d}_{k}$ to
$c^{d}_{k-1}\times\frac{h_{k-1}w_{k-1}}{h_{k}w_{k}}$. The purpose of expanding
the number of patches is to organize every $\frac{h_{k-1}w_{k-1}}{h_{k}w_{k}}$
patches into $c_{k-1}^{d}$ groups. Within each group, a pixel shuffle
operation is then applied to cyclically select elements, thereby creating a
new patch that is with $\frac{h_{k-1}w_{k-1}}{h_{k}w_{k}}$ times larger size.
As a result, the module outputs $c^{d}_{k-1}$ new patches, each with a larger
size of $h_{k-1}\times w_{k-1}$. In Stage 1, the 2D patch division module
divides the output into three patches, each with size $H\times W$ for
reconstruction.

The MambaJSCC encoder and decoder are jointly trained with the goal of
minimizing the distortion between the input image and the reconstructed image.
Therefore, the loss function $L$ can be derived as follows:

$\displaystyle L=d(\mathbf{s},\mathbf{\hat{}s}),$ (4)

where $d(\cdot)$ represents the distortion function.

##  III Structure of VSSM-CA Block

###  III-A The VSSM Module

Figure 2: The computational flow of V-S6 module.

The overall architecture of the VSSM-CA block is shown in Fig.1(b). The block
consists of two main components: the VSSM module and the CSI embedding module.
For the VSSM module, it first normalizes the input patches using the layer-
norm method and then feeds them into two branches. In the first branch,
depicted in Fig.1(b), an FC layer, a convolution layer and a SiLU activation
function process the patches into higher-dimensional patches. Each of the
$i$-th patch is denoted as $\mathbf{z}_{i}\in\mathbb{R}^{h_{z}\times
w_{z}},i=1,2,...,c_{z}$, where $c_{z}$ is the number of patches. These $c_{z}$
patches are then fed into $c_{z}$ V-S6 modules.

The computational flow of the V-S6 module is described in Fig. 2. Each vector
$\mathbf{z}_{i}$ is flattened into four 1D vectors
$\mathbf{v}_{i}^{j}\in\mathbb{R}^{D\times 1},j=1,2,3,4$ as follows:

$\displaystyle\mathbf{v}_{i}^{1}=\ vec(\mathbf{z}_{i}),$
$\displaystyle\mathbf{v}_{i}^{2}=\ vec(\mathbf{z}_{i}^{T}),$
$\displaystyle\mathbf{v}_{i}^{3}=\ R(vec(\mathbf{z}_{i})),\leavevmode\nobreak\
\leavevmode\nobreak\ $ $\displaystyle\mathbf{v}_{i}^{4}=\
R(vec(\mathbf{z}_{i}^{T})),$ (5)

where we use $D=h_{z}\times w_{z}$. The operation $vec(\cdot)$ denotes the
vectorization of a matrix. The operation $R(\cdot)$ is the reversal of the
order of elements in a vector or a matrix.

The four vectors are fed into four SSM modules, enabling the subsequent V-S6
module to acquire global information. For the $j$-th vector
$\mathbf{v}_{i}^{j}$, the $j$-th SSM module utilizes seven learnable matrices
$\mathbf{G}^{j}\in\mathbb{R}^{1\times N}$,
$\bm{\delta}^{j}\in\mathbb{R}^{N\times N}$ $\mathbf{H}_{1}^{j}$,
$\mathbf{H}_{2}^{j}$, $\mathbf{H}_{3}^{j}\in\mathbb{R}^{N\times D}$,
$\mathbf{A}^{j}\in\mathbb{R}^{N\times N}$ and
$\mathbf{D}^{j}\in\mathbb{R}^{1\times 1}$ to model the state space
$\mathbf{h}_{t}^{j}\in\mathbb{R}^{N\times 1}$. The module first computes the
matrices $\mathbf{\Delta}^{j}$, $\mathbf{B}^{j}$ and $\mathbf{C}^{j}$ as
follows:

$\mathbf{\Delta}^{j}=\mathbf{H}_{1}^{j}\mathbf{v}_{i}^{j}\mathbf{G}^{j}+\bm{\delta}^{j},\
\mathbf{B}^{j}=\mathbf{H}_{2}^{j}\mathbf{v}_{i}^{j},\
\mathbf{C}^{j}=(\mathbf{H}_{3}^{j}\mathbf{v}_{i}^{j})^{T},$ (6)

and then following the ZOH rule and the first-order Taylor series, the
matrices $\mathbf{\bar{A}}^{j}$ and $\mathbf{\bar{B}}^{j}$ are computed as
follows:

$\mathbf{\bar{A}}^{j}=\exp(\mathbf{\Delta}^{j}\mathbf{A}^{j}),\
\mathbf{\bar{B}}^{j}=\mathbf{\Delta}^{j}\mathbf{B}^{j}.$ (7)

Subsequently, the module iteratively computes the output for each element. For
the $t$-th element of $\mathbf{v}_{i}^{j}$, the outputs are:

$\displaystyle\mathbf{h}_{t}^{j}=\mathbf{\bar{A}}^{j}\mathbf{h}_{t-1}^{j}+\mathbf{\bar{B}}^{j}{v}_{i,t}^{j},\leavevmode\nobreak\
\leavevmode\nobreak\
{y}_{i,t}^{j}=\mathbf{C}^{j}\mathbf{h}_{t}^{j}+\mathbf{D}^{j}{v}_{i,t}^{j}.$
(8)

After the iterations, all the $y_{i,t}^{j}$ in this module are combined into a
2D matric $\mathbf{Y}_{i}^{j}\in\mathbb{R}^{h_{z}\times w_{z}}$. Then, we
merge the outputs as follows:

$\displaystyle\mathbf{Y}_{i}=\mathbf{Y}_{i}^{1}+(\mathbf{Y}_{i}^{2})^{T}+R(\mathbf{Y}_{i}^{3}+(\mathbf{Y}_{i}^{4})^{T}).$
(9)

Therefore, the $c_{z}$ V-S6 modules output $c_{z}$ matrices. The algorithm of
V-S6 module is summarized in Algorithm 1.

In parallel, the normalized input patches are processed by an FC layer and a
SiLU activation function in the second branch. The output of this branch is
multiplied in element-wise with the output of the first branch, and then the
result is projected back to the original dimension using another FC layer. A
residual connection adds this result to the original input, yielding the final
output.

Input: Input matrix $\mathbf{z}_{i}$;
Output: Output matrix $\mathbf{Y}_{i}$;

Algorithm 1 The computational flow of the $i$-th V-S6 module

1: Expand $\mathbf{z_{i}}$ to
$\mathbf{v}_{i}^{1},\mathbf{v}_{i}^{2},\mathbf{v}_{i}^{3},\mathbf{v}_{i}^{4}$
according to (III-A);

2: for $j=1$ to 4 do

3:
$\mathbf{\Delta}^{j}=\mathbf{H}_{1}^{j}\mathbf{v}_{i}^{j}\mathbf{G}^{j}+\bm{\delta}^{j},\
\mathbf{B}^{j}=\mathbf{H}_{2}^{j}\mathbf{v}_{i}^{j},\
\mathbf{C}^{j}=(\mathbf{H}_{3}^{j}\mathbf{v}_{i}^{j})^{T}$;

4: $\mathbf{\bar{A}}^{j}=\exp(\mathbf{\Delta}^{j}\mathbf{A}^{j}),\
\mathbf{\bar{B}}^{j}=\mathbf{\Delta}^{j}\mathbf{B}^{j}$;

5: $\mathbf{h}_{0}^{j}=\mathbf{0}$;

6: for $t=1$ to $N$ do

7:
$\mathbf{h}_{t}^{j}=\mathbf{\bar{A}}^{j}\mathbf{h}_{t-1}^{j}+\mathbf{\bar{B}}^{j}{v}_{i,t}^{j}$;

8:
${y}_{i,t}^{j}=\mathbf{C}^{j}\mathbf{h}_{t}^{j}+\mathbf{D}^{j}{v}_{i,t}^{j}$;

9: end for

10: Forming matrix $\mathbf{Y}_{i}^{j}$ using all $y_{i,t}^{j}$;

11: end for

12:
$\mathbf{Y}_{i}=\mathbf{Y}_{i}^{1}+(\mathbf{Y}_{i}^{2})^{T}+R(\mathbf{Y}_{i}^{3}+(\mathbf{Y}_{i}^{4})^{T})$.

###  III-B The CSI Embedding Method

The CSI embedding method, including CSI encoding modules and CSI embedding
modules, is responsible for embedding CSI into the VSSM module, enhancing the
single model’s ability to handle various channel conditions.

In the MambaJSCC, the CSI encoding module includes sine position coding, an FC
layer, a Swish activation function, and another FC layer. The module is
responsible for coding the CSI into a vector and injecting it into the VSSM-CA
blocks. All VSSM-CA blocks in the encoder share a single CSI encoding module,
and similarly, a distinct CSI encoding module is shared by all VSSM-CA blocks
within the decoder. Injecting the same CSI vector into each VSSM-CA block
ensures that each stage of the model obtains accurate SNR information, thereby
avoiding information loss that could occur with increased depth. This approach
also eliminates the need to propagate SNR information to deeper layers,
allowing the model to effectively encode the source information based on SNR,
which conserves model capacity. Encoding the CSI into a higher-dimensional
vector enhances the model’s ability to utilize it.

The CSI embedding module within the VSSM-CA block takes the CSI vector as
input and expands its length to match the number of patches using an FC layer.
As a result, each element of the expanded CSI vector can be added to each
patch. This addition operation is effective because it is analogous to
concatenating two features and processing them with a shared model. The design
also reduces delay and model size. The FC layer also allows different VSSM-CA
blocks to flexibly utilize the CSI.

##  IV EXPERIMENTS RESULTS

###  IV-A Experiment Setup

In the simulation, we use the DIV2K dataset, which contains 800 varied 2K
resolution images from a wide range of real-world scenes for training, along
with an additional 100 images for testing. During raining, images are randomly
cropped into $256\times 256$ patches.

The performance of the proposed MambaJSCC is evaluated using peak-signal-to-
noise ratio (PSNR). Computational complexity is assessed by the multiply-
accumulate operations (MACs), which count the number of multiplications and
additions required for inference. We also measure the inference delay for a
comprehensive evaluation. The model size is quantified by the number of
parameters. Both MACs and parameters size are calculated by the Pytorch-
OpCounter library, and the ID is measured on the NVIDIA A40 GPU.

We conduct a comparative analysis between MambaJSCC and SwinJSCC [9], with
main comparison schemes as follows: 1) MambaJSCC w/o CA: MambaJSCC without the
CSI encoding module and CSI embedding module; 2) SwinJSCC w/o SA&RA: SwinJSCC
without the Channel ModNet and Rate ModNet; 3) MambaJSCC with CSI embedding:
MambaJSCC with CSI embedding; 4) SwinJSCC w/ SA: SwinJSCC with the Channel
ModNet, which is an attention-based module for channel adaptation; 5)
MambaJSCC with Channel ModNet: this scheme uses the Channel ModNet instead of
CSI embedding in MambaJSCC.

Figure 3: Visual comparison between MambaJSCC w/o CA and SwinJSCC w/o SA&RA
with $N_{m}=N_{s}=6$ under the AWGN and the Rayleigh fading channel with
SNR=$5$ dB.

For the basic configurations of MambaJSCC, we adopt four stages in both the
encoder and decoder, with
$[n_{1}^{e},n_{2}^{e},n_{3}^{e},n_{4}^{e}]=[n_{1}^{d},n_{2}^{d},n_{3}^{d},n_{4}^{d}]=[2,2,6,2]$
and
$[c_{1}^{e},c_{2}^{e},c_{3}^{e},c_{4}^{e}]=[c_{1}^{d},c_{2}^{d},c_{3}^{d},c_{4}^{d}]=[128,192,256,320]$.
The CSI vector length is set to 128. To vary the model scales, we fix the
block number in Stage 3 at $n_{3}^{e}=n_{3}^{d}$ and adjust their values
$N_{m}$. Similarly, for SwinJSCC w/o SA&RA and SwinJSCC w/ SA, we use the same
method to modify the model size. The number of the Swin Transformer blocks in
Stage 3 of the encoder and decoder for both models is denoted as $N_{s}$. The
channel bandwidth ratio is set to $\frac{3}{128}$. We train all models using
the Adam optimizer over 600 epochs on an NVIDIA A40 GPU, with an initial
learning rate of $10^{-4}$ and a batch size of 4. The loss function is the
mean square error (MSE) loss.

###  IV-B Results of MambaJSCC w/o CA

Fig. 3 visualizes the reconstruction results of MambaJSCC w/o CA and SwinJSCC
w/o SA&RA with $N_{m}=N_{s}=6$ over different channels. We can observe that in
this example, the visual reconstruction quality of MambaJSCC w/o CA is
significantly better than that of SwinJSCC w/o SA&RA. For example, the edges
and colors of the flower in the images reconstructed by MambaJSCC w/o CA more
closely resemble the original image than those in images reconstructed by
SwinJSCC w/o SA&RA.

Fig. 4 displays the PSNR performance of MambaJSCC w/o CA and SwinJSCC w/o
SA&RA versus SNR under both the AWGN and the Rayleigh fading channels. The
legend provides the model size and computational complexity, represented by
parameters and MACs, respectively. We can see that MambaJSCC w/o CA requires
$15.30$ billion fewer MACs and $8.47$ million fewer parameters compared with
SwinJSCC w/o SA&RA. Meanwhile, we can see from Table I that the ID of
MambaJSCC w/o CA is 39.12 ms shorter than that of SwinJSCC w/o SA&RA. Even
with significant reductions in both MACs and parameters, and a lower ID, the
proposed MambaJSCC w/o CA still outperforms SwinJSCC w/o SA&RA in terms of
PSNR. For example, on the Rayleigh fading channel, Mamba JSCC w/o CA achieves
a $0.53$ dB gain at an SNR of $5$ dB. For the AWGN channel, Mamba JSCC w/o CA
achieves a $0.97$ dB gain at an SNR of $1$ dB.

Table I: MACs, ID, Parameters and PSNR at SNR=10 dB of models with different model sizes and computational complexity. Models | MACs (G) | ID (ms) | Parameters (M) | PSNR, the AWGN channel (dB) | PSNR, the Rayleigh fading channel (dB)
---|---|---|---|---|---
MambaJSCC | $N_{m}=6$ | 17.44339 | 31.95 | 9.88 | 27.85 | 26.63
w/o CA | $N_{m}=12$ | 22.36960 | 44.59 | 14.68 | 27.98 | 26.64
$N_{m}=18$ | 27.29581 | 56.65 | 19.48 | 28.21 | 26.79
SwinJSCC | $N_{s}=2$ | 26.27390 | 54.74 | 12.03 | 27.23 | 26.13
w/o SA&RA | $N_{s}=6$ | 32.74564 | 71.07 | 18.35 | 27.74 | 26.15

Furthermore, we investigate the potential performance gains of MambaJSCC w/o
CA compared to the SwinJSCC w/o SA&RA, considering similar model sizes and
computational complexities. We adjust $N_{m}$ and $N_{s}$, with the results
shown in Table I. We first fix $N_{s}=6$ for SwinJSCC w/o SA&RA, then increase
$N_{m}$ for MambaJSCC w/o CA to $12$ and $18$ to expand the model size and
computational complexity. The results show that when $N_{m}$ is increased to
$18$, MambaJSCC w/o CA achieves a similar model size and computational
complexity as SwinJSCC w/o SA&RA, with slightly more parameters and fewer
MACs. Under these conditions, the performance gain increases from $0.11$ dB to
$0.47$ dB at an SNR of $10$ dB under the AWGN channel, and from $0.48$ dB to
$0.64$ dB at an SNR of $10$ dB under the Rayleigh fading channel. Moreover, we
also reduce $N_{s}$ for SwinJSCC w/o SA&RA to $2$ to decrease its model size
and computational complexity. It can be observed that SwinJSCC w/o SA&RA with
$N_{s}=2$ has a similar model size and computational complexity to MambaJSCC
w/o CA with $N_{m}=12$, but with slightly more MACs and fewer parameters. In
this configuration, the performance gain increases from $0.11$ dB to $0.75$ dB
under the AWGN channel at an SNR of $10$ dB, and from $0.48$ dB to $0.51$ dB
under the Rayleigh fading channel at the same SNR. Experimental results
demonstrate that MambaJSCC w/o CA model can achieve significant performance
gains, even with similar model size and computational complexity.

###  IV-C Results of MambaJSCC with CSI embedding

Fig. 5 illustrates the performance of SwinJSCC w/ SA and MambaJSCC with
different channel adaptation methods versus SNR under the AWGN and the
Rayleigh fading channel. The $N_{m}$ is set to $12$ for the AWGN channel and
$18$ for the Rayleigh fading channel. $N_{s}$ is set to $6$. We can see that
MambaJSCC with CSI embedding significantly outperforms SwinJSCC w/ SA across
all SNRs in the both channels. For example, in the AWGN channel, a maximum
gain of $0.43$ dB is achieved at SNR=$1$ dB, and the minimum gain is $0.19$ dB
at SNR=$20$ dB. For the Rayleigh fading channel, the performance improvements
at different SNRs are consistent, with an average gain of $0.65$ dB. Moreover,
the parameter size and the computational complexity of MambaJSCC with CSI
embedding are significantly lower than those of SwinJSCC w/ SA. As shown in
Table II, compared to SwinJSCC w/ SA with $N_{s}=6$, MambaJSCC with CSI
embedding requires $12.32$ billion fewer MACs, $11.70$ million fewer
parameters for $N_{m}=12$, and $7.39$ billion fewer MACs and $6.90$ million
fewer parameters for $N_{m}=18$. Furthermore, ID is reduced by $29.48$ ms and
$17.30$ ms for $N_{m}=12$ and $N_{m}=18$ respectively.

Furthermore, we compare MambaJSCC with different channel adaptation modules.
It can be observed that under the AWGN channel, MambaJSCC with CSI embedding
and MambaJSCC w/o CA with Channel ModNet outperform MambaJSCC w/o CA,
achieving similar performance. For example, MambaJSCC with CSI embedding
achieves a $0.09$ dB higher PSNR at SNR=$1$ dB but performs $0.08$ dB lower as
the SNR increases to $20$ dB compared with MambaJSCC w/o CA with Channel
ModNet. For the Rayleigh fading channel, MambaJSCC with CSI embedding
outperforms MambaJSCC w/o CA with Channel ModNet. The maximum gain is $0.28$
dB at SNR=$20$ dB, and at SNR=$5$ dB, the gain is $0.16$ dB.

Figure 4: The PSNR performance of MambaJSCC w/o CA and SwinJSCC w/o SA&RA
versus SNR under the AWGN and the Rayleigh fading channels. $N_{m}$ and
$N_{s}$ are set to $6$.

Although the performance of the proposed CSI embedding is comparable to that
of the Channel ModNet in MambaJSCC, the advantage of the CSI embedding method
lies in its significantly low model size and computational complexity. As
shown in Table II, the CSI embedding method requires only
$0.290\leavevmode\nobreak\ (0.41)$ million of extra MACs,
$1.00\leavevmode\nobreak\ (1.12)$ ms for additional ID, and
$0.297\leavevmode\nobreak\ (0.297)$ million extra parameters at
$N_{m}=12\leavevmode\nobreak\ (18)$. In contrast, the Channel ModNet requires
$872.13$ million extra MACs and $9.86$ million extra parameters for both
configurations. These results indicate that the CSI embedding method
significantly reduces the parameter size and computational complexity compared
to the Channel ModNet, while maintaining comparable or superior performance.
This efficiency makes the CSI embedding a promising method for channel
adaptation in JSCC.

Table II: MACs, ID and Parameters of different models with diverse channel adaptation methods. Models |  MACs(G) |  ID(ms) |  Parameters(M)
---|---|---|---
MambaJSCC |  w/o CA |  22.36960 |  44.59 |  14.68
|  CSI embedding |  22.36989 |  45.59 |  14.97
$N_{m}=12$ |  w/o CA with |  23.24113 |  49.79 |  24.54
|  Channel ModNet |  |  |
MambaJSCC |  w/o CA |  27.29581 |  56.65 |  19.48
|  CSI embedding |  27.29622 |  57.77 |  19.77
$N_{m}=18$ |  w/o CA with |  28.16734 |  61.24 |  29.34
|  Channel ModNet |  |  |
SwinJSCC w/ SA, $N_{s}=6$ |  34.68512 |  75.07 |  26.67
Figure 5: The PSNR performance with different channel adaptation methods.

##  V CONCLUSION

In this paper, we propose MambaJSCC, a novel lightweight and efficient JSCC
architecture for wireless image transmission. MambaJSCC utilizes VSSM for
effective representation learning and employs a CSI embedding method,
integrating the CSI vector from CSI encoding module into each VSSM-CA block to
effectively handle various channel conditions. Experimental results confirm
that the proposed MambaJSCC achieves remarkable performance across a range of
channel conditions. Moreover, MambaJSCC requires fewer parameters, MACs and
exhibits a low ID, effectively overcoming the challenges of high computational
complexity and large model size associated with Transformer. Therefore,
MambaJSCC presents a promising solution for semantic communication systems
that demand real-time responsiveness, high energy efficiency, and operate
within storage constraints.

## References

  * [1] D. Gündüz, Z. Qin, I. E. Aguerri, H. S. Dhillon, Z. Yang, A. Yener, K. K. Wong, and C.-B. Chae, “Beyond Transmitting Bits: Context, Semantics, and Task-Oriented Communications,” _IEEE Journal on Selected Areas in Communications_ , vol. 41, no. 1, pp. 5–41, 2023.
  * [2] Z. Qin, F. Gao, B. Lin, X. Tao, G. Liu, and C. Pan, “A Generalized Semantic Communication System: From Sources to Channels,” _IEEE Wireless Communications_ , vol. 30, no. 3, pp. 18–26, 2023.
  * [3] J. Liu, S. Shao, W. Zhang, and H. V. Poor, “An Indirect Rate-Distortion Characterization for Semantic Sources: General Model and the Case of Gaussian Observation,” _IEEE Transactions on Communications_ , vol. 70, no. 9, pp. 5946–5959, 2022.
  * [4] J. Xu, B. Ai, N. Wang, and W. Chen, “Deep Joint Source-Channel Coding for CSI Feedback: An End-to-End Approach,” _IEEE Journal on Selected Areas in Communications_ , vol. 41, no. 1, pp. 260–273, 2023.
  * [5] T. Wu, Z. Chen, D. He, L. Qian, Y. Xu, M. Tao, and W. Zhang, “CDDM: Channel Denoising Diffusion Models for Wireless Semantic Communications,” _IEEE Transactions on Wireless Communications_ , pp. 1–1, 2024.
  * [6] S. Yao, K. Niu, S. Wang, and J. Dai, “Semantic Coding for Text Transmission: An Iterative Design,” _IEEE Transactions on Cognitive Communications and Networking_ , vol. 8, no. 4, pp. 1594–1603, 2022.
  * [7] S. Wang, J. Dai, Z. Liang, K. Niu, Z. Si, C. Dong, X. Qin, and P. Zhang, “Wireless Deep Video Semantic Transmission,” _IEEE Journal on Selected Areas in Communications_ , vol. 41, no. 1, pp. 214–229, 2023.
  * [8] E. Bourtsoulatze, D. Burth 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.
  * [9] K. Yang, S. Wang, J. Dai, X. Qin, K. Niu, and P. Zhang, “SwinJSCC: Taming Swin Transformer for Deep Joint Source-Channel Coding,” _arXiv preprint arXiv:2308.09361_ , 2023.
  * [10] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo, “Swin Transformer: Hierarchical Vision Transformer using Shifted Windows,” in _Proc. IEEE/CVF ICCV_ , 2021, pp. 9992–10 002.
  * [11] 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_ , vol. 32, no. 4, pp. 2315–2328, 2022.
  * [12] Y. Liu, Y. Tian, Y. Zhao, H. Yu, L. Xie, Y. Wang, Q. Ye, and Y. Liu, “Vmamba: Visual state space model,” _arXiv preprint arXiv:2401.10166_ , 2024.
  * [13] A. Gu and T. Dao, “Mamba: Linear-time sequence modeling with selective state spaces,” _arXiv preprint arXiv:2312.00752_ , 2023.