[ADD] einsum expression for transpose convolution's KFAC-reduce (#45)

* [ADD] Utility function to compute input sizes of a convolution

* [ADD] Einsum expression and functional for transpose input unfolding

* [DEL] Print statements

* [DOC] Polish docstrings

* [DOC] Minor polish

* [FIX] Too long lines

* [ADD] einsum expression for transpose convolution's KFAC-reduce

* [FIX] Long line

* [REF] Minor polish

* [FIX] flake8

docs/api/expressions.md

@@ -5,3 +5,4 @@
 :::einconv.expressions.convNd_kfc
 :::einconv.expressions.convNd_kfac_reduce
 :::einconv.expressions.conv_transposeNd_unfold
+:::einconv.expressions.conv_transposeNd_kfac_reduce

einconv/expressions/convNd_kfac_reduce.py

@@ -2,8 +2,10 @@
 
 KFAC-reduce was introduced by:
 
-- Eschenhagen, R. (2022). Kronecker-factored approximate curvature for linear
- weight-sharing layers, Master thesis.
+- [Eschenhagen, R., Immer, A., Turner, R. E., Schneider, F., & Hennig, P.
+ (2023). Kronecker-factored approximate curvature for modern neural network
+ architectures. In Advances in Neural Information Processing Systems (NeurIPS)]\
+(https://arxiv.org/abs/2311.00636).
 """
 
 from typing import List, Tuple, Union
@@ -27,6 +29,26 @@ def einsum_expression(
 ) -> Tuple[str, List[Tensor], Tuple[int, ...]]:
  """Generate einsum expression of input-based KFAC-reduce factor for convolution.
 
+ Let $\\mathbf{X}\\in\\mathbb{R}^{C_\\text{in}\\times I_1\\times I_2\\times\\dots}$
+ denote the input of a convolution. The unfolded input $[[\\mathbf{X}]]$
+ has dimension $(C_\\text{in} \\cdot K_1 \\cdot K_2 \\cdots) \\times (O_1 \\cdot O_2
+ \\cdots)$ where $K_i$ and $O_i$ are the kernel and output sizes of the convolution.
+ The input-based KFAC-reduce factor is the batch-averaged outer product
+ of the column-averaged unfolded input,
+
+ $$
+ \\hat{\\mathbf{\\Omega}} =
+ \\frac{1}{B \\cdot (O_1 \\cdot O_2 \\cdots)^2} \\sum_{b=1}^B
+ ( [[\\mathbf{X}_b]]^\\top \\mathbf{1} )
+ ( [[\\mathbf{X}_b]]^\\top \\mathbf{1} )^\\top
+ \\in \\mathbb{R}^{(C_\\text{in} \\cdot K_1 \\cdot K_2 \\cdots) \\times
+ (C_\\text{in} \\cdot K_1 \\cdot K_2 \\cdots)}
+ \\,,
+ $$
+
+ where $B$ is the batch size and $\\mathbf{X}_b$ is the convolution's input from the
+ $b$th data point.
+
  Args:
  x: Convolution input. Has shape ``[batch_size, in_channels, *input_sizes]``
  where ``len(input_sizes) == N``.
@@ -46,8 +68,8 @@ def einsum_expression(
  Einsum equation
  Einsum operands in order un-grouped input, patterns, un-grouped input, \
  patterns, normalization scaling
- Output shape: ``[groups, in_channels //groups * tot_kernel_sizes,\
- in_channels //groups * tot_kernel_sizes]``
+ Output shape: ``[groups, in_channels // groups * tot_kernel_sizes,\
+ in_channels // groups * tot_kernel_sizes]``
  """
  N = x.dim() - 2
 
@@ -83,7 +105,7 @@ def einsum_expression(
  x_ungrouped = rearrange(x, "n (g c_in) ... -> n g c_in ...", g=groups)
  output_tot_size = Tensor([p.shape[1] for p in patterns]).int().prod()
  batch_size = x.shape[0]
- scale = Tensor([1.0 / (batch_size * output_tot_size**2)]).to(x.device).to(x.dtype)
+ scale = Tensor([1.0 / (batch_size * output_tot_size**2)]).to(x.device, x.dtype)
  operands = [x_ungrouped, *patterns, *patterns, x_ungrouped, scale]
 
  # construct output shape

einconv/expressions/convNd_kfc.py

@@ -2,8 +2,9 @@
 
 KFC was introduced by:
 
-- Grosse, R., & Martens, J. (2016). A Kronecker-factored approximate Fisher matrix
- for convolution layers. International Conference on Machine Learning (ICML).
+- [Grosse, R., & Martens, J. (2016). A Kronecker-factored approximate Fisher matrix
+ for convolution layers. International Conference on Machine Learning (ICML).]\
+(https://arxiv.org/abs/1602.01407)
 """
 
 from typing import List, Tuple, Union
@@ -27,6 +28,25 @@ def einsum_expression(
 ) -> Tuple[str, List[Tensor], Tuple[int, ...]]:
  """Generate einsum expression of input-based KFC factor for convolution.
 
+ Let $\\mathbf{X}\\in\\mathbb{R}^{C_\\text{in}\\times I_1\\times I_2\\times\\dots}$
+ denote the input of a convolution. The unfolded input $[[\\mathbf{X}]]$
+ has dimension $(C_\\text{in} \\cdot K_1 \\cdot K_2 \\cdots) \\times (O_1 \\cdot O_2
+ \\cdots)$ where $K_i$ and $O_i$ are the kernel and output sizes of the convolution.
+ The input-based KFC factor is the batch-averaged outer product of the unfolded
+ input,
+
+ $$
+ \\mathbf{\\Omega} =
+ \\frac{1}{B} \\sum_{b=1}^B
+ [[\\mathbf{X}_b]] [[\\mathbf{X}_b]]^\\top
+ \\in \\mathbb{R}^{(C_\\text{in} \\cdot K_1 \\cdot K_2 \\cdots) \\times
+ (C_\\text{in} \\cdot K_1 \\cdot K_2 \\cdots)}
+ \\,,
+ $$
+
+ where $B$ is the batch size and $\\mathbf{X}_b$ is the convolution's input from the
+ $b$th data point.
+
  Args:
  x: Convolution input. Has shape ``[batch_size, in_channels, *input_sizes]``
  where ``len(input_sizes) == N``.

einconv/expressions/conv_transposeNd_kfac_reduce.py

@@ -0,0 +1,173 @@
+"""Input-based factor of the KFAC-reduce approximation for transpose convolutions.
+
+KFAC-reduce was introduced by:
+
+- [Eschenhagen, R., Immer, A., Turner, R. E., Schneider, F., & Hennig, P.
+ (2023). Kronecker-factored approximate curvature for modern neural network
+ architectures. In Advances in Neural Information Processing Systems (NeurIPS)]\
+(https://arxiv.org/abs/2311.00636).
+"""
+
+from typing import List, Optional, Tuple, Union
+
+from einops import rearrange
+from torch import Tensor
+
+import einconv
+from einconv.expressions.utils import create_conv_index_patterns, translate_to_torch
+from einconv.utils import _tuple, get_conv_input_size
+
+
+def einsum_expression(
+ x: Tensor,
+ kernel_size: Union[int, Tuple[int, ...]],
+ stride: Union[int, Tuple[int, ...]] = 1,
+ padding: Union[int, Tuple[int, ...]] = 0,
+ output_padding: Union[int, Tuple[int, ...]] = 0,
+ output_size: Optional[Union[int, Tuple[int, ...]]] = None,
+ dilation: Union[int, Tuple[int, ...]] = 1,
+ groups: int = 1,
+ simplify: bool = True,
+) -> Tuple[str, List[Tensor], Tuple[int, ...]]:
+ """Generate einsum expr. of input-based KFAC-reduce factor for transp. convolution.
+
+ We describe the `N`d transpose convolution using its associated `N`d convolution
+ which maps an input of shape `[batch_size, in_channels, *input_sizes]` to an output
+ of shape `[batch_size, out_channels, *output_sizes]`. The transpose convolution's
+ input has shape `[batch_size, out_channels, *output_sizes]` and the output has shape
+ `[batch_size, in_channels, *input_sizes]`.
+
+ Let $\\mathbf{X}\\in\\mathbb{R}^{C_\\text{out}\\times O_1\\times O_2\\times\\dots}$
+ denote the input of a transpose convolution. The unfolded input $[[\\mathbf{X}
+ ]]_\\top$ has dimension $(C_\\text{out} \\cdot K_1 \\cdot K_2 \\cdots) \\times
+ (I_1 \\cdot I_2 \\cdots)$ where $K_i$ and $I_i$ are the kernel and input sizes of
+ the associated convolution. The input-based KFAC-reduce factor is the batch-averaged
+ outer product of the column-averaged unfolded input,
+
+ $$
+ \\hat{\\mathbf{\\Omega}} =
+ \\frac{1}{B \\cdot (I_1 \\cdot I_2 \\cdots)^2} \\sum_{b=1}^B
+ ( [[\\mathbf{X}_b]]^\\top_\\top \\mathbf{1} )
+ ( [[\\mathbf{X}_b]]^\\top_\\top \\mathbf{1} )^\\top
+ \\in \\mathbb{R}^{(C_\\text{out} \\cdot K_1 \\cdot K_2 \\cdots) \\times
+ (C_\\text{out} \\cdot K_1 \\cdot K_2 \\cdots)}
+ \\,,
+ $$
+
+ where $B$ is the batch size and $\\mathbf{X}_b$ is the transpose convolution's
+ input from the $b$th data point.
+
+ Args:
+ x: Input tensor of shape `[batch_size, out_channels, *output_sizes]`.
+ kernel_size: Size of the convolutional kernel. Can be a single integer (shared
+ along all spatial dimensions), or an `N`-tuple of integers.
+ stride: Stride of the associated convolution. Can be a single integer (shared
+ along all spatial dimensions), or an `N`-tuple of integers. Default: `1`.
+ padding: Padding of the associated convolution. Can be a single integer (shared
+ along all spatial dimensions), or an `N`-tuple of integers. Default: `0`.
+ output_padding: Number of unused pixels at the end of the spatial domain.
+ This is used to resolve the ambiguity that a convolution can map different
+ input sizes to the same output size if its stride is different from 1.
+ Instead of specifying this argument, you can directly specify the output
+ size of the transpose convolution (i.e. the input size of the associated
+ convolution via the `output_size` argument). Can be a single integer
+ (shared along all spatial dimensions), or an `N`-tuple. Default: `0`.
+ output_size: Size of the output of the transpose convolution (i.e. the input
+ size of the associated convolution). Specifying this argument will override
+ the `output_padding` argument. Can be a single integer (shared along all
+ spatial dimensions), or an `N`-tuple of integers. Default: `None`.
+ dilation: Dilation of the convolution. Can be a single integer (shared along
+ all spatial dimensions), or an `N`-tuple of integers. Default: `1`.
+ groups: In how many groups to split the channels. Default: `1`.
+ simplify: Whether to simplify the einsum expression. Default: `True`.
+
+ Returns:
+ Einsum equation
+ Einsum operands in order un-grouped input, patterns, un-grouped input, \
+ patterns, normalization scaling
+ Output shape: `[groups, out_channels // groups * tot_kernel_sizes,\
+ out_channels // groups * tot_kernel_sizes]`
+ """
+ N = x.dim() - 2
+
+ # construct einsum equation
+ x1_str = "n g c_out " + " ".join([f"o{i}" for i in range(N)])
+ x2_str = "n g c_out_ " + " ".join([f"o{i}_" for i in range(N)])
+ pattern1_strs: List[str] = [f"k{i} o{i} i{i}" for i in range(N)]
+ pattern2_strs: List[str] = [f"k{i}_ o{i}_ i{i}_" for i in range(N)]
+ scale_str = "s"
+ lhs = ",".join([x1_str, *pattern1_strs, *pattern2_strs, x2_str, scale_str])
+ rhs = (
+ "g c_out "
+ + " ".join([f"k{i}" for i in range(N)])
+ + " c_out_ "
+ + " ".join([f"k{i}_" for i in range(N)])
+ )
+ equation = "->".join([lhs, rhs])
+ equation = translate_to_torch(equation)
+
+ conv_output_size = x.shape[2:]
+ t_kernel_size = _tuple(kernel_size, N)
+ t_stride = _tuple(stride, N)
+ t_padding = _tuple(padding, N)
+ t_dilation = _tuple(dilation, N)
+
+ # infer output_padding from convolution's input size
+ if output_size is not None:
+ t_output_size = _tuple(output_size, N)
+ t_output_padding = tuple(
+ output_size - get_conv_input_size(out, K, S, P, 0, D)
+ for output_size, out, K, S, P, D in zip(
+ t_output_size,
+ conv_output_size,
+ t_kernel_size,
+ t_stride,
+ t_padding,
+ t_dilation,
+ )
+ )
+ else:
+ t_output_padding = _tuple(output_padding, N)
+
+ conv_input_size = tuple(
+ get_conv_input_size(out, K, S, P, output_padding, D)
+ for out, K, S, P, output_padding, D in zip(
+ conv_output_size,
+ t_kernel_size,
+ t_stride,
+ t_padding,
+ t_output_padding,
+ t_dilation,
+ )
+ )
+
+ # construct einsum operands
+ patterns = create_conv_index_patterns(
+ N,
+ input_size=conv_input_size,
+ kernel_size=t_kernel_size,
+ stride=t_stride,
+ padding=t_padding,
+ dilation=dilation,
+ device=x.device,
+ dtype=x.dtype,
+ )
+ x_ungrouped = rearrange(x, "n (g c_in) ... -> n g c_in ...", g=groups)
+ conv_input_tot_size = Tensor(conv_input_size).int().prod()
+ batch_size, out_channels = x.shape[:2]
+ scale = Tensor([1.0 / (batch_size * conv_input_tot_size**2)]).to(x.device, x.dtype)
+ operands = [x_ungrouped, *patterns, *patterns, x_ungrouped, scale]
+
+ # construct output shape
+ t_kernel_size = _tuple(kernel_size, N)
+ kernel_tot_size = int(Tensor(t_kernel_size).int().prod())
+ shape = (
+ groups,
+ out_channels // groups * kernel_tot_size,
+ out_channels // groups * kernel_tot_size,
+ )
+
+ if simplify:
+ equation, operands = einconv.simplify(equation, operands)
+
+ return equation, operands, shape

einconv/expressions/conv_transposeNd_unfold.py

@@ -30,6 +30,27 @@ def einsum_expression(
  has shape `[batch_size, out_channels, *output_sizes]` and the output has shape
  `[batch_size, in_channels, *input_sizes]`.
 
+ Let $\\mathbf{X}\\in\\mathbb{R}^{C_\\text{out}\\times O_1\\times O_2\\times\\dots}$
+ denote the input of a transpose convolution, $\\mathbf{W} \\in \\mathbb{R}^{
+ C_\\text{out} \\times C_\\text{in} \\times K_1\\times K_2\\times\\dots}$ its kernel
+ and $\\mathbf{Y}\\in\\mathbb{R}^{C_\\text{in}\\times I_1\\times I_2\\times\\dots}$
+ its output. The unfolded input $[[\\mathbf{X}]]_\\top$ has dimension
+ $(C_\\text{out} \\cdot K_1 \\cdot K_2 \\cdots) \\times (I_1 \\cdot I_2 \\cdots)$ and
+ can be used to express transpose convolution as matrix multiplication,
+
+ $$
+ \\mathrm{mat}(\\mathbf{Y})
+ =
+ \\mathrm{mat}(\\mathbf{W})
+ [[\\mathbf{X})]]_\\top
+ \\,,
+ $$
+
+ where $\\mathrm{mat}(\\mathbf{Y}) \\in \\mathbb{R}^{C_\\text{in}\\times (I_1\\cdot
+ I_2 \\cdots)}$ and $\\mathrm{mat}(\\mathbf{W}) \\in \\mathbb{R}^{C_\\text{in}\\times
+ (C_\\text{out} \\cdot K_1\\cdot K_2 \\cdots)}$ are matrix views of $\\mathbf{Y},
+ \\mathbf{W}$ (note that $\\mathbf{W}$ must also be transposed before matricizing).
+
  Args:
  x: Input to the `N`d transpose convolution. Has shape
  `[batch_size, in_channels, *input_sizes]` where `len(input_sizes) == N`.

mkdocs.yml

@@ -1,8 +1,10 @@
 site_name: Einconv
-site_url: https://example.com # TODO Fill in the link from the hosting platform
+site_url: https://einconv.readthedocs.io
 repo_url: https://github.com/f-dangel/einconv/
 repo_name: f-dangel/einconv
 site_author: Felix Dangel
+watch:
+ - einconv
 nav:
  - Getting Started: index.md
  - Tutorials: