Searcharxiv⌕ Search

arXiv subjects

Zhiheng Zhang

Publications and source records attributed to Zhiheng Zhang.

At least 19 recordsLinked to original sources

Causal Influence Maximization with Steady-State Guarantees

Classical influence maximization optimizes expected reach, which can be misaligned with welfare when exposures have heterogeneous, saturating, or adverse effects. We study seed selection for steady-state causal welfare under network interference. Under our exposure-separable outcome model and weak independent live-edge propagation, a deterministic surrogate built from expected activations and exposure counts approximates welfare with uniform $O(\varepsilon^2)$ error for a fixed graph and seed budget, where $\varepsilon$ bounds edge activation probabilities. Building on this reduction, our two-stage method, CIM, learns shape-constrained exposure-response functions from logged diffusion-outcome data and estimates expected activations and exposures by Monte Carlo simulation. For unregularized fits on a design that covers the evaluated exposures, and under noise and simulation conditions, we bound the welfare-estimation error uniformly over seed sets, separating structural, statistical, and simulation errors; with an approximation guarantee for the optimizer, the bound transfers to the selected seeds. Experiments on real network topologies with simulated outcomes and a synthetic benchmark show that CIM attains the highest mean welfare among the compared methods.

stat.ME↗

Learning to Fluctuate: Statistical Foundations for Causal Tabular Pretraining

Causal tabular foundation models amortize effect estimation across synthetic mechanisms, but latent-effect supervision rewards posterior shrinkage rather than encoding the repeated-sample response needed in a fixed deployment population. We introduce fluctuation-supervised pretraining (FSP): each synthetic table is labeled by its average treatment effect plus its efficient influence-function fluctuation; deployment remains a frozen forward pass. Along the path $T_{λ,P}=θ(P)+λP_nψ_P$, we prove an endpoint transition: every fixed $λ<1$ retains label ambiguity of order $(1-λ)^2/n$, whereas full fluctuation makes the Gaussian label observable and reduces optimal finite-stratum causal label-prediction risk to order $n^{-2}$. A finite-pretraining bound combines label, network, episode-sampling, and optimization errors; its sampling defect controls fixed-mechanism bias, mean squared error, variance, Gaussian approximation, and, with variance-head accuracy, studentized coverage. Complementary lower bounds separate local $n^{-1}$ ATE risk from the $\log N/M$ excess risk of generic finite-dictionary episode learning. Experiments trace the learned sampling response. Across 24 nonlinear continuous-covariate cells at trained context lengths, continuous-row FSP lowers checkpoint-mean macro RMSE by 7.0% versus S-learner and wins all 12 weak-overlap cells; validation-selected Summary FSP deploys $11.6\times$ faster per table in our warm one-thread benchmark. Under effect shift, matched Raw FSP lowers mean-checkpoint RMSE by 54.2% and teacher defect by 99.0% versus latent-effect supervision, and RMSE by 10.2% versus the released CausalPFN-S checkpoint. Known-effect semisynthesis tests coverage; two randomized-study evaluations show that lower RMSE can coexist with residual attenuation.

stat.ML↗

Fengshui: Demystifying Chiplet Ecosystem and Bespoke Neural Network Accelerator Codesign

Modern ML workloads, with stringent latency and energy constraints, are increasingly hard to run efficiently on homogeneous commodity hardware. We argue that operator-level disaggregation--tailoring microarchitecture, batching, and memory hierarchy to each operator--is essential to overcome these limitations, though the resulting highly bespoke accelerators incur prohibitive Non-Recurring Engineering (NRE) costs. Chiplet-based integration amortizes NRE across applications, but choosing which chiplets to build and how to compose them into accelerators is circularly dependent--a chiplet pool's value depends on the constructed accelerators, while accelerator quality is constrained by available chiplets. This paper introduces Fengshui, a chiplet ecosystem and accelerator co-design framework that jointly optimizes chiplet pool composition and bespoke application-specific integrated circuit (BASIC) design. Fengshui constructs BASICs through operator-level disaggregation, co-exploring chiplet and memory heterogeneity, tensor fusion, and pipeline/tensor/expert parallelism with place-and-route validation for physical implementability. With just 8 strategically selected chiplets, encompassing network switches, processing-in-memory units, and accelerators with diverse microarchitectures, Fengshui-generated BASICs achieve 48.5%, 88.1%, 93.0%, and 97.8% reductions in energy, energy-cost product (EC), energy-delay product (EDP), and energy-delay-cost product (EDPC) over homogeneous accelerators, while scoring within 4.1% of unconstrained heterogeneous designs across diverse neural networks. For datacenter MoE and dense LLM serving, Fengshui reduces prefill energy and EC by up to 16.8% and 28.7%, respectively; for edge autonomous vehicle perception, it achieves 12.0% energy and 23.6% EC reductions under real-time latency constraints.

cs.AR↗

Beyond Factual Knowledge: Benchmarking and Learning Step-Level Procedural Rule Reasoning in Large Language Models

Large language models (LLMs) excel at text understanding and generation, yet still struggle to reliably understand and apply externally provided procedural rules at scale. To evaluate this capability, we introduce RuleWorld, a large-scale benchmark that reformulates rules as globally reusable abstract units rather than instance-specific facts. In RuleWorld, several scenarios, including single-rule, parallel multi-rule, and multi-hop reasoning, are settled for comprehensive evaluation. We further propose DynaRule, an end-to-end framework that injects the given rules into the KV cache and turns retrieval into an internal, learnable, step-wise process. Specifically, DynaRule employs Stacked Step-Level Attention Training with a special token to enable dynamic rule re-attention and updating during inference. In this way, the model can re-attend to the most relevant rules at each step, dynamically replacing outdated ones to support more stable multi-step reasoning. Experiments on RuleWorld show that existing LLMs face challenges under large rule pools, while DynaRule improves average QA accuracy by up to 19 points and achieves over 85% Recall@1 at 10K rules, outperforming strong baselines by large margins. We make our code and dataset available here: https://github.com/SharkSpicy-NLP/Beyond-Factual-Knowledge.

cs.CL↗

A Switch-Centric In-Network Architecture for Accelerating LLM Inference in Shared-Memory Network

Tensor parallelism (TP) has become a key technique for latency-sensitive LLM inference, but it introduces frequent, tightly synchronized All-Reduce operations that lie directly on the inference critical path, making inter-GPU communication a major performance bottleneck. In-network computing offers a promising approach to mitigating this bottleneck. However, we observe that the existing solution, NVLink SHARP (NVLS), is fundamentally constrained by its GPU-driven execution model that exposes in-network computing as GPU-issued, element-granular memory operations rather than a switch-orchestrated collective. This design forces redundant switch-GPU-switch data movement and cannot perform scale-aware quantized All-Reduce due to its lack of cross-element coordination. To address these limitations, we propose SCIN, the first switch-centric in-network computing architecture for multi-accelerator shared-memory networks. SCIN replaces GPU-driven collective orchestration with an in-switch accelerator (ISA) that directly initiates memory transactions and orchestrates the complete All-Reduce operation. This switch-driven execution eliminates redundant data movement inherent in NVLS, shortens synchronization paths, and fully offloads collective execution from GPUs. Furthermore, centralized ISA control enables cross-element coordination between quantized data and their scale factors. SCIN exploits this capability through a pipelined dequantization-reduction-requantization datapath, enabling scale-aware in-network quantized All-Reduce and nearly halving communication volume with negligible accuracy degradation. We implement SCIN on a multi-FPGA prototype to validate its feasibility. For a simulated 8-GPU system, SCIN accelerates small- and large-message All-Reduce by up to 1.8x and 2.6x over NVLS, respectively, translating into up to 1.12x TPOT and 1.42x TTFT speedups.

cs.AR↗

OoO-Spec: Out-of-Order Semantic Speculation for Fast Tool Calling

LLMs generate tool calls token by token, even though the function choice and argument values can often be predicted in parallel from the request and tool schema. ToolSpec reduces this cost by drafting schema tokens and retrieving earlier calls, but cannot propose request-specific values absent from either source. We present OoO-Spec, which computes these missing semantics out of order. At request arrival, a Qwen3-0.6B sidecar predicts the function choice and all schema-defined argument slots in one parallel request-level wave while the target begins ToolSpec decoding. The runtime joins the slot values, renders the resulting call as text, and exposes it to subsequent candidate-construction rounds. The target polls without blocking, re-tokenizes a ready hint with its own tokenizer, and remains the sole verifier and commit authority. The sidecar is trained once with LoRA on Qwen2.5-32B teacher traces and used unchanged across Qwen2.5, Qwen3, and Llama targets, without target-specific drafter training. Across seven fully ranked targets and three benchmarks under greedy batch-one decoding, OoO-Spec is fastest among all evaluated methods in all 21 target-benchmark cells, reaching 2.46x-5.34x over autoregressive decoding with an unweighted mean of 3.89x, versus 2.95x for ToolSpec. It also outperforms every evaluated released learned drafter in each comparable cell. Across Qwen3-4B, 8B, 14B, and 32B targets, the same sidecar improves on ToolSpec by 34.1% on average. Its compact semantic payload averages 85 bytes per request excluding protocol metadata, supporting effective split-GPU overlap.

cs.CL↗

DY-LUT: Depth-Aware YCbCr Lookup Tables for Real-Time Underwater Image Enhancement

Underwater image enhancement is challenged by spatially non-uniform, wavelength-dependent attenuation. Propagation distance and wavelength govern this degradation, while YCbCr separates luminance from chrominance for restoration. We propose DY-LUT, a depth-aware YCbCr lookup-table framework for real-time enhancement. A dual-branch encoder predicts image-level fusion weights and a joint pair of pixel-wise degradation indices from image and depth features. These quantities condition learnable 4D LUTs, followed by lightweight local refinement. DY-LUT preserves traditional LUT efficiency while enabling depth-conditioned, spatially adaptive restoration. With externally supplied depth, its 3.56M-parameter enhancement network achieves competitive quality on UIEB-90 and LSUI and runs $9$--$304\times$ faster than representative high-capacity baselines. Adaptive inference further maintains real-time performance ($\sim7$ ms) for 4K UIQAD images. DY-LUT also benefits downstream detection and feature matching. Ablations show that YCbCr is a more effective basis than RGB for depth-conditioned lookup, while the jointly learned indices further improve adaptive querying. These results provide a physically grounded route to efficient UIE on practical platforms.

eess.IV↗

No Training, Better Flights: Test-Time Scaled VLMs for UAV Navigation

Test-time scaling offers a promising method to improve the inference performance of Vision-Language Models (VLMs) without additional training. Existing approaches to vision-language navigation (VLN) for Unmanned Aerial Vehicle (UAV) typically relies on a single inference pass, which can falter in complex environments by producing suboptimal or unsafe trajectories. In this paper, we explore a simple and effective approach to apply test-time scaling to VLN for UAV. We enhance navigation reasoning through an iterative refinement process that requires no extra model training, guiding the model to re-evaluate its initial navigation plan for better accuracy and safety. Our method first prompts the model to generate multiple parallel candidates and then performs a self-correction step, achieving deeper and more robust planning without changing the underlying model. To further strengthen decision-making, we design a multi-criteria scoring function to evaluate the refined candidates based on safety, goal alignment, and forward-progress. This simple yet powerful combination enables a frozen UAV navigation VLMs to self-correct and generate more accurate and reliable flight plans, achieving SOTA performance in this task.

cs.CV↗

Rubrics on Trial: Evolving Rubrics from a Single Query via Synthetic Pairwise Evidence

Rubrics provide structured, fine-grained signals for training and evaluating large language models (LLMs). Yet reliable query-specific rubrics are difficult to construct. Existing approaches often derive supervision from human-written rubrics, preference data, or sampled responses. Direct query-to-rubric generation avoids these resources, but provides no explicit check that a plausible rubric is useful. Such a rubric may fail to distinguish answer quality, reward an optional style, or penalize a valid alternative strategy. We introduce Rubrics on Trial, a query-only framework that evolves a rubric set from an empty set without external annotations or model training. It derives supervision solely from synthetic rubric-conditioned response pairs and validates each proposed rubric before adding it, screening out non-discriminative, over-specific, and style-only candidate rubrics. Experiments across five preference benchmark suites demonstrate the effectiveness of Rubrics on Trial, which achieves the best average accuracy and leads on six of seven evaluation sets.

cs.CL↗

SpaCellAgent: A Self-Evolving LLM-Based Multi-Agent Framework for Trajectory Analysis

Spatial and Single-cell transcriptomics are transformative in deciphering cellular dynamics. As the fundamental paradigm for reconstructing cell developmental paths, trajectory inference (TI) is critical. However, existing methods require extensive manual intervention and proficiency in heterogeneous tools, posing a significant barrier to efficient TI analysis. To bridge this gap, we propose SpaCellAgent, an autonomous large language model (LLM) multi-agent framework that automates end-to-end spatiotemporal analysis and narrative generation. SpaCellAgent utilizes a multi-agent architecture for strategic workflow planning, a dynamic tool-orchestration engine for adaptive algorithm selection, and a self-evolution module that iteratively refines performance through feedback. We evaluate SpaCellAgent on six heterogeneous datasets encompassing complex temporal developmental trajectories, diverse sequencing platforms, and spatially-resolved tissue architectures. SpaCellAgent consistently demonstrates over 40\% improvement in analytical efficiency while maintaining expert-aligned performance. By converting natural language specifications into optimized analytical workflows and fully automating the pipeline, SpaCellAgent democratizes advanced spatiotemporal modeling and establishes a scalable, agent-driven paradigm for computational biology. The code and materials are available at https://github.com/LittleXH-shw/SpaCellAgent.

cs.AI↗

Denoised Conformal Alignment for Reliable Selection of Conditional Average Treatment Effect Predictions

In selective deployment, practitioners act only on a model-chosen subset of individuals based on predicted conditional average treatment effects, but marginal conformal guarantees need not control reliability on that selected subset. We study reliable selection for black-box CATE predictors: selecting candidates whose CATE errors are below a tolerance while controlling the false discovery rate (FDR). Since CATE errors are unobservable, we construct doubly robust proxy errors from pseudo-outcomes; however, naive proxies can lose power under heteroskedasticity because variance overwhelms the reliability signal. We propose Denoised Conformal Alignment, which subtracts an estimated conditional variance component and combines conformal calibration with Benjamini--Hochberg selection. Our analysis shows that validity is governed by stability of proxy/oracle threshold labels, rather than pointwise perfection of the variance estimator. Experiments show substantially improved power while maintaining FDR control across challenging settings.

stat.ML↗

Wasserstein Policy Learning for Distributional Outcomes

Offline policy learning has received growing attention in causal inference. The primary objective is to learn a policy (individualized treatment rule) as a mapping from covariates to treatment that maximizes the empirical welfare defined as the mean of scalar-valued potential outcomes. In this paper, we study offline policy learning with distribution-valued outcomes, where each potential outcome is a probability measure on $\mathbb{R}$ and the reward is defined through a utility functional applied to the Wasserstein barycenter of induced outcome distributions. We establish statistical guarantees for the policy learning framework based on both Inverse Probability Weighting (IPW) and Doubly Robust (DR) estimators. By handling the challenging uniform deviation over the product of the combinatorial policy class and the infinite-dimensional quantile domain, we prove that the finite-sample regret has leading dependence $\widetilde{\mathcal{O}}(\sqrt{\mathrm{N\text{-}dim}(Π)/N})$. In the one-dimensional Wasserstein setting and under the stated regularity conditions, the leading regret rate is still governed by the policy-class complexity. Moreover, we provide a minimax lower bound establishing the sharpness of the leading dependence on $N$ and $\mathrm{N\text{-}dim}(Π)$.

stat.ME↗

MLUBench: A Benchmark for Lifelong Unlearning Evaluation in MLLMs

Multimodal large language models (MLLMs) are trained on massive multimodal data, making data unlearning increasingly important as data owners may request the removal of specific content. In practice, these requests often arrive sequentially over time, giving rise to the challenging problem of MLLM Lifelong Unlearning. However, most existing benchmarks are limited in scale and scope, failing to capture the complexities of MLLM lifelong unlearning. To fill this gap, we introduce the MLUBench, a large-scale and comprehensive benchmark featuring 127 entities across 9 classes under lifelong unlearning requests. We perform extensive experiments using MLUBench and reveal that existing unlearning methods suffer from severe, cumulative degradation. More critically, we further identify the unique challenge of this problem: unlike in unimodal models, MLLM lifelong unlearning is constrained by the need to preserve multimodal alignment. Continually unlearning from one modality could degrade the entire model. To alleviate this challenge, we propose LUMoE, an effective method. Experiments demonstrate that LUMoE significantly mitigates the degradation problem faced by baselines. The source code and the MLUBench dataset are open-sourced in https://github.com/lihe-maxsize/Lifelong_Unlearning_main.

cs.AI↗

Partial Identification under High-Dimensional Potential Outcomes and Confounders via Optimal Transport

Partial identification provides informative causal guarantees when point identification is impossible, but existing approaches based on optimal transport (OT) become computationally and statistically intractable in high-dimensional settings. This limitation is particularly severe when both potential outcomes and confounders are high-dimensional, where classical OT-based bounds suffer from the curse of dimensionality and unfavorable convergence rates. To address this challenge, we propose a novel estimator that decomposes the transport problem into a low-dimensional signal subspace and a high-dimensional residual subspace. Unlike existing projection-based methods that discard residual information, we recover the residual transport energy using the Sliced Wasserstein distance, which is computationally efficient and robust to high dimensions. We establish interpretable conditions controlling the approximation gap based on residual structure and provide a data-driven rule for signal dimension selection. Empirical results show that our estimator consistently outperforms projection-only baselines by recovering lost transport energy, yielding more informative causal bounds while remaining computationally tractable in high dimensions.

stat.ME↗

Causal Representation Learning with Optimal Compression under Complex Treatments

Estimating Individual Treatment Effects (ITE) in multi-treatment scenarios faces two critical challenges: the Hyperparameter Selection Dilemma for balancing weights and the Curse of Dimensionality in computational scalability. This paper derives a novel multi-treatment generalization bound and proposes a theoretical estimator for the optimal balancing weight $α$, eliminating expensive heuristic tuning. We investigate three balancing strategies: Pairwise, One-vs-All (OVA), and Treatment Aggregation. While OVA achieves superior precision in low-dimensional settings, our proposed Treatment Aggregation ensures both accuracy and O(1) scalability as the treatment space expands. Furthermore, we extend our framework to a generative architecture, Multi-Treatment CausalEGM, which preserves the Wasserstein geodesic structure of the treatment manifold. Experiments on semi-synthetic and image datasets demonstrate that our approach significantly outperforms traditional models in estimation accuracy and efficiency, particularly in large-scale intervention scenarios.

cs.LG↗

Group Permutation Testing in Linear Model: Sharp Validity, Power Improvement, and Extension Beyond Exchangeability

We consider finite-sample inference for a single regression coefficient in the fixed-design linear model $Y = Zβ+ bX + \varepsilon$, where $\varepsilon\in\mathbb{R}^n$ may exhibit complex dependence or heterogeneity. We develop a group permutation framework, yielding a unified and analyzable randomization structure for linear-model testing. Under exchangeable errors, we place permutation-augmented regression tests within this group-theoretic setting and show that a grouped version of PALMRT controls Type I error at level at most $2α$ for any permutation group; moreover, we provide an worst-case construction demonstrating that the factor $2$ is sharp and cannot be improved without additional assumptions. Second, we relate the Type II error to a design-dependent geometric separation. We formulate it as a combinatorial optimization problem over permutation groups and bound it under additional mild sub-Gaussian assumptions. For the Type II error upper bound control, we propose a constructive algorithm for the permutation strategy that is better (at least no worse) than the i.i.d. permutation, with simulations empirically indicating substantial power gains, especially under heavy-tailed designs. Finally, we extend group-based CPT and PALMRT beyond exchangeability by connecting rank-based randomization arguments to conformal inference. The resulting weighted group tests satisfy finite-sample Type I error bounds that degrade gracefully with a weighted average of total variation distances between $\varepsilon$ and its group-permuted versions, recovering exact validity when these discrepancies vanish and yielding quantitative robustness otherwise. Taken together, the group-permutation viewpoint provides a principled bridge from exact randomization validity to design-adaptive power and quantitative robustness under approximate symmetries.

stat.ME↗

AppAgent-Claw: CLI Is All You Need for GUI Automation

The OpenClaw platform provides a practical foundation for automation through its skill-oriented architecture, organizing external capabilities into lightweight, reusable components that can be invoked efficiently through a command-line interface (CLI). However, a significant bottleneck remains: many real-world tasks are confined to graphical user interfaces (GUIs) with no stable API available. While LLM-based GUI agents offer generality, their reliance on repeated live model inference makes them too slow, costly, and inconsistent to serve as efficient OpenClaw skills. In this paper, we present AppAgent-Claw, a demonstration-driven system that converts GUI workflows into reliable, reusable skills without runtime inference. By following a ``record-once, replay-many'' paradigm, the system captures rich contextual metadata to facilitate robust execution. It employs a layered localization strategy to handle visual shifts and a validation-coupled execution model to ensure intended on-screen effects. AppAgent-Claw provides a practical, efficient, and diagnosable solution for integrating GUI-bound tasks into the OpenClaw ecosystem.

cs.HC↗

Preventing Safety Drift in Large Language Models via Coupled Weight and Activation Constraints

Safety alignment in Large Language Models (LLMs) remains highly fragile during fine-tuning, where even benign adaptation can degrade pre-trained refusal behaviors and enable harmful responses. Existing defenses typically constrain either weights or activations in isolation, without considering their coupled effects on safety. In this paper, we first theoretically demonstrate that constraining either weights or activations alone is insufficient for safety preservation. To robustly preserve safety alignment, we propose Coupled Weight and Activation Constraints (CWAC), a novel approach that simultaneously enforces a precomputed safety subspace on weight updates and applies targeted regularization to safety-critical features identified by sparse autoencoders. Extensive experiments across four widely used LLMs and diverse downstream tasks show that CWAC consistently achieves the lowest harmful scores with minimal impact on fine-tuning accuracy, substantially outperforming strong baselines even under high harmful data ratios.

cs.AI↗