bessel_expansion

Bases: transformation

The bessel expansion function.

Applies Bessel polynomial expansion to input data.

Notes

Formally, the Bessel polynomials are an orthogonal sequence of polynomials with the following closed-form representation:

\[
    \begin{equation}
    B_n(x) = \sum_{k=0}^n \frac{(n+k)!}{(n-k)! k!} \left( \frac{x}{2}\right)^k.
    \end{equation}
\]

The Bessel polynomials can be recursively defined as follows:

Base cases \(n=0\) and \(n=1\):

\[
    \begin{equation}
    B_0(x) = 1 \text{, and } B_1(x) = x + 1.
    \end{equation}
\]

High-order cases with degree \(n \ge 2\):

\[
    \begin{equation}
    B_n(x) = (2n - 1) x B_{n-1}(x) + B_{n-2}(x).
    \end{equation}
\]

The Bessel polynomials can be used to define the data expansion functions as follows:

\[
    \begin{equation}
    \kappa(\mathbf{x} | d) = \left[ B_1(\mathbf{x}), B_2(\mathbf{x}), \cdots, B_d(\mathbf{x}) \right] \in R^D,
    \end{equation}
\]

where the output dimension \(D = md\).

Attributes:

Name	Type	Description
d	int	The degree of Bessel polynomial expansion.

Methods:

Name	Description
calculate_D	Calculates the output dimension after expansion.
forward	Performs Bessel polynomial expansion on the input tensor.

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py

class bessel_expansion(transformation):
    r"""
        The bessel expansion function.

        Applies Bessel polynomial expansion to input data.

        Notes
        ----------

        Formally, the Bessel polynomials are an orthogonal sequence of polynomials with the following closed-form representation:

        $$
            \begin{equation}
            B_n(x) = \sum_{k=0}^n \frac{(n+k)!}{(n-k)! k!} \left( \frac{x}{2}\right)^k.
            \end{equation}
        $$

        The Bessel polynomials can be recursively defined as follows:

        __Base cases $n=0$ and $n=1$:__

        $$
            \begin{equation}
            B_0(x) = 1 \text{, and } B_1(x) = x + 1.
            \end{equation}
        $$

        __High-order cases with degree $n \ge 2$:__

        $$
            \begin{equation}
            B_n(x) = (2n - 1) x B_{n-1}(x) + B_{n-2}(x).
            \end{equation}
        $$

        The Bessel polynomials can be used to define the data expansion functions as follows:

        $$
            \begin{equation}
            \kappa(\mathbf{x} | d) = \left[ B_1(\mathbf{x}), B_2(\mathbf{x}), \cdots, B_d(\mathbf{x}) \right] \in R^D,
            \end{equation}
        $$

        where the output dimension $D = md$.

        Attributes
        ----------
        d : int
            The degree of Bessel polynomial expansion.

        Methods
        -------
        calculate_D(m: int)
            Calculates the output dimension after expansion.
        forward(x: torch.Tensor, device='cpu', *args, **kwargs)
            Performs Bessel polynomial expansion on the input tensor.
    """
    def __init__(self, name='bessel_polynomial_expansion', d: int = 2, *args, **kwargs):
        """
            Initializes the Bessel polynomial expansion transformation.

            Parameters
            ----------
            name : str, optional
                Name of the transformation. Defaults to 'bessel_polynomial_expansion'.
            d : int, optional
                The maximum order of Bessel polynomials for expansion. Defaults to 2.
            *args : tuple
                Additional positional arguments.
            **kwargs : dict
                Additional keyword arguments.
        """
        super().__init__(name=name, *args, **kwargs)
        self.d = d

    def calculate_D(self, m: int):
        """
            Calculates the output dimension after Bessel polynomial expansion.

            Parameters
            ----------
            m : int
                Input dimension.

            Returns
            -------
            int
                Output dimension after expansion.
        """
        return m * self.d

    def forward(self, x: torch.Tensor, device='cpu', *args, **kwargs):
        """
            Performs Bessel polynomial expansion on the input tensor.

            Parameters
            ----------
            x : torch.Tensor
                Input tensor of shape `(batch_size, input_dim)`.
            device : str, optional
                Device for computation ('cpu', 'cuda'). Defaults to 'cpu'.
            *args : tuple
                Additional positional arguments.
            **kwargs : dict
                Additional keyword arguments.

            Returns
            -------
            torch.Tensor
                Expanded tensor of shape `(batch_size, expanded_dim)`.

            Raises
            ------
            AssertionError
                If the output tensor shape does not match the expected dimensions.
        """
        b, m = x.shape
        x = self.pre_process(x=x, device=device)

        # base case: order 0
        expansion = torch.ones(size=[x.size(0), x.size(1), self.d + 1]).to(device)
        # base case: order 1
        if self.d > 0:
            expansion[:, :, 1] = x + 1
        # high-order cases
        for n in range(2, self.d + 1):
            expansion[:, :, n] = (2*n-1) * x * expansion[:, :, n-1].clone() + expansion[:, :, n-2].clone()

        expansion = expansion[:, :, 1:].contiguous().view(x.size(0), -1)

        assert expansion.shape == (b, self.calculate_D(m=m))
        return self.post_process(x=expansion, device=device)

__init__(name='bessel_polynomial_expansion', d=2, *args, **kwargs)

Initializes the Bessel polynomial expansion transformation.

Parameters:

Name	Type	Description	Default
name	str	Name of the transformation. Defaults to 'bessel_polynomial_expansion'.	'bessel_polynomial_expansion'
d	int	The maximum order of Bessel polynomials for expansion. Defaults to 2.	2
*args	tuple	Additional positional arguments.	()
**kwargs	dict	Additional keyword arguments.	{}

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py

def __init__(self, name='bessel_polynomial_expansion', d: int = 2, *args, **kwargs):
    """
        Initializes the Bessel polynomial expansion transformation.

        Parameters
        ----------
        name : str, optional
            Name of the transformation. Defaults to 'bessel_polynomial_expansion'.
        d : int, optional
            The maximum order of Bessel polynomials for expansion. Defaults to 2.
        *args : tuple
            Additional positional arguments.
        **kwargs : dict
            Additional keyword arguments.
    """
    super().__init__(name=name, *args, **kwargs)
    self.d = d

calculate_D(m)

Calculates the output dimension after Bessel polynomial expansion.

Parameters:

Name	Type	Description	Default
m	int	Input dimension.	required

Returns:

Type	Description
int	Output dimension after expansion.

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py

def calculate_D(self, m: int):
    """
        Calculates the output dimension after Bessel polynomial expansion.

        Parameters
        ----------
        m : int
            Input dimension.

        Returns
        -------
        int
            Output dimension after expansion.
    """
    return m * self.d

forward(x, device='cpu', *args, **kwargs)

Performs Bessel polynomial expansion on the input tensor.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor of shape (batch_size, input_dim).	required
device	str	Device for computation ('cpu', 'cuda'). Defaults to 'cpu'.	'cpu'
*args	tuple	Additional positional arguments.	()
**kwargs	dict	Additional keyword arguments.	{}

Returns:

Type	Description
Tensor	Expanded tensor of shape (batch_size, expanded_dim).

Raises:

Type	Description
AssertionError	If the output tensor shape does not match the expected dimensions.

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py

def forward(self, x: torch.Tensor, device='cpu', *args, **kwargs):
    """
        Performs Bessel polynomial expansion on the input tensor.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape `(batch_size, input_dim)`.
        device : str, optional
            Device for computation ('cpu', 'cuda'). Defaults to 'cpu'.
        *args : tuple
            Additional positional arguments.
        **kwargs : dict
            Additional keyword arguments.

        Returns
        -------
        torch.Tensor
            Expanded tensor of shape `(batch_size, expanded_dim)`.

        Raises
        ------
        AssertionError
            If the output tensor shape does not match the expected dimensions.
    """
    b, m = x.shape
    x = self.pre_process(x=x, device=device)

    # base case: order 0
    expansion = torch.ones(size=[x.size(0), x.size(1), self.d + 1]).to(device)
    # base case: order 1
    if self.d > 0:
        expansion[:, :, 1] = x + 1
    # high-order cases
    for n in range(2, self.d + 1):
        expansion[:, :, n] = (2*n-1) * x * expansion[:, :, n-1].clone() + expansion[:, :, n-2].clone()

    expansion = expansion[:, :, 1:].contiguous().view(x.size(0), -1)

    assert expansion.shape == (b, self.calculate_D(m=m))
    return self.post_process(x=expansion, device=device)