Skip to content

lphm_parameterized_concatenation_fusion

Bases: parameterized_concatenation_fusion

A parameterized concatenation fusion with low-rank hierarchical matrix (LPHM) parameter fabrication.

Notes

Formally, given input interdependence matrices \(\mathbf{A}_1, \mathbf{A}_2, \ldots, \mathbf{A}_k\), where each matrix \(\mathbf{A}_i \in R^{m \times n_i}\) has \(m\) rows and \(n_i\) columns, we define the fusion operator as follows:

\[
    \begin{equation}
    \begin{aligned}
    \mathbf{A} &= \text{fusion}(\mathbf{A}_1, \mathbf{A}_2, \cdots, \mathbf{A}_k) \\
    &= \left( \mathbf{A}_1 \sqcup \mathbf{A}_2 \sqcup \cdots \sqcup \mathbf{A}_k \right) \mathbf{W} \in R^{m \times n},
    \end{aligned}
    \end{equation}
\]

where \(\sqcup\) denotes the row-wise concatenation of the matrices.

Notation \(\mathbf{W} \in R^{(\sum_{i=1}^k n_i) \times n}\) denotes the parameter matrix fabricated from the learnable parameter vector \(\mathbf{w} \in R^{l}\), which can be represented as follows:

$$ \begin{equation} \psi(\mathbf{w}) = \mathbf{A} \otimes \mathbf{B} = \mathbf{A} \otimes ( \mathbf{S} \mathbf{T}^\top) \in R^{(\sum_{i=1}^k n_i) \times n}, \end{equation} $$ where \(\mathbf{A} \in R^{p \times q}\), \(\mathbf{S} \in R^{\frac{(\sum_{i=1}^k n_i)}{p} \times r}\) and \(\mathbf{T} \in R^{\frac{n}{q} \times r}\) are partitioned and reshaped from the parameter vector \(\mathbf{w}\).

The required length of parameter vector of this interdependence function is \(l = pq + r \times (\frac{(\sum_{i=1}^k n_i)}{p} + \frac{n}{q})\).

Attributes:

Name Type Description
r int

Rank for the low-rank approximation.

p int

Partition size for the hierarchical matrix.

q int

Block size for the hierarchical matrix.

Methods:

Name Description
__init__

Initializes the LPHM parameterized concatenation fusion function.

Source code in tinybig/fusion/parameterized_concatenation_fusion.py
class lphm_parameterized_concatenation_fusion(parameterized_concatenation_fusion):
    r"""
        A parameterized concatenation fusion with low-rank hierarchical matrix (LPHM) parameter fabrication.

        Notes
        ----------

        Formally, given input interdependence matrices $\mathbf{A}_1, \mathbf{A}_2, \ldots, \mathbf{A}_k$,
        where each matrix $\mathbf{A}_i \in R^{m \times n_i}$ has $m$ rows and $n_i$ columns,
        we define the fusion operator as follows:

        $$
            \begin{equation}
            \begin{aligned}
            \mathbf{A} &= \text{fusion}(\mathbf{A}_1, \mathbf{A}_2, \cdots, \mathbf{A}_k) \\
            &= \left( \mathbf{A}_1 \sqcup \mathbf{A}_2 \sqcup \cdots \sqcup \mathbf{A}_k \right) \mathbf{W} \in R^{m \times n},
            \end{aligned}
            \end{equation}
        $$

        where $\sqcup$ denotes the row-wise concatenation of the matrices.

        Notation $\mathbf{W} \in R^{(\sum_{i=1}^k n_i) \times n}$ denotes the parameter matrix fabricated from the learnable parameter vector $\mathbf{w} \in R^{l}$,
        which can be represented as follows:

        $$
            \begin{equation}
            \psi(\mathbf{w}) = \mathbf{A} \otimes \mathbf{B} = \mathbf{A} \otimes ( \mathbf{S} \mathbf{T}^\top) \in R^{(\sum_{i=1}^k n_i) \times n},
            \end{equation}
        $$
        where $\mathbf{A} \in R^{p \times q}$, $\mathbf{S} \in R^{\frac{(\sum_{i=1}^k n_i)}{p} \times r}$ and $\mathbf{T} \in R^{\frac{n}{q} \times r}$ are partitioned and reshaped from the parameter vector $\mathbf{w}$.

        The required length of parameter vector of this interdependence function is $l = pq + r \times (\frac{(\sum_{i=1}^k n_i)}{p} + \frac{n}{q})$.


        Attributes
        ----------
        r : int
            Rank for the low-rank approximation.
        p : int
            Partition size for the hierarchical matrix.
        q : int
            Block size for the hierarchical matrix.

        Methods
        -------
        __init__(...)
            Initializes the LPHM parameterized concatenation fusion function.
    """
    def __init__(self, r: int, p: int, q: int = None, name: str = 'lphm_parameterized_concatenation_fusion', *args, **kwargs):
        """
            Initializes the LPHM parameterized concatenation fusion function.

            Parameters
            ----------
            r : int
                Rank for the low-rank approximation.
            p : int
                Partition size for the hierarchical matrix.
            q : int, optional
                Block size for the hierarchical matrix. Defaults to `p`.
            name : str, optional
                Name of the fusion function. Defaults to "lphm_parameterized_concatenation_fusion".
            *args : tuple
                Additional positional arguments for the parent class.
            **kwargs : dict
                Additional keyword arguments for the parent class.

            Raises
            ------
            AssertionError
                If `n` is not divisible by `p` or if the sum of `dims` is not divisible by `q`.
        """
        super().__init__(name=name, *args, **kwargs)
        self.r = r
        self.p = p
        self.q = q if q is not None else p
        assert self.n is not None and self.n % self.p == 0
        assert self.dims is not None and sum(self.dims) % self.q == 0
        self.parameter_fabrication = lphm_reconciliation(p=self.p, q=self.q, r=self.r)

__init__(r, p, q=None, name='lphm_parameterized_concatenation_fusion', *args, **kwargs)

Initializes the LPHM parameterized concatenation fusion function.

Parameters:

Name Type Description Default
r int

Rank for the low-rank approximation.

required
p int

Partition size for the hierarchical matrix.

required
q int

Block size for the hierarchical matrix. Defaults to p.

None
name str

Name of the fusion function. Defaults to "lphm_parameterized_concatenation_fusion".

'lphm_parameterized_concatenation_fusion'
*args tuple

Additional positional arguments for the parent class.

()
**kwargs dict

Additional keyword arguments for the parent class.

{}

Raises:

Type Description
AssertionError

If n is not divisible by p or if the sum of dims is not divisible by q.

Source code in tinybig/fusion/parameterized_concatenation_fusion.py
def __init__(self, r: int, p: int, q: int = None, name: str = 'lphm_parameterized_concatenation_fusion', *args, **kwargs):
    """
        Initializes the LPHM parameterized concatenation fusion function.

        Parameters
        ----------
        r : int
            Rank for the low-rank approximation.
        p : int
            Partition size for the hierarchical matrix.
        q : int, optional
            Block size for the hierarchical matrix. Defaults to `p`.
        name : str, optional
            Name of the fusion function. Defaults to "lphm_parameterized_concatenation_fusion".
        *args : tuple
            Additional positional arguments for the parent class.
        **kwargs : dict
            Additional keyword arguments for the parent class.

        Raises
        ------
        AssertionError
            If `n` is not divisible by `p` or if the sum of `dims` is not divisible by `q`.
    """
    super().__init__(name=name, *args, **kwargs)
    self.r = r
    self.p = p
    self.q = q if q is not None else p
    assert self.n is not None and self.n % self.p == 0
    assert self.dims is not None and sum(self.dims) % self.q == 0
    self.parameter_fabrication = lphm_reconciliation(p=self.p, q=self.q, r=self.r)