arXiv subjects

Taiji Suzuki

Publications and source records attributed to Taiji Suzuki.

At least 19 recordsLinked to original sources

arXiv2026

Towards the Query Complexity of Sampling under Strong Log-Concavity and Log-Sobolev Inequalities

We study the query complexity of sampling from smooth distributions on $\mathbb R^d$ using first-order oracles. For $d$-dimensional strongly log-concave targets with condition number $κ=2$, we prove that any randomized sampler requires $\exp({Ω(\log^{1/3} d)})$ queries even from a constant warm start. This rules out polylogarithmic dimension dependence and strictly separates strongly log-concave sampling from Gaussian sampling. For targets with $L$-smooth potentials satisfying an $α$-log-Sobolev inequality with $κ= L/α$, we prove an $Ω(κ)$ query lower bound for $d\gtrsim\log^2κ$. This shows that log-Sobolev sampling is strictly harder than strongly log-concave sampling. Our proofs construct instance-dependent tests that target samples typically pass but outputs of algorithms making few queries typically fail. For strongly log-concave targets, we embed a sequence of hidden Gaussian into high-degree tensor polynomial with smooth gates, forcing any algorithm to uncover the layers sequentially. The log-Sobolev construction hides a constant fraction of the target mass in angular caps while controlling the log-Sobolev constant via delicate estimates.

cs.CC↗

arXiv2026

On Complexity of Gaussian Sampling

We study the number of exact matrix-vector queries needed to sample a centered $d$-dimensional Gaussian with condition number at most $κ$ to total variation error $\varepsilon$. The best known algorithm uses $O(\min\{d,\sqrtκ\log(\frac{d}{\varepsilon})\})$ queries, while the known lower bounds do not grow as $\varepsilon\to0$ and match this bound only for $κ\le d^{1/5-δ}$ at constant accuracy. We introduce a construction based on a hidden determinant sign and prove a lower bound of $Ω\!\left(\min\left\{d,\sqrtκ\log\left(1+\frac{d}{\sqrtκ\,\varepsilon}\right)\right\}\right)$ uniformly for all $d\ge1$, $κ\ge4$, and $0<\varepsilon\le\tfrac{1}{16}$. This bound matches the known upper bound up to constant factors whenever $κ\le d^{2-δ}$ for fixed $δ\in(0,2)$, or $\varepsilon\le d^{-c}$ for fixed $c>0$. We also extend a recent fixed-spectrum construction with a hidden eigenbasis from $κ=2$ to $2\leκ\le d^{1-δ}$, for any fixed $δ\in(0,1)$. This gives an $Ω_δ(\sqrtκ\log d)$ lower bound at constant accuracy that holds even when the spectrum of the precision matrix is known to the algorithm.

cs.CC↗

arXiv2026

On the Complexity of Finding Stationary Points of p-th Order Hölder Smooth Functions in One Dimension

We study the query complexity of finding an an $\varepsilon$-stationary points of a one-dimensional non-convex function whose $p$-th derivative is $β_p$-Hölder continuous. We determine the optimal complexity for deterministic and randomized algorithms and for oracles that return only derivatives, only function values, or full local information. Randomness provably helps when only derivatives are available, function values help in all cases, and higher-order derivatives never do. Together with known results in higher dimensions, this shows how these benefits depend on the dimension. Our deterministic lower bound also implies that the optimal accuracy exponents of objective-function-free adaptive regularization of order $p$ (OFFAR$p$) in every dimension $d\ge1$ and of deterministic gradient-only trapping in dimension $d=2$.

math.OC↗

arXiv2026

Finite-Sample Approximation of Hessian-Guided Perturbed Wasserstein Gradient Flows

Wasserstein gradient flow extends gradient descent to probability measures. Its Hessian-guided perturbed variant (PWGF) adds Gaussian perturbations to escape saddle points in nonconvex problems. We investigate when its approximation by finitely many interacting particles remains accurate over growing time horizons. Our analysis retains the curvature accumulated along the population-driven reference path: negative curvature can amplify approximation errors, while subsequent positive curvature can damp their influence. This captures favorable scenarios in which temporary instability is compatible with accurate tracking over growing horizons. Under regularity assumptions and a prescribed common perturbation schedule, we prove particle and objective-value tracking bounds on a high-probability event for reference paths satisfying explicit conditions on accumulated curvature. To handle state-dependent Gaussian jumps, we construct a population-first coupling that preserves the reference particles' conditional independence and reduces jump errors to covariance comparison. We verify the conditions in a variance-plus-cosine model, where curvature recovery yields a growing-horizon tracking guarantee. We also establish local attraction, transverse descent, and positive second variation in two regions of a regularized matrix-factorization model, motivating a positive-negative-positive curvature pattern.

cs.LG↗

arXiv2026

Two-Sample Testing via Generative Processes

Deciding whether two samples come from the same distribution is a classical problem in statistics, and generative transport offers a new way to approach it. We build a stochastic interpolant directly between the two samples and observe that, for a symmetric schedule, its law is invariant under the time reflection $t \mapsto 1-t$ whenever the two distributions coincide. We therefore test whether the marginals at times t and 1-t agree by computing their Jensen--Shannon divergence. Both marginals are explicit mixtures over all cross-pairs of observations, so nothing is learned, and permutation calibration gives an exact finite-sample level. For Gaussian noise, this divergence equals a time integral that pairs the reflection defects of the velocity field and of the score, so the test compares transport dynamics rather than endpoints alone. With a narrow-plus-broad noise design, the test attains the minimax separation rate n^{-2s/(4s+d)} over bounded, compactly supported densities whose difference has Sobolev smoothness s > 3d/4, with no lower bound on the densities. Fusing a dyadic grid of noise scales through their permutation ranks, without sample splitting, preserves exact level and adapts to unknown s at an iterated-logarithmic cost. Empirically, the test matches or outperforms state-of-the-art kernel two-sample tests.

stat.ML↗

arXiv2026

Reinforcement Learning for Hierarchical Reasoning Rewards: Minimax-Optimal Rates with Transformers

Reinforcement learning (RL) has become a standard tool for post-training language models on reasoning tasks, where the policy is updated by reward feedback while exploring the space of responses. Despite its empirical success, theoretical understanding of RL post-training remains limited, in particular of why on-policy exploration combined with a neural reward model is effective. In this paper, we address this question by modeling the reward as a hierarchical function on the response space: the reward consists of infinitely many local components, each of which becomes relevant only after the preceding ones have been resolved. We show that a natural Transformer-based actor--critic algorithm, which alternates between sampling from the current KL-regularized policy, fitting a Transformer critic to the observed rewards, and updating the policy, achieves the minimax optimal rates in the query budget and in the regularization strength up to logarithmic factors, and is minimax optimal for a fixed number of prompts. In contrast, we prove that sampling from the fixed reference distribution, as in offline reward modeling, can limit regret decay to a logarithmic rate. These results show that on-policy exploration progressively zooms in on the region where the reward is concentrated, and quantify its benefit for RL post-training.

cs.LG↗

arXiv2026

Steepest Guidance: A Practical and Principled Approach to Inference-Time Alignment of Flow and Diffusion-based Models

Inference-time alignment of flow and diffusion-based models is critical for achieving flexible generative modeling. Theoretically, Doob's $h$-transform provides an elegant solution to this problem, and most existing methods are based on this principle. However, in practice, estimating the optimal guidance derived from Doob's $h$-transform at inference time is challenging. To deal with this issue, we regard inference-time alignment as a sequential optimization problem in the space of probability measures and propose a novel framework called *Steepest Guidance*, based on the principle of maximizing local improvement in the objective. We provide a theoretical analysis of the proposed method and demonstrate its effectiveness through extensive experiments.

stat.ML↗

arXiv2026

Two-Sample Testing via Path-based Inference

Modern deep generative models are primarily studied for their ability to generate realistic samples, yet the generative dynamics they learn can also serve as objects of statistical inference. We develop this idea for two-sample testing, the problem of deciding whether the same distribution generated two finite datasets. Using stochastic interpolants, we connect both distributions to a shared Gaussian bottleneck, so that each half of the resulting path is a Gaussian channel acting on a single population. We prove that the null hypothesis holds if and only if the population denoiser, or equivalently, the velocity fields of the two halves, coincide at any single noise level, which amounts to a reflection symmetry of the path about the bottleneck. Deviations from this symmetry yield a continuum of two-sample witnesses, which we estimate via held-out regression risks on learned denoisers and velocities and aggregate along the path; under an information-theoretic weighting, the aggregated discrepancy equals the Jeffreys divergence between the noise-smoothed distributions. Calibrating the resulting statistics by permutation yields tests that are valid in finite samples for any trained networks and consistent when the fields are learned accurately. On a synthetic benchmark and three image benchmarks, the proposed tests improve power over the strongest baseline by up to 33 percentage points at an equal total sample budget, with the best choice of regression representation and path weighting depending on the data modality. These results show that generative paths provide a principled representation for statistical testing, extending stochastic-interpolant models beyond generation.

stat.ML↗

arXiv2026

Counterfactual Probing for Parallel Unmasking with Hidden Forest Structure

Masked generative models offer parallel token prediction, but accurate parallel sampling must account for dependencies among tokens. When dependencies are unknown, finding safe batches also costs model evaluations. We study whether total evaluations, including discovery, can be sublinear in sequence length $N$; sublinear sequential depth then follows. We consider discrete distributions with hidden forest structure, accessed through a fixed approximate conditional oracle. Under explicit regularity conditions and uniform Hellinger error bounds, for any fixed target accuracy $\varepsilon\in(0,1/8]$ and sufficiently large $N$, our sampler achieves seed-averaged total-variation error at most $\varepsilon$, with total masked-state submissions and sequential depth both bounded by $\widetilde{O}(N^C \varepsilon^{-a})$ for constants $0 0$. These guarantees use polynomial vocabulary size and an edge-response lower bound set by $N$ and $\varepsilon$. The sampler shares evaluations of hypothetical reveals across dependence tests to identify safe parallel batches without requiring full recovery of the hidden forest. A tunable parameter trades probing cost against irreversible commit rounds. In the same class, any admissible irreversible product-commit sampler attaining the same seed-averaged accuracy requires $Ω(N^c \varepsilon^b)$ counterfactual submissions or commit rounds in the worst case, for constants $c,b>0$.

cs.LG↗

arXiv2026

Learning without Overwriting: A Theory of Self-Distillation and Supervised Fine-Tuning in Continual Reasoning

On-policy self-distillation (OPSD) of large language models (LLMs) has demonstrated the ability to improve reasoning capabilities while preserving previously acquired knowledge. Despite substantial empirical success, the dynamics of OPSD in continual reasoning remain incompletely understood. Modeling LLM reasoning as search over a directed acyclic graph, we provide a unified theoretical analysis of both the dynamics of post-training---OPSD and supervised fine-tuning (SFT) in continual learning---and the impact of pre-training on subsequent performance. Our findings establish three key insights with an optimization guarantee: (i) OPSD with hints from correct outputs enables continual learning without forgetting by sparse yet effective gradient descent updates induced by the hint structure. (ii) SFT on correct reasoning paths can lead to catastrophic forgetting due to dense updates along the training paths, which overwrite the information previously acquired. (iii) Diversity in pre-training is crucial for enabling a post-trained model to reach a correct output when a rollout starts from an intermediate state. Our results, supported by theoretical analysis, show that reliable continual reasoning depends on how post-training updates interact with the reasoning structure established during pre-training.

cs.LG↗

arXiv2026

Diffusion Transformers are Provably Optimal In-context Generators

Generative foundation models are attracting interest for their ability to produce desired outputs from demonstrations given at inference time, without updating parameters. However, since a few demonstrations cannot uniquely identify the intended task, the challenge is how to learn and sample from an output distribution that reflects this task uncertainty. In this work, we theoretically analyze how a Diffusion Transformer (DiT), pretrained across diverse tasks, learns and generates predictive distributions for a new query from demonstrations. We first show that the natural target to generate from finite demonstrations is not an output derived from estimating a single task, but rather a predictive distribution that captures the task uncertainty remaining after observing the demonstrations. We then prove that a DiT can learn this predictive distribution through score estimation, using attention to aggregate information from demonstrations and diffusion to generate samples. Owing to this property, with sufficient pretraining resources and diffusion sampling steps, the resulting DiT achieves the minimax optimal rate over a Hölder class of test-time tasks. These results imply that DiT acts as a statistically grounded in-context generator capable of generating distributions adapted to new tasks while retaining the uncertainty inherent in finite demonstrations.

cs.LG↗

arXiv2026

Finding Stationary Points under Higher-Order Smoothness: Trapping Algorithms and First-Order Lower Bounds

We study the query complexity of finding approximate stationary points under higher-order smoothness using only function values or only gradients. Our results cover different dimensions, oracle models, randomization, and constraints. For constant dimension $d$ and $C^{p,β_p}$ objectives, with $ν:=p+β_p$, we give deterministic function-value gradient trapping algorithms with query complexities $O(\varepsilon^{-(d-1)})$ on $\mathbb R^d$ and $O(\varepsilon^{-(d-1)/ν})$ on $[0,1]^d$. High-order line integration enables gradient-only trapping. For two-dimensional $C^{1,1}$ objectives, randomized integration yields $O(\varepsilon^{-4/3})$ gradient queries on $\mathbb R^2$, breaking the deterministic gradient-only $Θ(\varepsilon^{-2})$ barrier. For general dimension, we introduce a trapping-or-descent framework that lazily reuses boundary models. For $C^{p,β_p}$ objectives with $p\ge2$, it achieves deterministic query complexities $O(d^p+d^{p-1+β_p/(ν-1)} \varepsilon^{-ν/(ν-1)})$ using function values and $O(d^{p-1}+d^{p-2+β_p/(ν-1)} \varepsilon^{-ν/(ν-1)})$ using gradients. These bounds apply to approximate stationary points on $\mathbb R^d$ and approximate KKT points on $[0,1]^d$. Their dimension-independent accuracy exponent matches the classical AR$p$ exponent $(p+1)/p$ when $β_p=1$. For sufficiently high dimension, we prove randomized lower bounds for gradient-only methods. For $C^{1,1}\cap C^{p,1}$ objectives, our lower bound is $Ω(\varepsilon^{-(3/2+1/(2p))})$, matching known upper bounds for $p=2$ and, up to logarithmic factors, for $p=3$. Without the $C^{1,1}$ assumption, we obtain an $Ω(\varepsilon^{-5/3})$ lower bound for every fixed $p\ge3$.

math.OC↗

arXiv2026

Zero Flux: Flow-Based Comparison of High-Dimensional Discrete Distributions

Comparing two high-dimensional discrete distributions has always been a challenging task due to the exponentially growing state space and complex changes in interactions. A recent work suggests comparing distributions through a vector field trained using flow matching between two continuous distributions. The resulting vector field at mid-point vanishes if and only if two distributions identical. However, such a flow-based criterion does not naturally apply to discrete distributions. We extend this principle to the discrete domain and introduce the \emph{Zero Flux} criterion, a discrepancy based on local probability fluxes. Under independent coupling, we show that all local probability fluxes vanish at the midpoint if and only if two distributions are the same. This discrepancy decomposes the joint distributional difference into smaller, local contributions and can be efficiently estimated from samples. We establish finite sample error bounds for our estimator. Experiments on synthetic and real categorical data demonstrate reliable recovery of sparse dependence signals and stable tracking of distribution shifts in high dimensions.

stat.ML↗

arXiv2026

Zero-Flow Two-Sample Tests

Motivated by the success of modern flow-based generative models in modeling complex data, we study two-sample testing through the lens of flow-based methods. We propose the Zero-Flow Two-Sample Test (ZF2ST), built on the zero-flow criterion, which characterizes distributional equality through a time-reversal antisymmetry of a learnable velocity field. We extend this criterion and further develop the Zero-Flow Discrepancy, an identifying discrepancy that controls the Wasserstein distance, and derive a variational representation in terms of a witness function. This representation naturally leads to a witness-based test whose power is governed by the signal-to-noise ratio (SNR), allowing direct power maximization for witness learning. ZF2ST learns the witness on one data split and performs testing on held-out samples, thereby maintaining Type-I error control and admitting a simple asymptotic null distribution. Experimentally, ZF2ST performs competitively across various synthetic and real-world benchmarks, while showing particularly strong performance in distinguishing image distributions from different sources.

cs.LG↗

arXiv2026

Wasserstein Saddle-Free Newton: Saddle Escape and Local Spectral Convergence

Extending saddle-free Newton methods to Wasserstein space requires controlling Hessian dependent operators along nonlinear transport trajectories and accounting for a Hessian spectrum that may accumulate at zero. We address these challenges through Wasserstein Saddle-Free Newton (WSFN), a second-order method for nonconvex optimization over Wasserstein space that preconditions the Wasserstein gradient by a regularized inverse square root Hessian operator. The dynamics exploit curvature magnitude while maintaining negative curvature repulsion and positive curvature attraction, avoiding key Wasserstein Newton limitations. Whereas previous Wasserstein saddle escape results rely on perturbed first-order dynamics, WSFN uses curvature both to modify the deterministic transport and to construct the perturbation mechanism. Our analysis combines stability estimates for Hessian operators with a perturbative saddle escape argument to establish polynomial time convergence to approximate second-order stationarity under regularity of landscape assumptions. The contribution of the saddle curvature parameter $δ$ improves from $\tilde O(δ^{-4})$ for prior perturbed first-order Wasserstein methods to $\tilde O(δ^{-3})$. For benign landscapes, this implies proximity to global minimizers. The cubic dependence corresponds to the scale of the objective decrease that can be guaranteed from negative curvature under Lipschitz Hessian regularity. Near a minimizer, the possible absence of a uniform spectral gap prevents a uniform contraction result and instead calls for a spectral characterization. We prove geometric contraction of the linearized dynamics on positive curvature spectral bands bounded away from zero and polynomial decay under spectral source conditions. Together, these results connect quantitative saddle escape with the local spectral behavior of second-order Wasserstein dynamics.

math.OC↗

arXiv2026

Test time training enhances in-context learning of nonlinear functions

Test-time training (TTT) enhances model performance by explicitly updating designated parameters prior to each prediction to adapt to the test data. While TTT has demonstrated considerable empirical success, its theoretical underpinnings remain limited, particularly for nonlinear models. In this paper, we investigate the combination of TTT with in-context learning (ICL), where the model is given a few examples from the target distribution at inference time. We analyze this framework in the setting of single-index models, where the feature vector is drawn from a hidden low-dimensional subspace. For single-layer transformers trained with gradient-based algorithms and adopting TTT, we establish an upper bound on the prediction risk. Our theory reveals that TTT enables the single-layer transformers to adapt to both the feature vector and the link function, which vary across tasks. This creates a sharp contrast with ICL alone, which is theoretically difficult to adapt to shifts in the link function. Moreover, we provide the convergence rate with respect to the data length, showing the predictive error can be driven arbitrarily close to the noise level as the context size and the network width grow.

stat.ML↗

arXiv2026

DPRM: A Plug-in Doob h transform-induced Token-Ordering Module for Discrete Diffusion Models

Discrete diffusion models admit many token orders, yet most systems rely on confidence-based decoding. Confidence is a strong and efficient heuristic, but it can be myopic because local certainty does not measure a position's effect on terminal quality. We introduce DPRM (Doob h-transform Process Reward Model), a plug-in token-ordering module induced by a terminal-reward-tilted trajectory distribution. The Doob transform converts terminal reward into a process reward for each candidate position. A compact estimator indexed by generation progress, confidence, and optional positional or task state makes this correction reusable across partial states. DPRM changes only token order; the host architecture, denoising objective, supervision, and token sampler remain fixed. We derive the exact reward-tilted law and bound its shortlist and online approximations. Across nine recent open-source hosts spanning language, multimodal generation, and scientific sequences, DPRM improves reasoning, numeric VQA, visual-codebook ordering, and preference-conditioned generation. Relative to confidence, it improves PUMA GSM8K by 14.6%, Omni-Diffusion CLIP-L/14 by 13.2%, Prism voted accuracy from 82.41% to 83.85%, and RealWorldQA numeric/count accuracy by 8.97 points. Matched scientific gains include 8.3% in DCM nonzero recovery, 15.0% in GenMol QED, and 53.3% in SDPO-DNA total utility. Entropy and random-order controls show that the gains do not come from favoring uncertain tokens indiscriminately. Token-order traces and shared-state interventions instead show how terminal reward can correct a locally confident but globally poor choice. Code is available at: https://github.com/DakeBU/DPRM-DLLM.

cs.LG↗

arXiv2026

The Geometry of Statistical Feature Learning in Mean-Field Langevin Dynamics

We introduce a geometric formulation of statistical feature learning for supervised regression. Feature learning is defined through a base--fiber decomposition: the base is the feature-side geometry produced by training, and the fiber is the learned feature space where estimation is performed. We prove this property for spherical mean-field Langevin dynamics, viewed as the Wasserstein gradient flow of a negative entropy-regularized empirical risk. In Gaussian multi-index models, the low-temperature stationary distribution concentrates near the hidden indices, forms a multi-spike structure, and yields parameter recovery with high probability, even though negative entropy regularization penalizes concentration. This concentration has a sharp transition at temperature $λ\asymp 1$. In Gaussian single-index models, the stationary measure satisfies a concentration property, with parity determining whether it lives on $S_2^{d-1}$ or $\mathbb{RP}^{d-1}$. The induced learned feature space aligns the regression signal and yields rates $d/N$ and $Md/N$, up to logarithmic factors.

math.ST↗