MAIT: Mamba-Attention Integrated Transformer

• The paper introduces MAIT as a hybrid encoder that interleaves multi-head self-attention and Mamba layers to manage variable system scales and long temporal horizons.
• MAIT optimizes multi-UAV trajectory planning by fusing global dependency modeling with efficient temporal dynamics for both discrete and continuous control.
• Empirical results demonstrate that MAIT outperforms traditional Transformers and graph-based approaches in energy efficiency and inference speed.

Searching arXiv for MAIT and closely related hybrid Mamba-attention architectures to ground the article in current literature. Mamba-Attention Integrated Transformer (MAIT) is a hybrid neural architecture introduced as the core trajectory-generation model of the ARMAIT framework for multi-UAV trajectory optimization. In its original formulation, MAIT is designed to combine the long-range dependency modeling capability of attention mechanisms with the efficient temporal dynamic representation of Mamba, so that a single policy network can handle variable system scale, long temporal horizons, heterogeneous entities, global coordination, and real-time efficiency in low-altitude economy networks (Jiang et al., 22 Aug 2025). The original MAIT is a serial hybrid encoder rather than a fused low-level operator: an embedding layer feeds a stack of attention layers and Mamba layers, each accompanied by MLP and RMSNorm sublayers, and a trajectory output layer maps the final representation to either discrete route choices or continuous control variables (Jiang et al., 22 Aug 2025).

1. Conceptual background and architectural rationale

MAIT belongs to a broader research program that treats Mamba neither as a universal replacement for Transformer attention nor as a mere auxiliary module, but as a complementary sequence operator whose strengths are most visible in long-context, high-resolution, or temporally extended settings. A recent survey explicitly identifies “The Combination of Transformer and Mamba” as a distinct line of development and groups together hybrid systems such as State Space Augmented Transformer, Block-State Transformer, MambaFormer, Jamba, Heracles, DiMBA, PoinTramba, and PointABM under that theme (Zou et al., 2024).

The compatibility between the two paradigms is not only empirical. “The Hidden Attention of Mamba Models” shows that a selective state-space layer can be rewritten as an input-dependent lower-triangular token-mixing matrix,

$$y = \tilde{\alpha} x,\qquad \tilde{\alpha}_{i,j} = C_i \left(\Pi_{k=j+1}^i \bar{A}_k \right) \bar{B}_j,$$

and argues that Mamba can be reformulated as an implicit form of causal self-attention rather than as a wholly unrelated mechanism (Ali et al., 2024). This interpretation does not make MAIT redundant; it instead clarifies why a hybrid can be coherent. Explicit attention supplies direct pairwise token interaction and normalized competition, whereas Mamba supplies selective recurrent memory propagation and favorable long-sequence scaling.

In the original MAIT motivation, this division of labor is framed operationally. Pure Transformer approaches improve variable-length modeling and global dependency capture, but their self-attention has quadratic complexity in sequence length and becomes expensive for large numbers of devices or UAVs; standard transformers are also described as weaker at continuously modeling physical state evolution such as energy, mobility, and resource-constraint dynamics. Pure Mamba or SSM-style models are efficient for long sequences and temporal dynamics, but do not provide the same explicit all-to-all relational interaction modeling as self-attention. MAIT is therefore introduced as a hybrid in which attention handles global dependencies and inter-node interactions, while Mamba handles efficient temporal dynamics and long-range sequence modeling (Jiang et al., 22 Aug 2025).

2. Canonical MAIT architecture

The original MAIT architecture has four named components: an embedding layer, attention layers, Mamba layers, and a trajectory output layer. Both the attention and Mamba layers include MLP sublayers and RMSNorm sublayers; the paper states that the MLP improves representation capacity and RMSNorm stabilizes training (Jiang et al., 22 Aug 2025).

The key integration mechanism is structural. MAIT is a transformer-like encoder stack in which each encoder layer is designated either as an attention layer or as a Mamba layer. During the forward pass, for each encoder layer $l$, the model computes multi-head self-attention if the layer is an attention layer and otherwise computes a Mamba-2 state-space transformation; RMS normalization and an MLP then follow. This means that the original MAIT is a serial hybrid encoder with configurable layer composition, not a parallel branch fusion module, not a gate between simultaneous attention and Mamba outputs, and not a single fused token mixer (Jiang et al., 22 Aug 2025).

This point matters because the term “integrated” can easily be overstated. The original paper does not provide an explicit algebraic fusion of the form

$$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$

nor any learned coefficient that mixes the two at token level. Integration is defined by layerwise coexistence inside one encoder stack, with the attention-to-Mamba ratio selected according to the task and system scale (Jiang et al., 22 Aug 2025).

The architecture is explicitly task-adaptive. In the paper’s continuous trajectory planning case, Agentic RAG recommends 10 encoder layers with 70% attention and 30% Mamba. In the discrete trajectory planning case, it recommends 15 encoder layers with 30% attention and 70% Mamba. The stated interpretation is scale-dependent: more attention for smaller systems where global relational precision is useful, and more Mamba for larger systems where efficiency matters more (Jiang et al., 22 Aug 2025).

3. Mathematical structure and spatiotemporal factorization

The original MAIT models the system state at time $t$ as

$$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$

Its embedding layer combines state embedding, temporal encoding, and identity embedding: $h_{i}^0 = x_{i, t} W_s + p_t + b_i.$ Here, $x_{i,t}$ is the raw state vector of the $i$-th input at time step $t$, $W_s$ is a learnable linear projection into model dimension, $l$0 is a learnable temporal encoding, and $l$1 is the identity embedding of input $l$2 (Jiang et al., 22 Aug 2025).

For attention layers, the previous hidden representation $l$3 is projected to query, key, and value: $l$4 with compatibility

$l$5

and self-attention output

$l$6

For multi-head attention with $l$7 heads, head outputs are concatenated and merged as

$l$8

In the MAIT interpretation given by the paper, this module captures global dependencies and dynamic interactions among input nodes, including coupling among UAV states, device states, and task or resource relations (Jiang et al., 22 Aug 2025).

For the Mamba side, the paper uses Mamba-2 and presents the standard continuous-time linear state-space equations

$l$9

followed by discretization

$$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$0

and discrete recurrence

$$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$1

The paper also gives the convolutional view over a sequence of length $$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$2: $$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$3 The intended role of the Mamba layer is to model continuity of device mobility and evolution of resource constraints while preserving linear-time sequence processing characteristics (Jiang et al., 22 Aug 2025).

The resulting factorization is explicitly spatiotemporal. The sequence dimension is primarily time, while the node or entity index captures UAVs and ground entities. Attention operates over the current set of input nodes; Mamba updates hidden state over time. The paper therefore treats MAIT as a model in which attention handles cross-entity interactions at a step, while Mamba handles temporal continuity and long-range sequential dynamics (Jiang et al., 22 Aug 2025).

4. Outputs, policy parameterization, and T-GRPO training

The final hidden representation at time $$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$4 is

$$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$5

which is decoded into multi-UAV actions

$$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$6

The paper defines two task-specific output heads. For discrete trajectory planning, each UAV chooses from a finite candidate set of stop points $$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$7, with output distribution

$$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$8

where $$\mathrm{Fuse}(\mathrm{Attention}(X), \mathrm{Mamba}(X)),$$9 is a mask matrix indicating feasible and infeasible stop points, and the action is

$t$0

For continuous trajectory planning, the output is

$t$1

where $t$2 is a task-specific scale matrix mapping normalized outputs to physically meaningful ranges. The examples given are $t$3, $t$4, and $t$5 (Jiang et al., 22 Aug 2025).

MAIT is not trained as a supervised sequence predictor in the original paper. It serves as the policy network in a reinforcement-learning formulation. The state, action, and trajectory are written as

$t$6

with return

$t$7

T-GRPO samples $t$8 independent rollouts from the same initial state and computes rollout returns $t$9. Relative-advantage normalization is

$$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$0

$$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$1

The importance ratio is

$$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$2

and the clipped PPO-style objective per step is

$$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$3

The overall objective is

$$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$4

with KL regularization

$$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$5

In this setup, MAIT is the policy $$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$6, and T-GRPO updates the parameters $$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$7 without a value network (Jiang et al., 22 Aug 2025).

5. Empirical behavior in multi-UAV trajectory optimization

The original MAIT paper reports two case studies. In continuous trajectory planning, the comparison is against Transformer, Linformer, Performer, MAIT 1 with 70% attention, and MAIT 2 with 30% attention. The reported results are:

• Transformer: 376.38 J, 2.25 s
• Linformer: 398.65 J, 1.54 s
• Performer: 338.43 J, 1.66 s
• MAIT 1: 317.64 J, 1.42 s
• MAIT 2: 340.57 J, 1.15 s

This supports the paper’s central scale-sensitive design claim. The more attention-heavy MAIT variant gives the best energy consumption in the smaller-scale case, while the more Mamba-heavy MAIT variant is fastest but slightly worse in energy (Jiang et al., 22 Aug 2025).

In discrete trajectory planning, MAIT is compared with GNN, GPN, and ATOM. The reported results are:

• GNN: 292.861 Wh, 2.42 s
• GPN: 278.88 Wh, 1.67 s
• ATOM: 260.64 Wh, 2.86 s
• MAIT: 249.57 Wh, 1.53 s

The paper interprets this as evidence that MAIT can outperform graph and attention-based prior trajectory designers in both energy and inference time, although these are broader baseline comparisons rather than pure ablations of the Mamba-attention composition itself (Jiang et al., 22 Aug 2025).

The paper also reports a generalization-style test under varying UAV storage capacities in which MAIT remains the best-performing method in energy cost. The attributed reasons are the hybrid self-attention and Mamba feature extraction and T-GRPO’s ability to adapt to dynamic elements such as remaining storage and battery capacity (Jiang et al., 22 Aug 2025).

These results should be read with a specific caveat. They validate the usefulness of the hybrid serial encoder in the stated UAV settings, but they do not isolate the exact contribution of layer ordering, normalization choices, or a principled attention-to-Mamba ratio selection rule. The empirical record therefore supports MAIT as a practical architecture template more strongly than it supports a complete theory of optimal hybrid composition.

6. Relationship to broader Transformer–Mamba hybrids

MAIT in the strict sense originates in ARMAIT, but its design sits within a larger family of hybrid attention–SSM systems that differ mainly in where and how the two operators interact.

System	Integration pattern	Domain
Dimba (Fei et al., 2024)	Interleaved attention, Mamba, cross-attention, and MLP inside a diffusion denoiser	Text-to-image generation
MAT (Zhang et al., 2024)	Multi-scale hybrid Mamba–Transformer modules for long-short range dependencies	Weather forecasting
Matten (Gao et al., 2024)	Spatial attention, temporal attention, and global-sequence bidirectional Mamba in series	Video generation
Tmamba (Zhu et al., 2024)	Parallel Mamba and linear-Transformer branches with T–M interaction and attention-level multimodal fusion	Multi-modality image fusion
A2Mamba (Lou et al., 22 Jul 2025)	Multi-scale attention maps injected into an attention-augmented SSM at hidden-state level	Visual recognition
TransMamba (Li et al., 31 Mar 2025)	Shared-parameter switching between attention and SSM at layer-specific TransPoints	Language modeling

Several of these systems clarify what MAIT is not. Dimba, for example, is explicitly a modular interleaving of attention and Mamba within a latent diffusion denoiser; in the reported models, the attention-to-Mamba ratio is $$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$8, and conditional text information is injected through cross-attention rather than by making Mamba itself the multimodal fusion mechanism (Fei et al., 2024). MAT, by contrast, assigns Mamba to long-range/global temporal modeling and Transformer attention to short-range/local dependency extraction, but leaves the exact fusion mechanism under-specified; it is best understood as a multi-scale complementary hybrid rather than a canonical block design (Zhang et al., 2024). Matten sharpens the division of labor for video by placing spatial attention, temporal attention, and global-sequence bidirectional Mamba in series, with the explicit claim that attention is superior for local video content and Mamba for global video content (Gao et al., 2024).

Other papers move beyond stack-level coexistence. Tmamba uses a dual-branch image-fusion network in which a linear Transformer branch carries channel information and a Mamba branch carries position information, with explicit T–M interaction between them and cross-modal attention at the attention level (Zhu et al., 2024). A2Mamba goes further by proposing a Multi-scale Attention-augmented State Space Model in which attention maps are injected into the state-space hidden-state path; its A2SSM performs a cross-attention-like aggregation of SSM hidden states using multi-scale attention maps, yielding a genuinely co-designed mixer rather than a simple stack (Lou et al., 22 Jul 2025). TransMamba unifies attention and Mamba through shared parameter matrices $$s_t = \left\{ s_{1,t}, s_{2,t}, \dots, s_{N,t} \right\}, \quad s_{i,t} \in \mathbb{R}^{d_{\text{in}}}.$$9 and $h_{i}^0 = x_{i, t} W_s + p_t + b_i.$0, introduces a parameter-free Memory Converter to map attention outputs into SSM-compatible states, and uses layer-specific TransPoint schedules so a single layer can switch from attention to SSM across token positions (Li et al., 31 Mar 2025).

Domain-specific work broadens the picture further. Mammo-Mamba modifies the deeper stages of MambaVision into sequences of SecMamba blocks and self-attention blocks governed by Sequential Mixture of Experts, while MambaCAFU places Mamba inside attention-guided fusion modules and uses a three-branch encoder that integrates CNNs, Transformers, and Mamba-based Attention Fusion (Bayatmakou et al., 23 Jul 2025, Bui et al., 4 Oct 2025). MAAT, an anomaly-detection model, retains association-discrepancy attention in the main path, adds a Mamba-based reconstruction refinement path, and fuses the two with a learned gate (Sellam et al., 11 Feb 2025). Differential Mamba shows a different form of “integration”: it does not add explicit attention blocks, but transfers a Transformer-derived differential design principle into Mamba’s implicit attention-like operator at the whole-block level (Schneider et al., 8 Jul 2025). Speech work offers another variant: replacing the MHSA sublayer inside Transformer- and Conformer-style speech encoders with bidirectional Mamba can outperform pure Mamba and even standard Conformer baselines on several tasks, indicating that the hybrid design space also includes scaffold-preserving replacement strategies (Zhang et al., 2024).

Taken together, these papers suggest that “Mamba-attention integration” is not a single architecture but a taxonomy. Serial layer composition, parallel interactive branches, stage-wise expert mixtures, hidden-state-level co-design, and dynamic switching all appear as viable patterns. The original MAIT corresponds to the most modular point in that design space: a configurable serial encoder whose layer types are chosen according to task scale.

7. Limitations, misconceptions, and open questions

The original MAIT formulation has several explicit limitations. It assumes global system information, addresses trajectory optimization rather than joint resource scheduling, does not provide an explicit MAIT fusion operator, and does not give a principled algorithm for determining exact layer ordering or the attention-to-Mamba ratio beyond Agentic RAG recommendation. It also leaves key implementation details unspecified, including hidden dimension $h_{i}^0 = x_{i, t} W_s + p_t + b_i.$1, number of heads $h_{i}^0 = x_{i, t} W_s + p_t + b_i.$2, MLP width, optimizer, learning rate, and exact layer ordering (Jiang et al., 22 Aug 2025).

More broadly, the surrounding hybrid literature remains uneven in mathematical specification. Dimba is valuable as a concrete large-scale image-generation instantiation of an alternating attention–Mamba backbone, but it does not print full block equations for attention, cross-attention, or Mamba, and it leaves several diffusion-training details unstated (Fei et al., 2024). MAT argues persuasively for long-short range decomposition, yet does not fully define the internal MAT block or its fusion operator (Zhang et al., 2024). A2Mamba explicitly criticizes prior hybrids for “simple stacking” and treats deep integration as an open problem, even while proposing a stronger hidden-state-level coupling (Lou et al., 22 Jul 2025). Differential Mamba also states that there is no rigorous theory for why differential design works in either Transformers or Mamba, despite strong empirical results (Schneider et al., 8 Jul 2025).

A recurring misconception is that any Mamba–attention hybrid must culminate in a single fused token mixer. The comparative literature does not support that claim. Some systems are intentionally modular. The original MAIT, Dimba, and Matten use serial decomposition; Tmamba uses parallel interactive branches; TransMamba uses within-layer switching; MAAT uses task-specialized dual paths with gated fusion (Jiang et al., 22 Aug 2025, Gao et al., 2024, Zhu et al., 2024, Li et al., 31 Mar 2025). This suggests that the decisive design variable is not the mere presence of both mechanisms, but where their respective inductive biases are inserted into the computation graph.

Another unresolved question concerns scheduling and scaling. TransMamba shows that hybridization can be improved by broad-range, fine-grained layer-specific TransPoint schedules rather than a globally shared switch, but its scheduling remains manual and scaling is only explored up to 1.5B (Li et al., 31 Mar 2025). In diffusion and forecasting, the evidence for attention–Mamba complementarities is strong, yet matched ablations over ratio, ordering, or complexity-quality trade-offs remain limited (Fei et al., 2024, Zhang et al., 2024). A plausible implication is that future MAIT research will be shaped less by the question of whether attention and Mamba can coexist—which is already well supported—and more by the questions of where to place the interface, how to transfer state across it, and how to preserve the efficiency benefits of Mamba without discarding the explicit relational bias of attention.