emerging_optimizers.utils#
- class emerging_optimizers.utils.SinkhornMapper(num_iters=20, eps=1e-08)[source]#
Applies the Sinkhorn-Knopp mapping to the input tensor.
The Sinkhorn-Knopp mapping is an iterative technique for normalizing the rows and columns of a matrix: Input -> [Exp] -> [Iterative Row/Col Normalization]
Supports batched inputs (3D+). The mapping operates on the last two dimensions.
- For an M×N matrix, the normalization targets are:
Square (M=N): row sums = 1.0, col sums = 1.0 (standard doubly-stochastic)
Wide (N>M): row sums = N/M, col sums = 1.0
Tall (M>N): row sums = 1.0, col sums = M/N
Based on Deepseek’s Manifold-Constrained Hyperconnections (https://arxiv.org/abs/2512.24880)
emerging_optimizers.utils.eig#
- emerging_optimizers.utils.eig.conjugate(a, p, diag=False)[source]#
Calculate similarity transformation
This function calculates \(B = P^T A P\). It assumes P is orthogonal so that \(P^{-1} = P^T\) and the similarity transformation exists.
- emerging_optimizers.utils.eig.eigh_with_fallback(x, force_double=False)[source]#
torch.linalg.eigh() function with double precision fallback
Unified wrapper over eigh() function with automatic fallback and force double precision options. Automatically falls back to double precision on failure and returns eigenvalues in descending order. Default 2nd argument of eigh UPLO is ‘L’.
- Parameters:
- Returns:
Eigenvalues and eigenvectors tuple (eigenvalues in descending order).
- Return type:
- emerging_optimizers.utils.eig.orthogonal_iteration(kronecker_factor, eigenbasis, power_iter_steps)[source]#
Refines an eigenbasis via power iteration with QR re-orthogonalization.
Performs
power_iter_stepsrounds ofQ = QR(kronecker_factor @ Q)starting fromeigenbasis. The columns ofeigenbasisare expected to already be aligned with the intended descending-eigenvalue ordering ofkronecker_factor(seeemerging_optimizers.soap.soap_utils.sort_eigenbasis_by_approx_eigvals()).- Parameters:
- Returns:
The refined eigenbasis.
- Return type:
emerging_optimizers.utils.modules#
- class emerging_optimizers.utils.modules.Conv1dFlatWeights(*args, **kwargs)[source]#
Conv1d with weights+bias stored in a single 2D tensor
There are conv1d used in some LLM, in mamba mixer for example. Because the weight is not 2d, we cannot apply many of the emerging optimizers originally introduced for 2d weights of Linear layers without bias. Since convolution can be viewed as a matrix multiplication with im2col (either implicit or explicit), we can flatten the weight into a single 2D tensor and then apply the emerging optimizers to it.
Bias is not commonly used in most LLM’s anymore, but they are often included in this type of conv1d. Since bias is mathematically the 0 order term of the polynomial, we can combine weight and bias into a single 2D tensor.
Arguments are the same as :
torch.nn.Conv1d.Note
This implementation potentially introduces a small overhead because of split weights can combining gradients of it. This should be trivial compared to computational cost of LLM training. If it becomes a concern, a kernel can be developed to eliminate the overhead.
Note
Similar flattening logic can be applied to N-D convolution. But since we don’t have use cases of them in LLM yet, they are not supported despite the __init__() function is generalized enough to support N-D convolution.
- extra_repr()[source]#
Return the extra representation of the module.
To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.
- Return type:
- forward(x)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.