SearcharxivSearch

arXiv subjects

Zichun Ye

Publications and source records attributed to Zichun Ye.

7 recordsLinked to original sources

Unlearning Offline Stochastic Multi-Armed Bandits

Machine unlearning aims to unlearn data points from a learned model, offering a principled way to process data-deletion requests and mitigate privacy risks without full retraining. Prior work has mainly studied unsupervised / supervised machine unlearning, leaving unlearning for sequential decision-making systems far less understood. We initiate the first study of a foundational sequential decision-making problem: offline stochastic multi-armed bandits (MAB). We formalize the privacy constraint for offline MAB and measure utility by the post-unlearning decision quality. We conduct a systematic study of both single- and multi-source unlearning scenarios under two data-generation models, the fixed-sample model and the distribution model. For these settings, our algorithmic design is built on two canonical base algorithms: Gaussian mechanism and rollback, and we propose adaptive algorithms that switch between them according to the data regime and privacy constraint. We further introduce a mixing procedure that elucidates the rationale behind these baselines. We provide performance guarantees across the above settings and establish lower bounds under both dataset models. Experiments validate the predicted tradeoffs and demonstrate the effectiveness of the proposed methods.

cs.LG

DVM: A Bytecode Virtual Machine Approach for Dynamic Tensor Computation

Dynamism is common in AI computation, e.g., the dynamic tensor shapes and the dynamic control flows in models. Due to the long compilation time, existing runtime compilation damages the model efficiency, while the offline compilers either suffer from the long compilation time and device memory footprint to cover all the possible execution instances of a dynamic model, or sacrifice optimization opportunities for usability. In this paper, we rethink the feasibility of runtime compilation for dynamic models and identify that the key for it to work is to speed up the compilation or hide the compilation overhead. To do this, we propose a real-time compiler, DVM. In DVM, we design a runtime operator compiler based on a bytecode virtual machine to perform effective and efficient compilation for each dynamic operator instance given its input. Specifically, instead of compiling programs into machine code, we encode the operator program into bytecode on the CPU and decode the bytecode into virtual instructions for direct execution on the NPU. Based on the runtime operator compiler, we further propose an operator fuser, which performs symbol-deduction-based fusion on static graphs and runtime fusion on dynamic graphs. Both pattern- and stacking-based fusion are supported to increase fusion opportunities. Evaluation on operators, subgraphs, and models shows that, compared with TorchInductor, PyTorch-eager and MindSpore-graph-O0, we are up to 11.77$\times$ better in terms of the operator/model efficiency and up to 5 orders of magnitude faster in terms of the maximum compilation time.

cs.PL

AKG kernel Agent: A Multi-Agent Framework for Cross-Platform Kernel Synthesis

Modern AI models demand high-performance computation kernels. The growing complexity of LLMs, multimodal architectures, and recommendation systems, combined with techniques like sparsity and quantization, creates significant computational challenges. Moreover, frequent hardware updates and diverse chip architectures further complicate this landscape, requiring tailored kernel implementations for each platform. However, manual optimization cannot keep pace with these demands, creating a critical bottleneck in AI system development. Recent advances in LLM code generation capabilities have opened new possibilities for automating kernel development. In this work, we propose AKG kernel agent (AI-driven Kernel Generator), a multi-agent system that automates kernel generation, migration, and performance tuning. AKG kernel agent is designed to support multiple domain-specific languages (DSLs), including Triton, TileLang, CPP, and CUDA-C, enabling it to target different hardware backends while maintaining correctness and portability. The system's modular design allows rapid integration of new DSLs and hardware targets. When evaluated on KernelBench using Triton DSL across GPU and NPU backends, AKG kernel agent achieves an average speedup of 1.46$\times$ over PyTorch Eager baselines implementations, demonstrating its effectiveness in accelerating kernel development for modern AI workloads.

cs.AI

Near-Optimal Regret for Efficient Stochastic Combinatorial Semi-Bandits

The combinatorial multi-armed bandit (CMAB) is a cornerstone of sequential decision-making framework, dominated by two algorithmic families: UCB-based and adversarial methods such as follow the regularized leader (FTRL) and online mirror descent (OMD). However, prominent UCB-based approaches like CUCB suffer from additional regret factor $\log T$ that is detrimental over long horizons, while adversarial methods such as EXP3.M and HYBRID impose significant computational overhead. To resolve this trade-off, we introduce the Combinatorial Minimax Optimal Strategy in the Stochastic setting (CMOSS). CMOSS is a computationally efficient algorithm that achieves an instance-independent regret of $O\big( (\log k)\sqrt{kmT}\big )$ when $k\leq \frac{m}{2}$ and $O\big((m-k)\sqrt{\log k\log(m-k)T}\big )$ when $k>\frac{m}{2}$ under semi-bandit feedback, where $m$ is the number of arms and $k$ is the maximum cardinality of a feasible action. Crucially, this result eliminates the dependency on $\log T$ and matches the established lower bounds of $\Omega\big(\sqrt{kmT}\big)$ when $k\leq \frac{m}{2}$ and $\Omega\big((m-k)\sqrt{\log (\frac{m}{m-k}) T}\big)$ when $k>\frac{m}{2}$ up to logarithmic terms of $k$ and $m$. We then extend our analysis to show that CMOSS is also applicable to cascading feedback. Experiments on synthetic and real-world datasets validate that CMOSS consistently outperforms benchmark algorithms in both regret and runtime efficiency.

cs.LG

Tight Gap-Dependent Memory-Regret Trade-Off for Single-Pass Streaming Stochastic Multi-Armed Bandits

We study the problem of minimizing gap-dependent regret for single-pass streaming stochastic multi-armed bandits (MAB). In this problem, the $n$ arms are present in a stream, and at most $m 0}Δ_i^{1 - 2α}\Big)$ and another applicable for $m<\frac{2}{3}n$ with regret at most $O_α\Big(\frac{T^{\frac{1}{α+1}}}{m^{\frac{1}{α+1}}}\displaystyle\sum_{i:Δ_i > 0}Δ_i^{1 - 2α}\Big)$. We also prove matching lower bounds for both cases by showing that for any constant $α\ge 1$ and any $m\leq k < n$, there exists a set of hard instances on which the regret of any algorithm is $Ω_α\Big(\frac{(k-m+1) T^{\frac{1}{α+1}}}{k^{1 + \frac{1}{α+1}}} \sum_{i:Δ_i > 0}Δ_i^{1-2α}\Big)$. This is the first tight gap-dependent regret bound for streaming MAB. Prior to our work, an $O\Big(\sum_{i\colonΔ>0} \frac{\sqrt{T}\log T}{Δ_i}\Big)$ upper bound for the special case of $α=1$ and $m=O(1)$ was established by Agarwal, Khanna and Patil (COLT'22). In contrast, our results provide the correct order of regret as $Θ\Big(\frac{1}{\sqrt{m}}\sum_{i\colonΔ>0}\frac{\sqrt{T}}{Δ_i}\Big)$.

cs.LG

Understanding Memory-Regret Trade-Off for Streaming Stochastic Multi-Armed Bandits

We study the stochastic multi-armed bandit problem in the $P$-pass streaming model. In this problem, the $n$ arms are present in a stream and at most $m<n$ arms and their statistics can be stored in the memory. We give a complete characterization of the optimal regret in terms of $m, n$ and $P$. Specifically, we design an algorithm with $\tilde O\left((n-m)^{1+\frac{2^{P}-2}{2^{P+1}-1}} n^{\frac{2-2^{P+1}}{2^{P+1}-1}} T^{\frac{2^P}{2^{P+1}-1}}\right)$ regret and complement it with an $\tilde Ω\left((n-m)^{1+\frac{2^{P}-2}{2^{P+1}-1}} n^{\frac{2-2^{P+1}}{2^{P+1}-1}} T^{\frac{2^P}{2^{P+1}-1}}\right)$ lower bound when the number of rounds $T$ is sufficiently large. Our results are tight up to a logarithmic factor in $n$ and $P$.

cs.LG

Models of Gradient Type with Sub-Quadratic Actions

We consider models of gradient type, which are the densities of a collection of real-valued random variables $ϕ:=\{ϕ_x: x \in Λ\}$ given by $Z^{-1}\exp({-\sum\nolimits_{j \sim k}V(ϕ_j-ϕ_k)})$. We focus our study on the case that $V(\nablaϕ) = [1+(\nablaϕ)^2]^α$ with $0 < α< 1/2$, which is a non-convex potential. We introduce an auxiliary field $t_{jk}$ for each edge and represent the model as the marginal of a model with log-concave density. Based on this method, we prove that finite moments of the fields $\left<[v \cdot ϕ]^p \right>$ are bounded uniformly in the volume. This leads to the existence of infinite volume measures. Also, every translation invariant, ergodic infinite volume Gibbs measure for the potential $V$ above scales to a Gaussian free field.

math.PR