Repository: TRAIS-Lab/dattri (language: python)

Issue:
API Design Proposal for LiSSA algorithm ihvp
## Background

With the `hvp` functions implemented before, further complete the LiSSA algorithm for ihvp calculation. Compared with CG, LiSSA algorithm reduces the number of hvp calculation and is more suitable for large datasets.

## Algorithm description

LiSSA algorithm approximates the ihvp function by averaging multiple samples. The samples are estimated by recursion based on Taylor expansion.

## API Design
The API design basically follows the existing implementations of `hvp` and `ihvp_cg`. There are also two versions which calculate the hvp for fixed or non-fixed x's.

The first input of the LiSSA algorithms is a function for estimation, such as $L(\cdot, \cdot)$. The input list should be a list of the form $[(z_0, \theta), \dots, (z_n,\ theta)]$.

```python
def ihvp_lissa(func: Callable,
               argnums: int = 0,
               num_repeat: int = 10, 
               recursion_depth: int = 5000,
               mode: str = "rev-rev") -> Callable:
    """LiSSA ihvp algorithm function.

    Standing for the inverse-hessian-vector product, returns a function that,
    when given vectors, computes the product of inverse-hessian and vector.

    LiSSA algorithm approximates the ihvp function by averaging multiple samples.
    The samples are estimated by recursion based on Taylor expansion.

    Args:
        func (Callable): A Python function that takes one or more arguments.
            Must return a single-element Tensor. The hessian will
            be estimated on this function.
        argnums (int): An integer default to 0. Specifies which argument of func
            to compute inverse hessian with respect to.
        num_repeat (int): An integer default 10. Specifies the number of samples
            of the hvp approximation to average on.
        recursion_depth (int): A integer default to 5000. Specifies the number of
            recursions used to estimate each ihvp sample.
        mode (str): The auto diff mode, which can have one of the following values:
            - rev-rev: calculate the hessian with two reverse-mode auto-diff. It has
                       better compatibility while cost more memory.
            - rev-fwd: calculate the hessian with the composing of reverse-mode and
                       forward-mode. It's more memory-efficient but may not be supported
                       by some operator.

    Returns:
        A function that takes a list of tuples of Tensor `x` and a vector `v` and returns
        the IHVP of the Hessian of `func` and `v`.
    """
```

```python
def ihvp_at_x_lissa(func: Callable,
                    input_list: List[Tuple],
                    argnums: int = 0,
                    num_repeat: int =10, 
                    recursion_depth: int =5000,
                    mode: str = "rev-rev") -> Callable:
    """LiSSA ihvp algorithm function (with fixed x).

    Standing for the inverse-hessian-vector product, returns a function that,
    when given vectors, computes the product of inverse-hessian and vector.

    LiSSA algorithm approximates the ihvp function by averaging multiple samples.
    The samples are estimated by recursion based on Taylor expansion.

    Args:
        func (Callable): A Python function that takes one or more arguments.
            Must return a single-element Tensor. The hessian will
            be estimated on this function.
        input_list (List[Tuple]): List of arguments for multiple calls of `func`. Each tuple
            inside the list should be a pair of valid arguments
        argnums (int): An integer default to 0. Specifies which argument of func
            to compute inverse hessian with respect to.
        num_repeat (int): An integer default 10. Specifies the number of samples
            of the hvp approximation to average on.
        recursion_depth (int): A integer default to 5000. Specifies the number of
            recursions used to estimate each ihvp sample.
        mode (str): The auto diff mode, which can have one of the following values:
            - rev-rev: calculate the hessian with two reverse-mode auto-diff. It has
                       better compatibility while cost more memory.
            - rev-fwd: calculate the hessian with the composing of reverse-mode and
                       forward-mode. It's more memory-efficient but may not be supported
                       by some operator.

    Returns:
        A function that takes a vector `v` and returns the IHVP of the Hessian
        of `func` and `v`.
    """
```

## Demonstration

```python
def mse_loss(x, y, theta):
    return (x * theta - y)**2

ihvp = ihvp_lissa(mse_loss, argnums=2)

theta = torch.randn(1)
input_list = [
    (torch.randn(1), torch.randn(1), theta),
    (torch.randn(1), torch.randn(1), theta)
]
vec = torch.randn(5, 2)
ihvp(input_list, vec)
```



Discussion / hints:
Could you provide more detail about the difference between `num_samples` and `recursion_depth`? I guess they are `r` and `t` in https://arxiv.org/pdf/1703.04730.pdf right? if so, maybe we can rename `num_samples` -> `num_repeat` and explain it as `"Specifies the number of repeat of the procedure to average on."`

minor issue:
1. LISSA -> LiSSA 

Others LGTM
I see. Thanks for the advice!
@TheaperDeng So how should we deal with hyperparameters, i.e. the damping and scaling factors?
> @TheaperDeng So how should we deal with hyperparameters, i.e. the damping and scaling factors?

I think you can currently make these 2 as TODO and fix them to the most common default value
@TheaperDeng Another minor concern. If I understand it correctly, it seems in the former implementations, the ihvp is given by $v \cdot H^{-1}$, which means the product is from the left?

Fix patch:
diff --git a/dattri/func/ihvp.py b/dattri/func/ihvp.py
index b5185cd7..3f743b7b 100644
--- a/dattri/func/ihvp.py
+++ b/dattri/func/ihvp.py
@@ -16,7 +16,8 @@
 
 if TYPE_CHECKING:
     from collections.abc import Callable
-    from typing import Tuple, Union
+    from typing import Optional, Tuple, Union
+
 
 import torch
 from torch import Tensor
@@ -679,5 +680,236 @@ def _ihvp_at_x_arnoldi(v: Tensor) -> Tensor:
     return _ihvp_at_x_arnoldi
 
 
+def _check_input_size(*x, in_dims: Optional[Tuple] = None) -> int:
+    """Check and return the size of input data.
+
+    Args:
+        *x: List of arguments to check. Each argument shoule be either:
+            1. A tensor with a batch size dimension. Each data point i
+            will take the i-th element along this dimension.
+            2. A tensor without a batch size dimension. Each data point will
+            share this tensor.
+        in_dims (Optional[Tuple]): A tuple with the same shape as *x, indicating
+            which dimension should be considered as batch size dimension. Take the
+            first dimension as batch size dimension by default.
+
+    Returns:
+        An integer indicating the number of input data.
+
+    Raises:
+        IHVPUsageError: if the input size is ambiguous or mismatches.
+    """
+    if in_dims is None:
+        in_dims = (0,) * len(x)
+
+    # Check batch size mismatch
+    batch_size = None
+    for i, (x_in, dim) in enumerate(zip(x, in_dims)):
+        if dim is None:
+            continue
+
+        if batch_size is None:
+            batch_size = x_in.shape[dim]
+        elif batch_size != x_in.shape[dim]:
+            message = (f"Input batch size mismatch! Expected {batch_size},"
+                       f"found {x_in.shape[dim]} for input tensor {i}.")
+            raise IHVPUsageError(message)
+
+    if batch_size is None:
+        message = "All inputs are identical for LiSSA ihvp!"
+        raise IHVPUsageError(message)
+
+    return batch_size
+
+
+def _sample_random_batch(*x,
+                         num_samples: int,
+                         in_dims: Optional[Tuple] = None,
+                         batch_size: int = 1) -> Tuple[torch.Tensor, ...]:
+    """Randomly sample a batch of `batch_size` from the input data, without replacement.
+
+    Args:
+        *x: List of arguments to check. Each argument shoule be either:
+            1. A tensor with a batch size dimension. Each data point i
+            will take the i-th element along this dimension.
+            2. A tensor without a batch size dimension. Each data point will
+            share this tensor.
+        num_samples (int): An integer, indicating the total number of samples.
+        batch_size (int): An integer default to 1, indicating the batch size.
+        in_dims (Optional[Tuple]): A tuple with the same shape as *x, indicating
+            which dimension should be considered as batch size dimension. Take the
+            first dimension as batch size dimension by default.
+
+    Returns:
+        An tuple of tensors corresponding to a batch of input data.
+    """
+    if in_dims is None:
+        in_dims = (0,) * len(x)
+
+    # Randomly sample and collate a batch
+    sampled_indices = torch.randperm(num_samples)[:batch_size]
+
+    return tuple(
+        x_in.index_select(dim, sampled_indices)
+        if dim is not None else x_in
+        for x_in, dim in zip(x, in_dims)
+    )
+
+
+def ihvp_lissa(func: Callable,
+               argnums: int = 0,
+               batch_size: int = 1,
+               num_repeat: int = 1,
+               recursion_depth: int = 5000,
+               damping: int = 0.0,
+               scaling: int = 50.0,
+               mode: str = "rev-rev") -> Callable:
+    """IHVP via LiSSA algorithm.
+
+    Standing for the inverse-hessian-vector product, returns a function that,
+    when given vectors, computes the product of inverse-hessian and vector.
+
+    LiSSA algorithm approximates the ihvp function by averaging multiple samples.
+    The samples are estimated by recursion based on Taylor expansion.
+
+    Args:
+        func (Callable): A Python function that takes one or more arguments.
+            Must return a single-element Tensor. The hessian will
+            be estimated on this function.
+        argnums (int): An integer default to 0. Specifies which argument of func
+            to compute inverse hessian with respect to.
+        batch_size (int): An integer default to 1. Specifies the batch size used
+            for LiSSA inner loop update.
+        num_repeat (int): An integer default to 1. Specifies the number of samples
+            of the hvp approximation to average on.
+        recursion_depth (int): A integer default to 5000. Specifies the number of
+            recursions used to estimate each ihvp sample.
+        damping (int): Damping factor used for non-convexity in LiSSA ihvp calculation.
+        scaling (int): Scaling factor used for convergence in LiSSA ihvp calculation.
+        mode (str): The auto diff mode, which can have one of the following values:
+            - rev-rev: calculate the hessian with two reverse-mode auto-diff. It has
+                       better compatibility while cost more memory.
+            - rev-fwd: calculate the hessian with the composing of reverse-mode and
+                       forward-mode. It's more memory-efficient but may not be supported
+                       by some operator.
+
+    Returns:
+        A function that takes a list of tuples of Tensor `x` and a vector `v` and
+        returns the IHVP of the Hessian of `func` and `v`.
+    """
+    hvp_func = hvp(func, argnums=argnums, mode=mode)
+
+    def _ihvp_lissa_func(x: Tuple[torch.Tensor, ...],
+                         v: Tensor,
+                         in_dims: Optional[Tuple] = None) -> Tensor:
+        """The IHVP function via LiSSA algorithm.
+
+        Args:
+            x (Tuple[torch.Tensor, ...]): The function will computed the
+                inverse hessian matrix with respect to these arguments.
+            v (Tensor): The vector to be produced on the inverse hessian matrix.
+            in_dims (Optional[Tuple]): A tuple with the same shape as *x. Indicating
+                which dimension should be considered as batch size dimension.
+
+        Returns:
+            The IHVP value.
+        """
+        num_samples = _check_input_size(*x, in_dims=in_dims)
+        if v.ndim == 1:
+            v = v.unsqueeze(0)
+
+        def _lissa_loop(vec: torch.Tensor) -> torch.Tensor:
+            ihvp_estimations = []
+            for _ in range(num_repeat):
+                curr_estimate = vec.detach().clone()  # No gradient on v
+                for _ in range(recursion_depth):
+                    sampled_input = _sample_random_batch(*x,
+                                                         batch_size=batch_size,
+                                                         num_samples=num_samples,
+                                                         in_dims=in_dims)
+                    hvp = hvp_func(sampled_input, curr_estimate)
+                    curr_estimate = (vec
+                                     + (1 - damping) * curr_estimate
+                                     - hvp / scaling)
+
+                ihvp_estimations.append(curr_estimate / scaling)
+
+            return torch.mean(torch.stack(ihvp_estimations), dim=0)
+
+        return torch.vmap(_lissa_loop, randomness="different")(v)
+
+    return _ihvp_lissa_func
+
+
+def ihvp_at_x_lissa(func: Callable,
+                    *x,
+                    in_dims: Optional[Tuple] = None,
+                    argnums: int = 0,
+                    batch_size: int = 1,
+                    num_repeat: int = 1,
+                    recursion_depth: int = 5000,
+                    damping: int = 0.0,
+                    scaling: int = 50.0,
+                    mode: str = "rev-rev") -> Callable:
+    """IHVP with fixed func inputs via LiSSA algorithm.
+
+    Standing for the inverse-hessian-vector product, returns a function that,
+    when given vectors, computes the product of inverse-hessian and vector.
+
+    LiSSA algorithm approximates the ihvp function by averaging multiple samples.
+    The samples are estimated by recursion based on Taylor expansion.
+
+    Args:
+        func (Callable): A Python function that takes one or more arguments.
+            Must return a single-element Tensor. The hessian will
+            be estimated on this function.
+        *x: List of arguments for `func`.
+        in_dims (Optional[Tuple]): A tuple with the same shape as *x, indicating
+            which dimension should be considered as batch size dimension. Take the
+            first dimension as batch size dimension by default.
+        argnums (int): An integer default to 0. Specifies which argument of func
+            to compute inverse hessian with respect to.
+        batch_size (int): An integer default to 1. Specifies the batch size used
+            for LiSSA inner loop update.
+        num_repeat (int): An integer default to 1. Specifies the number of samples
+            of the hvp approximation to average on.
+        recursion_depth (int): A integer default to 5000. Specifies the number of
+            recursions used to estimate each ihvp sample.
+        damping (int): Damping factor used for non-convexity in LiSSA ihvp calculation.
+        scaling (int): Scaling factor used for convergence in LiSSA ihvp calculation.
+        mode (str): The auto diff mode, which can have one of the following values:
+            - rev-rev: calculate the hessian with two reverse-mode auto-diff. It has
+                       better compatibility while cost more memory.
+            - rev-fwd: calculate the hessian with the composing of reverse-mode and
+                       forward-mode. It's more memory-efficient but may not be supported
+                       by some operator.
+
+    Returns:
+        A function that takes a vector `v` and returns the IHVP of the Hessian
+        of `func` and `v`.
+    """
+
+    def _ihvp_at_x_lissa_func(v: Tensor) -> Tensor:
+        """The IHVP function with fixed func inputs using LiSSA.
+
+        Args:
+            v (Tensor): The vector to be produced on the inverse hessian matrix.
+
+        Returns:
+            The IHVP value.
+        """
+        ihvp_lissa_func = ihvp_lissa(func,
+                                     argnums,
+                                     num_repeat,
+                                     batch_size,
+                                     recursion_depth,
+                                     damping,
+                                     scaling,
+                                     mode)
+        return ihvp_lissa_func(x, v, in_dims=in_dims)
+
+    return _ihvp_at_x_lissa_func
+
+
 class IHVPUsageError(Exception):
     """The usage exception class for ihvp module."""


Test patch:
diff --git a/test/dattri/func/test_ihvp.py b/test/dattri/func/test_ihvp.py
index 06d25597..f1fbf872 100644
--- a/test/dattri/func/test_ihvp.py
+++ b/test/dattri/func/test_ihvp.py
@@ -10,7 +10,9 @@
     ihvp_at_x_arnoldi,
     ihvp_at_x_cg,
     ihvp_at_x_explicit,
+    ihvp_at_x_lissa,
     ihvp_cg,
+    ihvp_lissa,
 )
 from dattri.func.utils import flatten_func, flatten_params
 
@@ -253,3 +255,89 @@ def f(params):
             rtol=1e-02,
             atol=1e-05,
         )
+
+    def test_ihvp_lissa(self):
+        """Test ihvp_lissa/ihvp_lissa_at_x."""
+
+        def mse_loss(xs, ys, theta):
+            return torch.mean((xs @ theta - ys)**2)
+
+        # Create data
+        data_size = (500, 2)
+        theta = torch.randn(2)
+        xs = torch.randn(data_size)
+        noise = torch.normal(mean=torch.tensor(0),
+                             std=torch.tensor(0.05),
+                             size=(data_size[0],))
+        ys = xs @ theta + noise
+        vec = torch.randn(1, 2)
+
+        ihvp_lissa_func = ihvp_lissa(mse_loss,
+                                     argnums=2,
+                                     num_repeat=100,
+                                     recursion_depth=100)
+
+        ihvp_lissa_at_x_func = ihvp_at_x_lissa(mse_loss,
+                                               *(xs, ys, theta),
+                                               in_dims=(0, 0, None),
+                                               argnums=2,
+                                               num_repeat=100,
+                                               recursion_depth=100)
+
+        ihvp_explicit_at_x_func = ihvp_at_x_explicit(mse_loss,
+                                                     *(xs, ys, theta),
+                                                     argnums=2)
+
+        # Set a larger tolerance for LiSSA
+        assert torch.allclose(ihvp_lissa_at_x_func(vec),
+                              ihvp_explicit_at_x_func(vec),
+                              atol=0.08)
+        assert torch.allclose(ihvp_lissa_func((xs, ys, theta),
+                                              vec,
+                                              in_dims=(0, 0, None)),
+                              ihvp_explicit_at_x_func(vec),
+                              atol=0.08)
+
+    def test_ihvp_lissa_batch_size(self):
+        """Test ihvp_lissa/ihvp_lissa_at_x with different batch sizes."""
+
+        def mse_loss(xs, ys, theta):
+            return torch.mean((xs @ theta - ys)**2)
+
+        # Create data
+        data_size = (500, 2)
+        theta = torch.randn(2)
+        xs = torch.randn(data_size)
+        noise = torch.normal(mean=torch.tensor(0),
+                            std=torch.tensor(0.05),
+                            size=(data_size[0],))
+        ys = xs @ theta + noise
+        vec = torch.randn(1, 2)
+
+        ihvp_lissa_func = ihvp_lissa(mse_loss,
+                                     argnums=2,
+                                     batch_size=4,
+                                     num_repeat=10,
+                                     recursion_depth=100)
+
+        ihvp_lissa_at_x_func = ihvp_at_x_lissa(mse_loss,
+                                               *(xs, ys, theta),
+                                               in_dims=(0, 0, None),
+                                               argnums=2,
+                                               batch_size=4,
+                                               num_repeat=10,
+                                               recursion_depth=100)
+
+        ihvp_explicit_at_x_func = ihvp_at_x_explicit(mse_loss,
+                                                     *(xs, ys, theta),
+                                                     argnums=2)
+
+        # Set a larger tolerance for LiSSA
+        assert torch.allclose(ihvp_lissa_at_x_func(vec),
+                              ihvp_explicit_at_x_func(vec),
+                              atol=0.08)
+        assert torch.allclose(ihvp_lissa_func((xs, ys, theta),
+                                              vec,
+                                              in_dims=(0, 0, None)),
+                              ihvp_explicit_at_x_func(vec),
+                              atol=0.08)