Attention-Based Cross-Stitching

• The paper demonstrates that attention-based cross-stitching significantly improves feature fusion using dynamic attention and gating, yielding measurable performance gains.
• The methodology employs spatially adaptive, bidirectional feature transfer between neural representations across modalities such as text and images.
• Empirical results reveal enhanced robustness and scalability in multi-modal and multi-task scenarios, outperforming static fusion approaches.

Attention-based cross-stitching is a feature fusion and sharing mechanism grounded in dynamic attention and gating, extending the original cross-stitch concept to enable context-dependent, bidirectional, and spatially adaptive information exchange between neural networks. This methodology addresses key bottlenecks in multi-modal and multi-task learning by replacing static, global mixing weights with cross-attention modules and adaptive gating, improving the transfer of complementary signals between modalities such as text and knowledge graphs or between multiple visual task decoders. Empirical studies demonstrate notable improvements over both no-sharing and fixed cross-stitch baselines by leveraging dynamically learned affinities and soft gating for feature interpolation (Dai et al., 2022, Kim et al., 2022).

1. Foundations and Motivation

Early cross-stitch approaches (as per Misra et al., CVPR 2016) operate by learning a matrix of static scalar mixing weights per channel between task-specific networks. While this facilitates basic feature sharing, it is limited by two main factors: (1) mixing weights are global for each feature channel and cannot adapt spatially to context or instance; (2) parameter complexity grows quadratically with the number of tasks, and manual layer-wise insertion is required. Attention-based cross-stitching was developed to overcome these problems by (a) employing attention modules that compute spatially and contextually varying affinities, and (b) introducing soft, learnable gating for bidirectional control of signal injection. This yields models that borrow information in a data-dependent manner and flexibly mediate between distinct neural representations (Dai et al., 2022, Kim et al., 2022).

2. Mathematical Formulation: Architectures and Operations

2.1. Bi-encoder Attention-based Cross-Stitching

In the XBE architecture for distantly-supervised relation extraction, attention-based cross-stitching fuses a text encoder (e.g., BERT) and a knowledge graph (KG) encoder (Transformer over triples) at several intermediate layers (Dai et al., 2022). At each cross-stitch layer, the key operations are:

• Cross-attention computation: At layer $i$, given text states $S_i \in \mathbb{R}^{N\times d}$ and KG states $T_i \in \mathbb{R}^{3 \times d}$, two attention matrices are formed:
  • $$A^{t2s} = \text{softmax}_{\text{col}}\left((W_p^{t2s} T_i) S_i^T\right) \in \mathbb{R}^{3 \times N}$$
  • $$A^{s2t} = \text{softmax}_{\text{row}}\left(S_i (W_p^{s2t} T_i)^T\right) \in \mathbb{R}^{N\times 3}$$
• Signal aggregation and projection:
  • KG→text update: $U_i^{t2s} = A^{t2s} T_i$, projected via $$T_i^{t2s} = W_{g2}^{t2s}\mathrm{ReLU}(W_{g1}^{t2s} (U_i^{t2s})^T)$$
  • text→KG update: $V_i^{s2t} = (A^{s2t})^T S_i$, projected via $$S_i^{s2t} = W_{g2}^{s2t}\mathrm{ReLU}(W_{g1}^{s2t} (V_i^{s2t})^T)$$
• Dynamic gating:
  • $G_i^{t2s} = \sigma(T_i^{t2s})$, $G_i^{s2t} = \sigma(S_i^{s2t})$
• Feature interpolation:
  • $S_i' = G_i^{t2s} \odot S_i + \lambda_t T_i^{t2s}$
  • $T_i' = G_i^{s2t} \odot T_i + \lambda_s S_i^{s2t}$

Updated embeddings ($S_i'$, $T_i'$) feed forward to the next encoder layer, enabling full bidirectional and context-dependent feature injection controlled by the gate outputs.

2.2. Sequential Cross Attention in Multi-task Vision

In multi-task visual scene understanding, the method replaces cross-stitch units with attention modules for two principal axes: tasks (CTAM) and scales (CSAM) (Kim et al., 2022):

• Cross-Task Attention Module (CTAM): For each target task $i$ and scale $k$, let $F_{k,i}$ be the feature map. CTAM projects these to queries ($Q_{k,i}$), keys ($K_{k,\lnot i}$), and values ($V_{k,\lnot i}$) using $1\times 1$ convolutions. The output at each pixel aggregates source task features via

$$A_{k,i}(x) = \sum_{y \in \Omega} \text{softmax}_y (Q_{k,i}(x) K_{k,\lnot i}(y)) V_{k,\lnot i}(y)$$

yielding pixel-wise, spatially adaptive fusion.

• Cross-Scale Attention Module (CSAM): For each task, features at fine scales query and aggregate coarser scale features in a similar fashion, enabling hierarchical, context-aware fusion without prohibitive computational costs.

This two-stage decomposition achieves near full-attention expressivity at substantially reduced complexity, and enables discriminative, content-aware feature mixing.

3. Comparison to Classical Cross-Stitching and Design Rationale

Traditional cross-stitch units implement fixed, location-invariant channel mixing matrices, resulting in non-adaptive, static feature sharing (Kim et al., 2022). In contrast, attention-based cross-stitching leverages (i) spatial adaptivity—mixing weights vary over spatial positions and are learned through attention; (ii) context sensitivity—affinities are a function of both the query and source context; and (iii) dynamic gating, so each layer and feature can softly regulate cross-modal or cross-task injection.

Empirically, fixed gates (e.g., gate value 0.5) lead to significant performance drops; dynamic gates learned via the network improve task-specific and overall performance by adapting to the input distribution and layer semantics (Dai et al., 2022).

4. Implementation Strategies and Hyperparameters

The method supports a wide range of backbone encoders and task settings. Architecturally, attention-based cross-stitch modules can be inserted at several intermediate layers, with empirical findings supporting placement in middle layers of deep transformers (e.g., between layers 1–6 of BERT). Scalar hyperparameters $\lambda_t$, $\lambda_s$ regulate the relative magnitude of injected updates and are tuned to harmonize the scale of features from disparate sources (e.g., text versus KG) (Dai et al., 2022).

In vision, projections to compressed latent spaces ($d_k, d_v$) via $1\times1$ convolutions, and attention heads tailored to the target memory/computation budget, control complexity. Using HRNet-18 backbone, Swin-Transformer self-attention, and attention-based feature propagation, sequential cross attention enables efficient, scalable multi-task fusion across four spatial scales (Kim et al., 2022).

5. Empirical Performance and Analysis

Attention-based cross-stitch bi-encoders and sequential cross-attention networks demonstrate superior empirical performance:

• In relation extraction, XBE delivers substantial AUC improvements on Medline21 (61.9 vs. 55.3, +6.6) and NYT10 (70.5 vs. 63.2, +7.3) over no-sharing baselines. Ablation removing the cross-stitch module or restricting it to unidirectional attention yields marked drops (up to 3.2 AUC) (Dai et al., 2022).
• In multi-task dense prediction, sequential cross attention attains joint-task gains ($A_m$) of +12.07% over single-task baselines on NYUD-v2 (mIoU=41.33%, depth RMSE=0.604), outperforming fixed cross-stitch and other recent multi-task fusion architectures (Kim et al., 2022).

The interplay of attention-based fusion and dynamic gating is responsible for improved robustness to noisy supervision and greater selectivity in information transfer, resulting in state-of-the-art results.

6. Computational Complexity and Scalability

Applying full cross-attention over all possible task-scale combinations scales quadratically, leading to prohibitive compute and parameter counts as the number of tasks or spatial scales increases. The sequential decomposition (first over tasks at each scale with CTAM, then over scales within each task with CSAM) reduces operations to $\mathcal{O}(KM^2 + MK^2)$ (for $M$ tasks, $K$ scales), as opposed to $\mathcal{O}((MK)^2)$ for naïve full attention. Projection to small latent dimensions (e.g., $d_k=64$) further reduces runtime and memory while retaining representational expressivity (Kim et al., 2022).

7. Extensions and Prospective Directions

Possible future advancements include neural architecture search to determine optimal insertion points or attention projection dimensions, extension to larger numbers of tasks (e.g., object detection, edge detection), and dynamic routing based on input complexity (e.g., bypassing cross-scale attention for fine scales). Attention-based cross-stitching frameworks are also adaptable to new modalities and regime shifts, offering a general, flexible foundation for scalable, adaptive feature sharing in multi-modal, multi-task, and cross-domain neural architectures (Kim et al., 2022).

---

Key References:

• "Cross-stitching Text and Knowledge Graph Encoders for Distantly Supervised Relation Extraction" (Dai et al., 2022)
• "Sequential Cross Attention Based Multi-task Learning" (Kim et al., 2022)