Digital Twin-Assisted Fast Beam Alignment for mmWave MIMO Networks

Digital Twin-Assisted Fast Beam Alignment for mmWave MIMO Networks

arXiv:2611.01107v1 [cs.IT] 31 Oct 2026

Muhammad Kamran Saeed and Ashfaq Khokhar, Fellow, IEEE

Abstract—Millimeter-wave (mmWave) communication relies on directional beam sweeping for initial access, which incurs substantial latency and signaling overhead as antenna arrays grow large and deployments become denser. Machine-learning-based beam prediction can reduce this overhead, but its effectiveness depends on training data that accurately reflects the deployment environment, a requirement that is difficult to meet in practice. This paper presents a latency-aware beam management architecture that integrates digital-twin-assisted prediction, regionbased localization, weighted ensemble learning, and explainabilitydriven beam pruning into a unified framework. A digital twin (DT) generates site-specific ray-traced channel knowledge offline, enabling environment-aware beam prediction without exhaustive real-world measurements and supporting both initial access and beam tracking. The proposed ensemble learning scheme combines a base classifier and a specialist classifier through reliabilityweighted fusion to improve prediction accuracy and robustness, while explainable artificial intelligence (XAI)-based attribution pruning reduces the beam codebook to a small subset of the most informative beams for beam allocation. To capture the distributional mismatch between the DT and real-world environments: the network was trained on imperfect digital-twin data but evaluated on real-world surrobase data. Results show that the proposed method achieves 96.80% top-1 and 98.36% top-3 beam prediction accuracy, while reducing over-the-air sounding overhead by nearly an order of magnitude relative to exhaustive sweeping, outperforming convolutional and deep neural network baselines from prior work. Index Terms—Beam management, beam prediction, millimeterwave (mmWave) communications, digital twin, explainable artificial intelligence (XAI), TreeSHAP, ensemble land earning, lowlatency initial access.

I. I NTRODUCTION Millimeter-wave (mmWave) communication is a key enabler of next-generation wireless systems, owing to the large available bandwidth that supports extremely high data rates and low-latency transmission. The short wavelength at mmWave frequencies further enables the use of large-scale antenna arrays, allowing highly directional beamforming that enhances spatial reuse and interference management in dense deployments. Despite these advantages, mmWave communications face significant challenges stemming from severe path loss and high sensitivity to blockage, which complicate reliable coverage and link maintenance, particularly in dynamic and non-line-ofsight (NLOS) environments [1]. To address these issues, current This research was partly funded by Palmer Department Chair Endowment Funds at Iowa State University. M. K. Saeed is with the Dept. of Electrical and Computer Engineering, Iowa State University, Ames, IA, USA (email: muhkamransaeed@gmail.com). A. Khokhar is with the College of Engineering, Kansas State University, Manhattan, KS, USA (email: akhokhar@k-state.edu).

mmWave systems rely on directional beam sweeping for initial access (IA). In this process, the base station periodically transmits synchronization signals over a predefined set of beams, while the user equipment scans its receive beams to detect and evaluate potential links. Once a suitable transmit–receive beam pair is identified, beam refinement and random access signaling are performed to establish the connection. In 5G new radio (NR), IA is realized through synchronization signal blocks (SSBs) mapped to spatial beams. Up to 64 SSBs may be transmitted within a single beam-sweeping burst spanning a 5 ms window, with a default periodicity of 20 ms. Because each SSB corresponds to a distinct beam, the user equipment (UE) may need to wait for an entire burst to assess all candidates, and must re-scan whenever mobility or blockage disrupts beam alignment, which occurs frequently in dynamic environments. Consequently, IA can incur considerable latency and overhead, especially in dense networks, high-mobility scenarios, and systems with large antenna arrays, where the channel coherence time may be shorter than the SSB periodicity [1], [2]. The IA procedure can be viewed as comprising two core components. The first is beam sweeping and measurement, in which synchronization signals are transmitted across multiple beams and the receiver measures the resulting signal strengths. The second is beam selection (or prediction), in which these measurements are processed to determine the optimal transmit– receive beam pair for link establishment, with beam prediction guiding the final access decision [3]. Among these two stages, beam sweeping and measurement accounts for most of the access delay because the total duration increases linearly with the number of beams that must be scanned. To address this bottleneck, machine learning (ML)–based beam prediction reduces beam training overhead by learning the relationship between partial measurements or side information and optimal beam directions, thereby limiting the need for exhaustive beam sweeping during initial access and beam tracking [3]–[5]. Prior studies show that ML models can accurately identify candidate beams using only a small subset of sensing beams or other low-dimensional observations, significantly shortening training time while maintaining near-optimal performance [3], [6]. Consequently, ML-driven beam management is increasingly recognized as a key enabler for reducing signaling overhead and latency in mmWave and 5G-Advanced systems [2]. Although ML–based beam prediction has shown strong potential for reducing beam training overhead, its effectiveness depends heavily on how well the training data reflects the target deployment environment. Prior studies indicate that models

trained on synthetic or limited datasets can experience noticeable performance degradation in real-world scenarios because the underlying data distribution may not accurately capture actual propagation conditions [4], [7]. Therefore, reliable MLbased beam prediction requires either training data that closely matches the deployment setting or adaptation mechanisms that can incorporate limited real-world measurements. One promising way to address this challenge is through DT–based approaches, which use virtual replicas of the physical environment to generate site-specific channel and beam information offline. By producing ray-tracing data that closely reflects the characteristics of the target environment, digital twins (DTs) enable beam prediction models to be trained on more representative data and to infer suitable beams directly, thereby substantially reducing the need for sequential beam sweeping during initial access and tracking [4]. DT–based beam management is highly effective when the virtual replica accurately captures the geometry, materials, and dynamics of the physical environment, thereby enabling reliable beam prediction with minimal real-world interaction [4], [8], [9]. In practice, however, achieving such a nearperfect DT is challenging, since modeling errors, environmental dynamics, and incomplete information inevitably create mismatches between the digital and physical worlds [4], [7]. These mismatches can undermine the robustness of digitaltwin-aided beam prediction, making mechanisms such as model adaptation and robustness enhancement essential for mitigating performance degradation [7], [8]. To address this limitation, we propose a region-based beam training framework that partitions the coverage area into localized regions. Rather than performing beam training across the entire spatial domain, the proposed method restricts beam exploration to region-specific subsets, thereby confining the impact of digital-twin inaccuracies to local areas. This regionalization supports more adaptive and robust beam training by accommodating localized mismatches between the DT and the physical environment, while also reducing unnecessary training overhead. To further reduce beam training overhead, we incorporate explainable artificial intelligence (XAI) into the beam management framework to enable more selective and informed beam exploration instead of exhaustive sweeping. Recent studies indicate that XAI can identify the most informative sensing beams or spatial directions, allowing redundant measurements to be pruned while preserving reliable beam alignment performance [7], [10]. In this paper, we propose a latency-aware beam management architecture that combines four complementary components to reduce beam training overhead while maintaining reliable beam alignment. First, a DT generates site-specific channel knowledge that supports environment-aware beam prediction without requiring exhaustive real-world measurements. Second, to mitigate the impact of digital-twin imperfections, a regionbased framework partitions the coverage area into localized regions and confines beam training within each region, thereby limiting the propagation of digital-physical mismatch across the network. Third, within each region, a base–specialist classifier

adapts inference cost to sample difficulty: a lightweight base resolves confident cases on its own, while uncertain samples are escalated to a specialist and fused with the base through a reliability-weighted ensemble that combines complementary base predictors to improve accuracy beyond any individual classifier and strengthen robustness under noisy, mismatched digital-twin conditions. Fourth, XAI is integrated into the regional beam training process to identify and prioritize the most informative beams within each region, enabling targeted exploration and further reducing training overhead. By jointly combining digital-twin-assisted prediction, region-based localization, ensemble-based regional prediction, and XAI-guided beam prioritization, the proposed framework improves robustness to environmental mismatch and provides an efficient, scalable solution for low-latency beam management in mmWave systems. A. Literature Review 1) Learning-Based Beam Prediction from In-Band Measurements: A large body of work reduces beam-training overhead by inferring narrow beams from a small set of measurements. Ma et al. [11] map coarse wide-beam measurements to refined narrow beams through data-driven calibration, avoiding exhaustive codebook scanning while preserving alignment accuracy. Liu et al. [12] extend this idea to near-field extremely largescale multiple-input multiple-output (XL-MIMO), mapping farfield wide-beam measurements to near-field beam decisions without exhaustive near-field codebook search. Deng et al. [13] instead exploit sub-6 GHz channels, predicting the mmWave beamspace with a convolutional neural network (CNN) and performing hybrid beamforming via a multi-layer graph neural network (GNN) and linear minimum mean-square error (LMMSE) combining. Chen et al. [14] propose a unified XLMIMO framework that integrates beam training, refinement, and tracking across near- and far-field regimes under practical hybrid architectures. These methods, however, require measurements for every prediction and train a single model over the entire coverage area. 2) Position- and Sensing-Aided Beam Prediction: Side information acquired outside the communication band offers another route to reducing beam-training overhead. Heng et al. [15] predict the optimal access point (AP) and beam directly from user equipment global positioning system (GPS) coordinates, shrinking the joint AP–beam search space on ray-traced urban data. Nugroho et al. [16] extend position-aided prediction to unmanned aerial vehicle (UAV) links, jointly forecasting current and future beams from sequential position data. Mollah et al. [17] fuse multi-modal out-of-band sensing data to predict the optimal beam index in vehicle-to-infrastructure (V2I) and vehicle-to-vehicle (V2V) links, while Rasheed et al. [18] combine transformer-based trajectory prediction with federated multi-agent deep reinforcement learning (DRL) to sustain alignment under high mobility. However, Morais et al. [19] show that position-aided models validated on synthetic data degrade under real commercial off-the-shelf (COTS) GPS

measurements, exposing a gap between simulated and deployed performance. 3) Environment- and Region-Aware Beam Management: Another line of work exploits spatial structure to make beam training more selective. Zhang et al. [20] leverage the spatial non-stationarity of near-field XL-MIMO channels through a two-stage strategy that first identifies the user’s visibility region and associated subarray, then performs fine alignment within that region using a learned hierarchical codebook. Alawadi et al. [21] predict emerging traffic hotspots with multiple learning models and proactively steer beams toward anticipated highdemand areas. In these works, regional context guides the beam search but is not used to specialize the predictor to local propagation conditions. 4) DT-Aided Beam Management: DTs reduce the cost of acquiring real-world training data. Zeulin et al. [22] train a lightweight ML beam predictor for time-sensitive industrial Internet of Things (IoT) using a map-based channel model derived from a real-world-calibrated, ray-traced 3D environment, narrowing the beam search space and reducing training overhead. Jiang et al. [4] use 3D environment models and ray tracing to generate synthetic channel data for training beam predictors, showing in a position-aided mmWave case study that DT-trained models achieve near-optimal performance while reducing real-world beam sweeping. Alkhateeb et al. [23] learn site-specific line-of-sight (LOS) and NLOS beam codebooks from DT-generated channels without costly real-world interaction, and identify ray-tracing accuracy as the dominant factor governing DT codebook performance. Zhang et al. [8] use virtual environment replicas for sample-efficient beam pattern design for interference nulling without explicit channel state information (CSI), extending the role of DTs beyond beam prediction. Bonora et al. [24] apply DT-guided iterative beamforming optimization to interference mitigation in amplify-andforward relay-assisted MIMO systems and validate it on a custom mmWave testbed. Despite this progress, how DT fidelity affects beam prediction accuracy remains largely unexplored, even though simulation–reality mismatch can substantially degrade performance [19]. 5) Explainable AI for Beam Management: Explainable artificial intelligence (XAI) has recently moved from interpreting models to actively reducing overhead. Khan et al. [6], [7] use XAI to identify the most informative sensing beams and prune redundant measurements without degrading alignment reliability. In a vehicle-to-everything (V2X) setting, Khan et al. [25] use SHapley Additive exPlanations (SHAP)-based feature importance to shrink the state space and model complexity of multi-agent DRL with minimal performance loss. These studies apply XAI-guided reduction to a single global model and do not exploit region-specific specialization. 6) Summary and Positioning: In summary, existing approaches reduce beam-training overhead through measurementto-beam mapping, out-of-band side information, regional context, DT-generated data, or XAI-guided pruning, but largely in isolation, with a single global model, and with limited attention to DT fidelity and inference complexity. This paper unifies these

directions in a DT-enabled, XAI-guided cascade framework that (i) partitions the coverage area into hexagonal regions using primary and auxiliary tilings, (ii) routes samples to a lightweight base model or a specialist predictor based on prediction confidence, with reliability-weighted ensemble fusion, (iii) constructs region-specific beam subsets via TreeSHAPbased pruning, and (iv) quantifies how DT fidelity, modeled by the number of retained propagation paths, affects prediction accuracy and computational cost. B. Contributions The main contributions of this paper are summarized as follows: • We propose a latency-aware beam management architecture that leverages a DT to generate site-specific channel knowledge and support environment-aware beam prediction, thereby reducing the need for exhaustive beam sweeping during initial access and beam tracking. • We develop a region-based beam training framework that partitions the coverage area into localized regions and restricts beam exploration to the relevant region, mitigating digital-twin mismatch and reducing unnecessary global training overhead. • We design a cascaded base–specialist ensemble learning architecture that combines two complementary base predictors with confidence-based routing and reliabilityweighted fusion: confident samples are resolved by the lightweight base alone, while uncertain samples are jointly predicted by the base and specialist with model-specific weights assigned according to their reliability. This design improves beam prediction accuracy beyond what any individual classifier achieves and strengthens the robustness of the regional beam selection stage under noisy and mismatched digital-twin conditions. • We integrate XAI into the regional beam management process to identify and prioritize the most informative sensing beams within each region, enabling targeted beam exploration and further reducing training overhead while preserving reliable beam alignment. • By jointly combining digital-twin-assisted prediction, regional localization, and XAI-guided beam prioritization, the proposed framework provides a robust and scalable solution for low-latency beam management in practical mmWave systems where the DT is imperfect and realworld interaction must be minimized. II. S YSTEM M ODEL We consider the downlink of a mmWave cellular system in which a single base station (BS), equipped with a uniform planar array (UPA) of N antenna elements, serves a set of K single-antenna user equipments (UEs) indexed by K = {1, 2, . . . , K}. The system operates at carrier frequency fc with signal bandwidth W , and since mmWave deployments typically satisfy W ≪ fc , the array steering response is treated as frequency-flat across the operating band; this assumption

where θn ∈ [0, 2π) is the phase shift √ introduced at the n-th antenna port. The normalization 1/ N ensures the unit-norm constraint ∥v∥2 = 1, thereby fixing the radiated power independently of the chosen beam direction. Because optimizing the continuous phase variables {θn } over RN is computationally prohibitive and incompatible with practical initial-access protocols (e.g., 5G NR SSB-based beam sweeping), the BS instead employs a finite-cardinality beam codebook of pre-designed analog beamformers, vc ∈ CN ×1 ,

V = {v1 , v2 , . . . , vC } , Fig. 1: Downlink of a mmWave system with active users in the user grid. is standard in the mmWave beam management literature and allows the precoder to be designed at the carrier wavelength without incurring beam-squint distortion. Owing to the high path loss, limited scattering, and quasi-optical propagation characteristics of the mmWave band, the BS-UE channel exhibits pronounced spatial sparsity and is most accurately captured by a geometric ray-based representation. Accordingly, the downlink channel vector gk ∈ CN ×1 from the BS to the k-th UE is modeled as a superposition of Lk dominant propagation paths, i.e., r Lk N X βk,l a(ϕk,l , ψk,l ) , (1) gk = Lk l=1

where βk,l ∈ C denotes the complex-valued small-scale fading coefficient of the l-th ray, while ϕk,l ∈ [−π, π] and ψk,l ∈ [−π/2, π/2] represent its azimuth and elevation angles of departure (AoDs) observed at the BS, respectively. The transmit array steering vector a(ϕ, ψ) ∈ CN ×1 associated with AoD pair (ϕ, ψ) depends on the geometric layout of the UPA. Assuming N = Nh Nv elements arranged in Nh columns and Nv rows with horizontal and vertical inter-element spacings dH and dV , the steering vector is given by iT 1 h j κ̃ ξ1 (ϕ,ψ) a(ϕ, ψ) = √ e , . . . , ej κ̃ ξN (ϕ,ψ) , (2) N where κ̃ = 2π/λc denotes the wavenumber corresponding to carrier wavelength λc , and the spatial phase contribution of the n-th antenna element is ξn (ϕ, ψ) = αn dH sin(ϕ) cos(ψ) + γn dV sin(ψ),

(3)

with (αn , γn ) denoting the horizontal and vertical position indices of element n within the planar grid are spaced halfwavelength, i.e., dH = dV = λc /2. To strike a favorable tradeoff between hardware cost, power dissipation, and beamforming flexibility, the BS adopts an analog-only transmit architecture in which all N antenna elements are driven by a single radiofrequency (RF) chain through a network of digitally controlled phase shifters. Under this front-end, the precoder v ∈ CN ×1 is constrained to have unit-modulus entries and is expressed as, T 1  jθ1 jθ2 e , e , . . . , ejθN , (4) v= √ N

(5)

where C denotes the codebook size. Every codeword inherits the constant-modulus constraint imposed by the analog architecture, namely ∥vc ∥2 = 1,

∀c ∈ {1, 2, . . . , C}.

(6)

The codebook may equivalently be stacked into matrix form as V = [v1 , v2 , . . . , vC ] ∈ CN ×C , which serves as the sweeping dictionary used during beam training. Common constructions include DFT, oversampled DFT (O-DFT), and hierarchical multi-resolution codebooks. DFT codebooks give broad coverage with C = N orthogonal beams but lack the resolution needed for precise alignment in narrow-beam, high-frequency regimes; O-DFT codebooks improve resolution at the cost of higher training overhead from evaluating more beams. The proposed framework is agnostic to codebook design, requiring only that V satisfy (6). During beam training, the BS sequentially probes the angular domain by transmitting known pilot symbols through each codeword in V. When the c-th codeword vc is activated, the baseband signal observed at the k-th UE is √ (c) y k = ϱ gkH vc x + zk , (7) where ϱ is the BS transmit power, x ∈ C denotes the unitenergy pilot symbol with E{|x|2 } = 1, and zk ∼ CN (0, σ 2 ) models the additive white Gaussian noise (AWGN) at the UE receiver with variance σ 2 . The corresponding instantaneous received signal-to-noise ratio (SNR) at user k under beam vc is therefore 2 ϱ gkH vc (c) . (8) γk = σ2 III. P ROBLEM F ORMULATION The system model introduced in Section II reduces, at its core, to the per-user beam selection problem c⋆k = arg

max c∈{1,2,...,C}

2

gkH vc ,

∀k ∈ K.

(9)

Solving (9) requires evaluating the effective beamforming gain |gkH vc |2 over the codebook V for every active user. In practice, however, the instantaneous channel gk is not available a priori, since acquiring it would itself demand a full sweep across V. This circular dependency between channel knowledge and beam selection lies at the heart of the mmWave beam management problem.

A. Limitations of Exhaustive Beam Sweeping The conventional approach, exhaustive beam sweeping (EBS), has the BS sequentially probe every codeword in V, with the UE measuring received signal strength per beam and reporting the strongest index. While EBS guarantees identification of c⋆k in static channels, its overhead scales with C, which can approach or exceed the channel coherence interval in dynamic environments, causing outdated beam selections and conflicting with the low-latency demands of 5G NR and 6G. It also wastes energy and feedback overhead by activating the RF chain across all C codewords, even though only a few beams typically offer meaningful gain. These limitations motivate approaches that exploit prior environmental information to prune the search space before over-the-air measurement. B. Site-Specific Prior via a DT In typical deployments, the propagation environment surrounding the BS, including building geometry, material properties, street layouts, and reflector locations, evolves on time scales far slower than the user channels themselves. This observation enables the construction of a DT: a site-specific computational replica of the physical environment from which synthetic channel realizations ĝk can be generated, for instance via ray-tracing over a three-dimensional model of the cell. Let V̂k ⊆ V denote a candidate beam subset predicted from DTderived priors for user k. The beam selection task can then be relaxed from (9) to the reduced-search problem ĉk = arg max c : vc ∈V̂k

2

gkH vc ,

D. Design Objectives of the Proposed Framework

2

∆k = ∥gk − ĝk ∥ ,

(11)

which is generally non-zero due to dynamic obstacles, calibration drift, and unmodeled scatterers. When ∆k is large, V̂k may fail to contain vc⋆k , leading to a beam-misalignment penalty 2

− gkH vĉk

2

≥ 0.

(12)

A central challenge is therefore to design a beam management framework that exploits DT priors aggressively enough to shrink |V̂k |, yet remains resilient when ∆k is non-negligible. C. Problem Statement Let Ttrain denote the total beam-training time consumed for serving all K users in a given training cycle, and let ηk be the per-user misalignment penalty in (12). We formulate the latency-aware beam management problem as   minimize Ttrain {V̂k } (13) {V̂k }k∈K

subject to

Pr(ηk ≤ ϵ) ≥ 1 − δ, V̂k ⊆ V,

where ϵ > 0 specifies the maximum tolerable beamforming gain loss, and δ ∈ (0, 1) controls the outage probability associated with that loss. The constraint (14) ensures that the proposed framework retains alignment reliability even under DT mismatch, while (13) captures the overarching goal of minimizing training latency. The optimization in (13)–(15) is intractable in its general form, for three reasons. First, the true channel gk governing ηk is unknown at the time the subset V̂k must be committed to. Second, the relationship between DT-generated priors and the optimal beam c⋆k is highly non-linear and environmentdependent, precluding closed-form solutions. Third, even if a data-driven predictor for V̂k were available, its black-box nature offers no principled mechanism for further pruning redundant beams within a candidate subset.

(10)

with overhead |V̂k | ≪ C. The fidelity of (10) hinges on the agreement between the DT and the true environment, captured by a model-mismatch error

ηk = gkH vc⋆k

Fig. 2: User-to-region assignment under the overlapping region partition.

∀k ∈ K,

∀k ∈ K,

(14) (15)

The above difficulties motivate three coupled design objectives that jointly guide the framework developed in the remainder of this paper. First, site-specific prediction exploits DT-generated channel knowledge to construct environmentaware beam predictions V̂k , substantially reducing the candidate set relative to the full codebook V and avoiding exhaustive real-world measurements. Second, robustness to DT mismatch is achieved through regionalization, where the coverage area is partitioned into localized regions {R1 , R2 , . . . , RR } and beam exploration for each user is restricted to its corresponding region. This localized scope limits the impact of DT inaccuracies to the affected region and improves resilience to local environmental discrepancies. Third, interpretable beamlevel pruning uses XAI to attribute prediction confidence to individual sensing beams, rank codewords according to their relevance, and retain only those whose contribution exceeds a learned threshold. This further reduces the regional candidate set V̂k to an importance-ordered probing sequence, enabling targeted exploration while reducing Ttrain and maintaining reliable beam alignment. Collectively, these objectives translate the optimization in (13)–(15) into a tractable three-stage pipeline comprising DT-assisted beam prediction, region-based localization, and XAI-guided beam prioritization. The detailed architecture implementing these objectives is presented in Section IV.

IV. R EGION -BASED U SER A SSOCIATION The intractability of the latency-aware beam management problem (13)–(15) stems chiefly from the dependence of V̂k on the unknown channel gk and on the fidelity of the DT. A direct attempt to predict per-user beam subsets from DT priors would be unduly sensitive to the inevitable mismatch term ∆k in (11), since even small errors in the inferred user position can translate into substantial deviations between predicted and actual dominant AoDs. To circumvent this brittleness, we discretize the coverage area into a finite collection of regions and lift the beam-subset design problem from the per-user level to the per-region level. Each region is then equipped with its own beam subset, and a user is associated with a region rather than with a precise spatial point. This indirection trades the unforgiving requirement of exact localization for the far milder requirement of approximate region membership, providing the structural robustness that DT-driven prediction lacks on its own. A. Regional Partitioning of the Coverage Area Let A ⊂ R2 denote the two-dimensional ground projection of the cell coverage area. We partition A using a regular hexagonal lattice of side length a, chosen so that the hexagon dimension is commensurate with the angular resolution of the BS beam codebook. The use of a hexagonal tiling, as opposed to a rectangular or Voronoi partition, is motivated by its near-isotropic spatial coverage, its uniform center-to-vertex distance, and its prevalence in cellular geometry; together these properties yield region shapes whose interior points share statistically similar AoD distributions with respect to the BS, which is precisely the property that justifies sharing a common beam subset across the region. We denote the resulting collection of hexagonal cells, henceforth referred to as the primary tiling, by R = {R1 , R2 , . . . , RR },

Rr ⊂ A,

R [

Rr = A,

r=1

(16) where adjacent hexagons share only their boundary edges so that the interiors form a disjoint partition of A. Each region Rr is identified by the planar coordinates of its center, or ∈ R2 , with the set of all primary centers denoted O = {o1 , . . . , oR }. Associated with Rr is a region-specific beam subset Br ⊆ V,

|Br | ≪ C,

(17)

which contains the codewords identified as informative for users located anywhere in Rr . The construction of Br from DT-generated channel data is performed offline through an explainability-driven beam ranking procedure detailed in Section V; for the purposes of the present section, we treat Br as a given attribute of each region. B. Boundary Ambiguity and the Need for an Auxiliary Tiling Although the primary tiling cleanly partitions A, it suffers from a well-known artifact common to all finite-grained spatial discretizations: users whose estimated location falls near the boundary of a hexagon are vulnerable to misclassification. e k ∈ R2 denote the estimated planar location of user Let p

k, obtained from any of the localization modalities reported in the literature, including LiDAR sensing, uplink-pilot-based angle/range estimation at the BS, or GPS, and assumed to satisfy e k = pk + ek , p ∥ek ∥ ≤ ε, (18) where pk is the true position and ε bounds the localization error. Specifying the particular localization mechanism is beyond the scope of this paper. If pk lies within a distance ε of the edge of e k may straddle the boundary some hexagon, the noisy estimate p and be assigned to the wrong region. The consequences are non-trivial: because Br is tuned to the angular statistics of Rr , applying Br′ for a neighbor r′ ̸= r can yield a beam subset that omits the user’s true dominant beam. In other words, while the regionalization contains DT mismatch to local neighborhoods, it does not by itself eliminate the sensitivity of edge-located users to localization noise. To resolve this boundary ambiguity, we introduce a second e congruent to R but spatially offset, such hexagonal tiling R, that the regions in which users are most vulnerable under one tiling become the regions in which they are most secure under the other. Specifically, the auxiliary tiling is obtained by translating the primary tiling by a shift vector  (19) τ = a [cos θs , sin θs ]T , θs ∈ / mπ 3 : m∈Z , where ∥τ ∥ = a equals the hexagon side length and the orientation θs is constrained to avoid the six directions parallel to the principal edges of the hexagonal lattice. This constraint is essential: any shift along an edge direction would realign the auxiliary lattice with the primary lattice and reproduce the same boundary structure, defeating the purpose of the second tiling. Choosing θs off the edge axes ensures that the centers e lie squarely in the interior of regions of R, so that every of R boundary of R is geometrically covered by an interior region e and vice versa, as illustrated in Fig. 2. The resulting of R, auxiliary collection is e = {R e1, R e2, . . . , R e R }, R

e cr = c r + τ ,

(20)

e r likewise endowed with its own with each auxiliary region R e beam subset Br ⊆ V obtained through the same DT-XAI procedure used for the primary tiling. C. User-to-Region Association Rule Given the two tilings, every user is by construction interior to either a primary region, or an auxiliary region. To exploit this property, we associate each user with the nearer center across the two tilings, on the rationale that a user closer to a region center is deeper in that region’s interior and therefore more robust to localization noise of magnitude ε. Formally, define the nearest primary and auxiliary centers for user k as rkP = arg rkA = arg

min r∈{1,...,R}

min e r∈{1,...,R}

∥e pk − or ∥ ,

(21)

er ∥ , ∥e pk − o

(22)

with corresponding center distances e k − orkP , dP k = p

ek − o erkA . dA k = p

The assigned region for user k is then ( A R rkP , if dP k ≤ dk , Ωk = e rA , otherwise, R k

(23)

(24)

and the corresponding beam subset to be probed by user k is ( B rkP , Ωk = R rkP , (25) B(Ωk ) = e rA . Be rkA , Ωk = R k The geometric guarantee underpinning this rule is that, for any point p ∈ A, √
3 P A min d (p), d (p) ≤ a, (26) 2 where dP (p) and dA (p) denote the distances from p to the √ nearest primary and auxiliary centers, respectively, and 3 2 a is the apothem (center-to-edge distance) of the hexagon. Consequently, every user is associated with at least one region whose center lies no farther than one apothem from the user, ensuring that the localization-error ball of radius ε around e k remains p well-contained in the interior of Ωk whenever √ A ε < 23 a − max(dP k , dk ). D. Implications for Beam Training Once Ωk is determined via (24), the beam-selection problem (9) reduces, for user k, to the localized search ĉk = arg

max

c : vc ∈ B(Ωk )

2

gkH vc ,

(27)

which requires probing only |B(Ωk )| codewords instead of the full C. Since |B(Ωk )| ≪ C by design, the proposed regionbased association achieves a substantial reduction in training latency relative to exhaustive sweeping, while the dual-tiling construction ensures that the reduction does not come at the cost of robustness for users located near regional boundaries. The mechanism by which each Br and Ber is constructed from the DT-generated channel knowledge, using XAI to rank and select the most informative beams within each region, is the subject of Section V. V. XAI-G UIDED B EAM S UBSET C ONSTRUCTION Having established in Section IV how each user is associated e the remaining task is to specify with a region Ωk ∈ R ∪ R, how the region-specific beam subsets Br and Ber are obtained from the DT-generated channel knowledge. The naive choice is keeping the top-ranked beams by average gain and ignores the joint statistical role that beams play when used together as input to a predictor, and tends to retain mutually redundant codewords pointing in similar angular directions. We instead pose subset construction as a feature selection problem driven by an interpretable, classifier-grounded importance measure. Concretely, for every region we train a two-stage cascaded classifier whose task is to predict the optimal beam from a

Fig. 3: The proposed stages of the offline training and online inference framework

vector of beam-power observations, and we use TreeSHAP to attribute the classifier’s decisions back to individual codewords. Beams whose attribution falls below a relevance threshold are pruned, yielding compact subsets that retain the codewords actually responsible for reliable prediction. A. Per-Region Dataset Generation from the DT Consider a fixed region Rr (the treatment for an auxiliary e r is identical). Using the DT of the deployment site, region R r we sample Mr virtual user locations {pr,m }M m=1 uniformly distributed within Rr . For each such location, the DT generates a site-specific channel realization via ray-tracing, from which the RSSI is recorded on each beam of the full codebook V.

Stacking the per-beam RSSI values yields the feature vector T  xr,m = RSSI1r,m , RSSI2r,m , . . . , RSSIC ∈ RC , (28) r,m where RSSIcr,m denotes the received signal strength on beam c at location pr,m . The classification label for sample m is the index of the beam yielding the highest RSSI, yr,m = arg

max c∈{1,...,C}

[xr,m ]c ,

(29)

where [ · ]c extracts the c-th entry. The resulting per-region dataset Mr (30) Dr = { (xr,m , yr,m ) }m=1 serves as the training corpus for the cascaded classifier associated with Rr . We emphasize that Dr is generated entirely offline from the DT and incurs no over-the-air measurement cost. B. Cascaded Base–Specialist Classifier Within each region, the optimal-beam prediction task is solved by a two-stage cascade consisting of a lightweight base classifier and a higher-capacity specialist classifier, trained jointly per region but invoked sequentially at inference time. The rationale follows the conditional-computation principle: most input samples are easy to classify and can be resolved cheaply, whereas a residual subset of harder cases warrants a more expressive model. Concentrating computational effort where uncertainty is high reduces inference latency without sacrificing accuracy in the hard regime. a) Base model.: The base frG : RC → ∆C−1 is realized as a gradient-boosted decision tree ensemble (LightGBM) trained on Dr to output a probability distribution over candidate beam indices, G C−1 pG , r (x) = fr (x) ∈ ∆

(31)

where ∆C−1 denotes the (C −1)-dimensional probability simplex. The base’s top-1 prediction ŷrG (x) and its associated confidence ρG r (x) are ŷrG (x) = arg max [pG r (x)]c ,

G ρG r (x) = max [pr (x)]c . c c (32) b) Cascade routing.: A confidence threshold τ ∈ (0, 1), fixed by validation on held-out DT samples, determines whether the base’s prediction is committed to or whether the sample is escalated. Define the gating indicator  G  Iesc (33) r (x) = 1 ρr (x) < τ ,

so that Iesc r (x) = 1 marks x as a hard sample requiring the specialist. c) Specialist model.: The specialist frS : RC → ∆C−1 is realized as an XGBoost ensemble trained on the full dataset Dr , with higher sample weights assigned to instances flagged as hard under the base model’s escalation rule. This weighted training strategy emphasizes challenging samples while preserving knowledge of the overall data distribution, enabling the specialist to improve predictions in ambiguous regions of the feature space.

d) Weighted ensemble on escalated samples: When the base is confident, its top-1 prediction is committed directly, as the sample lies in the easy regime where the base is reliable. When the sample is escalated, however, discarding the base’s probability vector in favor of the specialist alone is suboptimal: although the specialist is trained well on hard cases, the base has been exposed to the full training distribution and still carries useful soft information on a portion of the escalated samples. We therefore combine the two models’ probability outputs through a weighted average rather than routing to the specialist exclusively. Let wG , wS ≥ 0 denote the base and specialist weights, with wS > wG chosen to reflect the specialist’s stronger competence on hard inputs. The ensemble probability vector on an escalated sample is pE r (x) =

S wG pG r (x) + wS pr (x) ∈ ∆C−1 , wG + wS

and the corresponding predicted beam index is   ŷrE (x) = arg max pE r (x) c . c

The combined per-region predictor is therefore ( G ŷr (x), Iesc r (x) = 0, ŷr (x) = ŷrE (x), Iesc r (x) = 1.

(34)

(35)

(36)

C. TreeSHAP-Based Importance Attribution Once the cascade is trained, we quantify the contribution of each beam-power feature to the classifier’s predictions using TreeSHAP [26], a polynomial-time exact algorithm for computing Shapley values over tree ensembles. For a classifier f : RC → ∆C−1 and a sample x, TreeSHAP returns a matrix Φ(x) ∈ RC×C whose entry ϕc,c′ (x) measures the marginal contribution of feature c (the c-th beam) to the predicted logodds of class c′ (the c′ -th beam being optimal). The global importance of feature c for region r is obtained by averaging absolute attributions across all samples and predicted classes: Irq (c) =

Mr X C 1 X ϕqc,c′ (xr,m ) , Mr C m=1 ′

(37)

c =1

where q ∈ {G, S} denotes the base or specialist model, respectively, and ϕqc,c′ (xr,m ) is the SHAP attribution of feature c to class c′ for sample xr,m under model q. D. Region Beam Subset Construction The importance vectors Irq (·), where q ∈ {B, S}, capture complementary aspects of beam informativeness: the base reflects beams useful for the easy majority of users in Rr , while the specialist reflects beams useful for the ambiguous minority. Retaining only the base’s top features risks pruning codewords critical for hard cases; retaining only the specialist’s top features risks pruning codewords that anchor the easy decisions. We therefore construct the region beam subset as the union of top-importance features identified by both models. For each region Rr , where Kr denote the number of top-ranked beams retained per model. Let πrG (·) and πrS (·) denote permutations of

{1, . . . , C} that sort beams by IrG and IrS in decreasing order, respectively. The top-Kr beam sets selected by the base and specialist are TrG =

 G πr (1), πrG (2), . . . , πrG (Kr ) ,  TrS = πrS (1), πrS (2), . . . , πrS (Kr ) .

The final region beam subset is formed by the union  Br = vc : c ∈ TrG ∪ TrS ⊆ V,

(38) (39)

(40)

with effective cardinality |Br | = TrG ∪ TrS ≤ 2Kr ≪ C.

(41)

The parameter Kr directly controls the overhead–accuracy tradeoff: a smaller Kr reduces the number of beam measurements a user must perform but risks omitting codewords that are informative for either the easy or hard prediction regime, while a larger Kr improves coverage at the cost of increased sounding overhead. The union ensures that both the easy-case and hard-case codewords are preserved; overlap between TrG and TrS means the true subset size is typically strictly less than 2Kr . The identical procedure applied to the auxiliary tiling e r ∈ R. e yields Ber for every R E. Retraining on the Reduced Beam Subset The cascade described in Section V-B was deliberately trained on all C beam-power features so that the TreeSHAP analysis could identify, without bias, which codewords carry decision-relevant information for each region. We then retrain the cascade from scratch on the reduced feature space restricted to the surviving beams Br , yielding a deploymentready predictor whose input dimensionality matches the sparse sweep performed online. Concretely, let Pr ∈ {0, 1}|Br |×C be the row-selection matrix that extracts the entries of x ∈ RC corresponding to the beams retained in Br , ordered canonically by codeword index. The reduced feature vector for sample m of region r is then er,m = Pr xr,m ∈ R|Br | , x

(42)

and the corresponding pruned dataset is r er = { ( x er,m , yr,m ) }M D m=1 ,

(43)

where the labels yr,m remain unchanged from (29) because the ground-truth optimal beam index is independent of which er , we train features the classifier is allowed to observe. Using D a fresh base–specialist cascade, ferG : R|Br | → ∆C−1 ,

ferS : R|Br | → ∆C−1 ,

(44)

following the same architecture, hyperparameter family, and confidence-threshold routing rule adopted in the initial training stage, as discussed in Section V-B.

F. Inference-Time Beam Training At deployment, the BS stores three artifacts per region: the beam subsets {Br , Ber }, the retrained cascaded classifiers {(ferG , ferS )} and their auxiliary counterparts, and the regioncenter coordinates required for assignment. When a user k arrives, the following sequence is executed. a) Step 1: region assignment.: The user’s approximate e k is acquired through the localization modality delocation p scribed in Section IV, and the assigned region Ωk and its corresponding subset Br are determined via the nearest-center rule (24). b) Step 2: sparse beam sweep.: The BS transmits beams only on the codewords contained in Br , and the user reports the resulting reduced power vector  T ek = |gkH vc |2 : vc ∈ Br x ∈ R|Br | , (45) ordered canonically to match the input convention used during retraining. c) Step 3: cascade inference: The trained classifier associated with region Rr is then applied as follows: when xk ) ≥ τ , the base model’s top-1 prediction is committed ρeG r (e directly; when ρeG xk ) < τ , the base and specialist probability r (e vectors are instead combined under weights wG and wS before taking the argmax.The resulting index identifies the predicted optimal codeword ĉk ∈ {1, . . . , C}, which is committed for the ensuing data-plane transmission to user k. Compared with exhaustive sweeping, the proposed architecture replaces C beams transmissions per user by |B(Ωk )| ≪ C. VI. S IMULATION S ETUP AND R ESULTS We evaluate the proposed framework on the DeepMIMO ray-tracing dataset using the O1_28 outdoor street-canyon scenario, which captures a dense urban environment populated by buildings and street infrastructure typical of cellular mmWave deployments. The base station, mounted on transmitter set TX 7, is equipped with a uniform planar array of N = 64 antenna elements arranged as Nh × Nv = 8 × 8 with halfwavelength inter-element spacing in both dimensions, and the array boresight is oriented toward the served user grid. The system operates at a carrier frequency fc = 28 GHz with a bandwidth of W = 500 MHz. The service area spans 600 m × 300 m and is discretized into a uniform user grid with an inter-user spacing of approximately 0.37 m. Beam sweeping is performed using an oversampled DFT (O-DFT) codebook of size C = 96, providing finer angular resolution than the N = 64 baseline DFT. A. DT and Real-World Datasets To evaluate robustness under the DT mismatch term ∆k introduced in (11), we construct a family of coupled datasets per region from the same user locations, by varying the number of propagation paths retained from the ray tracer. The realworld surrogate, treated as ground truth, retains the full set of 25 DeepMIMO paths per location. To emulate realistic DT imperfections, we generate eight reduced-path datasets,

(a) Top-1

(b) Top-2

(c) Top-3

Fig. 4: Performance comparison for Top-1, Top-2, and Top-3 beam prediction accuracy. {3, 4, 5, 7, 10, 12, 15, 20}, plus a ninth dataset retaining all 25 paths as an idealized upper bound where the twin exactly replicates the real channel. Discarded paths model information loss from unmodeled scatterers, calibration drift, and dynamic obstacles. The number of retained paths thus acts as a tunable DT-fidelity parameter: small values (3–4) correspond to an impoverished twin, while large values (15–20) correspond to a twin in close agreement with the real-world channel. Throughout the simulation, classifier training is performed exclusively on a chosen DT dataset, while all reported inference results are obtained on the real-world surrogate, so that every accuracy figure reflects performance against the true channel rather than the data used for training. Each region contains Mr ≈ 7100 user samples, partitioned 80%/20% into training and test splits at the location level: DT samples from training locations fit the per-region classifiers, while real-world surrogate samples from test locations are reserved for evaluation. No real-world surrogate data is used during training, except in the 25-path setting, included as an idealized upper bound corresponding to a perfect DT. The coverage area is partitioned into ten overlapping hexagonal regions from the primary and auxiliary tilings of Section IV (five regions per tiling). The global codebook of C = 96 codewords is filtered per region prior to training, retaining only beams that appear among the top-3 strongest for at least one user in the region; the rest carry negligible energy and are excluded, leaving an effective candidate pool of roughly 40 beams per region before any XAIdriven pruning. The base classifier frG is implemented as a LightGBM gradient-boosted ensemble with 25 trees, leaf count 8, learning rate 0.2, column and row subsampling ratios of 0.8, and ℓ1 /ℓ2 regularization coefficients of 0.1 each. The specialist classifier frS is implemented as an XGBoost ensemble with 100 trees, maximum depth 8, learning rate 0.3, and column and row subsampling ratios of 0.8. The cascade routing threshold is τ = 0.9, the ensemble weights are wG = 0.5 and wS = 1.5, and the retained-beam parameter is swept over Kr ∈ {2, 3, 4, 5, 6} to characterize the tradeoff between sounding overhead and prediction accuracy. Training proceeds in two stages per region. In the first stage, the initial base–specialist cascade is fit on the full Cr -dimensional feature vector using one-third of the regional training data, which is sufficient for stable TreeSHAP attribution while keeping the SHAP computation tractable. The resulting per-beam importances IrG and IrS feed the union-based

subset construction (40) and produce Br . In the second stage, a fresh cascade is retrained from scratch on the reduced feature space of Br using the full regional training data, yielding the deployment-ready predictor (ferG , ferS ). We benchmark the proposed framework against four reference methods. The first is an Argmax baseline that, given a ek , predicts the beam corresponding to its largest feature vector x entry without any learning; this serves as a non-parametric lower bound that captures what is achievable by trusting the DT-generated RSSI directly. The second is a Random Forest classifier with 100 estimators trained on the same DT data and feature space as the proposed cascade. The remaining two baselines, taken from Khan et al. [6], are a convolutional neural network (CNN) and a deep neural network (DNN). For consistency, all learned baselines are trained on the DT dataset and evaluated on the real-world surrogate dataset, share the same per-region candidate pool, and produce a probability distribution over the regional codewords from which the topk predictions are extracted. Prediction performance is reported using top-k accuracy for k ∈ {1, 2, 3}, defined as the fraction of test samples for which the predicted beam ĉk falls within the k strongest codewords of the true codebook ordering. The top-1 metric thus reflects whether the predictor recovers the optimal beam exactly, while top-2 and top-3 quantify how close the predicted beam lies to the optimum in rank terms, with a top-2 or top-3 prediction incurring only a small beamforming-gain loss relative to the best codeword in the codebook. B. Performance Comparison with the Benchmark Schemes Fig. 4 reports the top-k beam prediction accuracy of the proposed framework and four baselines across the nine DTfidelity settings, with Kr = 5, so that each region operates on at most ten codewords. For top-1 accuracy, which measures whether the predicted beam exactly matches the optimal beam, the Argmax baseline remains constant at 21.08%, reflecting the intrinsic difficulty of the regional beam-selection task. All learned methods substantially outperform this baseline even at three paths, showing that DT-trained models can extract useful channel structure despite limited DT fidelity. The tree-based methods exhibit a clear monotonic improvement as the number of paths increases, followed by saturation. The proposed framework reaches 95.89% at saturation and outperforms XGBoostspecialist by 0.15 percentage points. In contrast, the DNN and CNN show non-monotonic behavior: the DNN increases to

(a) Top-1 Across Beams

(b) Top-2 Across Beams

(c) Top-3 Across Beams

Fig. 5: Top-k beam prediction accuracy across different beam configurations. 70.53% at five paths, drops to 63.32% at seven paths, recovers to 70.51% at ten paths, and then fluctuates between approximately 64–70%. The CNN follows a similar irregular trend, remaining largely within the 60–76% range. These fluctuations suggest that the neural models are more sensitive to distribution differences between DT-generated training data and real-world test data. The top-2 and top-3 results follow the same overall behavior while providing a more relaxed measure of prediction quality. The proposed framework increases from 95.89% top-1 to 96.88% top-2 and 97.98% top-3 at saturation. The relatively small 1 percentage-point increase from top-1 to top-2 indicates that most incorrect top-1 predictions are already very close to the optimal beam. The additional 1.10 percentage-point gain from top-2 to top-3 further shows that the remaining errors are concentrated within the immediate neighborhood of the optimal beam. The Argmax baseline similarly increases from 21.08% top-1 to 40.09% top-2 and 55.54% top-3, confirming the intrinsic difficulty of the task without learned DT-based prediction. Among the learned models, the tree-based methods again improve monotonically with DT fidelity and then saturate, whereas the DNN and CNN remain non-monotonic. For example, the DNN top-2 accuracy reaches 83.88% at five paths, falls to 77.98% at seven paths, and subsequently fluctuates around 80–84%. At top-3, the proposed framework achieves 97.98%, exceeding XGBoost-specialist at 97.79%, LightGBMbase at 96.00%, and Random Forest at 95.29%. Overall, the progression from 95.89% to 96.88% and 97.98% demonstrates that even when the optimal beam is not selected exactly, the proposed framework almost always predicts a beam among the strongest nearby codewords. With |Br | ≤ 10, this provides nearoptimal beam alignment while substantially reducing the beamsounding overhead compared with exhaustive search over the full C = 96 codebook. C. Accuracy across different beam configurations Fig. 5 characterizes the overhead–accuracy trade-off by sweeping the retained-beam parameter over Kr ∈ {2, 3, 4, 5, 6}, equivalently |Br | ∈ {4, 6, 8, 10, 12}, at the saturation operating point of ten retained DT paths. Considering first the top-1 accuracy, the proposed framework improves monotonically with the sounding budget, climbing from 79.62% at |Br | = 4 to 95.98% at |Br | = 12, and

dominates every alternative for all beams values. Random Forest also improves monotonically, but plateaus more slowly and reaches only 90.03% at the largest budget. LightGBMbase and XGBoost-specialist likewise grow monotonically, with XGBoost-specialist tracking the proposed framework closely across the entire sweep (e.g., 95.97% versus 95.98% at |Br | = 12) while LightGBM-base trails both by two to four percentage points from |Br | = 6 onward. The DNN and CNN baselines behave non-monotonically, saturating between |Br | = 8 and |Br | = 10 at identical accuracies of 70.51% and 75.76% respectively, reflecting their sensitivity to the precise composition of the input feature space. Notably, the proposed framework already attains 92.54% top-1 accuracy at |Br | = 8, exceeding every non-tree baseline observed at any sounding budget in this column and recovering more than 96% of the accuracy achieved at |Br | = 12, which confirms that the union-based subset extracts most of the achievable accuracy at substantially reduced sounding cost relative to the canonical |Br | = 10 setting. The top-2 accuracy follows the same trend, with the proposed framework increasing from 86.70% at |Br | = 4 to 97.32% at |Br | = 12. All tree-based methods improve monotonically, while the DNN slightly decreases from 83.65% at |Br | ∈ {8, 10} to 82.36% at |Br | = 12, and the CNN largely saturates at 87.17% for |Br | ∈ {8, 10}, reaching only 88.05% at the largest budget. At |Br | = 8, the proposed framework achieves 94.93% top-2 accuracy, indicating that the predicted beam is among the two strongest codewords in nearly 95% of cases. From |Br | = 10 onward, its accuracy remains above 96%, which the neural baselines do not achieve at any tested budget. Top-3 accuracy similarly approaches saturation as the sounding budget increases. The proposed framework reaches 96.91% at |Br | = 8 and increases to 97.66% and 98.64% at |Br | = 10 and 12, respectively. XGBoost-specialist follows closely at 97.60% and 98.52%, while Random Forest remains approximately two percentage points lower. LightGBM-base also improves monotonically, reaching 97.88% at |Br | = 12. In contrast, DNN and CNN saturate earlier at 91.17% and 92.70%, respectively, for |Br | ∈ {8, 10}. Overall, the proposed framework achieves near-saturated top-3 accuracy with only eight codewords, while increasing the budget beyond ten provides limited additional benefit.

(a) Sample routing

(b) Top-1 accuracy

(c) Computational Complexity

Fig. 6: Effect of the base confidence threshold t on the proposed cascade. (a) Percentage of samples retained by the base (blue) and escalated to the specialist (orange). (b) Top-1 beam-prediction accuracy of the proposed cascade against the standalone base (Base Model) and specialist (Specialist Model) baselines. (c) Computational complexity

D. Sensitivity to the base Confidence Threshold t To validate the proposed reliability-weighted cascade and identify a robust operating point, we sweep the base confidence threshold t over [0, 1] in increments of 0.01. The threshold t governs the routing decision: a sample is handled by the lightweight base alone if its confidence is at least t, and otherwise escalated to the heavier specialist, with the two predictions fused via the reliability-weighted ensemble (Section V). For this study, the cascade is trained on 12-path channel realizations using the Top-8 feature beams, and evaluated on a more demanding 25-path test set. Fig. 6a characterizes how the routing decision evolves with the threshold t. The blue curve indicates the percentage of test samples retained by the base, i.e., samples for which the base confidence is at least t, while the orange curve indicates the complementary fraction escalated to the specialist. For low thresholds (t ≲ 0.35), the base confidence exceeds the bar on virtually all samples and no escalation occurs; the cascade therefore reduces to the base alone. As t increases, progressively more samples whose confidence lies below the threshold are routed to the specialist. The escalation rate climbs gradually through the mid range (0.35 ≲ t ≲ 0.85) and rises sharply only as t approaches unity. In particular, at the operating point t⋆ = 0.78 only 19.83% of samples are escalated to the specialist, while the remaining 80.17% are handled exclusively by the base. At the limit t = 1, the base prediction is rejected for every sample, and all inference passes through the specialist. This confirms the efficiency of the cascade: the bulk of inference is served by the lightweight base, with the specialist invoked only for the minority of low-confidence samples. Fig. 6b compares the Top-1 accuracy of the proposed cascade against the standalone base (91.80%) and specialist (94.45%) baselines. For t ≲ 0.30, the cascade coincides with the base, since no samples are escalated. Beyond this, accuracy rises as more hard samples are corrected via reliability-weighted fusion, with a steep gain in 0.44 ≤ t ≤ 0.55 marking the transition from base-dominated to specialist-informed behavior. The cascade crosses the specialist baseline near t⋆ = 0.78, adopted as the operating point since it matches specialist-level accuracy

while escalating fewer than one in five samples. Accuracy peaks at 94.62% near t ≈ 0.86 (+2.82 percentage points over the base, +0.17 percentage points over the specialist), achieved at an escalation fraction well below 50% – evidence that the cascade leverages the complementary strengths of both models rather than merely interpolating between them. Beyond the peak (t ≳ 0.88), accuracy decays gradually toward the specialist baseline, consistent with Fig. 6a: as t → 1, the base is discarded even on samples where it is reliable, and the cascade reduces to specialist-only inference. Fig. 6c reports FLOPS/sample as a function of t. The standalone base and specialist incur constant costs of approximately 2.20 × 104 and 1.57 × 105 FLOPS/sample, respectively – the specialist is over 7× more expensive per prediction. The cascade’s cost, FLOPSprop (t) = FLOPSbase + r(t) FLOPSspec , increases monotonically with the escalation rate r(t) from Fig. 6a: every sample incurs the base cost, and only the escalated fraction additionally incurs the specialist cost. For t ≲ 0.35, escalation is negligible, and the cascade coincides with the base; as t → 1, all samples are escalated, and the cascade cost exceeds the standalone specialist, since the base pass is retained on top of full specialist inference. At t⋆ = 0.78, the cascade costs approximately 5.33 × 104 FLOPS/sample – only 33.9% of the specialist’s cost, a 2.95× reduction. Together with Fig. 6b, this establishes the central efficiency claim: the cascade matches the specialist’s 94.45% Top-1 accuracy at roughly one-third of its FLOPS budget, since the base handles the 80.17% of samples on which it is already confident, leaving only the remaining 19.83% to incur the specialist cost. Even at the peak-accuracy setting t ≈ 0.86 (94.62%, +0.17% over the specialist), the cost is approximately 7.20 × 104 FLOPS – 45.9% of the specialist’s cost, a 2.18× reduction. The cascade thus dominates the standalone specialist on both axes across 0.30 ≲ t ≲ 0.90: equal-or-higher accuracy at a fraction of the compute. This trade-off follows directly from reliability-weighted routing: the near-free base correctly disposes of most samples, so the expensive specialist is invoked only where it yields a real accuracy gain, avoiding the waste of running it on samples the base would already classify correctly.

Taken together, Figs. 6a–6c establish two key findings. First, the proposed reliability-weighted cascade matches or exceeds the accuracy of the standalone specialist across all three Top-K metrics, despite the specialist being more than 7× more expensive per prediction than the base. Second, the operating point is attained while routing only approximately 19.83% of samples to the specialist, meaning that the base alone serves more than four out of five inference requests. These observations indicate that the proposed cascade delivers specialist-level (or better) accuracy at roughly one third of the computational footprint of always invoking the specialist. VII. C ONCLUSION This paper presented a latency-aware beam management framework for mmWave MIMO systems that combines digitaltwin-based channel prediction, region-based beam localization, and explainability-driven beam pruning within a unified, cascaded classification architecture. By exploiting offline digitaltwin knowledge and a base–specialist cascade with weighted ensembling, the framework resolves the majority of prediction cases efficiently while reserving heavier computation for harder cases, substantially reducing beam sounding overhead without compromising alignment accuracy. Simulation results under a digital-twin mismatch protocol demonstrate that the proposed framework achieves consistently high beam prediction accuracy, outperforming existing baselines and remaining robust across varying levels of digital-twin fidelity and beam-subset budgets. These results confirm the framework’s practicality as a scalable, low-latency solution for beam management in nextgeneration mmWave and 6G networks. R EFERENCES [1] M. Giordani, M. Polese, A. Roy, D. Castor, and M. Zorzi, “A tutorial on beam management for 3GPP NR at mmWave frequencies,” IEEE Communications Surveys & Tutorials, vol. 21, no. 1, pp. 173–196, Firstquarter 2019. [2] Q. Xue, C. Ji, S. Ma, J. Guo, Y. Xu, Q. Chen, and W. Zhang, “A survey of beam management for mmWave and THz communications towards 6G,” IEEE Communications Surveys & Tutorials, vol. 26, no. 3, pp. 1520– 1559, 2024. [3] Q. Xue, J. Guo, B. Zhou, Y. Xu, Z. Li, and S. Ma, “AI/ML for beam management in 5G-Advanced: A standardization perspective,” IEEE Vehicular Technology Magazine, vol. 19, no. 4, pp. 64–72, Dec. 2024. [4] S. Jiang and A. Alkhateeb, “Digital twin based beam prediction: Can we train in the digital world and deploy in reality?” in Proc. IEEE Int. Conf. Communications Workshops (ICC Workshops), Rome, Italy, 2023, pp. 36–41. [5] I. Ioannou, M. Georgiades, P. Nagaradjane, A. Khalifeh, C. Christophorou, M. Raspopoulos, and V. Vassiliou, “A coupled multi-stage hybrid framework for ber prediction and beam angle optimization in massive mimo systems: Combining classical regression with coupled deep learning approaches,” Network, vol. 6, no. 2, p. 35, 2026. [6] N. Khan, A. Abdallah, A. Celik, A. M. Eltawil, and S. Coleri, “Explainable and robust millimeter wave beam alignment for AI-native 6G networks,” in Proc. IEEE Int. Conf. Communications (ICC), 2025, pp. 753–758. [7] N. Khan, A. Abdallah, A. Celik, and A. M. Eltawil, “Digital twin-assisted explainable AI for robust beam prediction in mmWave MIMO systems,” IEEE Transactions on Wireless Communications, vol. 25, pp. 2435–2451, 2026.

[8] Y. Zhang, T. Osman, and A. Alkhateeb, “A digital twin assisted framework for interference nulling in millimeter wave MIMO systems,” in Proc. IEEE Int. Conf. Communications Workshops (ICC Workshops), Rome, Italy, 2023, pp. 248–253. [9] J. Zhang, S. Xu, C. Li, Y. Huang, and L. Yang, “Efficient beam selection for isac in cell-free massive mimo via digital twin-assisted deep reinforcement learning,” IEEE Transactions on Wireless Communications, vol. 25, pp. 9875–9890, 2026. [10] N. Khan, S. Coleri, A. Abdallah, A. Celik, and A. M. Eltawil, “Explainable and robust artificial intelligence for trustworthy resource management in 6G networks,” IEEE Communications Magazine, vol. 62, no. 4, pp. 50–56, 2024. [11] K. Ma, D. He, H. Sun, Z. Wang, and S. Chen, “Deep learning assisted calibrated beam training for millimeter-wave communication systems,” IEEE Transactions on Communications, vol. 69, no. 10, pp. 6706–6721, Oct. 2021. [12] W. Liu, H. Ren, C. Pan, and J. Wang, “Deep learning based beam training for extremely large-scale massive MIMO in near-field domain,” IEEE Communications Letters, vol. 27, no. 1, pp. 310–314, Jan. 2023. [13] W. Deng, M. Li, M.-M. Zhao, M.-J. Zhao, and O. Simeone, “Csi transfer from sub-6g to mmwave: Reduced-overhead multi-user hybrid beamforming,” IEEE Journal on Selected Areas in Communications, vol. 43, no. 3, pp. 973–987, 2025. [14] K. Chen, C. Qi, C.-X. Wang, and G. Y. Li, “Beam training and tracking for extremely large-scale MIMO communications,” IEEE Transactions on Wireless Communications, vol. 23, no. 5, pp. 5048–5062, May 2024. [15] Y. Heng and J. G. Andrews, “Machine learning-assisted beam alignment for mmwave systems,” IEEE Transactions on Cognitive Communications and Networking, vol. 7, no. 4, pp. 1142–1155, 2021. [16] V. Ardianto Nugroho and B. Moo Lee, “Gps-aided deep learning for beam prediction and tracking in uav mmwave communication,” IEEE Access, vol. 13, pp. 117 065–117 077, 2025. [17] M. B. Mollah, H. Wang, M. A. Karim, and H. Fang, “Multi-modality sensing in mmwave beamforming for connected vehicles using deep learning,” IEEE Transactions on Cognitive Communications and Networking, vol. 12, pp. 327–341, 2026. [18] I. Rasheed and H. Mostafa, “Deepbeam: A multi-agent deep reinforcement learning framework for predictive mmwave beam management in dynamic v2x networks,” IEEE Transactions on Vehicular Technology, vol. 74, no. 11, pp. 17 854–17 864, 2025. [19] J. Morais, A. Bchboodi, H. Pezeshki, and A. Alkhateeb, “Position-aided beam prediction in the real world: How useful gps locations actually are?” in ICC 2023 - IEEE International Conference on Communications, 2023, pp. 1824–1829. [20] X. Zhang, H. Zhang, C. Li, Y. Huang, and X. You, “Environment-specific beam training for extremely large-scale MIMO systems via contrastive learning,” IEEE Communications Letters, vol. 28, no. 1, pp. 167–171, Jan. 2024. [21] A. M. G. Alawadi, H. Seyedarabi, R. Afrouzian, and J. Musevi Niya, “A hotspot prediction scheme based on ensemble learning with adaptive beamforming for virtual small cell in 5G networks,” IEEE Access, vol. 13, pp. 109 435–109 450, 2025. [22] N. Zeulin, A. Ponomarenko-Timofeev, O. Galinina, and S. Andreev, “Mlassisted beam selection via digital twins for time-sensitive industrial iot,” IEEE Internet of Things Magazine, vol. 5, no. 1, pp. 36–40, 2022. [23] H. Luo and A. Alkhateeb, “Digital twin aided millimeter wave mimo: Site-specific beam codebook learning,” ICC 2026 - IEEE International Conference on Communications, pp. 1–6, 2025. [24] A. Bonora, A. V. Guglielmi, D. Scazzoli, M. Giordani, M. Magarini, V. Teeda, and S. Tomasin, “Digital twin-based beamforming for interference mitigation in AF relay MIMO systems,” arXiv preprint arXiv:2602.22991, Feb. 2026. [25] N. Khan, A. Abdallah, A. Celik, A. M. Eltawil, and S. Coleri, “Explainable AI-aided feature selection and model reduction for DRL-based V2X resource allocation,” 2025, pp. 1–1, early access. [26] S. M. Lundberg, G. Erion, H. Chen, A. DeGrave, J. M. Prutkin, B. Nair, R. Katz, J. Himmelfarb, N. Bansal, and S.-I. Lee, “From local explanations to global understanding with explainable ai for trees,” Nature Machine Intelligence, vol. 2, pp. 56–67, 2020.

Related documents

• Digital Twin-Assisted Fast Beam Alignment for mmWave MIMO Networks#1536916
• SpaceBeam#1341220
• Knowledge Distillation for mmWave Beam Prediction Using Sub-6 GHz Channels#798758
• Knowledge Distillation for mmWave Beam Prediction Using Sub-6 GHz Channels#797344
• Beamforming Design Via GNN in mmWave Cell-Free Massive MIMO Using Sub-6 GHz CSI#799717