arXiv subjects

Yihan He

Publications and source records attributed to Yihan He.

At least 19 recordsLinked to original sources

arXiv2026

GRACE: Generative Recommender Acceleration Engine for Real-Time Ads Retrieval

Productionizing generative recommenders for high-volume, real-time ads retrieval creates two serving challenges: eligibility, ensuring that each generated ad is eligible for the request under the advertiser's audience targeting rules, and compute, which requires meeting strict latency and GPU cost requirements while remaining capable of generating thousands of ads per request with wide-beam decoding. This paper presents GRACE, a serving system for ads generative retrieval that addresses both challenges. For eligibility, GRACE introduces Generative Target Matching (GTM), which extends catalog-valid constrained decoding with personalized filtering over Semantic ID (SID) prefixes using bitmask and Bloom filter matchers derived from targeting rules. SID-level GTM improves final ad-level target matching pass rate from 23.55% to 40.42% over constrained decoding alone. For compute-cost and latency, GRACE targets encoder-decoder Transformers, which are more lightweight than LLMs. It redesigns the decoder around the wide-beam, short-sequence regime, covering attention kernels, KV cache, and beam search optimizations. On NVIDIA GH200, compared with the faster of FlashAttention-2 and FlashAttention-3 baselines, GRACE improves cross-attention latency by 68.0 times and self-attention latency by 23.4-25.8 times across decode steps. Together, these changes reduce decoder latency by 11.1 times, keeping ads generative retrieval within latency and compute requirements.

cs.IR↗

arXiv2026

A Theory on Flow Matching with Neural Networks

In this work, we develop theoretical foundation for flow matching with neural-network-parameterized conditional velocity fields. We establish convergence guarantees for gradient descent in the over-parameterized 2-layered ReLU neural network regime. We derive generalization bounds for the conditional velocity-field matching objective. Building on these results, we provide Wasserstein-distance guarantees for the samples generated by the induced flow. Our analysis is based on generalization bound for multi-task representation learning with unbounded losses, which may be of independent interest beyond flow-based generative modeling. These theoretical results are validated through extensive experiments on both synthetic and real-world image benchmarks.

cs.LG↗

arXiv2026

T2SGrid: Temporal-to-Spatial Gridification for Video Temporal Grounding

Video Temporal Grounding (VTG) aims to localize the video segment that corresponds to a natural language query, which requires a comprehensive understanding of complex temporal dynamics. Existing Vision-LMMs typically perceive temporal dynamics via positional encoding, text-based timestamps, or visual frame numbering. However, these approaches exhibit notable limitations: assigning each frame a text-based timestamp token introduces additional computational overhead and leads to sparsity in visual attention, positional encoding struggles to capture absolute temporal information, and visual frame numbering often compromises spatial detail. To address these issues, we propose Temporal to Spatial Gridification (T2SGrid), a novel framework that reformulates video temporal understanding as a spatial understanding task. The core idea of T2SGrid is to process video content in clips rather than individual frames. we employ a overlapping sliding windows mechanism to segment the video into temporal clips. Within each window, frames are arranged chronologically in a row-major order into a composite grid image, effectively transforming temporal sequences into structured 2D layouts. The gridification not only encodes temporal information but also enhances local attention within each grid. Furthermore, T2SGrid enables the use of composite text timestamps to establish global temporal awareness. Experiments on standard VTG benchmarks demonstrate that T2SGrid achieves superior performance.

cs.CV↗

arXiv2025

KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta

Making deep learning recommendation model (DLRM) training and inference fast and efficient is important. However, this presents three key system challenges - model architecture diversity, kernel primitive diversity, and hardware generation and architecture heterogeneity. This paper presents KernelEvolve-an agentic kernel coding framework-to tackle heterogeneity at-scale for DLRM. KernelEvolve is designed to take kernel specifications as input and automate the process of kernel generation and optimization for recommendation model across heterogeneous hardware architectures. KernelEvolve does so by operating at multiple programming abstractions, from Triton and CuTe DSL to low-level hardware agnostic languages, spanning the full hardware-software optimization stack. The kernel optimization process is described as graph-based search with selection policy, universal operator, fitness function, and termination rule, dynamically adapts to runtime execution context through retrieval-augmented prompt synthesis. We designed, implemented, and deployed KernelEvolve to optimize a wide variety of production recommendation models across generations of NVIDIA and AMD GPUs, as well as Meta's AI accelerators. We validate KernelEvolve on the publicly-available KernelBench suite, achieving 100% pass rate on all 250 problems across three difficulty levels, and 160 PyTorch ATen operators across three heterogeneous hardware platforms, demonstrating 100% correctness. KernelEvolve reduces development time from weeks to hours and achieves substantial performance improvements over PyTorch baselines across diverse production use cases and for heterogeneous AI systems at-scale. Beyond performance efficiency improvements, KernelEvolve significantly mitigates the programmability barrier for new AI hardware by enabling automated kernel generation for in-house developed AI hardware.

cs.LG↗

arXiv2025

Solving Boolean Satisfiability Problems Using A Hypergraph-based Probabilistic Computer

Boolean Satisfiability (SAT) problems are critical in fields such as artificial intelligence and cryptography, where efficient solutions are essential. Conventional probabilistic solvers often encounter scalability issues due to complex logic synthesis steps. In this work, we present a novel approach for solving the 3-SAT Boolean satisfiability problem using hypergraph-based probabilistic computers obtained through direct mapping. This method directly translates 3-SAT logical expressions into hypergraph structures, thereby circumventing conventional logic decomposition and synthesis procedures, and offering a more streamlined solver architecture. For representative uf100-430 instances, the proposed approach reduces the node count from 631 to 100 and the edge count from ~2,423 to ~1,013. Under identical simulated annealing conditions, the conventional simple undirected graph (SUG)-based solver achieves a 0% success rate across the tested instances, whereas the hypergraph-based solver attains an average success rate of ~77.6%. In addition, the hypergraph-based method reaches an average minimum energy of ~0.24, close to the theoretical ground state, while the SUG-based architecture remains trapped at substantially higher energy levels (~9.12 on average). The direct hypergraph mapping can further be extended to k-SAT formulations, providing a scalable framework for more complex satisfiability problems in probabilistic computing.

physics.comp-ph↗

arXiv2025

The First Hardware Demonstration of a Universal Programmable RRAM-based Probabilistic Computer for Molecular Docking

Molecular docking is a critical computational strategy in drug design and discovery, but the complex diversity of biomolecular structures and flexible binding conformations create an enormous search space that challenges conventional computing methods. Although quantum computing holds promise for these challenges, it remains constrained by scalability, hardware limitations, and precision issues. Here, we report a prototype of a probabilistic computer (p-computer) that efficiently and accurately solves complex molecular docking for the first time, overcoming previously encountered challenges. At the core of the system is a p-computing chip based upon our artificial tunable probabilistic bits (p-bits), which are compatible with computing-in-memory schemes, based upon 180 nm CMOS technology and BEOL HfO2 RRAM. We successfully demonstrated the superior performance of the p-computer in practical ligand-protein docking scenarios. A 42-node molecular docking problem of lipoprotein with LolA-LolCDE complex-a key point in developing antibiotics against Gram-negative bacteria, was successfully solved. Our results align well with the Protein-Ligand Interaction Profiler tool. This work marks the first application of p-computing in molecular docking-based computational biology, which has great potential to overcome the limitations in success rate and efficiency of current technologies in addressing complex bioinformatics problems.

physics.comp-ph↗

arXiv2025

A Novel P-bit-based Probabilistic Computing Approach for Solving the 3-D Protein Folding Problem

In the post-Moore era, the need for efficient solutions to non-deterministic polynomial-time (NP) problems is becoming more pressing. In this context, the Ising model implemented by the probabilistic computing systems with probabilistic bits (p-bits) has attracted attention due to the widespread availability of p-bits and support for large-scale simulations. This study marks the first work to apply probabilistic computing to tackle protein folding, a significant NP-complete problem challenge in biology. We represent proteins as sequences of hydrophobic (H) and polar (P) beads within a three-dimensional (3-D) grid and introduce a novel many-body interaction-based encoding method to map the problem onto an Ising model. Our simulations show that this approach significantly simplifies the energy landscape for short peptide sequences of six amino acids, halving the number of energy levels. Furthermore, the proposed mapping method achieves approximately 100 times acceleration for sequences consisting of ten amino acids in identifying the correct folding configuration. We predicted the optimal folding configuration for a peptide sequence of 36 amino acids by identifying the ground state. These findings highlight the unique potential of the proposed encoding method for solving protein folding and, importantly, provide new tools for solving similar NP-complete problems in biology by probabilistic computing approach.

physics.app-ph↗

arXiv2025

Robust Inference for the Direct Average Treatment Effect with Treatment Assignment Interference

This paper develops methods for uncertainty quantification in causal inference settings with random network interference. We study the large-sample distributional properties of the classical difference-in-means Hajek treatment effect estimator, and propose a robust inference procedure for the (conditional) direct average treatment effect. Our framework allows for cross-unit interference in both the outcome equation and the treatment assignment mechanism. Drawing from statistical physics, we introduce a novel Ising model to capture complex dependencies in treatment assignment, and derive three results. First, we establish a Berry-Esseen-type distributional approximation that holds pointwise in the degree of interference induced by the Ising model. This approximation recovers existing results in the absence of treatment interference, and highlights the fragility of inference procedures that do not account for the presence of interference in treatment assignment. Second, we establish a uniform distributional approximation for the Hajek estimator and use it to develop robust inference procedures that remain valid uniformly over all interference regimes allowed by the model. Third, we propose a novel resampling method to implement the robust inference procedure and validate its performance through Monte Carlo simulations. A key technical innovation is the introduction of a conditional i.i.d. Gaussianization that may have broader applications. We also discuss extensions and generalizations of our results.

stat.ME↗

arXiv2025

Transformers versus the EM Algorithm in Multi-class Clustering

LLMs demonstrate significant inference capacities in complicated machine learning tasks, using the Transformer model as its backbone. Motivated by the limited understanding of such models on the unsupervised learning problems, we study the learning guarantees of Transformers in performing multi-class clustering of the Gaussian Mixture Models. We develop a theory drawing strong connections between the Softmax Attention layers and the workflow of the EM algorithm on clustering the mixture of Gaussians. Our theory provides approximation bounds for the Expectation and Maximization steps by proving the universal approximation abilities of multivariate mappings by Softmax functions. In addition to the approximation guarantees, we also show that with a sufficient number of pre-training samples and an initialization, Transformers can achieve the minimax optimal rate for the problem considered. Our extensive simulations empirically verified our theory by revealing the strong learning capacities of Transformers even beyond the assumptions in the theory, shedding light on the powerful inference capacities of LLMs.

stat.ML↗

arXiv2025

Transformers and Their Roles as Time Series Foundation Models

We give a comprehensive analysis of transformers as time series foundation models, focusing on their approximation and generalization capabilities. First, we demonstrate that there exist transformers that fit an autoregressive model on input univariate time series via gradient descent. We then analyze MOIRAI, a multivariate time series foundation model capable of handling an arbitrary number of covariates. We prove that it is capable of automatically fitting autoregressive models with an arbitrary number of covariates, offering insights into its design and empirical success. For generalization, we establish bounds for pretraining when the data satisfies Dobrushin's condition. Experiments support our theoretical findings, highlighting the efficacy of transformers as time series foundation models.

cs.LG↗

arXiv2025

Transformers Simulate MLE for Sequence Generation in Bayesian Networks

Transformers have achieved significant success in various fields, notably excelling in tasks involving sequential data like natural language processing. Despite these achievements, the theoretical understanding of transformers' capabilities remains limited. In this paper, we investigate the theoretical capabilities of transformers to autoregressively generate sequences in Bayesian networks based on in-context maximum likelihood estimation (MLE). Specifically, we consider a setting where a context is formed by a set of independent sequences generated according to a Bayesian network. We demonstrate that there exists a simple transformer model that can (i) estimate the conditional probabilities of the Bayesian network according to the context, and (ii) autoregressively generate a new sample according to the Bayesian network with estimated conditional probabilities. We further demonstrate in extensive experiments that such a transformer does not only exist in theory, but can also be effectively obtained through training. Our analysis highlights the potential of transformers to learn complex probabilistic models and contributes to a better understanding of large language models as a powerful class of sequence generators.

stat.ML↗

arXiv2025

Learning Spectral Methods by Transformers

Transformers demonstrate significant advantages as the building block of modern LLMs. In this work, we study the capacities of Transformers in performing unsupervised learning. We show that multi-layered Transformers, given a sufficiently large set of pre-training instances, are able to learn the algorithms themselves and perform statistical estimation tasks given new instances. This learning paradigm is distinct from the in-context learning setup and is similar to the learning procedure of human brains where skills are learned through past experience. Theoretically, we prove that pre-trained Transformers can learn the spectral methods and use the classification of bi-class Gaussian mixture model as an example. Our proof is constructive using algorithmic design techniques. Our results are built upon the similarities of multi-layered Transformer architecture with the iterative recovery algorithms used in practice. Empirically, we verify the strong capacity of the multi-layered (pre-trained) Transformer on unsupervised learning through the lens of both the PCA and the Clustering tasks performed on the synthetic and real-world datasets.

stat.ML↗

arXiv2024

Linear Convergence of the Proximal Gradient Method for Composite Optimization Under the Polyak-{\L}ojasiewicz Inequality and Its Variant

We study the linear convergence rates of the proximal gradient method for composite functions satisfying two classes of Polyak-{\L}ojasiewicz (PL) inequality: the PL inequality, the variant of PL inequality defined by the proximal map-based residual. Using the performance estimation problem, we either provide new explicit linear convergence rates or improve existing complexity bounds for minimizing composite functions under the two classes of PL inequality. Finally, we illustrate numerically the effects of our theoretical results.

math.OC↗

arXiv2024

One-Layer Transformer Provably Learns One-Nearest Neighbor In Context

Transformers have achieved great success in recent years. Interestingly, transformers have shown particularly strong in-context learning capability -- even without fine-tuning, they are still able to solve unseen tasks well purely based on task-specific prompts. In this paper, we study the capability of one-layer transformers in learning one of the most classical nonparametric estimators, the one-nearest neighbor prediction rule. Under a theoretical framework where the prompt contains a sequence of labeled training data and unlabeled test data, we show that, although the loss function is nonconvex when trained with gradient descent, a single softmax attention layer can successfully learn to behave like a one-nearest neighbor classifier. Our result gives a concrete example of how transformers can be trained to implement nonparametric machine learning algorithms, and sheds light on the role of softmax attention in transformer models.

cs.LG↗

arXiv2024

Global Convergence in Training Large-Scale Transformers

Despite the widespread success of Transformers across various domains, their optimization guarantees in large-scale model settings are not well-understood. This paper rigorously analyzes the convergence properties of gradient flow in training Transformers with weight decay regularization. First, we construct the mean-field limit of large-scale Transformers, showing that as the model width and depth go to infinity, gradient flow converges to the Wasserstein gradient flow, which is represented by a partial differential equation. Then, we demonstrate that the gradient flow reaches a global minimum consistent with the PDE solution when the weight decay regularization parameter is sufficiently small. Our analysis is based on a series of novel mean-field techniques that adapt to Transformers. Compared with existing tools for deep networks (Lu et al., 2020) that demand homogeneity and global Lipschitz smoothness, we utilize a refined analysis assuming only $\textit{partial homogeneity}$ and $\textit{local Lipschitz smoothness}$. These new techniques may be of independent interest.

stat.ML↗

arXiv2023

Ground-State Probabilistic Logic with the Simplest Binary Energy Landscape for Probabilistic Computing

We investigate the ground-state probabilistic logic based on a binary energy landscape (GSPL-BEL) model, implementing the many-body interactions within Ising model cells. The GSPL-BEL model offers a simplified binary energy landscape, enabling the conversion of traditional CMOS-based logic into a probabilistic graphical representation based on desired truth tables. Stochastic Ising cells, coupled with generic probabilistic devices exhibiting sigmoidal electrical responses, serve as the building blocks of the GSPL-BEL. Multi-body interactions are realized through cascaded CMOS-based XNOR gates and a passive resistor network. Through circuit simulations of three-node, four-node, and five-node systems, the functionality of the GSPL-BEL model is verified in forward, reverse, and partial-reverse operating modes, and applied to various arithmetic tasks. The many-body effect provides additional degrees of freedom in describing the system's energy function, resulting in distinct energy levels for valid and invalid states. This observation is supported by the binarized probability distribution observed in the steady state of the probabilistic circuits. Furthermore, compared to conventional combinatorial logic circuits, the GSPL-BEL-based circuit design requires a minimal number of probabilistic devices, as demonstrated in the invertible multiplier/integer factorizer circuit. These findings highlight the potential of the GSPL-BEL model for future high-performance logic circuit designs leveraging probabilistic devices.

physics.app-ph↗

arXiv2023

Hidden Clique Inference in Random Ising Model I: the planted random field Curie-Weiss model

We study the problem of testing and recovering the hidden $k$-clique Ferromagnetic correlation in the planted Random Field Curie-Weiss model (a.k.a. the pRFCW model). The pRFCW model is a random effect Ising model that exhibits richer phase diagrams both statistically and physically than the standard Curie-Weiss model. Using an alternative characterization of parameter regimes as 'temperatures' and the mean values as 'outer magnetic fields,' we establish the minimax optimal detection rates and recovery rates. The results consist of $7$ distinctive phases for testing and $3$ phases for exact recovery. Our results also imply that the randomness of the outer magnetic field contributes to countable possible convergence rates, which are not observed in the fixed field model. As a byproduct of the proof techniques, we provide two new mathematical results: (1) A family of tail bounds for the average magnetization of the Random Field Curie-Weiss model (a.k.a. the RFCW model) across all temperatures and arbitrary outer fields. (2) A sharp estimate of the information divergence between RFCW models. These play pivotal roles in establishing the major theoretical results in this paper. Additionally, we show that the mathematical structure involved in the pRFCW hidden clique inference problem resembles a 'sparse PCA-like' problem for discrete data. The richer statistical phases than the long-studied Gaussian counterpart shed new light on the theoretical insight of sparse PCA for discrete data.

math.ST↗

arXiv2023

Hidden Clique Inference in Random Ising Model II: the planted Sherrington-Kirkpatrick model

We study the problem of testing and recovering $k$-clique Ferromagnetic mean shift in the planted Sherrington-Kirkpatrick model (i.e., a type of spin glass model) with $n$ spins. The planted SK model -- a stylized mixture of an uncountable number of Ising models -- allows us to study the fundamental limits of correlation analysis for dependent random variables under misspecification. Our paper makes three major contributions: (i) We identify the phase diagrams of the testing problem by providing minimax optimal rates for multiple different parameter regimes. We also provide minimax optimal rates for exact recovery in the high/critical and low temperature regimes. (ii) We prove a universality result implying that all the obtained rates still hold with non-Gaussian couplings. (iii) To achieve the major results, we also establish a family of novel concentration bounds and central limiting theorems for the averaging statistics in the local and global phases of the planted SK model. These technical results shed new insights into the planted spin glass models. The pSK model also exhibits close connections with a binary variant of the single spike Gaussian sparse principle component analysis model by replacing the background identity precision matrix with a Wigner random matrix.

math.ST↗