Lower-Order-Invariant High-Order Representation Loss

Attach the loss to d selected channels, token groups, or attention heads immediately before the task head. For each minibatch, compute a scalar a_i^(b) from group i using a learned linear projection followed by a sigmoid, then form R_i^(b)=2 sigmoid(a_i^(b))-1. Enumerate all nonempty proper subsets S of [d] for d<=8, or sample a fixed collection of subsets for larger d. Compute H_hat_S=B^{-1} sum_b prod_{i in S} R_i^(b), set the lower-order penalty to the sum of squared moments, and set the d-way target penalty to (H_hat_[d]-tau)^2. Warm up with lambda_high=0, then ramp both coefficients over 5-10% of training to prevent representation collapse; normalize each moment by its minibatch standard error or maintain exponential moving averages. The mathematical target is zero lower-order moments and nonzero full interaction; the paper's 2^{-d} and 1/3 results motivate testing whether strong full dependence survives despite lower-order invariance, but the moment estimator is empirical. First experiment: train a small Transformer or MLP on a synthetic task combining a parity feature with a spurious pairwise feature, comparing the baseline, lower-order penalty alone, and complete loss. Evaluate in-distribution accuracy, accuracy after permuting pairwise correlations, estimated moments, and calibration. Success means lower-order moments approach zero, parity accuracy remains high, and out-of-distribution accuracy improves without substantial optimization instability.