Skip to content

lorr_reconciliation

Bases: fabrication

The low-rank parameter reconciliation function.

It performs the low-rank parameter reconciliation, and returns the low-rank reconciled parameter matrix of shape (n, D). This class inherits from the reconciliation class (i.e., the fabrication class in the module directory).

...

Notes

Formally, given the parameter vector \(\mathbf{w} \in {R}^{l}\) and a rank hyper-parameter \(r\), low-rank parameter reconciliation function partitions \(\mathbf{w}\) into two sub-vectors and subsequently reshapes them into two matrices \(\mathbf{A} \in {R}^{n \times r}\) and \(\mathbf{B} \in {R}^{D \times r}\), each possessing a rank of \(r\). These two sub-matrices \(\mathbf{A}\) and \(\mathbf{B}\) help define the low-rank reconciliation function as follows: $$ \begin{equation} \psi(\mathbf{w}) = \mathbf{A} \mathbf{B}^\top \in {R}^{n \times D}. \end{equation} $$ In implementation, the matrix rank \(r\) is defined as the hyper-parameter, which in turn determines the desired parameter length \(l\) in accordance with the stated constraints. This necessitates imposing certain limitations on these dimension and rank parameters represented as follows: $$ \begin{equation} l = (n + D) \times r. \end{equation} $$

Attributes:

Name Type Description
name str, default = 'low_rank_reconciliation'

Name of the low-rank parameter reconciliation function
r int, default = 2

Submatrix rank parameter.

Methods:

Name Description
__init__

It initializes the low-rank parameter reconciliation function.
calculate_l

It calculates the length of required parameters for the reconciliation function.
forward

It implements the abstract forward method declared in the base reconciliation class.
Source code in tinybig/reconciliation/lowrank_reconciliation.py 
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
class lorr_reconciliation(fabrication):
    r"""
    The low-rank parameter reconciliation function.

    It performs the low-rank parameter reconciliation, and returns the low-rank reconciled parameter matrix of shape (n, D).
    This class inherits from the reconciliation class (i.e., the fabrication class in the module directory).

    ...

    Notes
    ----------
    Formally, given the parameter vector $\mathbf{w} \in {R}^{l}$ and a rank hyper-parameter $r$,
    low-rank parameter reconciliation function partitions $\mathbf{w}$ into two sub-vectors and subsequently
    reshapes them into two matrices $\mathbf{A} \in {R}^{n \times r}$ and $\mathbf{B} \in {R}^{D \times r}$,
    each possessing a rank of $r$.
    These two sub-matrices $\mathbf{A}$ and $\mathbf{B}$ help define the low-rank reconciliation function as follows:
    $$
        \begin{equation}
            \psi(\mathbf{w}) = \mathbf{A} \mathbf{B}^\top \in {R}^{n \times D}.
        \end{equation}
    $$
    In implementation, the matrix rank $r$ is defined as the hyper-parameter, which in turn determines the
    desired parameter length $l$ in accordance with the stated constraints.
    This necessitates imposing certain limitations on these dimension and rank parameters represented as follows:
    $$
        \begin{equation}
            l = (n + D) \times r.
        \end{equation}
    $$

    Attributes
    ----------
    name: str, default = 'low_rank_reconciliation'
        Name of the low-rank parameter reconciliation function
    r: int, default = 2
        Submatrix rank parameter.

    Methods
    ----------
    __init__
        It initializes the low-rank parameter reconciliation function.

    calculate_l
        It calculates the length of required parameters for the reconciliation function.

    forward
        It implements the abstract forward method declared in the base reconciliation class.
    """
    def __init__(self, name: str = 'low_rank_reconciliation', r: int = 2, *args, **kwargs):
        """
        The initialization method of the low-rank parameter reconciliation function.

        It initializes a low-rank parameter reconciliation function object.
        This method will also call the initialization method of the base class as well.

        Parameters
        ----------
        name: str, default = 'low_rank_reconciliation'
            Name of the low-rank parameter reconciliation function.
        r: int, default = 2
            Matrix rank parameter of the low-rank parameter reconciliation.

        Returns
        ----------
        fabrication
            The low-rank parameter reconciliation function object.
        """
        super().__init__(name=name, *args, **kwargs)
        self.r = r

    def calculate_l(self, n: int, D: int):
        r"""
        The required parameter number calculation method.

        It calculates the number of required learnable parameters, i.e., $l$, of the parameter reconciliation function
        based on the intermediate and output space dimensions, $n$ and $D$, and the rank parameters $r$,
        which can be represented as follows:
        $$
            \begin{equation}
                l = (n + D) \times r.
            \end{equation}
        $$

        Parameters
        ----------
        n: int
            The dimension of the output space.
        D: int
            The dimension of the intermediate expansion space.

        Returns
        -------
        int
            The number of required learnable parameters.
        """
        return self.r * (n + D)

    def forward(self, n: int, D: int, w: torch.nn.Parameter, device='cpu', *args, **kwargs):
        r"""
        The forward method of the parameter reconciliation function.

        It applies the low-rank parameter reconciliation operation to the input parameter vector $\mathbf{w}$,
        and returns the reconciled parameter matrix of shape (n, D) subject to rank parameters $r$ as follows:
        $$
            \begin{equation}
                \psi(\mathbf{w}) = \mathbf{A} \mathbf{B}^\top \in {R}^{n \times D}.
            \end{equation}
        $$
        where $\mathbf{A} \in {R}^{n \times r}$ and $\mathbf{B} \in {R}^{D \times r}$ are two low-rank matrices of rank
        $r$ obtained by partitioning $\mathbf{w}$ into two sub-vectors and subsequently reshaping them into matrices.

        Parameters
        ----------
        n: int
            The dimension of the output space.
        D: int
            The dimension of the intermediate expansion space.
        w: torch.nn.Parameter, default = None
            The learnable parameters of the model.
        device: str, default = 'cpu'
            Device to perform the parameter reconciliation.

        Returns
        ----------
        torch.Tensor
            The reconciled parameter matrix of shape (n, D).
        """
        assert w.ndim == 2 and w.size(1) == self.calculate_l(n=n, D=D)
        A, B = torch.split(w, [self.r*n, self.r*D], dim=1)
        return torch.matmul(A.view(n, self.r), B.view(self.r, D))

__init__(name='low_rank_reconciliation', r=2, *args, **kwargs)

The initialization method of the low-rank parameter reconciliation function.

It initializes a low-rank parameter reconciliation function object. This method will also call the initialization method of the base class as well.

Parameters:

Name Type Description Default
name str

Name of the low-rank parameter reconciliation function.

 'low_rank_reconciliation'
r int

Matrix rank parameter of the low-rank parameter reconciliation.

 2

Returns:

Type Description
fabrication

The low-rank parameter reconciliation function object.
Source code in tinybig/reconciliation/lowrank_reconciliation.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def __init__(self, name: str = 'low_rank_reconciliation', r: int = 2, *args, **kwargs):
    """
    The initialization method of the low-rank parameter reconciliation function.

    It initializes a low-rank parameter reconciliation function object.
    This method will also call the initialization method of the base class as well.

    Parameters
    ----------
    name: str, default = 'low_rank_reconciliation'
        Name of the low-rank parameter reconciliation function.
    r: int, default = 2
        Matrix rank parameter of the low-rank parameter reconciliation.

    Returns
    ----------
    fabrication
        The low-rank parameter reconciliation function object.
    """
    super().__init__(name=name, *args, **kwargs)
    self.r = r

calculate_l(n, D)

The required parameter number calculation method.

It calculates the number of required learnable parameters, i.e., \(l\), of the parameter reconciliation function based on the intermediate and output space dimensions, \(n\) and \(D\), and the rank parameters \(r\), which can be represented as follows: $$ \begin{equation} l = (n + D) \times r. \end{equation} $$

Parameters:

Name Type Description Default
n int

The dimension of the output space.

 required
D int

The dimension of the intermediate expansion space.

 required

Returns:

Type Description
int

The number of required learnable parameters.
Source code in tinybig/reconciliation/lowrank_reconciliation.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def calculate_l(self, n: int, D: int):
    r"""
    The required parameter number calculation method.

    It calculates the number of required learnable parameters, i.e., $l$, of the parameter reconciliation function
    based on the intermediate and output space dimensions, $n$ and $D$, and the rank parameters $r$,
    which can be represented as follows:
    $$
        \begin{equation}
            l = (n + D) \times r.
        \end{equation}
    $$

    Parameters
    ----------
    n: int
        The dimension of the output space.
    D: int
        The dimension of the intermediate expansion space.

    Returns
    -------
    int
        The number of required learnable parameters.
    """
    return self.r * (n + D)

forward(n, D, w, device='cpu', *args, **kwargs)

The forward method of the parameter reconciliation function.

It applies the low-rank parameter reconciliation operation to the input parameter vector \(\mathbf{w}\), and returns the reconciled parameter matrix of shape (n, D) subject to rank parameters \(r\) as follows: $$ \begin{equation} \psi(\mathbf{w}) = \mathbf{A} \mathbf{B}^\top \in {R}^{n \times D}. \end{equation} $$ where \(\mathbf{A} \in {R}^{n \times r}\) and \(\mathbf{B} \in {R}^{D \times r}\) are two low-rank matrices of rank \(r\) obtained by partitioning \(\mathbf{w}\) into two sub-vectors and subsequently reshaping them into matrices.

Parameters:

Name Type Description Default
n int

The dimension of the output space.

 required
D int

The dimension of the intermediate expansion space.

 required
w Parameter

The learnable parameters of the model.

 required
device

Device to perform the parameter reconciliation.

 'cpu'

Returns:

Type Description
Tensor

The reconciled parameter matrix of shape (n, D).
Source code in tinybig/reconciliation/lowrank_reconciliation.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def forward(self, n: int, D: int, w: torch.nn.Parameter, device='cpu', *args, **kwargs):
    r"""
    The forward method of the parameter reconciliation function.

    It applies the low-rank parameter reconciliation operation to the input parameter vector $\mathbf{w}$,
    and returns the reconciled parameter matrix of shape (n, D) subject to rank parameters $r$ as follows:
    $$
        \begin{equation}
            \psi(\mathbf{w}) = \mathbf{A} \mathbf{B}^\top \in {R}^{n \times D}.
        \end{equation}
    $$
    where $\mathbf{A} \in {R}^{n \times r}$ and $\mathbf{B} \in {R}^{D \times r}$ are two low-rank matrices of rank
    $r$ obtained by partitioning $\mathbf{w}$ into two sub-vectors and subsequently reshaping them into matrices.

    Parameters
    ----------
    n: int
        The dimension of the output space.
    D: int
        The dimension of the intermediate expansion space.
    w: torch.nn.Parameter, default = None
        The learnable parameters of the model.
    device: str, default = 'cpu'
        Device to perform the parameter reconciliation.

    Returns
    ----------
    torch.Tensor
        The reconciled parameter matrix of shape (n, D).
    """
    assert w.ndim == 2 and w.size(1) == self.calculate_l(n=n, D=D)
    A, B = torch.split(w, [self.r*n, self.r*D], dim=1)
    return torch.matmul(A.view(n, self.r), B.view(self.r, D))