Jucys–Murphy spectral token mixer

Integrate this module as a replacement for one self-attention block in a small encoder, initially with sequence length n fixed or padded to a fixed maximum. The input is X with shape [batch,n,d]. First apply LayerNorm and a linear projection W_v. For each k=2,...,n, construct Y_k=(1/(k-1))\sum_{i<k}P_{ik}Y, where P_{ik} is implemented by exchanging columns i and k of Y; then update Y <- a_k Y + b_k Y_k. This sequential implementation equals the product F_theta because the Jucys–Murphy operators commute, although floating-point ordering can still cause tiny numerical differences. Use scalar parameters a_k initialized to 1 and b_k initialized to 0, constrain |b_k| with tanh or weight decay, and initialize the residual scale gamma to 1e-2. Apply W_o to the final Y and add it to the input. For variable sequence lengths, mask nonexistent positions and renormalize each J_k by the number of valid i<k terms. The commutation relation and product form come from the paper; the normalization, residual connection, and parameterization are engineering adaptations. The first experiment should compare a 4-layer Set Transformer, standard self-attention, and the Jucys–Murphy mixer on permutation-sensitive set classification and synthetic hierarchical aggregation tasks using 1k–4k token sequences. Measure validation accuracy, robustness to unseen permutations, parameter count, and wall-clock cost. Success is higher accuracy or better permutation generalization than an equal-parameter MLP mixer, without destabilizing training; a secondary target is matching attention accuracy at substantially lower quadratic compute.