Class FactorizationSpec

axlearn/common/base_layer.py:131–155

A FactorizationSpec describes how to factorize a parameter's gradient. Used by AdaFactor optimizer for memory-efficient second-moment estimation by factorizing the second-moment matrix into row and column statistics instead of storing the full matrix. Attributes: axes: A list o

@dataclasses.dataclass
class FactorizationSpec:
    """A FactorizationSpec describes how to factorize a parameter's gradient.

    Used by AdaFactor optimizer for memory-efficient second-moment estimation by factorizing
    the second-moment matrix into row and column statistics instead of storing the full matrix.

    Attributes:
        axes: A list of None/str corresponding to the axes of the parameter shape.
            Each element is either:
            - None: no factorization along this axis.
            - str: the factorization axis name (typically "row" or "col").

            For AdaFactor, either:
            - All axes are None (no factorization, used for small parameters)
            - Exactly two axes are "row" and "col" (factorized, used for large matrices)

            Example:
                For a weight matrix of shape [1024, 4096]:
                axes=["row", "col"] enables factorization

                For a 3D tensor of shape [8, 1024, 4096]:
                axes=[None, "row", "col"] factorizes only the last two dimensions
    """

    axes: Sequence[Optional[str]]


# Ideally this would be a recursive type: