Goto

Collaborating Authors

 Statistical Learning


Wasserstein Transfer Learning

Neural Information Processing Systems

Transfer learning is a powerful paradigm for leveraging knowledge from source domains to enhance learning in a target domain. However, traditional transfer learning approaches often focus on scalar or multivariate data within Euclidean spaces, limiting their applicability to complex data structures such as probability distributions. To address this limitation, we introduce a novel transfer learning framework for regression models whose outputs are probability distributions residing in the Wasserstein space. When the informative subset of transferable source domains is known, we propose an estimator with provable asymptotic convergence rates, quantifying the impact of domain similarity on transfer efficiency. For cases where the informative subset is unknown, we develop a data-driven transfer learning procedure designed to mitigate negative transfer. The proposed methods are supported by rigorous theoretical analysis and are validated through extensive simulations and real-world applications. The code is available at https://github.com/h7nian/WaTL.


Identifying multi-compartment Hodgkin-Huxley models with high-density extracellular voltage recordings

Neural Information Processing Systems

Multi-compartment Hodgkin-Huxley models are biophysical models of how electrical signals propagate throughout a neuron, and they form the basis of our knowledge of neural computation at the cellular level. However, these models have many free parameters that must be estimated for each cell, and existing fitting methods rely on intracellular voltage measurements that are highly challenging to obtain in vivo. Recent advances in neural recording technology with high-density probes and arrays enable dense sampling of extracellular voltage from many sites surrounding a neuron, allowing indirect measurement of many compartments of a cell simultaneously. Here, we propose a method for inferring the underlying membrane voltage, biophysical parameters, and the neuron's position relative to the probe, using extracellular measurements alone. We use an Extended Kalman Filter to infer membrane voltage and channel states using efficient, differentiable simulators. Then, we learn the model parameters by maximizing the marginal likelihood using gradient-based methods. We demonstrate the performance of this approach using simulated data and real neuron morphologies.


A data and task-constrained mechanistic model of the mouse outer retina shows robustness to contrast variations

Neural Information Processing Systems

Visual processing starts in the outer retina where photoreceptors transform light into electrochemical signals. These signals are modulated by inhibition from horizontal cells and sent to the inner retina via excitatory bipolar cells. The outer retina is thought to play an important role in contrast invariant coding of visual information, but how the different cell types implement this computation together remains incompletely understood. To understand the role of each cell type, we developed a fully-differentiable biophysical model of a circular patch of mouse outer retina. The model includes 200 cone photoreceptors with a realistic phototransduction cascade and ribbon synapses as well as horizontal and bipolar cells, all with celltype specific ion channels. Going beyond decades of work constraining biophysical models of neurons only by experimental data, we used a dual approach, constraining some parameters of the model with available measurements and others by a visual task: (1) We fit the parameters of the cone models to whole cell patch-clamp measurements of photocurrents and two-photon glutamate imaging measurements of synaptic release.


ROOT: Rethinking Offline Optimization as Distributional Translation via Probabilistic Bridge

Neural Information Processing Systems

This paper studies the black-box optimization task which aims to find the maxima of a black-box function using a static set of its observed input-output pairs. This is often achieved via learning and optimizing a surrogate function with that offline data. Alternatively, it can also be framed as an inverse modeling task that maps a desired performance to potential input candidates that achieve it. Both approaches are constrained by the limited amount of offline data. To mitigate this limitation, we introduce a new perspective that casts offline optimization as a distributional translation task.


Theoretical Investigation of Adafactor for Non-Convex Smooth Optimization

Neural Information Processing Systems

Adafactor is an early memory-efficient optimization algorithm proposed as an alternative to Adam. By eliminating first-order momentum and employing a rank-1 matrix factorization to approximate the second-moment matrix, Adafactor achieves near-zero memory overhead compared to traditional gradient descent methods. Despite its practical suitability for large-scale training tasks where memory efficiency is critical, its theoretical convergence analysis remains unexplored, largely due to the challenges posed by its matrix factorization and update clipping mechanisms. In this work, we provide a convergence analysis of Adafactor for non-convex smooth optimization. We establish optimal convergence rates (up to logarithmic factors) for finding stationary points in both deterministic and stochastic settings, the latter under sub-Gaussian noise. Central to our analysis is viewing Adafactor as an approximation of Adam, and the use of a new proxy step-size to approximate the unique adaptive step-size induced by Adafactor's matrix factorization and update clipping, along with an induction argument to control the gradient magnitude. Our findings may theoretically suggest that involving rank-1 matrix approximation of the second-moment matrix in Adam does not fundamentally hinder the convergence.


Riemannian Proximal Sampler for High-accuracy Sampling on Manifolds

Neural Information Processing Systems

We introduce the Riemannian Proximal Sampler, a method for sampling from densities defined on Riemannian manifolds. The performance of this sampler critically depends on two key oracles: the Manifold Brownian Increments (MBI) oracle and the Riemannian Heat-kernel (RHK) oracle. We establish high-accuracy sampling guarantees for the Riemannian Proximal Sampler, showing that generating samples with ฮต-accuracy requires O(log(1/ฮต)) iterations in Kullback-Leibler divergence assuming access to exact oracles and O(log2(1/ฮต))iterations in the total variation metric assuming access to sufficiently accurate inexact oracles.


Tropical Attention: Neural Algorithmic Reasoning for Combinatorial Algorithms

Neural Information Processing Systems

Can algebraic geometry enhance the sharpness, robustness, and interpretability of modern neural reasoning models by equipping them with a mathematically grounded inductive bias? To answer this, we introduce Tropical Attention, an attention mechanism grounded in tropical geometry that lifts the attention kernel into tropical projective space, where reasoning is piecewise-linear and 1-Lipschitz, thus preserving the polyhedral decision structure inherent to combinatorial reasoning. We prove that Multi-Head Tropical Attention (MHTA) stacks universally approximate tropical circuits and realize tropical transitive closure through composition, achieving polynomial resource bounds without invoking recurrent mechanisms. These guarantees explain why the induced polyhedral decision boundaries remain sharp and scale-invariant, rather than smoothed by Softmax. Empirically, we show that Tropical Attention delivers stronger out-of-distribution generalization in both length and value, with high robustness against perturbative noise, and substantially faster inference with fewer parameters compared to Softmax-based and recurrent attention baselines, respectively. For the first time, we push the domain of neural algorithmic reasoning beyond PTIME problems to NP-hard/complete problems, paving the way toward sharper and more expressive Large Reasoning Models (LRMs) capable of tackling complex combinatorial challenges in Phylogenetics, Cryptography, Particle Physics, and Mathematical Discovery.


Overleaf Example

Neural Information Processing Systems

A foundation model for medical time series, pretrained on ethically approved clinical datasets, can substantially reduce annotation burdens, minimize the need for task-specific tuning, and promote reliable transferability across healthcare institutions, data modalities, and clinical tasks, especially in data-scarce or privacysensitive environments. However, existing generalist time series foundation models struggle to handle medical time series data due to their inherent challenges, including irregular intervals, heterogeneous sampling rates, and frequent missing values. To address these challenges, we introduce MIRA, a unified foundation model specifically designed for medical time series forecasting. MIRA incorporates a Continuous-Time Rotary Positional Encoding that enables fine-grained modeling of variable time intervals, a frequency-specific mixture-of-experts layer that routes computation across latent frequency regimes to further promote temporal specialization, and a Continuous Dynamics Extrapolation Block based on Neural ODE that models the continuous trajectory of latent states, enabling accurate forecasting at arbitrary target timestamps. Pretrained on a large-scale and diverse medical corpus comprising over 454 billion time points collect from publicly available datasets, MIRA achieves reductions in forecasting errors by an average of 8% and 6% in out-of-distribution and in-distribution scenarios, respectively, when compared to other zero-shot and fine-tuned baselines. We also introduce a comprehensive benchmark spanning multiple downstream clinical tasks, establishing a foundation for future research in medical time series modeling. Our code is available at Microsoft/MIRA.



SignFlow Bipartite Subgraph Network For Large-Scale Graph Link Sign Prediction

Neural Information Processing Systems

Link sign prediction in signed bipartite graphs, which are extensively utilized across diverse domains such as social networks and recommendation systems, has recently emerged as a pivotal challenge. However, significant space and time complexities associated with the scalability of bipartite graphs pose substantial challenges, particularly in large-scale environments. To address these issues, this paper introduces the SignFlow Bipartite Subgraph Network (SBSN), balancing sublinear training memory growth through a heuristic subgraph extraction method integrated with a novel message passing module, with optimal inference efficiency achieved via the node feature distillation module. Our subgraph sampling approach reduces the graph size by focusing on neighborhoods around target links and employs an optimized directed message passing mechanism to aggregate critical structural patterns. This mechanism allows SBSN to efficiently learn rich local structural patterns essential for accurate sign prediction. Furthermore, to overcome the inefficiency of subgraph sampling-based models during inference, SBSN incorporates a node feature distillation module after the first training stage.