# Exact Tensorized Distribution Loss

- ID: 2946
- Canonical URL: https://synthcore.org/idea/2946/exact-tensorized-distribution-loss
- API JSON: https://synthcore.org/api/idea/2946.json
- API Markdown: https://synthcore.org/api/idea/2946.md
- Verification status: unverified
- Source: [arXiv:2608.28556](https://arxiv.org/abs/2608.28556)
- Category: training
- Solves: stability, sample-efficiency, scalability
- ML areas: loss, vae, diffusion
- Math tags: information-theory, probability, functional-analysis
- Ratings: usefulness 6/10; difficulty 3/10; novelty 6/10

## Idea description

When a neural model defines a product distribution over coordinates, compute the discrepancy to a target product distribution from per-coordinate divergences using the exact tensorization rule instead of sampling full vectors and estimating a joint divergence. Implement KL, chi-squared, and squared-Hellinger variants as drop-in losses, with an optional learned choice among these mathematically tensorizable families.

## Mathematical statement

The paper formalizes tensorization by requiring continuous maps \(\tau_n\) satisfying \(D_f(P_1\otimes\cdots\otimes P_n\|Q_1\otimes\cdots\otimes Q_n)=\tau_n(D_f(P_1\|Q_1),\ldots,D_f(P_n\|Q_n))\), where \(P_i,Q_i\) are probability measures, \(\otimes\) is the product measure, \(D_f\) is the Csiszár \(f\)-divergence, and \(\tau_n\) depends only on marginal divergence values. The paper states that every admissible tensorization formula has a multi-affine form and identifies KL, \(\chi^2\), and squared Hellinger as tensorizable divergences. For implementation, use the exact identities \(D_{\mathrm{KL}}(\otimes_iP_i\|\otimes_iQ_i)=\sum_iD_{\mathrm{KL}}(P_i\|Q_i)\), \(1+D_{\chi^2}(\otimes_iP_i\|\otimes_iQ_i)=\prod_i(1+D_{\chi^2}(P_i\|Q_i))\), and \(1-H^2(\otimes_iP_i,\otimes_iQ_i)=\prod_i(1-H^2(P_i,Q_i))\). Here \(H^2(P,Q)=1-\int\sqrt{dP\,dQ}\). Multiplicative identities should be evaluated in log-space.

## Key formulas

- $$D_f(P_1\otimes\cdots\otimes P_n\|Q_1\otimes\cdots\otimes Q_n)=\tau_n\bigl(D_f(P_1\|Q_1),\dots,D_f(P_n\|Q_n)\bigr).$$
- $$D_{\mathrm{KL}}\left(\bigotimes_{i=1}^nP_i\middle\|\bigotimes_{i=1}^nQ_i\right)=\sum_{i=1}^nD_{\mathrm{KL}}(P_i\|Q_i).$$
- $$1+D_{\chi^2}\left(\bigotimes_{i=1}^nP_i\middle\|\bigotimes_{i=1}^nQ_i\right)=\prod_{i=1}^n\left(1+D_{\chi^2}(P_i\|Q_i)\right).$$
- $$1-H^2\left(\bigotimes_{i=1}^nP_i,\bigotimes_{i=1}^nQ_i\right)=\prod_{i=1}^n\left(1-H^2(P_i,Q_i)\right),\qquad H^2(P,Q)=1-\int\sqrt{dP\,dQ}.$$

## Implementation notes

(1) Integration point: modify the output-distribution loss of a VAE decoder, diffusion denoiser with factorized likelihood, masked-token model, or world model whose prediction is \(P_\theta(x_{1:n}\mid c)=\prod_iP_{\theta,i}(x_i\mid c)\), with a target distribution \(Q(x_{1:n}\mid c)=\prod_iQ_i(x_i\mid c)\). Do not apply this to an actually correlated target unless correlation is deliberately approximated away. (2) Pseudocode: for each batch context \(c\), obtain predicted categorical probabilities \(p_i\) and target probabilities \(q_i\); compute \(d_i^{KL}=\sum_kq_{ik}(\log q_{ik}-\log p_{ik})\), \(d_i^{chi}=\sum_k(q_{ik}-p_{ik})^2/p_{ik}\), or \(h_i=\sum_k\sqrt{p_{ik}q_{ik}}\); then set `loss_KL = mean(sum_i d_i_KL)`, `loss_chi = mean(exp(sum_i log1p(d_i_chi))-1)`, or `loss_H = mean(1-exp(sum_i log(clamp(h_i,min=eps))))`. Average over contexts and backpropagate. Use float64 or compensated summation for long products, and clamp probabilities only for numerical safety. (3) Computed from the mathematics: the aggregation rules and log-domain evaluation are exact under the product assumption. Estimated empirically: calibration, target factorization error, and whether gradients have lower variance than a sampled joint-divergence estimator. (4) First experiment: train a small 6-layer Transformer or VAE on synthetic length-64 categorical sequences and MNIST, comparing standard per-coordinate cross-entropy, sampled joint Hellinger or chi-squared estimation, and exact tensorized losses. Equalize optimizer, parameter count, training steps, and total forward FLOPs; report wall-clock time, gradient variance, and held-out NLL or distributional discrepancy. (5) Pre-register mechanism tests: on synthetic independent coordinates, exact and Monte Carlo losses should agree within 1% relative error at 64 coordinates; gradient coefficient of variation should be at least 25% lower than joint sampling at the same sample budget; and chi-squared/Hellinger losses should remain finite for sequence lengths where naive products overflow or underflow. Measure each over 10 seeds. Ablate only tensorized aggregation by replacing it with joint sampling or an arithmetic mean of marginal divergences. The transfer is falsified if exact aggregation fails to reduce estimator variance, or if its product-rule value disagrees with direct enumeration by more than 1% when factorization holds exactly.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.