The framework, called RL-Axial and described in the journal Cognitive Computation, attacks a stubborn inefficiency at the heart of modern segmentation models. Convolutional encoder-decoder architectures such as U-Net excel at capturing local texture through skip connections and multi-scale aggregation, but their limited receptive fields leave them weak at modeling global context. Transformer-style models fix that weakness with self-attention, which lets every pixel consult every other pixel, yet the computational cost of that conversation grows quadratically with image size, a serious problem for the gigapixel tiles typical of aerial surveys. Worse, transformers apply the same attention budget everywhere: easy regions of sky or uniform pavement receive the same expensive global computation as tangled street intersections full of small, ambiguous objects. The result is redundant computation on simple areas and, paradoxically, under-modeling of the hard cases that matter most. RL-Axial’s central idea is to allocate global modeling capacity on demand rather than imposing a single, fixed computation path on every sample.

The architecture rests on two pillars. The first is a Cross-Modal Gated Fusion head, or CMGF, which handles the delicate business of merging RGB and DSM features. Naive fusion strategies, whether early stacking of co-registered inputs or simple feature concatenation, are fragile in practice: resolution mismatches, imperfect geometric correction, and different noise statistics between sensors introduce subtle misalignments that static fusion can amplify into boundary blurring and semantic drift. CMGF instead constructs an explicit cross-modal state by concatenating the RGB features, the DSM features, and their difference, a term that makes the module directly sensitive to inconsistencies between the two modalities. A small predictor built from a 3×3 convolution, group normalization, ReLU, and a 1×1 convolution then produces fusion logits, and a softmax over the two modalities yields per-pixel weight maps. The fused feature is a weighted sum, so every single pixel can lean toward appearance or geometry as the evidence dictates, and unreliable modality responses are suppressed locally rather than globally.

The second pillar is the element that gives the framework its name: a controllable axial-attention operator with four distinct execution modes. Full two-dimensional self-attention scales as the square of the number of pixels, but axial attention decomposes the problem into row-wise and column-wise one-dimensional attention, reducing complexity to a far more manageable order that suits dense prediction. RL-Axial goes further by making the operator switchable: for any given image patch, attention can be turned off entirely, applied along rows only, applied along columns only, or applied along both axes. The off mode is not a mere convenience; it provides a stable fallback during training and inference, ensuring the network never depends on attention to function. Each axial branch follows a projection-residual-normalization pattern, with its output added back to the input and normalized position-wise, a design that helps avoid the distribution shifts attention branches can introduce.

Deciding which mode to use is framed as a reinforcement learning problem. A lightweight policy network, a two-layer multilayer perceptron, receives a globally average-pooled state vector from each patch’s fused representation and maps it to a categorical distribution over the four modes. Crucially, states are never averaged across the mini-batch, so different patches in the same batch can select different modes independently. Training the selector uses what the authors call a group-normalized policy gradient, an objective that borrows the relative-reward normalization idea popularized by group-relative methods but deliberately avoids the probability-ratio clipping and reference-policy regularization of PPO and GRPO. For each patch, a reference prediction is generated with attention off, then K candidate actions are sampled and evaluated. Each candidate’s reward is the difference in soft-IoU between its prediction and the reference, and these within-sample rewards are normalized into advantages. An entropy bonus, linearly annealed over training, encourages broad exploration early and more deterministic choices later.

Training proceeds in two carefully separated stages. In the first, axial attention is disabled and the dual modality encoders, the CMGF module, the decoder, and the segmentation head are trained jointly with standard cross-entropy loss using stochastic gradient descent. In the second stage, the encoders and fusion head are frozen, and the policy, axial operator, decoder, and segmentation head continue training with Adam. Gradient paths are explicitly partitioned: the policy loss, built from detached rewards, updates only the selector parameters, while the segmentation loss updates only the axial operator and decoder, never differentiating through the discrete action choice. This separation keeps the non-differentiable scheduling decision cleanly isolated from supervised learning. At inference time, per-sample decisions do not force batch size one; samples selecting the same mode are grouped, the corresponding axial branch is executed for each group, and the original order is restored before decoding, preserving patch-wise adaptivity within ordinary batched tensor operations.

Evaluated on the ISPRS Vaihingen and Potsdam benchmarks, the standard proving grounds for urban remote sensing segmentation, RL-Axial delivered strong single-run results. On Vaihingen, aerial imagery at 9-centimeter ground sampling distance paired with a DSM, the model reached 84.58 percent mean intersection over union, along with 92.73 percent overall accuracy and 91.50 percent mean F1. On the larger Potsdam dataset, with 6000-by-6000-pixel orthophotos at 5-centimeter resolution, it achieved 86.73 mIoU, 91.73 overall accuracy, and 92.74 mean F1. Those mIoU figures sit numerically 0.35 and 0.53 points above the literature-reported values for FTransUNet, a strong multimodal transformer baseline. The authors are notably careful here: because baseline numbers come from their original publications with source-specific evaluation protocols, the comparison is descriptive rather than a controlled head-to-head test, and no claims of statistical superiority are made. Qualitatively, the predictions show cleaner region layouts, sharper boundaries at class transitions, fewer boundary breaks, and more coherent object shapes, particularly in the dense, complex urban layouts of Potsdam.

The ablation studies reveal exactly where the gains come from, and the honesty of the decomposition is refreshing. An optical-only backbone reaches 82.14 mIoU; naive RGB-DSM fusion adds 1.23 points, confirming the value of elevation even without learned gating. Replacing naive fusion with CMGF contributes a further 0.54 points, the best fixed axial mode adds 0.40, and adaptive scheduling contributes the final 0.27. In a controlled comparison where every fixed attention branch and every alternative selector shared the same backbone, initialization, and training budget, the learned selector beat the best fixed native mode, a fixed guided focal-axial adaptation, random scheduling, a supervised action classifier, and a Gumbel-Softmax selector, while an oracle that exhaustively evaluates all four actions per patch sets an upper bound of 85.07 mIoU, leaving a 0.49-point gap. The reward design matters too: optimizing the change in mIoU directly outperformed rewards based on F1 or overall accuracy, evidence that aligning the learning signal with the target metric pays off. Group size ablations showed performance rising from K equals 2 to a peak at K equals 8, the value adopted.

Perhaps the most revealing result is what the learned policy actually does. Rather than collapsing onto a single favorite mode, it uses all four: row-plus-column attention for 30 percent of test patches, the off path for 29 percent, row-only for 23 percent, and column-only for 18 percent. Profiling each branch on a single RTX 3090 GPU, the authors estimate that this distribution yields roughly 40.125 gigaflops and 28.8 milliseconds per patch, a branch-only reduction of 7.55 percent in FLOPs and 10.51 percent in latency compared with always running full row-plus-column attention, though they stress these figures exclude selector overhead and are not end-to-end speedup claims. The broader significance lies in the paradigm: axial attention itself is not new to remote sensing, and the authors explicitly disclaim that novelty. What is new is treating global context as a schedulable resource, learned through direct feedback from the segmentation metric itself. As earth observation constellations multiply and the appetite for fine-grained urban mapping grows, frameworks that spend computation where it matters, and skip it where it does not, may define the next generation of geospatial AI.

Subject of Research: Reinforcement-learning-based adaptive axial attention scheduling for multimodal RGB and digital surface model segmentation in high-resolution remote sensing

Article Title: RL-Axial: Learning When to Compute Global Context for Multimodal Segmentation

Article References: Xing, D., Yang, H., Zhang, J., Liu, P., & Wang, Y. (2026). RL-Axial: Learning When to Compute Global Context for Multimodal Segmentation. Cognitive Computation, 18(1), Article 113. https://doi.org/10.1007/s12559-026-10661-z

Image Credits: AI Generated

DOI: 10.1007/s12559-026-10661-z

Keywords: remote sensing, semantic segmentation, multimodal fusion, reinforcement learning, axial attention, digital surface model, ISPRS Vaihingen, ISPRS Potsdam, policy gradient, dynamic computation, urban mapping, transformer