Transformer Architecture with Minimal Inference Latency for Multi-Modal Wireless Networks

1. Problem Statement

Next-generation wireless networks (e.g., 6G) require situational awareness to execute critical tasks like beamforming and proactive handover in dynamic environments (e.g., vehicular networks). These tasks rely on multi-modal data (LiDAR, cameras, radar, GPS, RSSI) fused via Transformer models.

However, existing Transformer-based approaches face a critical bottleneck:

• Quadratic Complexity: The self-attention mechanism scales as $O(N^2)$, where $N$ is the number of tokens. Multi-modal inputs significantly increase $N$, leading to high inference latency and memory usage.
• Real-Time Failure: High latency causes predictions to become obsolete before execution. For example, in a proactive handover scenario, if the model predicts a blockage after the line-of-sight (LoS) is already blocked, the link fails.
• Inefficiency of Current Solutions: Existing token pruning/merging methods (e.g., ToMe, MoD) are designed for single-modal tasks (vision or language) and do not optimize the token count jointly with model performance under a specific computational budget for multi-modal wireless scenarios.

Goal: Design a framework that processes only the most important tokens to minimize inference latency and FLOPs while maintaining task accuracy under a target computational budget.

---

2. Methodology

The authors propose a novel Fast Multi-Modal Transformer Inference Framework consisting of three core components:

A. Modality-Specific Tokenization

To bridge the gap between different data types, the framework uses lightweight tokenizers to project all modalities into a shared embedding dimension ($d$):

• Images & Point Clouds: Processed via tiny CNNs (ResNet blocks + 1D Conv) to extract semantic features and project them to tokens.
• Radar: FMCW radar cubes are downsampled and projected via linear layers.
• GPS & RSSI: Scalar values are projected into single tokens via linear layers.
• Result: All modalities are concatenated into a unified token sequence $X$, allowing the Transformer to process them in the same space.

B. Token-Level Routing (Token Router)

Instead of processing all tokens, the framework employs a Token Router in every Transformer block:

• Mechanism: A lightweight router (two linear layers with GELU activation) predicts an importance score ($s_l$) for each token.
• Selection: Only the Top-$K_l$ tokens are passed to the subsequent Transformer block (Self-Attention + MLP).
• Bypass: Non-selected tokens bypass the block entirely, reducing the attention complexity from $O(N^2)$ to $O(K_l^2)$.
• Training: The router is trained end-to-end using a gradient path generated by weighting the outputs of selected vs. bypassed tokens.

C. Differentiable Keep Ratio Optimization

To determine the optimal number of tokens ($K_l$) for each layer without manual hyperparameter tuning:

• Learnable Parameter: A continuous keep ratio ($r_l \in [0, 1]$) is defined for each layer, representing the fraction of tokens to keep.
• Differentiable Approximation: Since the Top-$K$ operation is non-differentiable, the authors use a piecewise relaxation. They round $r_l$ to the two nearest integers ($K_{up}$ and $K_{down}$), perform two forward passes, and linearly combine the outputs based on the distance to the integers. This creates a valid gradient path for $r_l$.
• Constraint Optimization: An optimization problem is formulated to minimize task loss subject to a target FLOPs budget ($\gamma$). A Mean Squared Error (MSE) penalty is added to the loss function to enforce the average keep ratio ($r_{avg}$) to match a target ratio ($\gamma'$), effectively controlling the computational budget.

---

3. Key Contributions

1. Novel Framework: The first multi-modal Transformer inference framework specifically designed to balance performance and latency for wireless communication tasks by processing only important tokens.
2. Joint Optimization: Formulates and solves an optimization problem to jointly learn the model parameters, token routers, and the optimal number of tokens per layer under a strict computational budget.
3. Differentiable Token Counting: Introduces a trainable keep ratio with a piecewise relaxation technique, enabling end-to-end optimization of token counts without heuristic rules or discrete approximations like Gumbel-Softmax.
4. Real-World Validation: Developed a new real-world multi-modal handover dataset collected on a testbed (Tokyo Institute of Science), validating the framework's ability to perform proactive handovers before signal blockage occurs.

---

4. Experimental Results

The framework was evaluated on two public mmWave beamforming datasets (DeepSense 6G and CARLA-MATLAB) and the authors' custom real-world handover dataset.

A. Efficiency Gains (DeepSense 6G)

Compared to a baseline processing all tokens (MHA):

• Inference Latency: Reduced by 86.2% (at $\gamma'=30\%$).
• GPU Memory: Reduced by 35%.
• FLOPs: Reduced by 80%.
• Accuracy: Negligible loss (only 0.12% drop in Top-1 accuracy).
• Note: At higher keep ratios (e.g., 50-70%), the proposed method actually achieved higher accuracy than the full-baseline, suggesting the router acts as a regularizer by removing noise.

B. Real-World Handover Performance

• Proactive Handover: In a testbed scenario with a moving vehicle and blockage:
  • Baseline (MHA): Inference took >300ms. The decision was made after the blockage occurred, resulting in a 70% degradation in Received Signal Strength (RSSI).
  • Proposed Framework: Inference was 83% faster. It successfully predicted the blockage and initiated handover before the link was lost, maintaining stable RSSI.
• Efficiency: At $\gamma'=40\%$, latency was reduced by 83.2% with only a 0.61% accuracy loss compared to the baseline.

C. Ablation Studies

• Tokenizers: Essential for reducing modality gaps; freezing them caused significant accuracy drops.
• Token Routers: Critical for identifying important tokens; random selection led to poor performance.
• Learnable Keep Ratios: Uniform token reduction (heuristic) was inferior to layer-specific learned ratios. Early layers were found to require more tokens than deeper layers.

---

5. Significance

This paper addresses a critical barrier to deploying AI in 6G wireless networks: latency. By demonstrating that Transformers can be made 86% faster with negligible accuracy loss through intelligent token selection, the authors enable:

• Real-time Situational Awareness: Systems can now react to dynamic changes (vehicles, blockages) within the coherence time of the channel.
• Edge Deployment: The drastic reduction in FLOPs and memory makes it feasible to run complex multi-modal models on resource-constrained edge devices (e.g., RSUs, vehicles).
• Proactive Network Management: The ability to predict blockages before they happen allows for seamless, zero-interruption handovers, a key requirement for ultra-reliable low-latency communication (URLLC).

The work establishes a new paradigm for efficient multi-modal learning in wireless systems, moving beyond "process everything" to "process what matters."