MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

Muon has recently emerged as a promising alternative to AdamW for language model pretraining by orthogonalizing momentum matrices using Newton-Schulz iterations. Although Muon mitigates gradient anisotropy, it does not explicitly account for the curvature geometry of the loss landscape and may therefore remain sensitive to curvature anisotropy. We bridge this gap by proposing MALT (Muon Augmented by Lightweight Two-sided Preconditioning), which uses lightweight diagonal preconditioners to reduce the sensitivity of Muon to curvature anisotropy. Specifically, MALT uses two-sided diagonal preconditioners with low memory and computational overhead to approximately capture the curvature geometry of the loss landscape. It orthogonalizes the preconditioned momentum using Newton-Schulz iterations and maps the result back to define the update direction, while norm grafting controls the update magnitude. To improve the robustness of MALT to stochastic gradient noise, we further propose MALTER (MALT with Adaptive stEpsize Rescaling). Convergence guarantees are provided for MALT in the stochastic non-convex setting. Experiments on GPT-2 Small, Medium, and Large pretraining show that the proposed methods outperform Muon while maintaining nearly the same memory footprint and wall-clock time.

Learning to Beamform for Cooperative Localization and Communication: A Link Heterogeneous GNN-Based Approach

Integrated sensing and communication (ISAC) has emerged as a key enabler for next-generation wireless networks, supporting advanced applications such as high-precision localization and environment reconstruction. Cooperative ISAC (CoISAC) further enhances these capabilities by enabling multiple base stations (BSs) to jointly optimize communication and sensing performance through coordination. However, CoISAC beamforming design faces significant challenges due to system heterogeneity, large-scale problem complexity, and sensitivity to parameter estimation errors. Traditional deep learning-based techniques fail to exploit the unique structural characteristics of CoISAC systems, thereby limiting their ability to enhance system performance. To address these challenges, we propose a Link-Heterogeneous Graph Neural Network (LHGNN) for joint beamforming in CoISAC systems. Unlike conventional approaches, LHGNN models communication and sensing links as heterogeneous nodes and their interactions as edges, enabling the capture of the heterogeneous nature and intricate interactions of CoISAC systems. Furthermore, a graph attention mechanism is incorporated to dynamically adjust node and link importance, improving robustness to channel and position estimation errors. Numerical results demonstrate that the proposed attention-enhanced LHGNN achieves superior communication rates while maintaining sensing accuracy under power constraints. The proposed method also exhibits strong robustness to communication channel and position estimation error.