@@ -167,3 +167,100 @@ def get_occ(natom):
167167 def get_occ (natom ):
168168 return new_occ [natom ]
169169 return get_occ
170+
171+ class SafeEigh (torch .autograd .Function ):
172+ """
173+ A custom autograd function for eigendecomposition of real symmetric matrices.
174+ It handles degenerate eigenvalues by masking out the infinite gradients
175+ caused by the term 1/(lambda_i - lambda_j) when lambda_i approx lambda_j.
176+
177+ Reference:
178+ Derivatives of Partial Eigendecomposition of a Real Symmetric Matrix
179+ for Degenerate Cases (Kasim et al., 2020), Equation (27).
180+ """
181+
182+ @staticmethod
183+ def forward (ctx , a ):
184+ """
185+ Forward pass: Standard eigendecomposition.
186+
187+ Args:
188+ a: Input symmetric matrix. Shape: (..., N, N), supports batching.
189+ Returns:
190+ e: Eigenvalues. Shape: (..., N).
191+ v: Eigenvectors. Shape: (..., N, N).
192+ """
193+ # Ensure the input is float/complex as required by eigh
194+ # Note: 'U' (Upper) or 'L' (Lower) doesn't matter much for valid symmetric inputs
195+ e , v = torch .linalg .eigh (a )
196+
197+ # Save tensors for the backward pass
198+ ctx .save_for_backward (e , v )
199+ return e , v
200+
201+ @staticmethod
202+ def backward (ctx , grad_e , grad_v ):
203+ """
204+ Backward pass: Computes gradient with respect to input matrix 'a'.
205+
206+ This implementation specifically handles the degeneracy issue where
207+ eigenvalues are identical or very close, which would normally cause
208+ NaNs or Infs in the gradient of eigenvectors.
209+ """
210+ e , v = ctx .saved_tensors
211+
212+ # 1. Handle cases where gradients might be None
213+ # (e.g., if eigenvalues or eigenvectors are not used in the loss function)
214+ if grad_e is None :
215+ grad_e = torch .zeros_like (e )
216+ if grad_v is None :
217+ grad_v = torch .zeros_like (v )
218+
219+ # 2. Construct the pairwise difference matrix of eigenvalues
220+ # Shape of e: (Batch, N)
221+ # Use unsqueeze to broadcast: (Batch, N, 1) - (Batch, 1, N) -> (Batch, N, N)
222+ # e_diff[..., i, j] = e[..., i] - e[..., j] (column - row)
223+ e_diff = e .unsqueeze (- 2 ) - e .unsqueeze (- 1 )
224+
225+ # 3. Handle Degeneracy (Masking)
226+ # Define a small threshold to detect degeneracy
227+ epsilon = 1e-8
228+
229+ # Create a mask where |lambda_i - lambda_j| > epsilon
230+ mask = torch .abs (e_diff ) > epsilon
231+
232+ # Construct the F matrix: F_ij = 1 / (lambda_j - lambda_i)
233+ # Note: We use the transposed definition implicit in the matrix formula below.
234+ # Here we initialize f_matrix with zeros, effectively ignoring degenerate terms.
235+ f_matrix = torch .zeros_like (e_diff )
236+
237+ # Only compute division for non-degenerate pairs
238+ # This prevents division by zero and corresponds to setting the gradient
239+ # contribution of degenerate subspaces to zero (Gauge Invariance).
240+ f_matrix [mask ] = 1.0 / e_diff [mask ]
241+
242+ # 4. Compute the gradient w.r.t. the input matrix 'a'
243+ # Formula: grad_a = v @ (diag(grad_e) + F * (v^T @ grad_v)) @ v^T
244+
245+ # Projection of gradients onto the eigenvector basis: v^T @ grad_v
246+ # transpose(-2, -1) handles the last two dimensions for batch processing
247+ vt = v .transpose (- 2 , - 1 )
248+ v_t_grad_v = vt @ grad_v
249+
250+ # The middle term: diag(grad_e) + F * (v^T @ grad_v)
251+ # torch.diag_embed creates a diagonal matrix from the eigenvalue gradients
252+ # f_matrix * v_t_grad_v performs element-wise multiplication (Hadamard product)
253+ mid_term = torch .diag_embed (grad_e ) + f_matrix * v_t_grad_v
254+
255+ # Transform back to the original basis
256+ grad_a = v @ mid_term @ vt
257+
258+ # 5. Enforce symmetry
259+ # Since the input 'a' is symmetric, its gradient must also be symmetric.
260+ grad_a = 0.5 * (grad_a + grad_a .transpose (- 2 , - 1 ))
261+
262+ return grad_a
263+
264+ # Wrapper function for easy usage
265+ def safe_eigh (input_tensor ):
266+ return SafeEigh .apply (input_tensor )
0 commit comments