How Mamba State Space Models Are Reshaping Multimodal AI Fusion

The problem that pushed researchers past plain attention

Attention became the default tool for sequence and spatial reasoning for a good reason. It lets every element look at every other element and decide, dynamically, what matters. That flexibility is also its cost. Comparing every pair of positions scales quadratically with sequence length, and once you are working with hyperspectral imagery carrying over a hundred bands, or high resolution feature maps, that cost adds up fast, especially on hardware that has to run in the field or on a tight budget.

State space models take a different approach, one borrowed from classical control theory. Instead of asking a token to look back at everything that came before it, a state space model keeps a compressed hidden state and updates it step by step as it moves through a sequence. Recent selective variants let that update depend on the actual input at each step rather than following a fixed rule, which is what allows a state space block to behave a bit like attention, deciding what to keep and what to forget, while keeping the underlying computation close to linear in the length of the sequence.

The case study we are using here did not set out to write a new architecture paper. It set out to solve a very specific applied problem, matching satellite hyperspectral imagery with satellite LiDAR data to map eight tree species across a mountainous subtropical forest in Anhui province. But solving that problem forced a set of architectural decisions that turn out to be a genuinely useful lens for understanding when a state space block is the right tool.

Why this particular fusion problem is harder than it sounds

Two data sources rarely agree with each other cleanly. In this study, the hyperspectral imagery comes from the ZY1 02E satellite sensor and offers continuous spatial coverage with 151 usable spectral bands after noisy shortwave infrared bands were removed. The structural data comes from GEDI, a spaceborne LiDAR system that samples the forest through narrow footprints rather than a continuous sweep. One modality is dense and spectrally rich. The other is sparse and structurally rich. Forcing them into a shared representation without losing what makes each one useful is the entire design challenge.

The researchers also point out a subtler issue with how hyperspectral data is usually compressed before it reaches a model. Standard dimensionality reduction methods such as PCA tend to keep whatever explains the most variance, which is not the same thing as keeping whatever is most discriminative for telling species apart. A faint pigment signal that barely moves the variance needle might be exactly the signal a botanist would use to separate two similar looking species, and a generic compression step can throw it away before the classifier ever sees it.

How the fusion pipeline is actually built

The full model, MUFNet, is organized around three stages, and it is worth walking through them in order because each one solves a distinct piece of the mismatch problem described above.

Stage one, a task aware spectral projector

Rather than reducing the 151 hyperspectral bands with an offline method like PCA, the authors train a learnable projector called the manifold orthogonal spectral projector, or MOSP. It is a simple one by one convolution, but it is trained end to end alongside the classifier and constrained to stay close to orthogonal, which keeps the compressed 30 dimensional representation from collapsing redundant information into itself. A shared positional convolution is then applied to both the hyperspectral and LiDAR streams using the same weights, which gives both modalities a common spatial coordinate system before anything else happens.

Stage two, gating followed by a state space refinement

This is the part most relevant to anyone studying Mamba style architectures directly. The fusion stage itself has two steps. First, a pixelwise gate called early stage adaptive gating, or EAG, produces a soft weighting between the hyperspectral and LiDAR features at every pixel. It is a practical fix for a real world problem. Hyperspectral signal quality drops in shadowed terrain, while LiDAR structural signal drops where the canopy is sparse, so the gate lets each pixel lean on whichever modality is more trustworthy at that location.

Second comes the module the authors call MCI-SSM, mid stage cross modal interaction via state space modeling. It runs cross attention first, using hyperspectral features as queries and LiDAR features as keys and values, then feeds the attended output into a lightweight Mamba style state space block for further refinement. The ordering matters. Cross attention establishes which spatial locations in one modality correspond to which locations in the other. The state space block then sweeps across the resulting feature map to capture long range and direction dependent structure, the kind of elongated pattern you would expect from an irregular tree crown or a ridge line, which a small convolution kernel would struggle to see in one pass.

Here \(S(\cdot)\) is the state space refinement itself, applied after the attended features are concatenated with the original hyperspectral stream and projected back to the working channel dimension. The practical effect is that the model gets two different kinds of reasoning stacked on top of each other rather than relying on either one alone.

Stage three, uncertainty aware aggregation and routing

The final stage does not just average the early, mid, and late stage predictions. It first estimates an aleatoric uncertainty value for each prediction branch, using the log variance of a small evidential head, then feeds both the pooled features and the uncertainty scores into a lightweight routing network that decides how much weight each stage deserves for a given input. The authors call this meta cognitive dynamic routing, and the name is a fair description of what it does. The model is, in a limited but real sense, weighing its own confidence at each stage before deciding whose vote counts more.

What the numbers actually show

The researchers compared MUFNet against six other multimodal fusion models on the same labeled dataset, which covers 245763 pixels across eight tree species collected from field surveys conducted across four separate visits in 2023 and 2024.

One comparison stands out for anyone specifically interested in state space architectures. MSFMamba is itself a Mamba based fusion model built for multisource remote sensing classification, so it is the most direct architectural cousin in this table. MUFNet still beats it by close to five points of overall accuracy and nearly nine points of kappa, while using roughly twenty four times fewer parameters. That gap is not really an argument that one state space design beats another in the abstract. It is a reminder that where you place the state space block, and what you pair it with, matters as much as the fact that you used one at all.

The ablation results make the contribution of each stage explicit rather than leaving it to intuition. Using only the mid stage branch, essentially the cross attention plus state space module on its own, produced 90.07 percent overall accuracy and 87.05 percent kappa. Adding the late stage evidential aggregation on top pushed that to 92.16 percent and 89.48 percent. The full three stage model reached 93.34 percent and 91.17 percent. A separate component removal test found that dropping the LiDAR input entirely cost 2.87 points of accuracy, dropping MOSP cost 1.36 points, dropping the Mamba block cost 2.73 points, and dropping the routing network cost 1.11 points. None of these numbers is dramatic on its own, which is exactly the point. The performance comes from several modest, well targeted design choices stacking together rather than from one dominant trick.

Calibration, not just accuracy

A model that is right ninety three percent of the time but confidently wrong the rest of the time is a liability in any application where a human downstream has to trust the output. The authors report an expected calibration error of 0.0073 for MUFNet, meaning its predicted confidence scores track its actual accuracy closely across the full range of confidence levels. They also show a risk coverage curve where the lowest uncertainty predictions carry the lowest error rate and that error rate climbs steadily as more of the less confident predictions are included. That is exactly the behavior you want from an uncertainty estimate, and it is a direct product of building the log variance heads into the architecture rather than bolting on a post hoc calibration step afterward.

Where this generalizes beyond forestry

Strip away the tree species labels and the GEDI footprints, and what remains is a fairly general recipe. Two modalities disagree with each other in coverage and reliability across space or time. A learnable, task specific projector compresses the high dimensional one without throwing away subtle signal. A pixelwise or tokenwise gate handles the cases where one modality is simply missing or degraded. Cross attention aligns what corresponds between the two. A state space block then sweeps across the aligned result to catch dependencies that span a wide area or follow a particular direction. An uncertainty aware head at the end lets the model, and the people relying on it, know when to trust a given output less.

That pattern maps fairly directly onto other domains already dealing with the same tension between attention and state space computation, including video paired with audio, wearable sensor streams paired with sparse clinical events, and robotics platforms fusing camera frames with lower frequency lidar or radar sweeps. The specific numbers in this paper will not transfer, but the architectural sequencing, gate first, attend second, sweep with a state space block third, is a genuinely portable idea.

A minimal, runnable implementation

The code below is a compact, dependency light reimplementation of the ideas above. It is not a copy of the authors’ training pipeline, which was not published alongside the paper, but it follows the same architectural sequence, a spectral projector with an orthogonality penalty, a pixelwise gate, cross attention followed by a simplified selective state space block, and an uncertainty aware aggregation head with a small routing network. It runs end to end on random dummy tensors so you can see the full forward and backward pass without needing real satellite data.

# mamba_fusion_lite.py
# A compact, dependency light reimplementation of a Mamba style multimodal
# fusion pipeline inspired by the three stage MUFNet design described above.
# This is an independent educational implementation, not the authors' code.

import torch
import torch.nn as nn
import torch.nn.functional as F


class SpectralProjector(nn.Module):
    """Learnable 1x1 spectral projector with an orthogonality penalty,
    playing the role of MOSP in the source architecture."""

    def __init__(self, in_bands: int, proj_dim: int = 30):
        super().__init__()
        self.proj = nn.Conv2d(in_bands, proj_dim, kernel_size=1)
        self.proj_dim = proj_dim

    def forward(self, x):
        return self.proj(x)

    def orthogonal_loss(self):
        w = self.proj.weight.view(self.proj_dim, -1)
        gram = w @ w.t()
        eye = torch.eye(self.proj_dim, device=w.device)
        return F.mse_loss(gram, eye)


class SharedPositionalConv(nn.Module):
    """Same depthwise conv and group norm applied to both streams to build
    a shared spatial coordinate system before fusion."""

    def __init__(self, channels: int):
        super().__init__()
        self.dw = nn.Conv2d(channels, channels, kernel_size=3, padding=1, groups=channels)
        self.gn = nn.GroupNorm(4, channels)

    def forward(self, x):
        return F.gelu(self.gn(self.dw(x)))


class AdaptiveGate(nn.Module):
    """Pixelwise adaptive gate between two modalities, playing the role of
    EAG in the source architecture."""

    def __init__(self, channels: int):
        super().__init__()
        self.h_pref = nn.Conv2d(channels, 1, kernel_size=1)
        self.l_pref = nn.Conv2d(channels, 1, kernel_size=1)
        self.mix = nn.Conv2d(channels, channels, kernel_size=1)

    def forward(self, h, l):
        h_score = torch.sigmoid(self.h_pref(h))
        l_score = torch.sigmoid(self.l_pref(l))
        stacked = torch.cat([h_score, l_score], dim=1)
        alpha = F.softmax(stacked, dim=1)[:, 0:1]
        fused = self.mix(alpha * h + (1 - alpha) * l)
        entropy = -(alpha * torch.log(alpha + 1e-6) +
                    (1 - alpha) * torch.log(1 - alpha + 1e-6)).mean()
        return fused, entropy


class SimpleSelectiveSSM(nn.Module):
    """A small, readable selective state space block. It scans a flattened
    spatial sequence position by position, updating a hidden state with
    input dependent gating, in the spirit of a Mamba style block, without
    depending on an external CUDA specific kernel library."""

    def __init__(self, channels: int, state_dim: int = 16):
        super().__init__()
        self.channels = channels
        self.state_dim = state_dim
        self.in_proj = nn.Linear(channels, state_dim)
        self.gate_proj = nn.Linear(channels, state_dim)
        self.out_proj = nn.Linear(state_dim, channels)
        self.decay = nn.Parameter(torch.zeros(state_dim))

    def forward(self, x):
        # x has shape B, N, C where N is a flattened spatial sequence
        b, n, c = x.shape
        u = self.in_proj(x)
        gate = torch.sigmoid(self.gate_proj(x))
        decay = torch.sigmoid(self.decay).view(1, 1, -1)
        state = torch.zeros(b, self.state_dim, device=x.device, dtype=x.dtype)
        outputs = []
        for t in range(n):
            state = decay.squeeze(1) * state + gate[:, t, :] * u[:, t, :]
            outputs.append(state)
        seq = torch.stack(outputs, dim=1)
        return self.out_proj(seq)


class CrossModalStateSpaceFusion(nn.Module):
    """Structure guided cross attention followed by a selective state space
    refinement, playing the role of MCI-SSM in the source architecture."""

    def __init__(self, channels: int, heads: int = 4):
        super().__init__()
        self.attn = nn.MultiheadAttention(channels, heads, batch_first=True)
        self.ssm = SimpleSelectiveSSM(channels * 2, state_dim=channels)
        self.proj = nn.Linear(channels * 2, channels)
        self.norm = nn.LayerNorm(channels)

    def forward(self, h_seq, l_seq):
        attn_out, _ = self.attn(h_seq, l_seq, l_seq)
        concat = torch.cat([h_seq, attn_out], dim=-1)
        refined = self.ssm(concat)
        return self.norm(self.proj(refined) + h_seq)


class UncertaintyHead(nn.Module):
    """Predicts classwise logits and a log variance term, playing the role
    of the UAH block inside LHEA in the source architecture."""

    def __init__(self, in_dim: int, n_classes: int, temperature: float = 1.0):
        super().__init__()
        self.mu = nn.Linear(in_dim, n_classes)
        self.log_var = nn.Linear(in_dim, n_classes)
        self.t = temperature

    def forward(self, x):
        mu = self.mu(x) / self.t
        log_var = torch.clamp(self.log_var(x), min=-4.0, max=2.0)
        return mu, log_var


class MambaFusionLite(nn.Module):
    """End to end model tying every block above together into a compact
    three stage fusion network."""

    def __init__(self, hsi_bands: int, lidar_channels: int,
                 n_classes: int, proj_dim: int = 30):
        super().__init__()
        self.hsi_proj = SpectralProjector(hsi_bands, proj_dim)
        self.lidar_proj = nn.Conv2d(lidar_channels, proj_dim, kernel_size=1)
        self.shared_pos = SharedPositionalConv(proj_dim)
        self.gate = AdaptiveGate(proj_dim)
        self.cross_ssm = CrossModalStateSpaceFusion(proj_dim)
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.early_head = UncertaintyHead(proj_dim, n_classes)
        self.mid_head = UncertaintyHead(proj_dim, n_classes)
        self.router = nn.Sequential(
            nn.Linear(proj_dim * 2 + 2, 64),
            nn.ReLU(),
            nn.Linear(64, 2),
        )

    def forward(self, hsi, lidar):
        h = self.shared_pos(self.hsi_proj(hsi))
        l = self.shared_pos(self.lidar_proj(lidar))

        fused_early, gate_entropy = self.gate(h, l)

        b, c, hh, ww = h.shape
        h_seq = h.flatten(2).transpose(1, 2)
        l_seq = l.flatten(2).transpose(1, 2)
        mid_seq = self.cross_ssm(h_seq, l_seq)
        fused_mid = mid_seq.transpose(1, 2).reshape(b, c, hh, ww)

        e_vec = self.pool(fused_early).flatten(1)
        m_vec = self.pool(fused_mid).flatten(1)

        mu_e, logvar_e = self.early_head(e_vec)
        mu_m, logvar_m = self.mid_head(m_vec)

        s_e = logvar_e.mean(dim=1, keepdim=True)
        s_m = logvar_m.mean(dim=1, keepdim=True)
        route_in = torch.cat([e_vec, m_vec, s_e, s_m], dim=1)
        weights = F.softmax(self.router(route_in), dim=1)

        final_logits = weights[:, 0:1] * mu_e + weights[:, 1:2] * mu_m
        return {
            "logits": final_logits,
            "gate_entropy": gate_entropy,
            "orth_loss": self.hsi_proj.orthogonal_loss(),
            "uncertainty": 0.5 * (s_e + s_m),
        }


def fusion_loss(outputs, targets, orth_weight=5e-3, entropy_weight=5e-3):
    ce = F.cross_entropy(outputs["logits"], targets)
    total = ce + orth_weight * outputs["orth_loss"] - entropy_weight * outputs["gate_entropy"]
    return total, ce


def train_one_epoch(model, optimizer, hsi, lidar, labels, batch_size=8):
    model.train()
    n = hsi.shape[0]
    perm = torch.randperm(n)
    running_loss = 0.0
    for i in range(0, n, batch_size):
        idx = perm[i:i + batch_size]
        optimizer.zero_grad()
        out = model(hsi[idx], lidar[idx])
        loss, _ = fusion_loss(out, labels[idx])
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * idx.shape[0]
    return running_loss / n


def evaluate(model, hsi, lidar, labels):
    model.eval()
    with torch.no_grad():
        out = model(hsi, lidar)
        preds = out["logits"].argmax(dim=1)
        acc = (preds == labels).float().mean().item()
        mean_uncertainty = out["uncertainty"].mean().item()
    return acc, mean_uncertainty


if __name__ == "__main__":
    # Smoke test on random dummy data, matching the shapes described in
    # the study, 151 hyperspectral bands and a small LiDAR feature stack,
    # over a small patch size and a handful of species classes.
    torch.manual_seed(0)
    n_samples, hsi_bands, lidar_channels, patch, n_classes = 32, 151, 15, 11, 8

    hsi_dummy = torch.randn(n_samples, hsi_bands, patch, patch)
    lidar_dummy = torch.randn(n_samples, lidar_channels, patch, patch)
    labels_dummy = torch.randint(0, n_classes, (n_samples,))

    model = MambaFusionLite(hsi_bands, lidar_channels, n_classes)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-7)

    for epoch in range(3):
        loss = train_one_epoch(model, optimizer, hsi_dummy, lidar_dummy, labels_dummy)
        acc, unc = evaluate(model, hsi_dummy, lidar_dummy, labels_dummy)
        print(f"epoch {epoch} loss {loss:.4f} train acc {acc:.3f} mean log var {unc:.3f}")

Conclusion

The headline result here, 93.34 percent overall accuracy on an eight species satellite classification task, is a solid number, but the more durable contribution is architectural. The study shows a specific, load bearing role for a state space block inside a larger fusion pipeline, sitting after cross attention rather than competing with it, tasked specifically with catching long range and direction dependent structure that a small convolution or a single attention pass would miss. That is a more precise claim than the general excitement around Mamba style architectures usually offers, and precision is exactly what makes it useful to other builders.

The conceptual shift worth carrying away is that efficiency and capability are not in tension here so much as they are the product of good sequencing. Nobody swapped attention out for a state space block and called it a day. The gate decides which modality to trust. Attention decides what corresponds to what. The state space block decides how far that correspondence should reach across space. Each piece does one job, and the ablation numbers show that removing any single piece costs something specific and measurable rather than something vague.

That sequencing logic is not tied to satellites or trees. Anywhere two data sources disagree in coverage and reliability, and anywhere a model needs to reason about structure that spans a wide area or follows a direction, the same pattern of gate, align, then sweep is worth testing before reaching for a bigger, more generic backbone. It is a cheap experiment relative to the parameter and compute savings on offer here, 0.400 million parameters and 0.086 GFLOPs is a tiny footprint by current standards.

The honest limitations matter too. A single region, eight species, and no published cross region test mean the specific numbers should be read as a strong proof of concept rather than a settled benchmark. Anyone adapting this pattern to a new domain should expect to retune the gating and routing components rather than assuming the same weights or thresholds will transfer directly.

Where this goes next is fairly predictable. Expect more fusion papers across medical imaging, video, and robotics to adopt some version of the gate, align, then sweep pattern, because the ablation evidence for it is now hard to argue with. The bigger open question is whether the same efficiency gains hold once the label space grows from eight classes to hundreds, and whether the uncertainty calibration remains this tight once a model is deployed somewhere far from its training region. Those are the experiments worth watching for next.

Frequently asked questions

Related reading