Two-Stage Causal Beam Selection

• The paper demonstrates that causal discovery identifies key RSSI measurements which drive optimal beam prediction in mmWave systems.
• It employs a two-stage process where DirectLiNGAM first learns a causal graph and then selects a minimal set of sensing beams for DNN classification.
• Empirical results on the Boston5G scenario show top-1 accuracy of ~87% and a 59.4% reduction in beam sweeping overhead.

Searching arXiv for the cited paper and closely related beam-selection works. The Two-Stage Causal Beam Selection Algorithm is a causally-aware deep learning procedure for beam management in mmWave multiple-input multiple-output systems, in which causal discovery first learns a directed acyclic graph over received power measurements and the optimal beam, and causal feature selection then restricts beam prediction to a minimal set of causally relevant sensing beams. In the formulation introduced for reliable initial access in AI-driven beam management, the algorithm is designed for narrowband mmWave downlink MIMO with analog-only beamforming, and addresses the claim that conventional deep learning beam alignment methods often rely on statistical correlations rather than causal structure, with consequences for interpretability, generalization, and beam sweeping overhead (Khan et al., 22 Aug 2025).

1. Problem Setting and Motivation

Future 6G systems are expected to rely heavily on millimeter-wave frequencies to achieve ultra-high data rates. In this regime, path loss is severe, channels are sparse in the angular domain, and highly directional beams are required. Beam management therefore becomes a central control problem, especially during initial access, where the base station must identify a viable beam pair for a user equipment before data transmission. In conventional mmWave operation, this is typically implemented through beam sweeping over a beam codebook, with the user measuring received signal strength indicators and selecting the best beam (Khan et al., 22 Aug 2025).

The main difficulty is that practical codebooks are large. With $N_{\mathrm{BS}}$ antennas and angular oversampling, the narrow-beam codebook can easily reach $Y \sim 128$ or more beams, so exhaustive sweeping incurs high latency and consumes time-frequency resources. This is compounded by channel sparsity, mobility, blockage, measurement noise, hardware impairments, incomplete sensing, and non-stationary environments. The desired alternative is AI-driven beam prediction using a small set of sensing beams that can still predict the best narrow beam, or the top-$k$ beams, with high accuracy.

Within this setting, the causal beam selection proposal is motivated by four limitations attributed to conventional deep learning beam alignment. First, feature relevance is usually inferred from statistical correlation rather than interventionally meaningful dependence, so models may exploit spurious correlations and degrade under domain shift. Second, post-hoc explainable AI tools such as SHAP provide model-dependent importance scores and require a trained predictor. Third, statistical relevance may not be stable across environments. Fourth, feature-selection overhead can itself be large; the paper explicitly contrasts SHAP/XAI feature selection, which requires a pre-trained model and is computationally expensive, with a causal-discovery-based alternative (Khan et al., 22 Aug 2025).

2. System Model and Causal Formulation

The algorithm is defined for a narrowband mmWave downlink with a base station equipped with a uniform linear array of $N_{\mathrm{BS}}$ antennas, $N_{\mathrm{U}}$ single-antenna users, and analog-only beamforming with a single RF chain. During beam sweeping, the base station transmits pilots over a beam codebook $\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$, and the optimal beam index for user $u$ is the beam maximizing received power: $$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$

To reduce overhead, the system does not use all narrow beams as inputs. Instead, it sweeps a small set of wide sensing beams from an $N_{\mathrm{BS}}$-DFT codebook, with $M_{\mathrm{w}} \ll Y$. For each user $Y \sim 128$0, the resulting observation vector is

$Y \sim 128$1

where the entries are RSSIs measured under the sensing beams. The supervised learning dataset is

$Y \sim 128$2

with labels obtained by noise-free exhaustive search over an O-DFT codebook of size $Y \sim 128$3 (Khan et al., 22 Aug 2025).

A standard deep neural beam classifier $Y \sim 128$4 maps $Y \sim 128$5 to a softmax over $Y \sim 128$6 narrow beams and is trained with cross-entropy loss. The causal reformulation augments this with a directed acyclic graph $Y \sim 128$7 whose nodes are the RSSI variables $Y \sim 128$8 and the target $Y \sim 128$9. A causal edge $k$0 is interpreted as meaning that changing RSSI measurement $k$1 by intervention would change the distribution of the optimal beam. The objective is then to identify the parents, or more broadly the Markov blanket, of $k$2, since these form a minimal set of causally relevant predictors.

The central modeling assumption is that the target beam index is not a cause of any RSSI feature, so edges $k$3 are excluded. The paper treats this as a physically grounded directionality: channel and beamforming generate the measured powers, while the optimal beam index is a consequence of that generative process.

3. Algorithmic Structure

The proposed method has two stages. Stage 1 performs causal discovery via DirectLiNGAM, learning a Bayesian graph over the sensing-beam RSSIs and the target beam index. Stage 2 performs causal feature selection using that graph, and trains the beam classifier only on the resulting reduced feature set (Khan et al., 22 Aug 2025).

For causal discovery, the variables are stacked as

$k$4

and modeled by the Direct Linear Non-Gaussian Acyclic Model

$k$5

where $k$6 is a strictly lower-triangular effect matrix and $k$7 contains mutually independent, non-Gaussian disturbances. A nonzero entry $k$8 denotes a directed edge $k$9. DirectLiNGAM infers a causal ordering from regression residual independence, then estimates $N_{\mathrm{BS}}$0 by least squares on predecessors in that ordering.

Once $N_{\mathrm{BS}}$1 is learned, the direct parents of the target beam index are

$N_{\mathrm{BS}}$2

These are the core causal predictors. To include strong indirect influencers, the method also computes, for each feature $N_{\mathrm{BS}}$3,

$N_{\mathrm{BS}}$4

$N_{\mathrm{BS}}$5

and ranks features by $N_{\mathrm{BS}}$6 in descending order. The selected feature set $N_{\mathrm{BS}}$7 is initialized with $N_{\mathrm{BS}}$8, then augmented by highly connected nodes until the desired reduced dimension $N_{\mathrm{BS}}$9 is reached.

Stage	Operation	Output
1	Train DirectLiNGAM, derive causal ordering, estimate $N_{\mathrm{U}}$0	DAG / Bayesian graph
2	Select $N_{\mathrm{U}}$1, extend by $N_{\mathrm{U}}$2, train DNN on $N_{\mathrm{U}}$3	Reduced sensing-beam set and classifier

The reduced input for user $N_{\mathrm{U}}$4 is

$N_{\mathrm{U}}$5

with $N_{\mathrm{U}}$6. The same DNN architecture is then trained on the restricted dataset. The paper emphasizes that DirectLiNGAM-based selection does not require a pre-trained model, unlike SHAP, and that both sensing overhead and model complexity shrink with $N_{\mathrm{U}}$7.

4. Complexity, Overhead, and Empirical Performance

Operationally, the method has an offline phase and an online phase. Offline, it collects $N_{\mathrm{U}}$8, runs DirectLiNGAM, extracts $N_{\mathrm{U}}$9, and trains a DNN on the reduced feature space. Online, the base station sweeps only the $\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$0 sensing beams in $\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$1, users report the corresponding RSSIs, and the DNN predicts top-$\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$2 narrow beams from the O-DFT codebook; the base station may then sweep the predicted narrow beams for refinement (Khan et al., 22 Aug 2025).

The resulting overhead scales linearly in $\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$3. Beam sweeping complexity is $\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$4 sensing beams, and if top-$\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$5 refinement is used, sweeping overhead is

$\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$6

with feedback overhead

$\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$7

Model complexity is also reduced because the DNN parameter count depends on $\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$8. The paper explicitly reports DirectLiNGAM-based beam selection runtime of approximately $\mathbf{W}=\{\mathbf{w}_1,\dots,\mathbf{w}_Y\}$9 s, versus approximately $u$0 s for SHAP-based feature selection, corresponding to a $u$1 reduction in input selection time (Khan et al., 22 Aug 2025).

The evaluation uses the Boston5G scenario from DeepMIMO. The base station is a 32-element ULA at 28 GHz with 500 MHz bandwidth and transmit power 30 dBm; users are single-antenna terminals in a $u$2 area. Wide sensing uses $u$3 beams from an $u$4-DFT codebook, and narrow prediction uses an O-DFT codebook with oversampling factor 4, giving $u$5 narrow beams. The DNN uses hidden layers of 64, 64, and 128 ReLU neurons, Adam with learning rate $u$6, 100 epochs, and a dataset split of $u$7 train, $u$8 validation, and $u$9 test, with $$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$0 (Khan et al., 22 Aug 2025).

The learned DAG identifies direct parents of the target node at RSSI$$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$1, RSSI$$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$2, RSSI$$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$3, RSSI$$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$4, RSSI$$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$5, and RSSI$$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$6, with additional features selected through connectivity strength. Using only $$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$7 selected sensing beams, the proposed method achieves top-1 accuracy of approximately $$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$8 and top-2 accuracy of approximately $$y_u^*=\arg\max_{y_u\in\{1,2,\dots,Y\}} \left|\mathbf{h}_u^H \mathbf{w}_{y_u}\right|^2 .$$9. The reported beam sweeping overhead reduction is $N_{\mathrm{BS}}$0, corresponding to reducing sensing beams roughly from 32 to approximately 13. In effective spectral efficiency, the method with $N_{\mathrm{BS}}$1 beams outperforms the SVD-based solution and matches SHAP-based feature selection, while exhaustive search has the lowest spectral efficiency because of large initial-access overhead. The spectral-efficiency curve exhibits a unique maximizer $N_{\mathrm{BS}}$2, increasing for $N_{\mathrm{BS}}$3 and decreasing when initial-access overhead dominates (Khan et al., 22 Aug 2025).

5. Interpretation, Significance, and Terminological Distinctions

The learned causal graph is intended to separate causally relevant beams from redundant or spurious beams. In the paper’s interpretation, beams near the dominant angle-of-departure region, such as RSSI$N_{\mathrm{BS}}$4–RSSI$N_{\mathrm{BS}}$5, show strong causal links to the optimal narrow-beam index, while other beams may be merely correlated or noisy. This underpins the claim that removing causally irrelevant features can reduce noise and overfitting while preserving alignment accuracy (Khan et al., 22 Aug 2025).

A central misconception in the broader literature is that the word causal always denotes the same concept. In the causal-discovery formulation above, causality refers to structural relations in a DAG and to feature relevance under interventions. In several adjacent beam-management papers, by contrast, causal means that beam decisions are made sequentially from past and current observations only, with no look-ahead. Posterior-matching initial alignment, heteroscedastic Track-and-Stop beam alignment, and optimized two-stage search all use this sequential sense of causality (Ronquillo et al., 2019, Wei et al., 2022, Li et al., 2018). The distinction matters because the first sense concerns structural explainability and generalization, whereas the second concerns online decision admissibility.

The practical significance of the DAG-based approach lies in three related properties emphasized by the paper: reduced sensing codebooks, interpretable beam policies, and improved robustness under domain shift when the causal mechanism is preserved. A plausible implication is that the method is especially attractive when operators require both initial-access efficiency and an auditable explanation of why specific sensing beams drive the final narrow-beam prediction.

6. Relation to Other Two-Stage Beam Selection Paradigms

The broader arXiv literature contains several two-stage beam-selection designs, but they instantiate different notions of structure. In adaptive initial alignment, "Sequential Learning of CSI for MmWave Initial Alignment" organizes two sequential algorithms around a learning stage for $N_{\mathrm{BS}}$6 and a posterior-matching beam-refinement stage over a hierarchical codebook; the causal aspect is explicitly no-look-ahead posterior updating (Ronquillo et al., 2019). In pure-exploration bandit beam alignment, the Two-Phase Heteroscedastic Track-and-Stop algorithm first identifies a best beam set and then refines within a neighboring local set, exploiting both correlation among nearby beams and heteroscedastic Gaussian rewards (Wei et al., 2022).

A different two-stage principle appears in "Explore and Eliminate: Optimized Two-Stage Search for Millimeter-Wave Beam Alignment", where all beam pairs are measured once with small energy, the $N_{\mathrm{BS}}$7 worst are eliminated, and surviving beams are remeasured and coherently combined; this is an explore-then-eliminate design rather than a causal-discovery design (Li et al., 2018). In distributed antenna systems, the two-stage schedule-and-align framework performs coarse scanning over the full angular range, obtains beam indices and peak-to-background ratios from feedback, then reconfigures search angles, steps, and power for a refined second-stage search (Wei et al., 2019).

Deep-learning-based probing methods also employ a two-stage decomposition. "Entropy-based Probing Beam Selection and Beam Prediction via Deep Learning" first selects probing beams from a location-specific codebook designed by an entropy-based criterion and then identifies the optimal beam using additional probing beams with the highest predicted power values (Meng et al., 2024). In abruptly changing environments, K-sequential halving and exhaustive search uses sequential halving until $N_{\mathrm{BS}}$8 beams remain and then exhaustive search on that restricted set, explicitly targeting the case where a near-optimal beam becomes optimal during selection (Ghatak, 2023).

Taken together, these works show that two-stage may denote causal discovery followed by feature selection, posterior learning followed by refinement, beam-set search followed by local identification, coarse scan followed by fine scan, or entropy-based coarse probing followed by targeted verification. This suggests that the term identifies a design pattern rather than a single algorithmic family, while the 2025 causal beam selection method is distinctive in making the first stage a structural causal-discovery problem rather than a coarse search heuristic.

7. Assumptions, Limitations, and Extensions

The causal beam selection formulation is explicitly tied to a narrowband mmWave downlink, analog-only beamforming with a single RF chain, single-antenna users, and coherence-time stationarity during beam sweeping and reporting. It also inherits the assumptions of DirectLiNGAM: linear relations in $N_{\mathrm{BS}}$9, non-Gaussian and independent disturbances, and acyclicity of the learned graph. The DAG is learned offline for a given environment, so significant environmental change may require retraining (Khan et al., 22 Aug 2025).

The reported validation is scenario-specific. The evaluation uses Boston5G in DeepMIMO with a particular geometry, 32 sensing beams, 128 narrow beams, and a specific DNN architecture, so generalization to indoor, rural, or highly dynamic environments is described as promising but not proven. The paper also notes that multi-user interference and coordinated beams are not explicitly modeled, and that although DirectLiNGAM is scalable, very large arrays and codebooks may still create complexity issues.

The extensions proposed in the source are correspondingly broad: online or incremental causal discovery, multi-cell and multi-user coordination, integration of the causal graph with reinforcement learning, extension to hybrid analog-digital beamforming and multi-stream MIMO, and semantic-level causal modeling that couples physical-layer beam selection to higher-level KPIs such as QoS and mobility (Khan et al., 22 Aug 2025). Within that trajectory, the two-stage causal beam selection algorithm marks a specific synthesis: causal discovery is used not as post-hoc explanation, but as the mechanism that determines which sensing beams are worth measuring at all.