# Investigating the Effectiveness of Cross-Attention to Unlock Zero-Shot
Editing of Text-to-Video Diffusion Models

Saman Motamed1 Wouter Van Gansbeke1 Luc Van Gool1,2,3
1 INSAIT, Sofia University 2 ETH Zurich 3 KU Leuven

###### Abstract

With recent advances in image and video diffusion models for content creation,
a plethora of techniques have been proposed for customizing their generated
content. In particular, manipulating the cross-attention layers of Text-to-
Image (T2I) diffusion models has shown great promise in controlling the shape
and location of objects in the scene. Transferring image-editing techniques to
the video domain, however, is extremely challenging as object motion and
temporal consistency are difficult to capture accurately. In this work, we
take a first look at the role of cross-attention in Text-to-Video (T2V)
diffusion models for zero-shot video editing. While one-shot models have shown
potential in controlling motion and camera movement, we demonstrate zero-shot
control over object shape, position and movement in T2V models. We show that
despite the limitations of current T2V models, cross-attention guidance can be
a promising approach for editing videos. Code: https://github.com/sam-
motamed/Video-Editing-X-Attention.git

##  1 Introduction

Figure 1: This figure shows an overview of backward guidance in T2V models. On
the left, we show the generated frames of the T2V model after t steps, given
an initial input latent $z_{t}$ and the text prompt “A burger floats on the
water”. To edit the video and move the burger from the top-left of the screen
to the bottom-left in a straight line, we generate
$\mathcal{A}_{tar_{\text{fi}}}$ for each frame fi reflecting this edit.
Following the scheme in Section 4.2, we update the latent through the
denoising process based on objective $E$. At time step 0, $z_{0}$ generates
the video on the right which reflects the intended edit.

Text-to-Video diffusion models [18, 3, 2, 41, 46] have been fast advancing in
generating temporally consistent scenes with plausible object interactions.
There has been a series of works that have focused on editing T2V models to
enable greater control over video generation. More successful editing methods
have made use of small sets of reference videos to learn an object’s motion or
camera movement. They subsequently transfer that specific movement or camera
motion to a new object and scene [56, 21, 52] by training parts of the video
diffusion model or performing Low Rank Adaptation (LoRA) [19]. While these
methods can be effective, they require additional data and compute with
limited flexibility, which limits their adoption in practice.

Several works [45, 40, 47, 49] have shown the promise of attention maps in
object discovery and segmentation. In the domain of text-to-image models,
cross-attention and its role in controlling the scene layout has also been
well studied. In particular, cross-attention is responsible for determining
the objects’ shape and size in the image. Cross-attention facilitates
maintaining semantic consistency between the text and the generated image. By
attending to relevant textual features, the model ensures that the generated
visual content aligns with the overall semantics of the input description. One
of the works that exploited cross-attentions to enable editing images was
Prompt-to-Prompt [16]. This work showed that the shape of an object a can be
replaced with the shape of another object b by replacing a’s cross-attentions
with those of b. Training-Free Layout Control [7] was another work that
proposed an energy-based objective to control the position of objects in the
generated image. Given a user-specified bounding box, the energy function
encourages the cross-attention maps of a token to form within the bounding box
and hence position the object within the bounding box. Diffusion self-guidance
[12] generalized the Training-Free Layout Control method such that editing the
scene could be done by using the cross-attention maps alone, without the need
for external inputs (e.g., bounding box), in a zero-shot manner. This was
achieved by applying transformations (e.g., relocating and resizing) to the
original cross-attentions of a token and using the resulting cross-attentions
as the target of the objective.

With the success of the above methods for editing images generated by T2I
models [5, 16, 12, 30] or adapting T2I models for video editing [27], we ask
the question; “Do such approaches for editing images transfer to the video
domain?”. In particular, we are interested in exploring the effectiveness of
cross-attention layers for editing the subject’s size, positioning, and motion
in videos.

In this paper, we build upon the achievements of prior image-editing
techniques by extending them to the video domain. More specifically, our
contributions are threefold:

  * •

We take a first look at cross-attention layers in T2V diffusion models and
their role in editing videos.

  * •

We explore two possible ways to use cross-attentions in editing videos; namely
forward and backward guidance.

  * •

We investigate the limitations of current T2V models that hinder the
capabilities of video editing methods.

##  2 Related Works

#### Denoising Diffusion Models.

The denoising diffusion paradigm [42, 1, 43] emerged as a new method to
generate images with high photo-realism and diversity. It has rapidly advanced
text-conditioned image generation [36, 31, 55, 39, 54, 14, 35, 34, 17, 11],
which is important for gaining control over its generated content. Due to
their versatility and representation learning capabilities [20, 53, 8], they
have also been successfully adapted for specialized tasks such as
classification [10, 26], depth prediction [23] and segmentation [44, 6, 22].

#### Personalizing Image Generation.

Personalizing [55, 29, 38, 13, 50, 25] and editing [5, 16, 12, 30, 32] T2I
models has become a research focus to enable user-intuitive control for
creating content with these generative models. In particular, the cross-
attention layers of diffusion models have been studied for their role in
determining a scene’s layout and their ability to enable zero-shot editing of
generated images. Similar to [5], we split cross-attention-based editing of
T2I and T2V models into two categories of 1) forward and 2) backward guidance.

In forward guidance, cross-attention manipulation occurs directly during the
denoising process via a forward pass through the model. A notable example of
forward guidance is Prompt-to-Prompt [16], which proposes replacing the
token’s cross-attentions from a source prompt with those of a target prompt.
Figure 2 shows one such example in the video domain where the cross-attentions
of “car”, from the source prompt “car drives on the road”, are replaced with
cross-attentions of “truck”, from the target prompt “truck drives on the
road”. To enable more precise modifications to a specific source token, while
preserving the overall scene, forward guidance requires source and target
prompts that differ by a single token, limiting its applicability.

In contrast to forward guidance that directly manipulates cross-attentions,
backward guidance biases the cross-attention through backpropagation. By
designing an energy-based loss that encourages some desired edit [5, 12], the
gradient of the loss is then used to update the input latent $z_{t}$ of the
model. Training-Free Layout Control [5] is an example of backward guidance
where the energy function encourages the cross-attentions of the user-
specified token to obtain higher values inside a user-defined bounding box. At
multiple time steps, the input latent is updated to realize this objective.
Similarly, Diffusion self-guidance [12] designed energy functions that
encourage the cross-attentions to take certain shapes or positions within the
image. This paper is inspired by the success of these two works in the image
domain. In Section 3.3, we show that forward guidance is too restrictive to
enable effective video editing. In Section 3.4, we show backward guidance’s
promise in enabling zero-shot editing of T2V models.

#### Text-to-Video Generation.

Diffusion models have been improving at high-quality video generation by
training conditional denoising networks (e.g, 3D U-Net [9], DiT [33]) to
denoise randomly sampled sequences of Gaussian noises [18, 3, 2, 41, 46, 28].
Some works take advantage of large, pre-trained text-to-image foundation
models to build text-to-video models. This is done by inflating the T2I model
with temporal layers, like Tune-A-Video [51], Text2Video-Zero [24] and
AnimateDiff [15].

#### Personalizing Video Generation.

Following the same desire to control image generation, a few works focused on
video editing and customizing the motion and camera movement in T2V models
[56, 21, 52, 48, 7]. Most current editing and customization methods work by
tuning parts of the network or performing LoRA [19] based on example videos
containing the desired effect. Such methods lack the flexibility of a zero-
shot approach and require additional training data and resources. For this
reason, we investigate the effectiveness of forward and backward guidance
using cross-attention for T2V models.

##  3 Method

###  3.1 How Do Video Diffusion Models Work?

Video diffusion models train a 3D denoising network, traditionally U-Nets but
more recently transformer-based [33] networks, to generate videos from
randomly sampled Gaussian noise. In this work, we use T2V models with 3D U-Net
backbone [46] which consists of down-blocks, middle-blocks, and up-blocks.
Each block has several convolution layers, spatial transformers, and temporal
transformers. During training on videos, the U-Net ($\epsilon_{\theta}$) and a
text encoder ($\tau_{\theta}$) are optimized with the following objective:

$\mathcal{L}=\mathbb{E}_{z_{0},y,\epsilon\sim\mathcal{N}(0,I),t\sim\mathcal{U}(0,T)}=\lVert\epsilon-\epsilon_{\theta}(\mathbf{z}_{t},t,\tau_{\theta}(y))\rVert^{2}_{2},$
(1)

where $z_{0}\in\mathbb{R}^{f\times b\times h\times w\times c}$ is the initial
latent input of the training videos ($b$ indicates the batch size, $f$ is the
number of frames, $h$, $w$ and $c$ are the height and width and channels
respectively) and $y$ is the text description of the video, with $\epsilon$
and $t$ being the added Gaussian noise to the videos and the time step. At
time step $t$, the noised latent is defined as:

$z_{t}=\sqrt{\bar{\alpha}_{t}}z_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon,$ (2)

where ${\alpha}_{t}$ controls the noise strength.

###  3.2 T2V Cross-attention

The cross-attention mechanism in the spatial transformers of the 3D U-Net
enables the model to capture spatial relationships between the video frames
and the input text. In this work, we focus on changing an object’s size,
location and motion given a latent input and text prompt to the T2V model. To
this end, we work with the cross-attention layers of the 3D U-Net where
$\\{\mathcal{A}_{i,t,.,.,k}\in\mathbb{R}^{H_{i}\times W_{i}\times|k|}$} is the
Softmax-normalized cross-attention map of the $i^{th}$ layer of the U-Net, at
time step $t$ for token $k$.

###  3.3 Forward T2V Guidance

Following the works that perform forward guidance in T2I models [16, 32, 4],
we implemented forward guidance in the T2V pipeline. Figure 2 is one example
where the cross-attentions of “car” are replaced with the cross-attentions of
“truck”. Below are the two main limitations with forward guidance that have
also been observed in the T2I domain.

  * •

Size and Shape Mismatch. Forward guidance is restrictive and can lead to
artifacts due to the difference in shape and size of the two objects. In the
example of Figure 2, since the truck is larger than the car, injecting the
cross-attentions of the truck to replace the car’s has led to artifacts around
the car without changing the car’s size to match the truck’s.

  * •

Cross-attention Overlap. The cross-attentions of different tokens can overlap.
We refer to the top row of Figure 3, where the shark is still visible in the
cross-attention maps of tokens “in” and “the”. For this reason, forward
guidance can work reasonably well where the two source and target sentences
only differ by one token (i.e., Prompt-to-Prompt’s setting). This overlap can
cause degradation in the image and video quality, especially when the text
inputs differ by more than one token.

We note that some of these artifacts are due to the current T2V models
generating noisy cross-attentions. We go over more details in Section 4.1
regarding this limitation.

Figure 2: We show an example of forward guidance by swapping the cross-
attention maps of “car” with cross-attention maps of the “truck”. The two
input texts only differ in one token (“truck” and “car”). While the car
follows the motion and location of the truck in the video, artifacts can be
seen around the car due to the mismatch in size and shape of the truck and
car. Figure 3: We compare the cross-attention maps for the same prompt to a
T2I and T2V model. The cross-attention maps are extracted and averaged at the
$16\times 16$ resolution from the mid-blocks and up-blocks of the U-Net. Open-
source T2I models currently produce much less noisy cross-attention maps
compared to T2V models. In Section 4.1, we give details on how the noisy
cross-attentions hinder backward guidance and propose a procedure for
bypassing this limitation for our experiments in this paper.

###  3.4 Backward T2V Guidance

Following Diffusion self-guidance [12] and Training-Free Layout Control [5],
we define an energy function $E$ to encourage specific shape, size and motion
properties on the cross-attentions of some user-specified token $k$. Figure 1
gives an overview of our backward guidance where
$\mathcal{A}_{orig_{\text{fi}}}$ is the cross-attention map of some user-
specified token $k$ (e.g., token corresponding to “burger”) in frame fi of the
video generated by the T2V model. we omit the layer number and the token $k$
in our notation of the cross-attention. $\mathcal{A}_{tar_{\text{fi}}}$ is the
target cross-attention that captures the properties of the editing task. In
Figure 1, the task is to move the burger from the top-left to the bottom-left
of the scene. We define the energy function $E$ below. To control the shape
and size of an object (indicated by token $k$) through its corresponding
cross-attention maps, we threshold the attention map to eliminate the effect
of background noise and overlapping attention from other tokens. This is
achieved by taking a soft threshold at the midpoint of the per-channel minimum
and maximum values:

$\text{shape}(k)=\mathcal{A}_{k}^{\text{threshold}}.$

Using the thresholded original cross-attention and the target cross-attention,
we define the energy function $E$ as:

$E=\text{shape}(A_{tar})-\text{shape}(A_{orig}).$ (3)

This objective is zero-shot since $\text{shape}(A_{tar})$ can be computed as
($M\times\text{shape}(A_{orig})$) where $M$ defines some transformation such
as resizing and relocating the original attention. At time step $t$, we update
the latent $z_{t}$ according to the gradient of the loss defined by the energy
function $E$. This is realized through the following equation:

$z_{t}\leftarrow z_{t}-\delta_{t}^{2}\eta\nabla_{Z_{t}}\sum
E(A_{tar},A_{orig}),$ (4)

where $\eta>0$ controls the strength of backward guidance and
$\delta_{t}=\sqrt{(1-\alpha_{t})/\alpha_{t}}$. Updating the latent $z$ in this
manner indirectly influences the cross-attentions. Please refer to Section 4.2
for more details on our experimental setup.

Figure 4: We show qualitative results for shrinking and enlarging objects
through backward guidance. The middle image of each row visualizes the first
frame of the original video. We enlarge and shrink the target cross-attentions
at four different levels (Big / Bigger and small / smaller) and update the
latent through backward guidance. The first frame for each edited video is
shown.

##  4 Experiments

###  4.1 Limitation of Current T2V Models

In Figure 3, we visualize the cross-attention maps for all tokens of the
prompt “a shark swims in the ocean” generated with Stable Diffusion [37] and
our T2V model [46]. The cross-attention maps in T2I models capture the tokens
much better than T2V models. We attribute this to deeper denoising networks of
T2I models, larger training datasets, and more cross-attention layers. Using
such noisy cross-attention maps hinders both forward and backward guidance. To
perform backward guidance more effectively, we opted to directly generate
$\text{shape}(A_{tar})$ for each frame. Instead of transforming
$\text{shape}(A_{orig})$ to calculate $\text{shape}(A_{tar})$ for each video
frame, we generate binary cross-attention maps for the token of interest.
Despite this backward guidance setup not being zero-shot, we rely on future
T2V models with better cross-attention maps to replace this manual effort.

Figure 1 shows an example of user-generated target cross-attentions. In this
example, instead of transforming the cross-attention maps of the “burger” to
calculate the target, we directly generate each frame’s cross-attention
according to our editing task. Here, the task is to move the burger from the
top-left of the scene to the bottom-left in a straight line. Hence, we
generate cross-attention maps for each frame. For frame 1,
$A_{tar_{\text{f1}}}$ is placed at the top-left of the scene and in the
following frames, the cross-attention map moves slightly down such that in the
last frame, $A_{tar_{\text{f16}}}$ is placed at the bottom-left.

###  4.2 Experiment Details

We use the ModelScope [46] T2V model in our experiments and generate 16 frame
videos with $256\times 256$ resolution. Image editing methods such as
Diffusion self-guidance [12] have used the extracted image features learned by
the denoising network to preserve the background details and appearance
features of the object being edited. In this work, we only focus on
controlling objects’ motion and size and leave background and appearance
consistency for future works. We experimented with text prompts that describe
a simple scene, to further control the limitation of current T2V models and
get less noisy cross-attention maps for the object we want to edit.

Our 3D U-Net has cross-attentions with resolutions $4\times 4$, $8\times 8$ ,
$16\times 16$ and $32\times 32$. We find the $8\times 8$ and $16\times 16$
cross-attentions to be the most important dimensions for effectively
minimizing the energy function and editing the scene. There are 10 such layers
in the down-blocks, mid-blocks, and up-blocks of the 3D U-Net (4 down-block, 2
mid-block, and 4 up-block layers). We found that mid-block’s cross-attentions
played a vital role in backward guidance. Excluding the two mid-block layers
resulted in failed edits whereas excluding either all of down-block’s or all
of up-block’s cross-attentions resulted in fewer failures.

We experimented with different schemes for updating the latent $z$ and found
the most effective strategy to be that of Diffusion self-guidance [12]. During
the first $N/4$ iterations, we update $z$ at each step. For the subsequent
$3N/4$ iterations, we update $z$ at every other step. The guidance scale
$\eta$ (eq. 4) also plays an important role in the method’s effectiveness.
Increasing $\eta$ too much leads to degradation in the generated frames.
Selecting a very low scale does not change the latent enough for effective
editing. We found that in our setting, a scale of $15<\eta<25$ provided a good
balance between guidance strength and synthesis quality.

##  5 Results

In this section, we show the capabilities of backward guidance for two
different tasks. Figure 4 shows qualitative results for changing an object’s
size through backward guidance. Figure 5 presents qualitative results of
backward guidance for controlling the motion of an object in a video. To edit
an object of interest, we generate binary cross-attention maps that capture
the target position for the object’s token. For “burger”, we placed the first
cross-attention at the top-left of the scene and slowly moved it down. For the
“ball”, we placed the cross-attention at the top-left and moved it towards the
bottom-right of the scene. Finally, we moved the “shark” from the top-right
towards the bottom-left of the scene. Each sequence of frames with the black
caption shows the original video without performing guidance. The sequence of
frames with blue instruction shows the video after updating the latent with
backward guidance. The object successfully follows the cross-attention at each
frame.

We also observe that the original video can be missing an object described in
the text. The example with prompt “A wolf howls to the moon” in Figure 5 is
missing the moon. Interestingly, backward guidance encourages the moon to be
present in the scene. Attend-and-Excite [4] achieves the same objective in the
T2I domain.

Figure 5: The figure visualizes the results of backward cross-attention
guidance. For each of the 4 examples, we show the output of the T2V model
given the prompt in black. The blue text describes the applied transformation
to the cross-attentions at each frame. We update the input latent accordingly.
The red bounding box highlights the edit’s success.

##  6 Observations

In this section, we go over a few interesting observations when experimenting
with backward guidance.

#### Perspective.

In our experiments, we used a fixed size for $A_{tar}$ for all 16 video
frames. However, if the object is moving away or toward the camera, we should
see a change in the object’s size. In Figure 5, we see the burger, ball, and
shark getting larger as they move closer to the camera while the moon remains
the same size as it is static in the sky. It is noteworthy that despite
updating the model’s input latent with fixed-size target cross-attentions, the
model consistently generates videos with reasonable perspective. However, this
comes at the expense of not strictly adhering to the exact size defined by
$A_{tar}$.

#### Motion Control.

To control the motion of objects, we interpolate the cross-attention maps
between the attention map $A_{tar_{\text{f1}}}$, placed at starting position
$a$ and the attention map of the last frame $A_{tar_{\text{f16}}}$ placed in
final position $b$. We observed that the model keeps the temporal consistency
at the expense of not following the exact start and end location defined by
the target cross-attention. For instance, in Figure 5 \- last row, we placed
the cross-attention of the “shark” at the top-right for the first frame and at
the bottom-left for the last ($16^{th}$) frame. However, after $t$ steps, the
shark is not at the bottom of the scene where $A_{tar_{\text{f16}}}$ was
positioned. To do so, the model needs to move the shark much faster to go from
$A_{tar_{\text{f1}}}$ to $A_{tar_{\text{f16}}}$ in a short number of frames.
We also note that compared to resizing an object, controlling its motion is
often prone to failures using backward guidance. This failure takes the form
of the object being statically positioned at $A_{tar_{\text{f1}}}$. We leave
further exploration of this mode of failure for future work.

##  7 Discussion

This study conducted an initial investigation into the significance of cross-
attention layers within the 3D U-Net framework of video diffusion models. More
specifically, focusing on their role in determining objects’ size, position
and motion in T2V models. We examined the efficacy of utilizing cross-
attention maps to manipulate object size and motion, employing both forward
and backward guidance. In Section 3.3, we showed that forward guidance in
videos faces the same limitations that were previously observed in the T2I
domain [7] which hinders its performance. In Section 5, we showed results for
editing the size and motion of an object through backward guidance. Our
findings emphasize the promise of backward guidance in enabling zero-shot
editing capabilities for video generation. Moreover, in Section 4.1, we
highlighted current limitations that impede the transition of cross-attention-
based editing methods from the image domain to videos. This analysis provides
insights into the challenges and opportunities inherent to adapting editing
techniques to be used in dynamic video content.

##  8 Impact and Future Directions

Enabling zero-shot editing capabilities for generative video models is a
valuable approach to enhance user control without the need for model fine-
tuning with additional data. While current video models face limitations in
quality, length, and cross-attention accuracy, we anticipate that editing
methodologies like ours will leverage future advancements in Text-to-Video
models, similar to the progress seen in the Text-to-Image domain.

In this study, we focused on manipulating objects’ size and motion with
backward guidance. However, practical applications for editing tools require
further exploration, particularly enabling editing of real videos. This needs
additional constraints such as controlling background alterations and
maintaining the fidelity of different objects to the original video. These
aspects remain open for future work.

## References

  * ddp [2020] Denoising diffusion probabilistic models.  _Advances in neural information processing systems_ , 33:6840–6851, 2020.
  * Bar-Tal et al. [2024] Omer Bar-Tal, Hila Chefer, Omer Tov, Charles Herrmann, Roni Paiss, Shiran Zada, Ariel Ephrat, Junhwa Hur, Yuanzhen Li, Tomer Michaeli, et al.  Lumiere: A space-time diffusion model for video generation.  _arXiv preprint arXiv:2401.12945_ , 2024.
  * Blattmann et al. [2023] Andreas Blattmann, Robin Rombach, Huan Ling, Tim Dockhorn, Seung Wook Kim, Sanja Fidler, and Karsten Kreis.  Align your latents: High-resolution video synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 22563–22575, 2023.
  * Chefer et al. [2023] Hila Chefer, Yuval Alaluf, Yael Vinker, Lior Wolf, and Daniel Cohen-Or.  Attend-and-excite: Attention-based semantic guidance for text-to-image diffusion models.  _ACM Transactions on Graphics (TOG)_ , 42(4):1–10, 2023.
  * Chen et al. [2024a] Minghao Chen, Iro Laina, and Andrea Vedaldi.  Training-free layout control with cross-attention guidance.  In _Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision_ , pages 5343–5353, 2024a.
  * Chen et al. [2023a] Ting Chen, Lala Li, Saurabh Saxena, Geoffrey Hinton, and David J Fleet.  A generalist framework for panoptic segmentation of images and videos.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 909–919, 2023a.
  * Chen et al. [2023b] Weifeng Chen, Jie Wu, Pan Xie, Hefeng Wu, Jiashi Li, Xin Xia, Xuefeng Xiao, and Liang Lin.  Control-a-video: Controllable text-to-video generation with diffusion models.  _arXiv preprint arXiv:2305.13840_ , 2023b.
  * Chen et al. [2024b] Xinlei Chen, Zhuang Liu, Saining Xie, and Kaiming He.  Deconstructing denoising diffusion models for self-supervised learning, 2024b.
  * Çiçek et al. [2016] Özgün Çiçek, Ahmed Abdulkadir, Soeren S Lienkamp, Thomas Brox, and Olaf Ronneberger.  3d u-net: learning dense volumetric segmentation from sparse annotation.  In _Medical Image Computing and Computer-Assisted Intervention–MICCAI 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19_ , pages 424–432. Springer, 2016.
  * Clark and Jaini [2024] Kevin Clark and Priyank Jaini.  Text-to-image diffusion models are zero shot classifiers.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Dhariwal and Nichol [2021] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  _Advances in neural information processing systems_ , 34:8780–8794, 2021.
  * Epstein et al. [2024] Dave Epstein, Allan Jabri, Ben Poole, Alexei Efros, and Aleksander Holynski.  Diffusion self-guidance for controllable image generation.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Gal et al. [2022] Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit H Bermano, Gal Chechik, and Daniel Cohen-Or.  An image is worth one word: Personalizing text-to-image generation using textual inversion.  _arXiv preprint arXiv:2208.01618_ , 2022.
  * Gu et al. [2022] Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo.  Vector quantized diffusion model for text-to-image synthesis.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 10696–10706, 2022.
  * Guo et al. [2023] Yuwei Guo, Ceyuan Yang, Anyi Rao, Yaohui Wang, Yu Qiao, Dahua Lin, and Bo Dai.  Animatediff: Animate your personalized text-to-image diffusion models without specific tuning.  _arXiv preprint arXiv:2307.04725_ , 2023.
  * Hertz et al. [2022] Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or.  Prompt-to-prompt image editing with cross attention control.  _arXiv preprint arXiv:2208.01626_ , 2022.
  * Ho and Salimans [2022] Jonathan Ho and Tim Salimans.  Classifier-free diffusion guidance.  _arXiv preprint arXiv:2207.12598_ , 2022.
  * Ho et al. [2022] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al.  Imagen video: High definition video generation with diffusion models.  _arXiv preprint arXiv:2210.02303_ , 2022.
  * Hu et al. [2021] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.  Lora: Low-rank adaptation of large language models.  _arXiv preprint arXiv:2106.09685_ , 2021.
  * Jaini et al. [2023] Priyank Jaini, Kevin Clark, and Robert Geirhos.  Intriguing properties of generative classifiers.  _arXiv preprint arXiv:2309.16779_ , 2023.
  * Jeong et al. [2023] Hyeonho Jeong, Geon Yeong Park, and Jong Chul Ye.  Vmc: Video motion customization using temporal attention adaption for text-to-video diffusion models.  _arXiv preprint arXiv:2312.00845_ , 2023.
  * Karazija et al. [2023] Laurynas Karazija, Iro Laina, Andrea Vedaldi, and Christian Rupprecht.  Diffusion models for zero-shot open-vocabulary segmentation, 2023.
  * Ke et al. [2024] Bingxin Ke, Anton Obukhov, Shengyu Huang, Nando Metzger, Rodrigo Caye Daudt, and Konrad Schindler.  Repurposing diffusion-based image generators for monocular depth estimation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2024.
  * Khachatryan et al. [2023] Levon Khachatryan, Andranik Movsisyan, Vahram Tadevosyan, Roberto Henschel, Zhangyang Wang, Shant Navasardyan, and Humphrey Shi.  Text2video-zero: Text-to-image diffusion models are zero-shot video generators.  _arXiv preprint arXiv:2303.13439_ , 2023.
  * Kumari et al. [2023] Nupur Kumari, Bingliang Zhang, Richard Zhang, Eli Shechtman, and Jun-Yan Zhu.  Multi-concept customization of text-to-image diffusion.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 1931–1941, 2023.
  * Li et al. [2023] Alexander C. Li, Mihir Prabhudesai, Shivam Duggal, Ellis Brown, and Deepak Pathak.  Your diffusion model is secretly a zero-shot classifier.  In _2023 IEEE/CVF International Conference on Computer Vision (ICCV)_. IEEE, 2023.
  * Liu et al. [2023] Shaoteng Liu, Yuechen Zhang, Wenbo Li, Zhe Lin, and Jiaya Jia.  Video-p2p: Video editing with cross-attention control, 2023.
  * Mei and Patel [2023] Kangfu Mei and Vishal Patel.  Vidm: Video implicit diffusion models.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , pages 9117–9125, 2023.
  * Motamed et al. [2023] Saman Motamed, Danda Pani Paudel, and Luc Van Gool.  Lego: Learning to disentangle and invert concepts beyond object appearance in text-to-image diffusion models.  _arXiv preprint arXiv:2311.13833_ , 2023.
  * Mou et al. [2023] Chong Mou, Xintao Wang, Jiechong Song, Ying Shan, and Jian Zhang.  Dragondiffusion: Enabling drag-style manipulation on diffusion models.  _arXiv preprint arXiv:2307.02421_ , 2023.
  * Nichol et al. [2021] Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen.  Glide: Towards photorealistic image generation and editing with text-guided diffusion models.  _arXiv preprint arXiv:2112.10741_ , 2021.
  * Patashnik et al. [2023] Or Patashnik, Daniel Garibi, Idan Azuri, Hadar Averbuch-Elor, and Daniel Cohen-Or.  Localizing object-level shape variations with text-to-image diffusion models, 2023.
  * Peebles and Xie [2023] William Peebles and Saining Xie.  Scalable diffusion models with transformers.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 4195–4205, 2023.
  * Ramesh et al. [2021] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever.  Zero-shot text-to-image generation.  In _International Conference on Machine Learning_ , pages 8821–8831. PMLR, 2021.
  * Ramesh et al. [2022] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen.  Hierarchical text-conditional image generation with clip latents.  _arXiv preprint arXiv:2204.06125_ , 1(2):3, 2022.
  * Rombach et al. [2022a] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 10684–10695, 2022a.
  * Rombach et al. [2022b] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 10684–10695, 2022b.
  * Ruiz et al. [2023] Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman.  Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 22500–22510, 2023.
  * Saharia et al. [2022] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al.  Photorealistic text-to-image diffusion models with deep language understanding.  _Advances in Neural Information Processing Systems_ , 35:36479–36494, 2022.
  * Siméoni et al. [2021] Oriane Siméoni, Gilles Puy, Huy V Vo, Simon Roburin, Spyros Gidaris, Andrei Bursuc, Patrick Pérez, Renaud Marlet, and Jean Ponce.  Localizing objects with self-supervised transformers and no labels.  In _BMVC_ , 2021.
  * Singer et al. [2022] Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang, Oron Ashual, Oran Gafni, et al.  Make-a-video: Text-to-video generation without text-video data.  _arXiv preprint arXiv:2209.14792_ , 2022.
  * Sohl-Dickstein et al. [2015] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International conference on machine learning_ , pages 2256–2265. PMLR, 2015.
  * Song et al. [2021] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  In _ICLR_ , 2021.
  * Van Gansbeke and De Brabandere [2024] Wouter Van Gansbeke and Bert De Brabandere.  A simple latent diffusion approach for panoptic segmentation and mask inpainting.  _arXiv preprint arXiv:2401.10227_ , 2024.
  * Van Gansbeke et al. [2022] Wouter Van Gansbeke, Simon Vandenhende, and Luc Van Gool.  Discovering object masks with transformers for unsupervised semantic segmentation.  _arXiv preprint arXiv:2206.06363_ , 2022.
  * Wang et al. [2023a] Jiuniu Wang, Hangjie Yuan, Dayou Chen, Yingya Zhang, Xiang Wang, and Shiwei Zhang.  Modelscope text-to-video technical report.  _arXiv preprint arXiv:2308.06571_ , 2023a.
  * Wang et al. [2023b] Xudong Wang, Rohit Girdhar, Stella X Yu, and Ishan Misra.  Cut and learn for unsupervised object detection and instance segmentation.  In _CVPR_ , 2023b.
  * Wang et al. [2024] Xiang Wang, Hangjie Yuan, Shiwei Zhang, Dayou Chen, Jiuniu Wang, Yingya Zhang, Yujun Shen, Deli Zhao, and Jingren Zhou.  Videocomposer: Compositional video synthesis with motion controllability.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Wang et al. [2022] Yangtao Wang, Xi Shen, Shell Xu Hu, Yuan Yuan, James L Crowley, and Dominique Vaufreydaz.  Self-supervised transformers for unsupervised object discovery using normalized cut.  In _CVPR_ , 2022.
  * Wu et al. [2022] Chen Henry Wu, Saman Motamed, Shaunak Srivastava, and Fernando D De la Torre.  Generative visual prompt: Unifying distributional control of pre-trained generative models.  _Advances in Neural Information Processing Systems_ , 35:22422–22437, 2022.
  * Wu et al. [2023] Jay Zhangjie Wu, Yixiao Ge, Xintao Wang, Stan Weixian Lei, Yuchao Gu, Yufei Shi, Wynne Hsu, Ying Shan, Xiaohu Qie, and Mike Zheng Shou.  Tune-a-video: One-shot tuning of image diffusion models for text-to-video generation.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 7623–7633, 2023.
  * Yang et al. [2024] Shiyuan Yang, Liang Hou, Haibin Huang, Chongyang Ma, Pengfei Wan, Di Zhang, Xiaodong Chen, and Jing Liao.  Direct-a-video: Customized video generation with user-directed camera movement and object motion, 2024.
  * Yang and Wang [2023] Xingyi Yang and Xinchao Wang.  Diffusion model as representation learner.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 18938–18949, 2023.
  * Yu et al. [2022] Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, et al.  Scaling autoregressive models for content-rich text-to-image generation.  _arXiv preprint arXiv:2206.10789_ , 2(3):5, 2022.
  * Zhang et al. [2023] Lvmin Zhang, Anyi Rao, and Maneesh Agrawala.  Adding conditional control to text-to-image diffusion models.  In _2023 IEEE/CVF International Conference on Computer Vision (ICCV)_. IEEE, 2023.
  * Zhao et al. [2023] Rui Zhao, Yuchao Gu, Jay Zhangjie Wu, David Junhao Zhang, Jiawei Liu, Weijia Wu, Jussi Keppo, and Mike Zheng Shou.  Motiondirector: Motion customization of text-to-video diffusion models.  _arXiv preprint arXiv:2310.08465_ , 2023.


# TransClaw U-Net: Claw U-Net with transformers for medical image segmentation

###### Abstract

In recent years, computer-aided diagnosis has become an increasingly popular
topic. Methods based on convolutional neural networks have achieved good
performance in medical image segmentation and classification. Due to the
limitations of the convolution operation, the long-term spatial features are
often not accurately obtained. Hence, we propose a TransClaw U-Net network
structure, which combines the convolution operation with the transformer
operation in the encoding part. The convolution part is applied for extracting
the shallow spatial features to facilitate the recovery of the image
resolution after upsampling. The transformer part is used to encode the
patches, and the self-attention mechanism is used to obtain global information
between sequences. The decoding part retains the bottom upsampling structure
for better detail segmentation performance. The experimental results on
Synapse Multi-organ Segmentation Datasets show that the performance of
TransClaw U-Net is better than other network structures. The ablation
experiments also prove the generalization performance of TransClaw U-Net.

Index Terms—  TransClaw U-Net, transformer, medical image segmentation, deep
learning

##  1 Introduction

With the continuous development of deep learning, computer vision has been
widely applied in many fields, including manufacturing, automatic driving,
computer-aided diagnosis and agriculture. Specially, computer aided diagnosis
has attracted the attention of many scholars at home and abroad. Many methods
based on convolutional neural network (CNN) have been used in image
classification and segmentation [1], and have achieved good performance. In
particular, the accurate medical image segmentation plays an important role in
clinical diagnosis [2]. Among various network structures, U-Shape based UNet
[3] framework has achieved great success in the field of medical image
segmentation. The encoder of UNet provides a high-level semantic feature map,
and the decoder provides a low-level detailed feature map. These two phases
are combined through skip connections [4]. Many improved network structures
based on U-Net have achieved better results in different medical datasets,
such as UNet++ [5], Resnet34-UNet [6], Channel-UNet [7], Attention-UNet [8],
and R2UNet [9], which can fully learn the rich detail information in the image
and restore it based on the location information. The success of these studies
in organ segmentation [10], lesion segmentation [11] and cardiac segmentation
[12] have fully proved the high efficiency and superiority of CNNs.

Although most of the current medical image segmentation tasks are based on
CNNs, it is still difficult to meet the requirements of accurate segmentation
in computer-aided diagnosis. The main reason is that CNN makes full use of the
end-to-end learning method, and it lacks the modeling of some long-range
ordered features in the images [13], which leads to the poor segmentation
performance of many regions. Each convolution kernel only pays attention to
the feature information of itself and its boundary, and lacks a large range of
feature fusion, thus affecting the overall effect. To overcome these
disadvantages of convolution operation, many scholars put forward some novel
solutions, including the self-attention mechanism [14], atrous convolution
[15] and image pyramid [16].These methods does not consider the long-range
dependencies of medical images.

In natural language processing (NLP), the prevalent architecture transformer
demonstrates strong learning capabilities for long-term features such as BERT
[17] and GPT [18] models in sequence modeling and transduction tasks [19]. For
the transformer operation, the distributed convolution operator is combined
with the multi-head attention mechanism, which can find the relationship among
different input parts [20]. Many researches have started to use the
transformer structure in computer vision. Dosovitskiy et al [21] proposed
visual transformer (ViT) technology for image recognition tasks. Liu et al
[19] proposed a hierarchical swin transformer structure, which uses the shift
window for feature calculation and achieves better results in image
classification. Recently, transformer has also been used in medical image
segmentation tasks and has achieved satisfactory results. Chen et al [13]
inserted transformer into the encoding layer of UNet structure, which encodes
patches based on feature maps so as to realize the characteristics of long-
term dependency. Cao et al [22] proposed a UNet-like pure transformer
architecture, which uses transformer with drifted windows for feature
extraction and patch-expanding transformer for restoration of image
resolution.

In this research, we propose the TransClaw U-Net architecture, which combines
the detailed segmentation ability of Claw U-Net with the long-term feature
learning ability of transformer structure, and apply it to medical image
segmentation. This network uses the convolution operation to obtain the
feature map, and then encodes the patches of high-resolution images using the
feature map. Transformer strengthens the global connection between the
encoding parts, which makes the image features sufficiently visible, and the
self-attention mechanism is more conducive to the detail segmentation. In the
decoding part, the bottom upsampling module of Claw U-Net is retained, and the
deep feature map is combined with the shallow layer to achieve accurate
positioning. Finally, the encoding part, the upsampling part and the decoding
part of the corresponding layers are combined to achieve the restoration of
image resolution. Experimental results on abdominal CT scans datasets show
that our network structure achieves more accurate organ segmentation.
Meanwhile, we have also verified that larger input image size, smaller patch
size and more skip-connections can improve the model performance, but they
require more computing resources.

##  2 RELATED WORKS

CNN and variants: As the basic framework of deep learning, CNN has gradually
replaced the traditional machine learning methods [23] and become the most
effective methods in the field of medical image segmentation due to its
powerful image learning ability. Since 2015, UNet structure has been proposed
and proved to have superior performance in medical image segmentation tasks.
More and more new network structures based on U-shape have been proposed,
which include Resnet34-UNet [6], Attention-UNet [7], Dense-UNet [10], and
UNet++ [5]. Some improved convolution operations are introduced to improve the
model performance, such as deformable convolution [24] and deep convolution
[25]. CNN and its variants have become the most primary back bone architecture
in the field of artificial intelligence due to their strong learning ability.

Self-attention/Transformers to complement CNNs: In recent years, there are a
lot of researches on how to integrate self-attention mechanisms or transformer
operations with CNNs. By encoding the images, the self-attention mechanism can
make up for the deficiency of CNNs’ long-range dependence and heterogeneous
interactions, thus improving the head network [26] and back bone network [27].
Wang et al [27] inserted new non-local operators into the non-local operators
of multiple intermediate convolution layers. By adding transformer structure
into UNet structure and introducing global self-attention mechanism, Chen et
al [13] have achieved good results in 2D CT image dataset. Based on the above
researches, we adopt the newly proposed Claw U-Net structure [4] as the
framework, and introduce the transformer structure to explore its effect in
medical image segmentation task.

Transformers: the transformer structure was first used in the NLP domain, and
has achieved excellent results in machine translation [20]. As it has strong
ability to model the long-term features of sequences, it can make up for some
deficiencies of CNNs in global features. Many scholars have started to combine
transformer structures with convolution operations to achieve better
performance. Recently, in imagenet classification tasks, ViT [21] has achieved
a great break-through by applying the global attention mechanism to images.
However, the network results of the UNet structure can be trained on small
datasets, the transformer structure requires a large amount of data to be pre-
trained. Deit [28] proposed several training strategies for applying ViT to
small datasets. In the structure proposed in this paper, the Claw U-Net
network structure is adopted as the skeleton, and the convolution operation is
first carried out in the encoding part. Transformer structure is introduced to
the obtained feature maps to encode them, and the corresponding features are
extracted through the multi-head attention mechanism. The decoding part
combines the bottom upsampling feature maps with the encoding feature maps and
the decoding feature maps to restore the spatial features of the whole image.

##  3 Method

In this section, we will introduce the overall structure of the TransClaw
U-Net. This network introduces the transformer structure in the encoding
layer, and uses the self-attention mechanism to achieve better global feature
learning. During pixel restoration, it combines encoding, bottom up-sampling,
and decoding part, so that detailed information can be maintained, and finally
the image segmentation task is realized.

###  3.1 Overall structure of TransClaw U-Net

Fig. 1: The architecture of TransClaw U-Net

The overall structure of TransClaw U-Net is shown in Fig. 1, which mainly
includes the encoding part, the bottom upsampling part and the decoding part.
These parts are connected by skip-connections. Compared with Claw U-Net,
TransClaw U-Net uses a combination of convolution and transformer mechanism in
the encoding part. In the encoding part, assuming that $I_{in}$ denotes the
input image, $I_{fea}$ denotes the feature images after convoluton operation,
$I_{deep}$ denotes the deepest image in the encoding part. For the convolution
parts, the operation can be formulated as:

$\mathrm{F}_{0}=\sigma\left(BN\left(\operatorname{Conv}\left(\mathrm{I}_{in}\right)\right)\right)$
(1)

where $\sigma$ represents the ReLU function, $BN$ represents the
batchnormalization function, and $conv$ represents thr convolution function.
Before each convolution parts, maxpooling operation is used to reduce the
feature map size by two times.

$\mathrm{F}_{1}=\operatorname{MaxPool}\left(\mathrm{I}_{\text{fea }}\right)$
(2)

Table 1: Comparison of model performance on the Synapse multi-organ CT
dataset. The best performers are highlighted in bold.

Model | DSC | HD | Aorta | Gallbladder | Kidney(L) | Kidney(R) | Liver | Pancreas | Spleen | Stomach
---|---|---|---|---|---|---|---|---|---|---
V-Net [29] | 68.81 | - | 75.34 | 51.87 | 77.10 | 80.75 | 87.84 | 40.05 | 80.56 | 56.98
DARR [30] | 69.77 | - | 74.74 | 53.77 | 72.31 | 73.24 | 94.08 | 54.18 | 89.90 | 45.96
U-Net [3] | 76.85 | 39.70 | 89.07 | 69.72 | 77.77 | 68.60 | 93.43 | 53.98 | 86.67 | 75.58
R50 U-Net [13] | 74.68 | 36.87 | 87.74 | 63.66 | 80.60 | 78.19 | 93.74 | 56.90 | 85.87 | 74.16
R50 Att-UNet [13] | 75.57 | 36.97 | 55.92 | 63.91 | 79.20 | 72.71 | 93.56 | 49.37 | 87.19 | 74.95
R50 ViT [13] | 71.29 | 32.87 | 73.73 | 55.13 | 75.80 | 72.20 | 91.51 | 45.99 | 81.99 | 73.95
TransUNet [13] | 77.48 | 31.69 | 87.23 | 63.13 | 81.87 | 77.02 | 94.08 | 55.86 | 85/08 | 75.62
TransClaw U-Net | 78.09 | 26.38 | 85.87 | 61.38 | 84.83 | 79.36 | 94.28 | 57.65 | 87.74 | 73.55

where $MaxPool$ represents the maxpooling function. The purpose of these three
convolution operations is to retain low-level feature maps, so that the
shallow features and contours in the medical image can be effectively combined
with the decoding part. To make TransClaw U-Net learn better long-term spatial
features, we introduce a transformer mechanism after the convolution parts
which can be summarized as:

$\mathrm{I}_{\text{deep }}=\operatorname{Trans}\left(\mathrm{I}_{\text{fea
}}\right)$ (3)

where $Trans$ represents the transformer operation, which will be described in
detail in section 3.2. After the transformer module, the size of the encoded
image is restored to $H/P\times W/P$ (where $H$,and $W$ is the height and
width of the input image, and $P$ is the patch size) and sent to the decoding
part and the bottom upsampling part.

In the decoding part, the feature maps can be restored to the original
resolution after several upsampling operations. The structure of TransClaw
U-Net combines the corresponding encoding part, the bottom upsampling part,
and the decoding part through skip-connections. The expression to skip the
link can be expressed as:

when $i=1,\cdots,N-1$,

$\mathrm{x}_{\mathrm{De}}^{\mathrm{i},2}=\left[\mathrm{C}\left(\mathcal{D}\left(X_{\mathrm{En}}^{\mathrm{k},0}\right)\right)_{k=1}^{i-1},\mathrm{C}\left(X_{\mathrm{Up}}^{\mathrm{i},1}\right),\mathrm{C}\left(U\left(X_{\mathrm{De}}^{\mathrm{k},2}\right)\right)_{k=i+1}^{N}\right]$
(4)

when $i=N$,

$\mathrm{x}_{\mathrm{De}}^{\mathrm{i},2}=\mathrm{x}_{\mathrm{En}}^{\mathrm{i},0}$
(5)

where $i$ represents the downsampling layer along with the encoder, $N$
represents the total layer of the encoder, $x_{De}^{i,2}$ represents the
decoding part of the $i-th$ layer, $x_{En}^{i,0}$ represents the encoding part
of the $i-th$ layer, $x_{Up}^{i,1}$ represents the decoding part of the $i-th$
layer. Let function $C(\cdot)$ denotes a convolution operation,
$\mathcal{D}(\cdot)$ and $U(\cdot)$ denote upsampling and downsampling
operation separately, $[\cdot]$ denotes connecting together.

Each upsampling operation will go through a $3\times 3$ convolution operation
before. After the image is restored to $H/2\times W/2$, no concatenation
operation is required for the last upsampling. After that, the image size goes
to the initial size.

###  3.2 Transformer in TransClaw U-Net

To preserve the initial features of the image and the contours of the organs,
several convolutions are performed before the transformer operation, and the
feature maps are used as input of transformer. Divide the image into $N$
patches of $P\times P$ size, which have no overlapping with each other,
forming a flatten 2D
sequence$\left\\{\mathbf{x}_{p}^{i}\in\mathbb{R}^{P^{2}\cdot C}\mid
i=1,\cdots,N-1\right\\}.$

Concat the vectors of $N$ patches that are reshaped together, and we can get
the $N\times\ P^{2}$ matrix, which is similar to the word vector input to the
transformer. $N$ is the number of patches, which determines the length of the
input sequence. It is worth noting that when the size of the patch changes,
the length of the $P^{2}$ dimensional vector will also change after each patch
is reshaped e.

The structure in transformer is similar to that in NLP, which consists of
$L-th$ layer transformer blocks. For the first layer of the encoder, assuming
that the input is $z_{l-1}$ and the output is $z_{l}$, the results after they
have passed through the Multihead Self-Attention (MSA) and Multi-Layer
Perceptron (MLP) blocks can be represented by formulas 3 and 4, respectively.

$\displaystyle\mathbf{z}_{\ell}^{\prime}$
$\displaystyle=\operatorname{MSA}\left(\mathrm{LN}\left(\mathbf{z}_{\ell-1}\right)\right)+\mathbf{z}_{\ell-1},$
$\displaystyle\ell$ $\displaystyle=1\ldots L$ (6)

$\displaystyle\mathbf{z}_{\ell}$
$\displaystyle=\operatorname{MLP}\left(\operatorname{LN}\left(\mathbf{z}_{\ell}^{\prime}\right)\right)+\mathbf{z}_{\ell}^{\prime},$
$\displaystyle\ell$ $\displaystyle=1\ldots L$ (7)

where $L(\cdot)$ represents the layer normalization, and $z_{l}$ represents
the images after encoding.

The Multi-Head Attention used in the transformer block is the core of the
encoding. The study found that it is better to train multiple value vectors in
parallel and then stitch them together to get the output than to use an
attention function to get the $d_{model}$ dimensional value vector.
Specifically, given an input matrix, we calculate multiple sets of $V$, $K$,
and $Q$ matrices based on different parameter matrices, and then calculate
multiple weighted $V$ matrices through multiple attention functions, and
finally concatenate these matrices to get the final output through a weight
matrix $W^{O}$. The above process can be expressed by the formula as follows,

$\text{ MultiHead }(Q,K,V)=\text{ Concat }\left(\text{ head
}_{1},\ldots,\text{ head }_{\mathrm{b}}\right)W^{O}$ (8)

$\text{ head }_{\mathrm{i}}=\text{ Attention
}\left(QW_{i}^{Q},KW_{i}^{K},VW_{i}^{V}\right)$ (9)

where $W_{i}^{Q}\in\mathbb{R}^{d_{\text{model }}\times
d_{k}},W_{i}^{K}\in\mathbb{R}^{d_{\text{model }}\times
d_{k}},W_{i}^{V}\in\mathbb{R}^{d_{\text{model }}\times
d_{v}},W^{O}\in\mathbb{R}^{hd_{v}\times d_{\text{model }}}$.

Since the dimensionality of the vector in each attention layer is equally
divided, the computational cost of this method is not much different than that
of using a single head attention.

##  4 Experiments

###  4.1 Datasets

Synapse Multi-organ Segmentation Dataset (Synapse): From the MICCAI 2015
Multi-Atlas Abdomen Labeling Challenge, 30 cases of abdominal CT scans were
used, and there were a total of 3779 axial contrast-enhanced abdominal
clinical CT images. 18 (2212 slices) of 30 cases are taken as the training
dataset, and 12 (1567 slices) as the testing dataset. Following [13] [22], We
used average Hausdor Distance (HD) and average Dice-Similarity coeffcient
(DSC) as evaluation indicators to evaluate the segmentation performance of 8
organs in CT images including aorta, gallbladder, spleen, kidney(L),
kidney(R), liver, pancreas, spleen and stomach.

Fig. 2: Comparison of segmentation results of different models on the Synapse
multi-organ CT dataset.

###  4.2 Implementation Details

In this experiment, the proposed TransClaw U-Net is implemented in Python
using the Pytorch deep learning framework and trained on the Synapse multi-
organ CT dataset. In addition, the Stochastic Gradient Descent (SGD) is used
with an initialized learning rate $1\times 10^{-2}$ and momentum $0.9$. The
weight decay is set to $1\times 10^{-4}$. Considering that the size of the
input image is $224\times 224$, we set the batch size to $24$, and train $150$
times. For a fair comparison, the parameters of all experiments are set to the
same situation. In the training process, the model with the best performance
on the validation set is selected as the final model. We use cross entropy as
the loss function for optimization. All the experiments are performed on
NVIDIA GeForce RTX 2080 Ti.

###  4.3 Comparison with other models on Synapse dataset

We test TransClaw U-Net on the Synapse multi-organ CT dataset and compare it
with several state-of-the-art models, including: 1) V-Net [29], 2) DARR [30],
3) U- Net [3], 4) AttnUNet [8] and 5)TransU-Net[13]. Table 1 shows the
comparison results of various networks on the testing set. It can be clearly
seen that our TransClaw U-Net has achieved the best results on the evaluation
indicators with DSC and HD of $78.09\%$ and $26.38\%$, respectively. The DSC
indicators of other networks are all lower than $77.5\%$ and the HD of other
networks are all above $31.5$. Compared with the newly proposed TransUNet, our
model improves the accuracy of DSC by $0.6\%$ and the accuracy of HD by
$5.3\%$, which proves the effect of the upsampling part of the bottom layer on
edge detail segmentation. It can also be seen that our model has the best
segmentation performance in three of the 8 organs, while TransUNet only has
one, which proves the effectiveness of our model.

Fig. 2 shows in detail the segmentation results of different network
structures on the Synapse dataset. It can be seen that Attention-UNet has a
problem of excessive segmentation of organs, and UNet’s segmentation is not
accurate. This is mainly due to the limitations of the convolution operation
and the lack of the ability to extract global features of the image. Compared
with TransUNet, our model has a better effect on the segmentation of organ
boundaries, especially between organs that occupy a small area. Claw U-Net’s
preservation of image detail features and transformer’s effective learning of
global context information enable TransClaw U-Net to achieve better accurate
medical image segmentation tasks.

###  4.4 Ablation Experiments

To discuss the impact of some input parameters used in the model on the
results, we conduct several ablation experiments. The parametres to be
explored include the size of the input image, the number of skip-connections,
and the patch size.

Table 2: Comparison of performance of different input image sizes.

Resolution | DSC | Aorta | Gallbladder | Kidney(L) | Kidney(R) | Liver | Pancreas | Spleen | Stomach
---|---|---|---|---|---|---|---|---|---
$224\times 224$ | 78.09 | 85.87 | 61.38 | 84.83 | 79.36 | 94.28 | 57.65 | 87.74 | 73.55
$512\times 512$ | 80.39 | 90.00 | 56.86 | 83.27 | 76.21 | 95.06 | 67.76 | 91.16 | 82.82

Table 3: Comparison of performance of different patch sizes.

Patch size | DSC | Aorta | Gallbladder | Kidney(L) | Kidney(R) | Liver | Pancreas | Spleen | Stomach
---|---|---|---|---|---|---|---|---|---
$16\times 16$ | 78.09 | 85.87 | 61.38 | 84.83 | 79.36 | 94.28 | 57.65 | 87.74 | 73.55
$20\times 20$ | 77.92 | 85.96 | 61.81 | 84.66 | 76.67 | 94.28 | 55.94 | 86.90 | 77.13
$24\times 24$ | 77.27 | 86.84 | 60.41 | 84.05 | 77.83 | 93.74 | 56.16 | 86.33 | 72.77

1) Effect of the input size: when the image size is larger, it often contains
more pixel information. We compared model performance when the input image
size is $224\times 224$ and $512\times 512$. Due to the limitation of
computing resources, the batch size is set to $24$ and $6$. It can be seen
from Table 2 that when the image resolution increases, the DSC will increase
significantly. It may be that as the image size increases, the patch size is
unchanged, the length of the sequence increases, so that more global and
detailed information in the image can be learned and converted into more
effective encoding results, thereby improving segmentation performance.
However, the increasing of image size needs more computing resources, longer
training time. Taking into account the calculation cost and other factors, the
input image size we finally adopted is $224\times 224$.

2) Effect of the number of skip-connections: by setting the number of skip-
connections to $0$, $1$, $2$, and $3$, we can explore its impact on model
performance. Fig. 3 shows the DSC index results of 8 organs segmentation under
different numbers of skip-connections. It can be seen that as the number of
skip-connections decreases, the effect of segmentation becomes worse. The
skip-connections can effectively combine the encoding part, the bottom
upsampling part and the decoding part, so that the model is more robust and
can achieve better performance. It is worth noting that when the skip-
connections reduce to $0$, for Arota and Pancreas, the DSC drops by $10.8\%$
and $10.16\%$ respectively.

3) Effect of the patch size and sequence length: the length of the transformer
encoding sequence depends on the size of the patch, and the sequence length
has a certain impact on the acquisition of image features. The use of
different patch sizes will have a greater impact on segmentation performance.
According to the existing knowledge, the product of the square of the patch
size and the length of the transformer sequence is a fixed value. When the
patch size is larger, the sequence is shorter, and the effective information
that can be learned between encoders is less, which affects the acquisition of
global spatial features. When the patch size is smaller, as the sequence
grows, the parameters that need to be learned grow rapidly, which greatly
occupies computing resources. Therefore, it is important to choose the
suitable patch size. In this experiment, we use $16\times 16$ patches. It can
be seen from Table 3 that when the patch size is $16$, the DSC is $78.09\%$,
and when the patch size is $24$, the DSC is $77.27\%$. The growth rate of DSC
is limited, but the occupied memory shows a multiple growth.

##  5 CONCLUSION

In this paper, we propose a TransClaw U-Net network structure, which combines
the convolution operation of CNN and the transformer mechanism, and has
achieved excellent performance on the Synapse multi-organ CT dataset. Compared
with the traditional CNN-network, the multi-head self-attention in the
transformer can effectively obtain the global context information and make up
for the short-comings of the convolution operation in this respect. In the
encoding part, these two parts are combined with each other to achieve better
results. The decoding part retains the upsampling part, which makes the
network segmentation of medical images more accurate. Ablation experimental
results show that TransClaw U-Net has excellent generalization performance.

Fig. 3: Comparison of performance of different numbers of skip-connections.

## References

  * [1] Jonathan Long, Evan Shelhamer, and Trevor Darrell,  “Fully convolutional networks for semantic segmentation,”  in Proceedings of the IEEE conference on computer vision and pattern recognition, 2015, pp. 3431–3440.
  * [2] Ali Hatamizadeh, Dong Yang, Holger Roth, and Daguang Xu,  “Unetr: Transformers for 3d medical image segmentation,”  arXiv preprint arXiv:2103.10504, 2021.
  * [3] Olaf Ronneberger, Philipp Fischer, and Thomas Brox,  “U-net: Convolutional networks for biomedical image segmentation,”  in International Conference on Medical image computing and computer-assisted intervention. Springer, 2015, pp. 234–241.
  * [4] Chang Yao, Jingyu Tang, Menghan Hu, Yue Wu, Wenyi Guo, Qingli Li, and Xiao-Ping Zhang,  “Claw u-net: A unet-based network with deep feature concatenation for scleral blood vessel segmentation,”  arXiv preprint arXiv:2010.10163, 2020 (This paper has been accepted in 2021 CICAI.).
  * [5] Zongwei Zhou, Md Mahfuzur Rahman Siddiquee, Nima Tajbakhsh, and Jianming Liang,  “Unet++: Redesigning skip connections to exploit multiscale features in image segmentation,”  IEEE transactions on medical imaging, vol. 39, no. 6, pp. 1856–1867, 2019.
  * [6] Xiao Xiao, Shen Lian, Zhiming Luo, and Shaozi Li,  “Weighted res-unet for high-quality retina vessel segmentation,”  in 2018 9th international conference on information technology in medicine and education (ITME). IEEE, 2018, pp. 327–331.
  * [7] Yilong Chen, Kai Wang, Xiangyun Liao, Yinling Qian, Qiong Wang, Zhiyong Yuan, and Pheng-Ann Heng,  “Channel-unet: a spatial channel-wise convolutional neural network for liver and tumors segmentation,”  Frontiers in genetics, vol. 10, pp. 1110, 2019.
  * [8] Ozan Oktay, Jo Schlemper, Loic Le Folgoc, Matthew Lee, Mattias Heinrich, Kazunari Misawa, Kensaku Mori, Steven McDonagh, Nils Y Hammerla, Bernhard Kainz, et al.,  “Attention u-net: Learning where to look for the pancreas,”  arXiv preprint arXiv:1804.03999, 2018.
  * [9] Md Zahangir Alom, Mahmudul Hasan, Chris Yakopcic, Tarek M Taha, and Vijayan K Asari,  “Recurrent residual convolutional neural network based on u-net (r2u-net) for medical image segmentation,”  arXiv preprint arXiv:1802.06955, 2018.
  * [10] Xiaomeng Li, Hao Chen, Xiaojuan Qi, Qi Dou, Chi-Wing Fu, and Pheng-Ann Heng,  “H-denseunet: hybrid densely connected unet for liver and tumor segmentation from ct volumes,”  IEEE transactions on medical imaging, vol. 37, no. 12, pp. 2663–2674, 2018.
  * [11] Yuyin Zhou, Lingxi Xie, Wei Shen, Yan Wang, Elliot K Fishman, and Alan L Yuille,  “A fixed-point model for pancreas segmentation in abdominal ct scans,”  in International conference on medical image computing and computer-assisted intervention. Springer, 2017, pp. 693–701.
  * [12] Lequan Yu, Jie-Zhi Cheng, Qi Dou, Xin Yang, Hao Chen, Jing Qin, and Pheng-Ann Heng,  “Automatic 3d cardiovascular mr segmentation with densely-connected volumetric convnets,”  in International conference on medical image computing and computer-assisted intervention. Springer, 2017, pp. 287–295.
  * [13] Jieneng Chen, Yongyi Lu, Qihang Yu, Xiangde Luo, Ehsan Adeli, Yan Wang, Le Lu, Alan L Yuille, and Yuyin Zhou,  “Transunet: Transformers make strong encoders for medical image segmentation,”  arXiv preprint arXiv:2102.04306, 2021.
  * [14] Jo Schlemper, Ozan Oktay, Michiel Schaap, Mattias Heinrich, Bernhard Kainz, Ben Glocker, and Daniel Rueckert,  “Attention gated networks: Learning to leverage salient regions in medical images,”  Medical image analysis, vol. 53, pp. 197–207, 2019.
  * [15] Zaiwang Gu, Jun Cheng, Huazhu Fu, Kang Zhou, Huaying Hao, Yitian Zhao, Tianyang Zhang, Shenghua Gao, and Jiang Liu,  “Ce-net: Context encoder network for 2d medical image segmentation,”  IEEE transactions on medical imaging, vol. 38, no. 10, pp. 2281–2292, 2019.
  * [16] Hengshuang Zhao, Jianping Shi, Xiaojuan Qi, Xiaogang Wang, and Jiaya Jia,  “Pyramid scene parsing network,”  in Proceedings of the IEEE conference on computer vision and pattern recognition, 2017, pp. 2881–2890.
  * [17] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova,  “Bert: Pre-training of deep bidirectional transformers for language understanding,”  arXiv preprint arXiv:1810.04805, 2018.
  * [18] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever,  “Improving language understanding by generative pre-training,”  2018\.
  * [19] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo,  “Swin transformer: Hierarchical vision transformer using shifted windows,”  arXiv preprint arXiv:2103.14030, 2021.
  * [20] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin,  “Attention is all you need,”  arXiv preprint arXiv:1706.03762, 2017.
  * [21] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.,  “An image is worth 16x16 words: Transformers for image recognition at scale,”  arXiv preprint arXiv:2010.11929, 2020.
  * [22] Hu Cao, Yueyue Wang, Joy Chen, Dongsheng Jiang, Xiaopeng Zhang, Qi Tian, and Manning Wang,  “Swin-unet: Unet-like pure transformer for medical image segmentation,”  arXiv preprint arXiv:2105.05537, 2021.
  * [23] Andy Tsai, Anthony Yezzi, William Wells, Clare Tempany, Dewey Tucker, Ayres Fan, W Eric Grimson, and Alan Willsky,  “A shape-based approach to the segmentation of medical imagery using level sets,”  IEEE transactions on medical imaging, vol. 22, no. 2, pp. 137–154, 2003.
  * [24] Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He,  “Aggregated residual transformations for deep neural networks,”  in Proceedings of the IEEE conference on computer vision and pattern recognition, 2017, pp. 1492–1500.
  * [25] Xizhou Zhu, Han Hu, Stephen Lin, and Jifeng Dai,  “Deformable convnets v2: More deformable, better results,”  in Proceedings of the IEEE conference on computer vision and pattern recognition, 2019, pp. 9308–9316.
  * [26] Han Hu, Jiayuan Gu, Zheng Zhang, Jifeng Dai, and Yichen Wei,  “Relation networks for object detection,”  in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2018, pp. 3588–3597.
  * [27] Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He,  “Non-local neural networks,”  in Proceedings of the IEEE conference on computer vision and pattern recognition, 2018, pp. 7794–7803.
  * [28] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou,  “Training data-efficient image transformers & distillation through attention,”  arXiv preprint arXiv:2012.12877, 2020.
  * [29] Fausto Milletari, Nassir Navab, and Seyed-Ahmad Ahmadi,  “V-net: Fully convolutional neural networks for volumetric medical image segmentation,”  in 2016 fourth international conference on 3D vision (3DV). IEEE, 2016, pp. 565–571.
  * [30] Shuhao Fu, Yongyi Lu, Yan Wang, Yuyin Zhou, Wei Shen, Elliot Fishman, and Alan Yuille,  “Domain adaptive relational reasoning for 3d multi-organ segmentation,”  in International Conference on Medical Image Computing and Computer-Assisted Intervention. Springer, 2020, pp. 656–666.