Exact Tensorized Distribution Loss

(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.