Condition-Wise Sinkhorn Drifting for One-Shot Learned Channel Simulation

Abstract: Learned communication systems may evaluate stochastic channel surrogates millions of times inside differentiable training loops, making diffusion-style reverse sampling expensive. This paper proposes condition-wise Sinkhorn drifting, a one-shot channel surrogate that preserves the transmitted symbol and transports only the conditional output laws (p(y\mid x)). We formulate a conditional Sinkhorn objective over repeated outputs at the same transmitted symbol and train the generator with finite-sample barycentric velocities followed by detached particle regression. Experiments on additive white Gaussian noise (AWGN), Rayleigh fading, solid-state power amplifier (SSPA) nonlinearity, and a compact tapped-delay-line (TDL) channel compare direct drifting, joint Sinkhorn drifting, condition-wise Sinkhorn drifting, conditional denoising diffusion probabilistic modeling (DDPM), denoising diffusion implicit modeling (DDIM), and Wasserstein generative adversarial network (WGAN) references. Within the evaluated one-shot drifting-family variants, condition-wise Sinkhorn is strongest under conditional diagnostics and symbolic-coding checks, while diffusion remains strongest on the hardest downstream symbol-error-rate (SER) curves. The resulting operating point is a condition-preserving one-shot simulator for settings where repeated channel calls make diffusion-style sampling too costly.

• The paper introduces condition-wise Sinkhorn drifting to directly match conditional distributions in a one-step forward pass.
• It leverages barycentric velocity fields from optimal transport to efficiently simulate complex channel behaviors.
• Empirical results show lower sliced Wasserstein distances and improved SER/BER performance compared to diffusion-based methods.

Condition-Wise Sinkhorn Drifting for One-Shot Learned Channel Simulation

Introduction and Context

The accurate simulation of communication channels is foundational to the design and optimization of modern learned communication systems, especially as these systems increasingly rely on end-to-end differentiable training loops. Channel models that provide realistic, sample-accurate, and efficiently computable conditional distributions $p(y|x)$ are essential for both performance benchmarking and inner-loop optimization. Traditional approaches using analytical models fall short when real-world channels deviate from idealized assumptions or contain complexities such as nonlinearities, fading, and memory effects. Recent work leveraging diffusion models has significantly advanced the fidelity of learned channel surrogates; however, the iterative, gradient-based sampling procedures inherent in diffusion and related approaches incur nontrivial inference-time costs, making them impractical for scenarios with heavy inner-loop invocation.

This paper introduces "condition-wise Sinkhorn drifting," a one-shot learned channel surrogate that explicitly enforces preservation of the conditioning input (transmitted symbol $x$) and efficiently transports the conditional output laws $p(y|x)$ by construction. Building upon the drifting generative modeling paradigm, the method reformulates the Sinkhorn/Wasserstein gradient flow objective for the conditional law, training a direct generator to match $p(y|x)$ for almost every $x$ in a one-step forward-pass manner. This contribution addresses the operational constraint where repeated channel calls dominate the computational budget, a scenario endemic to neural encoder-decoder pipelines and Monte Carlo-based training regimes.

Methodological Framework

One-Shot Drifting and Conditional Optimal Transport

Distinct from diffusion models, which realize sample refinement through a sequence of time-reversed denoising steps, drifting models relocate the iterative matching to the training phase: generator samples are directly nudged toward the target data manifold via a data-driven drift field, and the generator is trained to minimize the barycentric (Sinkhorn) distance to the true conditional distribution. Once trained, generation consists of a single network pass per sample, with no iterative correction at inference time.

The condition-wise Sinkhorn drifting formalism arises from reformulating the optimal transport problem at fixed input conditions. Let $p(dx, dy)$ denote the joint distribution of transmitted symbol and channel output, with generator law $q_\theta(dx, dy) = p(dx) q_{\theta,x}(dy)$. The Sinkhorn divergence is computed separately for each fiber (the conditional output law at fixed $x$) and then integrated over the marginal of $x$, yielding the objective:

$$S_{\text{cond}}(q_\theta, p) = \int S_\varepsilon(q_{\theta, x}, p_x) p(dx)$$

where $x$0 denotes the entropic Sinkhorn divergence. The generator $x$1 is trained using finite-sample barycentric velocity estimates, with the update projected onto the generator parameterization via detached (stop-gradient) particle regression. This method ensures that transport occurs only within the conditional output, preserving the essential structure for downstream communication-systems.

Implementation and Training

The practical implementation employs an MLP generator with conditioning on $x$2 and random latent $x$3, optimized via Adam with problem-specific architectural and optimizer hyperparameters. During minibatch training, repeated samples at each $x$4 are drawn from both the analytic/simulator channel and the current generator. The empirical conditional Sinkhorn coupling is solved per-condition, and the resulting barycentric field provides the update direction for each sample.

The method is compared to both joint/global Sinkhorn drifting (which matches the full joint law, allowing bleed across conditioning variables) and direct kernel-based drifting. The ablation between these drift-field choices isolates the effect of restricting transport to the conditional law, a central claim of the paper.

Empirical Evaluation

Benchmarking Across Channel Types

The study evaluates the method on four canonical channel models:

• AWGN: Standard additive white Gaussian noise.
• Rayleigh: Independent per-component fading followed by AWGN.
• SSPA: Solid-state power amplifier with nonlinear gain, operating in I/Q space.
• TDL: Short-block 3GPP TDL-D multipath channel (memory channel, complex).

For each, the method is benchmarked using generator-level metrics (direct-sliced Wasserstein distance, SWD, and anchor-conditioned moments), and downstream symbol/bit error rates (SER/BER) in an autoencoding setting (system trained via surrogate channel, evaluated on analytic truth).

Numerical Results

• Generator-Level Accuracy: Condition-wise Sinkhorn drifting obtains the lowest or nearly-lowest SWD among all one-shot generator variants across AWGN, Rayleigh, and TDL, indicating strong fidelity to the output distribution. On SSPA, direct drifting is marginally better under global SWD, but condition-wise Sinkhorn remains highly competitive.
• Conditional Law Preservation: Anchor-conditioned diagnostics (e.g., fixed-$x$5 SWD, Wasserstein-2) robustly favor condition-wise Sinkhorn, especially evident in SSPA where nonlinearity introduces pronounced conditional error in joint transport models.
• Downstream Performance: In symbolic autoencoder experiments, the condition-wise Sinkhorn surrogate yields the lowest learned-implant SER/BER among one-shot generative models in most channels. Diffusion-based surrogates outperform one-shot surrogates on the most challenging SER/BER curves, but at a steep inference-time cost.
• Latency and Training Cost: One-shot models (drifting, condition-wise Sinkhorn, WGAN) require microseconds per sample at inference, compared to tens to hundreds of microseconds for 100-step diffusion/denoising samplers. Sinkhorn drift fields introduce manageable training-time overhead, with per-condition Sinkhorn solving scaling efficiently due to the independence of fiber-wise couplings.

Key Claims

• Condition-wise Sinkhorn drifting achieves the best one-shot conditional-channel surrogate performance under conditional diagnostics and downstream symbolic-coding checks, except for the hardest channels where diffusion remains strongest on SER.
• Strict matching of the joint distribution across $x$6 is insufficient for channel simulation; explicit condition-wise transport is required to preserve $x$7.
• Discrepancies between global sample-matching metrics (e.g., SWD), anchor-conditioned metrics, and actual system-level performance motivate the inclusion of fixed-input and downstream diagnostics in evaluation protocols.

Theoretical and Practical Implications

Theoretically, condition-wise Sinkhorn drifting clarifies the optimal transport interpretation of channel simulation: accurate surrogates must realize equilibrium in the family of conditional output laws—not merely in joint distributional overlap. This insight generalizes beyond communication, informing any scenario where generator conditioning or side-information is central to the modeling objective. The use of Sinkhorn-based barycentric velocity fields aligns training with the trajectory of the entropic Wasserstein gradient flow, offering improved sample quality and stability relative to heuristic or adversarial matching.

Practically, the approach realizes significant acceleration for training protocols that embed channel simulation tightly in inner loops. The trade-off is clear: one-shot surrogates trade some ultimate fidelity (in the hardest-case SER curves) for orders-of-magnitude reduction in inference latency, a trade directly relevant for real-world learned communication system deployment and control settings.

For empirical datasets lacking repeated conditional samples (e.g., measured real-world channels), the method proposes an extension based on kernel-weighted local conditional Sinkhorn problems. Future work on robust neighborhood selection and kernel-induced bias control is required to operationalize this in the field.

Speculation on Future Directions

The separation between global and conditional matching exposes further opportunities for conditional optimal transport in generative modeling, especially in high-dimensional or structured-output scenarios. Extensions may involve integrating richer transport features (phase, amplitude, channel-state information) or adaptive kernel geometries, and developing rigorous methods for local conditional law estimation in small-sample or non-simulator regimes.

Advancements in the computational efficiency of Sinkhorn solvers and mini-batch approximations will further enhance scalability, enabling application to longer block codes, large MIMO systems, and joint source-channel learning. The approach also provides a principled bridge to consistency/distillation-based diffusion acceleration methods, which also pursue one-shot or few-shot generation via compressed matching of iterative models.

Conclusion

Condition-wise Sinkhorn drifting realizes a principled, one-shot generative approach to learned channel simulation, enforcing transport only within fixed-input conditional laws and yielding surrogates with superior conditional and downstream performance relative to existing drift-based baselines. The method effectively balances accuracy and inference latency, making it highly suitable for inner-loop and real-time learned communication settings. The divergence between global, conditional, and system-level metrics underscores the importance of comprehensive evaluation. The framework sets a foundation for further research on measured data adaptation and conditional optimal transport in generative modeling (2606.17893).