Equivariance constraint

Overview

An equivariance constraint trains a predictor f so that warping the input by a transform T gives the same result as warping the output of f: f(T(x)) should match T(f(x)). The transform is known and sampled at training time, so it supplies the target and no annotation is needed 2003.00196 §3.3 2203.14367 §3.3.

Three keypoint and region papers use this as a training loss. First Order Motion Model applies it to keypoint positions and to their local 2×2 Jacobians, with a random affine map plus a thin-plate-spline (TPS) warp on a 5×5 control grid 2003.00196 §3.3 2003.00196 app. B. Thin-Plate Spline Motion Model applies it to keypoint positions only, with a random TPS warp 2203.14367 §3.3. Motion Representations for Articulated Animation applies it to whole region transforms, each made of a centre and a linear part, rather than to single points 2104.11280 §3.5 2104.11280 eq. 10. The First Order paper notes that the idea comes from earlier unsupervised-keypoint work (Jakab et al., 2018; Zhang et al., 2018) and extends it with the Jacobian term 2003.00196 §3.3. Those papers have no card in this wiki, so nothing from them is used here.

The word is also used differently in two molecule and graph papers. DiGress's graph network is permutation equivariant by construction, and its loss is permutation invariant. No transform is sampled and no extra loss term is added 2209.14734 §3.3 2209.14734 eq. 3. GeoDiff builds roto-translation (SE(3)) equivariance into its reverse kernels, so the density of generated 3D conformations does not change under rotation or translation 2203.02923 §4.2. It also swaps the noise regression target for an equivariant one 2203.02923 §4.3. These are architectural or target changes, not the sampled-warp penalty above, so they are listed as their own variants below.

The informed-machine-learning survey gives the wider frame. Its knowledge term is a loss that takes the inputs and the model but no labels 1903.12394 eq. 1 1903.12394 §5.1, and it lists symmetries among the forms that prior knowledge can take 1903.12394 §4.2. The survey has no experiments of its own, so it does not show that any variant works 1903.12394 §1.

Cards: First Order Motion Model for Image Animation, Thin-Plate Spline Motion Model for Image Animation, Motion Representations for Articulated Animation, DiGress, GeoDiff, Informed Machine Learning survey.

Variants

The first three are extra loss terms on a keypoint or region detector, trained with a reconstruction loss. The fourth swaps the regression target of a diffusion model for an equivariant one. The fifth is a property of the network, with no sampled transform and no added loss.

Keypoint position equivariance

The detector runs on the original frame and on the warped frame. The keypoint positions found on the warped frame, mapped back through the warp, should match the positions found on the original frame. First Order Motion Model states this as an L1 distance and calls it the keypoint loss 2003.00196 eq. 10 2003.00196 eq. 11 2003.00196 §3.3. Card: First Order Motion Model for Image Animation.

Thin-Plate Spline Motion Model uses the same position idea with a random TPS warp. Its equivariance term has no Jacobian part 2203.14367 eq. 10 2203.14367 §3.3. Card: Thin-Plate Spline Motion Model for Image Animation.

Jacobian equivariance in product form

The same argument applies to the local 2×2 Jacobians. The paper writes the constraint as a condition that a product of three Jacobians, the inverse of the original one, the warp's Jacobian, and the warped-frame one, equals the identity 2003.00196 eq. 13. The paper says a direct L1 term on the Jacobian pulls its magnitude to zero, so the product form is used instead 2003.00196 §3.3. In the ablation on Tai-Chi-HD, the variant without this term reaches AKD 9.887 and MKR 0.052, against 6.862 and 0.036 for the full model 2003.00196 tab. 1. The paper's ablation has no row that removes only the keypoint position term, so its separate effect is not reported 2003.00196 tab. 1. Card: First Order Motion Model for Image Animation.

Region-transform equivariance

The region paper does not work with single points. Each of its K regions is a soft heatmap, and a centre and a linear part are measured from it (see the region variant in the Implementation sketch). The equivariance term compares these region transforms: the transform measured on a warped driving frame, mapped through the warp, should match the transform measured on the original frame 2104.11280 eq. 10 2104.11280 §3.5. The paper writes one such term and gives no weight for it 2104.11280 §3.5. The same random-warp recipe as the motion papers is used, with an affine map and a thin-plate-spline warp, but the spread values are not stated in the paper 2104.11280 §3.5. Its reported gains are author-reported: L1 0.047 against 0.056 for FOMM at K = 10 on TaiChiHD (256) 2104.11280 tab. 2. The paper's ablation removes the PCA measurement and the background model, but not the equivariance term, so the term's own effect is not reported 2104.11280 tab. 4. Card: Motion Representations for Articulated Animation.

Equivariant noise target (chain-rule form)

Diffusion models regress the noise added to the clean coordinates. That noise is not equivariant with the noisy coordinates, so a loss written with it would not be invariant as written 2203.02923 §4.3. GeoDiff replaces the noise with an equivariant target in one of two ways. GeoDiff-A rigidly aligns the clean coordinates to the noisy ones and takes the noise from the aligned coordinates. GeoDiff-C writes the noise as a weighted score over interatomic distances and pulls it back to atom displacements with the chain rule, which goes through the invariant distances 2203.02923 §4.3. The paper uses GeoDiff-C as its main model 2203.02923 §5.2. Card: GeoDiff.

Architectural equivariance

These two papers put the symmetry into the network and add no sampled transform. The training loss is then unchanged.

• Permutation equivariance (DiGress). The graph transformer is permutation equivariant and the loss is permutation invariant. Training therefore needs no random node ordering and no graph matching 2209.14734 §3.3 2209.14734 eq. 3. The paper proves these two properties; it adds no sampled transform and no loss term for them. Card: DiGress.
• SE(3) equivariance (GeoDiff). The noise network keeps node features invariant and coordinates equivariant, and stacking such layers gives an equivariant noise field 2203.02923 §4.2. With an invariant prior that has zero centre of mass, and equivariant reverse kernels, the marginal density over conformations is SE(3)-invariant 2203.02923 §4.2 2203.02923 app. A. Card: GeoDiff.

When to use which

Implementation sketch

The loss below follows the First Order paper's objective. The position and Jacobian terms use the same warp. Comments give the source in the paper. The TPS variant drops the Jacobian lines and keeps the position term only.

# Equivariance terms for a keypoint detector  [arxiv:2003.00196 §3.3]
# Input:  driving frame D, a random warp T (affine + TPS)  [arxiv:2003.00196 app. B]
# Output: scalar loss, added to the reconstruction loss with equal weight

def equivariance_loss(D, T):
    # Keypoints on the original frame: positions p and 2x2 Jacobians J
    p, J = detect(D)
    # Keypoints on the warped frame
    p_w, J_w = detect(T(D))
    # Positions: mapping the warped detections back should match p  [eq. 10] [eq. 11]
    loss_pos = l1(p, T_points(p_w))
    # Jacobians: the product of the three Jacobians should be identity  [eq. 13]
    product = inverse(J) @ T_jacobian(p) @ J_w
    loss_jac = l1(identity_2x2, product)
    return loss_pos + loss_jac

The reconstruction term is separate and is not sketched here 2003.00196 §3.3. The variants differ only in which terms are included.

The region variant keeps the same warp and the same loss shape, but replaces each point and Jacobian with a region's centre and linear part, measured from a heatmap by a weighted mean and an SVD of its covariance 2104.11280 eq. 5 2104.11280 eq. 6 2104.11280 eq. 7. Its term compares the region transform of the driving frame with the warped copy's transform 2104.11280 eq. 10. The paper writes this as one term; the released code adds a separate shift term and an affine term 2104.11280 §3.5. The architectural variants have no loss term to sketch. The GeoDiff chain-rule target replaces the noise regression target, and the paper gives its formula 2203.02923 §4.3.

Contradictions

• Spread of the random warp. The appendix gives the affine and TPS spreads as variances, 0.05 and 0.005 2003.00196 app. B. The released code draws them with those values as standard deviations model.py model.py. The card records both readings and does not pick one. The reported results cannot be tied to either reading from the paper alone. The TPS paper does not give these values; its released config uses the same numbers taichi-256.yaml. The region paper also leaves them unstated in the text, while its shipped config names sigma 0.05 and 0.005 taichi256.yaml. The same reading question applies.
• Loss weights. The First Order paper states that the reconstruction, keypoint and Jacobian terms carry equal weights in all experiments 2003.00196 §3.3. The shipped Tai-Chi-HD configuration sets weight 10 for each of these terms taichi-256.yaml. The relative weights agree; the absolute scale is not stated in the paper. The region paper gives no weights for its equivariance term, while its code uses 10 for both the shift and the affine part 2104.11280 §3.5 taichi256.yaml. The TPS paper writes its total as an unweighted sum 2203.14367 eq. 14, while its code multiplies each term by 10 taichi-256.yaml.
• Which frame gets the warp. The TPS paper's equivariance loss uses the source image S 2203.14367 eq. 10. Its code applies the random TPS to the driving frame model.py. The card records both.
• TPS kernel. The First Order paper gives no kernel formula for the TPS warp. The released code uses the L1 distance between points inside the thin-plate kernel model.py. An implementation that follows the paper alone cannot pick this form.
• Chain-rule versus alignment (GeoDiff). The paper says the chain-rule form GeoDiff-C is slightly better and uses it as the main model [arxiv:2203.02923 §5.2]. On GEOM-Drugs, GeoDiff-C has the higher COV-R mean, 89.13% against 88.36% for GeoDiff-A 2203.02923 tab. 1. On GEOM-QM9, GeoDiff-A has the higher COV-R mean and median, 90.54% and 94.61% against 90.07% and 93.39%, while GeoDiff-C has the lower MAT-R mean, 0.2090 Å against 0.2104 Å 2203.02923 tab. 5. Both sides are reported; the single-number tables show no consistent winner and no variance.

No contradictions between the two motion papers' reported results are known. DiGress, GeoDiff and the motion papers do not compete on the same task; the word "equivariance" means different things across them (see Overview).

Pitfalls

• The warp is part of the method. Its distribution (affine spread, TPS spread, control grid) sets how far the keypoints are pushed during training 2003.00196 app. B. Record it with the training run.
• Padding at the warp boundary. The First Order code uses reflection padding for the warped frame. The paper does not state this 2003.00196 §3.3.
• Matrix inversion. The product form needs the inverse of a driving Jacobian. The paper assumes the local transform is bijective. The code has no explicit guard against near-singular Jacobians 2003.00196 §3.1.
• Jacobian magnitude. A direct L1 term on the Jacobian drives it to zero, so the product form is required 2003.00196 §3.3.
• Singular region covariance. The region transform needs the inverse of a driving linear part and an SVD of a heatmap covariance. The released code uses no epsilon in either inverse, so a collapsed heatmap gives a singular matrix pixelwise_flow_predictor.py util.py. The paper does not discuss this case.
• Axis sign. An SVD returns each axis only up to sign. The region paper does not give a sign rule, and its code fixes one before the equivariance term is formed model.py. Skipping this step is untested in either the paper or the code.
• No isolated ablation of the constraint. The First Order ablation has no row without the keypoint position term 2003.00196 tab. 1. The TPS ablation adds components in turn and never removes the equivariance loss 2203.14367 tab. 4. The effect of the constraint alone is therefore not measured in either paper.
• Cost. The First Order paper reports about 100k iterations at batch 20 to 40, roughly two days on two GPUs for Tai-Chi-HD and VoxCeleb. This is author-reported and was not verified 2003.00196 app. B.
• Architectural equivariance is not exchangeability of the data. DiGress proves the network is permutation equivariant and the loss is permutation invariant 2209.14734 §3.3. That the generated distribution is exchangeable relies on an external theorem the paper cites from GeoDiff 2209.14734 app. D. The paper does not show exact likelihood invariance beyond these claims.
• Architectural cost. Each DiGress layer costs Θ(n²) in the number of nodes, and the spectral features need an eigendecomposition per graph 2209.14734 §3.2 2209.14734 app. B.2.
• GeoDiff-A is not in the released code. The repository has no alignment step, so only the chain-rule form can be run from it dualenc.py geometry.py.
• Coordinate convention must be one. The paper's noisy state is √ᾱ C⁰ + √(1−ᾱ) ε 2203.02923 app. A. The code works in a scaled state that differs by √ᾱ dualenc.py. Mixing the two gives a wrong noise target and a wrong sampler start.
• Centring affects the invariance proof. Algorithm 1 centres the state before every prediction 2203.02923 alg. 1. The released code centres after each update dualenc.py dualenc.py. The card reads this as equivalent but did not test it.
• The released denoiser takes no step index. The paper's noise network takes the step t 2203.02923 eq. 4. The code's forward pass does not use it dualenc.py. A time-conditioned reimplementation will not match the released weights.
• Division by edge length. The equivariant projection divides by each edge length with no epsilon geometry.py. Atoms that come very close give large values; the code clips the global force by default test.py.
• Sampling cost. One conformation needs T network evaluations, 5000 at the paper's default 2203.02923 tab. 4. The T = 1000 variant is weaker than T = 5000 2203.02923 tab. 6.
• Survey, not evidence. The survey gives no counts and no search protocol. Its "main paths" are the authors' sample, not a census of the field 1903.12394 §1 1903.12394 §3.1.2. Use it for vocabulary and the taxonomy, not for claims about any one variant.

Papers

• First Order Motion Model for Image Animation: trains a keypoint detector with a reconstruction loss and two equivariance terms (positions and Jacobians), without object annotations.
• Thin-Plate Spline Motion Model for Image Animation: replaces local affine motion with TPS fields and trains with a position-only equivariance term under a random TPS warp.
• Motion Representations for Articulated Animation: measures K heatmap regions by PCA and trains with an equivariance term on whole region transforms, under a random affine and TPS warp.
• DiGress: Discrete Denoising diffusion for graph generation: builds permutation equivariance into a graph transformer and uses a permutation-invariant loss, with no sampled transform.
• GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation: diffusion over 3D atom coordinates with SE(3)-equivariant reverse kernels and an equivariant noise target.
• Informed Machine Learning: A Taxonomy and Survey of Integrating Knowledge into Learning Systems: a taxonomy in which symmetries are one form of prior knowledge; a survey with no experiments.