Skip to content

FUTON, bases and combiners

The FUTON concept page derives the model; this page is the reference for inro.models.futon: FUTON and its encoding, FUTONEncoding, the bases and combiners, and the SharedRows and Taps features its bases return on grids and for local bases.

futon

FUTON: a fixed basis per axis, a tensor-network combiner and a head.

FUTON computes output_activation(head(combiner(basis(x)))) for coordinates (N, C) in [-1, 1]:

  1. A basis gives each axis c features (N, K_c). Every basis takes in_features, the number of axes C; num_components, the widths K_c, one for all axes or one per axis; normalize, whether to L2-normalize each feature vector (default True); and grid_size, the sizes of the grid linspace(-1, 1, size) all inputs lie on, if any, whose features it tabulates once.
  2. A combiner fuses the axes with a tensor network, stated once as an einsum equation over the axes' projections, into (N, R) features.
  3. A head maps them to the signal: "linear", "mlp" or "identity".

Points share the features of their shared coordinates. On a grid, a basis gives each axis as SharedRows, a table of its coordinates and every point's row in it, so a combiner projects S_c rows rather than N points, and it returns its features as a GridFeatures, its equation and operands, which the head's first linear layer reads off the whole grid in one einsum with its weight when the batch covers the grid. A kernel of compact support gives scattered points only their nonzero features, as Taps. Coordinates that need gradients are evaluated point by point.

SharedRows

Features of one axis on a grid: a table of S rows and each point's row.

Behaves as the (N, K_c) matrix table[rows], which dense returns. A product with a factor projects the S rows rather than N points, and the product shares the rows still.

Attributes:

  • table (Tensor) –

    Features (S, K_c) of the axis's S coordinates.

  • rows (Tensor) –

    Row of each of the N points.

unflatten

unflatten(dim: int, sizes: Sequence[int]) -> SharedRows

Unflatten a dimension of the features, as torch.Tensor.unflatten.

Source code in inro/models/futon.py
def unflatten(self, dim: int, sizes: Sequence[int]) -> "SharedRows":
    """Unflatten a dimension of the features, as :meth:`torch.Tensor.unflatten`."""
    return SharedRows(self.table.unflatten(dim, sizes), self.rows)

dense

dense() -> Tensor

Return the matrix table[rows], a row per point.

Source code in inro/models/futon.py
def dense(self) -> Tensor:
    """Return the matrix ``table[rows]``, a row per point."""
    return self.table[self.rows]

Taps

Sparse features of one axis: the L nonzero values of each coordinate.

Behaves as the (N, K_c) matrix whose row n holds values[n] in columns start[n] to start[n] + L - 1.

Attributes:

  • values (Tensor) –

    Nonzero features (N, L).

  • start (Tensor) –

    First column (N,) of each row's window.

  • compiled (bool) –

    Whether a compiled graph made the taps, whose product then gathers windows of rows. Set where the taps are made: a product traced after a graph break may no longer see torch.compiler.is_compiling().

CosineBasis

CosineBasis(in_features: int, num_components: Sizes, normalize: bool = True, grid_size: Sizes | None = None)

Cosines \(\cos(k \pi u)\) with \(u = (x + 1) / 2\), of frequencies \(k = 0, \dots, K_c - 1\): the paper's basis without its \(\sqrt{2}\), which the factors absorb.

Source code in inro/models/futon.py
def __init__(
    self,
    in_features: int,
    num_components: Sizes,
    normalize: bool = True,
    grid_size: Sizes | None = None,
):
    super().__init__()
    self.in_features = in_features
    self.num_components = _broadcast(num_components, in_features, "num_components")
    if min(self.num_components) < 2 * self.radius:
        raise ValueError(f"num_components must be at least {2 * self.radius}")
    self.normalize = normalize
    self.grid_size: list[int] | None = None
    if grid_size is not None:
        self.grid_size = _broadcast(grid_size, in_features, "grid_size")
    self._tabulate(torch.empty(0))

SincBasis

SincBasis(in_features: int, num_components: Sizes, normalize: bool = True, grid_size: Sizes | None = None)

Cardinal sines \(\operatorname{sinc}((x - \mu_k) / w)\) centered on \(\mu_k = -1 + k w\), spaced by \(w = 2 / (K_c - 1)\).

Each component is one at its own center and zero at the others, so normalize=False gives cardinal (Whittaker) interpolation.

Source code in inro/models/futon.py
def __init__(
    self,
    in_features: int,
    num_components: Sizes,
    normalize: bool = True,
    grid_size: Sizes | None = None,
):
    super().__init__()
    self.in_features = in_features
    self.num_components = _broadcast(num_components, in_features, "num_components")
    if min(self.num_components) < 2 * self.radius:
        raise ValueError(f"num_components must be at least {2 * self.radius}")
    self.normalize = normalize
    self.grid_size: list[int] | None = None
    if grid_size is not None:
        self.grid_size = _broadcast(grid_size, in_features, "grid_size")
    self._tabulate(torch.empty(0))

ChebyshevBasis

ChebyshevBasis(in_features: int, num_components: Sizes, normalize: bool = True, grid_size: Sizes | None = None)

Chebyshev polynomials \(T_k(x) = \cos(k \arccos x)\) of the first kind, orthogonal under the weight \((1 - x^2)^{-1/2}\).

Source code in inro/models/futon.py
def __init__(
    self,
    in_features: int,
    num_components: Sizes,
    normalize: bool = True,
    grid_size: Sizes | None = None,
):
    super().__init__()
    self.in_features = in_features
    self.num_components = _broadcast(num_components, in_features, "num_components")
    if min(self.num_components) < 2 * self.radius:
        raise ValueError(f"num_components must be at least {2 * self.radius}")
    self.normalize = normalize
    self.grid_size: list[int] | None = None
    if grid_size is not None:
        self.grid_size = _broadcast(grid_size, in_features, "grid_size")
    self._tabulate(torch.empty(0))

LegendreBasis

LegendreBasis(in_features: int, num_components: Sizes, normalize: bool = True, grid_size: Sizes | None = None)

Legendre polynomials, the standard, unscaled ones, as fixed combinations of Chebyshev polynomials with nonnegative coefficients that sum to one.

Source code in inro/models/futon.py
def __init__(
    self,
    in_features: int,
    num_components: Sizes,
    normalize: bool = True,
    grid_size: Sizes | None = None,
):
    super().__init__()
    self.in_features = in_features
    self.num_components = _broadcast(num_components, in_features, "num_components")
    if min(self.num_components) < 2 * self.radius:
        raise ValueError(f"num_components must be at least {2 * self.radius}")
    self.normalize = normalize
    self.grid_size: list[int] | None = None
    if grid_size is not None:
        self.grid_size = _broadcast(grid_size, in_features, "grid_size")
    self._tabulate(torch.empty(0))

TriangleBasis

TriangleBasis(in_features: int, num_components: Sizes, normalize: bool = True, grid_size: Sizes | None = None)

Triangles (hats) \(\max(0, 1 - |x - \mu_k| / w)\) centered on \(\mu_k = -1 + k w\), spaced by \(w = 2 / (K_c - 1)\): piecewise-linear interpolation, two components nonzero, K_c >= 2.

With normalize=False, a CP combiner and an MLP head give the TensoRF-CP parametrization: a linear map of these features samples a 1D feature grid.

References

Chen et al., "TensoRF: Tensorial Radiance Fields", ECCV 2022.

Source code in inro/models/futon.py
def __init__(
    self,
    in_features: int,
    num_components: Sizes,
    normalize: bool = True,
    grid_size: Sizes | None = None,
):
    super().__init__()
    self.in_features = in_features
    self.num_components = _broadcast(num_components, in_features, "num_components")
    if min(self.num_components) < 2 * self.radius:
        raise ValueError(f"num_components must be at least {2 * self.radius}")
    self.normalize = normalize
    self.grid_size: list[int] | None = None
    if grid_size is not None:
        self.grid_size = _broadcast(grid_size, in_features, "grid_size")
    self._tabulate(torch.empty(0))

LanczosBasis

LanczosBasis(in_features: int, num_components: Sizes, radius: int = 3, normalize: bool = True, grid_size: Sizes | None = None)

Lanczos kernels, windowed sines centered on \(\mu_k = -1 + k w\), spaced by \(w = 2 / (K_c - 1)\), K_c >= 2 * radius.

With offset \(t = (x - \mu_k) / w\) and radius \(a\), the kernel is \(\operatorname{sinc}(t) \operatorname{sinc}(t / a)\) for \(|t| < a\) and zero elsewhere: one at its own center, zero at the others, and at most 2 * a components nonzero.

Parameters:

  • radius (int, default: 3 ) –

    Kernel radius in grid steps: 2 gives Lanczos-2 (four taps), 3 Lanczos-3 (six taps).

Source code in inro/models/futon.py
def __init__(
    self,
    in_features: int,
    num_components: Sizes,
    radius: int = 3,
    normalize: bool = True,
    grid_size: Sizes | None = None,
):
    if radius < 1:
        raise ValueError(f"radius must be at least 1, got {radius}")
    self.radius = int(radius)  # first: the checks and the tables need it
    super().__init__(in_features, num_components, normalize, grid_size)

GridFeatures

GridFeatures(equation: str, operands: list[Operand], shape: tuple[int, ...])

A combiner's features on a grid, its network not yet contracted.

torch.nn.functional.linear, as a head's first layer calls it, contracts the network with the layer's weight in one einsum over the axes' tables when the grid has at most 64 / out_features cells per point, which spares every point's features and the scatter of their gradients. A sparser batch, and every other torch function, take the features point by point, dense.

Source code in inro/models/futon.py
def __init__(self, equation: str, operands: list[Operand], shape: tuple[int, ...]):
    self.equation, self.operands, self.shape = equation, operands, shape

dense

dense() -> Tensor

Return the features (N, prod(shape)), contracted point by point.

Source code in inro/models/futon.py
def dense(self) -> Tensor:
    """Return the features ``(N, prod(shape))``, contracted point by point."""
    return _fold(self.equation, self.operands).flatten(-len(self.shape))

CPCombiner

CPCombiner(in_features: Sequence[int], rank: int)

Canonical polyadic (CP) combiner.

Projects each axis to R channels with its factor \(\boldsymbol{U}^{(c)}\), factors[c] of shape (K_c, R), and multiplies the projections elementwise, ...r,...r->...r (Eqs. 10 and 11 of the paper). A linear head's weight is then the CP model's output factor \(\boldsymbol{V}\).

Parameters:

  • in_features (Sequence[int]) –

    Number of features K_c of each axis.

  • rank (int) –

    CP rank R, the number of the network's channels.

References

Kolda and Bader, "Tensor Decompositions and Applications", SIAM Review 2009.

Source code in inro/models/futon.py
def __init__(self, in_features: Sequence[int], rank: int):
    super().__init__(in_features)
    self.rank = rank
    self.factors = nn.ParameterList(torch.empty(k, rank) for k in in_features)
    self.equation = ",".join(["...r"] * len(self.factors)) + "->...r"
    self.reset_parameters()
    self.shape = (rank,)

reset_parameters

reset_parameters() -> None

Initialize the factors Kaiming-uniform, with fan-in K_c.

Source code in inro/models/futon.py
def reset_parameters(self) -> None:
    """Initialize the factors Kaiming-uniform, with fan-in ``K_c``."""
    for factor in self.factors:
        nn.init.kaiming_uniform_(factor.T)

TRCombiner

TRCombiner(in_features: Sequence[int], rank: int)

Tensor-ring (TR) combiner.

Axis c contracts its features with a core (R, K_c, R) into an R x R matrix, and the matrices multiply in order, ...ab,...bc->...ac on two axes. The product is flattened, R**2 wide, leaving the closure of the ring, such as a trace, to what follows. Extends the paper's CP model.

Parameters:

  • in_features (Sequence[int]) –

    Number of features K_c of each axis.

  • rank (int) –

    Tensor-ring rank R.

References

Zhao et al., "Tensor Ring Decomposition", arXiv 2016.

Source code in inro/models/futon.py
def __init__(self, in_features: Sequence[int], rank: int):
    super().__init__(in_features)
    self.rank = rank
    self.cores = nn.ParameterList(torch.empty(rank, k, rank) for k in in_features)
    bonds = string.ascii_lowercase[: len(self.cores) + 1]  # axis c joins c, c + 1
    links = ",".join(f"...{a}{b}" for a, b in zip(bonds, bonds[1:]))
    self.equation = f"{links}->...{bonds[0]}{bonds[-1]}"
    self.reset_parameters()
    self.shape = (rank, rank)

reset_parameters

reset_parameters() -> None

Initialize the cores with Kaiming-uniform weights.

Source code in inro/models/futon.py
def reset_parameters(self) -> None:
    """Initialize the cores with Kaiming-uniform weights."""
    for core in self.cores:
        nn.init.kaiming_uniform_(core)

TuckerCombiner

TuckerCombiner(in_features: Sequence[int], rank: Sizes)

Tucker combiner.

Projects axis c to rank[c] channels with its factor (K_c, rank[c]) and contracts the projections with a core (rank[0], ..., rank[C - 1], max(rank)), such as ...b,abz,...a->...z on two axes whose second rank is the larger. With a linear head the weight tensor is in Tucker format, and CP is the case of a superdiagonal core. Extends the paper's CP model.

Parameters:

  • in_features (Sequence[int]) –

    Number of features K_c of each axis.

  • rank (Sizes) –

    Multilinear rank, shared or per axis; the network has as many channels as its largest.

References

Tucker, "Some mathematical notes on three-mode factor analysis", Psychometrika 1966.

Source code in inro/models/futon.py
def __init__(self, in_features: Sequence[int], rank: Sizes):
    super().__init__(in_features)
    modes = len(self.in_features)
    self.rank = _broadcast(rank, modes, "rank")
    widths = zip(self.in_features, self.rank)
    self.factors = nn.ParameterList(torch.empty(k, r) for k, r in widths)
    self.core = nn.Parameter(torch.empty(*self.rank, max(self.rank)))
    # Per point, the core absorbs one projection at a time, largest rank first,
    # which keeps the intermediates small; z is the core's output.
    self.order = sorted(range(modes), key=self.rank.__getitem__, reverse=True)
    ranks = string.ascii_lowercase[:modes]
    first, *others = (f"...{ranks[c]}" for c in self.order)
    self.equation = ",".join([first, f"{ranks}z", *others]) + "->...z"
    self.reset_parameters()
    self.shape = (max(self.rank),)

reset_parameters

reset_parameters() -> None

Initialize the factors (fan-in K_c) and the core Kaiming-uniform.

Source code in inro/models/futon.py
def reset_parameters(self) -> None:
    """Initialize the factors (fan-in ``K_c``) and the core Kaiming-uniform."""
    for factor in self.factors:
        nn.init.kaiming_uniform_(factor.T)
    nn.init.kaiming_uniform_(self.core.view(-1, self.core.shape[-1]))

FUTONEncoding

FUTONEncoding(in_features: int, basis: ModuleSpec, combiner: ModuleSpec)

FUTON's encoding, combiner(basis(x)), which FUTON decodes.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • basis (ModuleSpec) –

    Basis spec, built as basis(in_features, **kwargs). A custom basis must expose num_components, one count per axis.

  • combiner (ModuleSpec) –

    Combiner spec, built as combiner(basis.num_components, **kwargs). A custom combiner must expose out_features.

Shape
  • Input: (N, C) coordinates in [-1, 1].
  • Output: (N, out_features).
Source code in inro/models/futon.py
def __init__(self, in_features: int, basis: ModuleSpec, combiner: ModuleSpec):
    super().__init__()
    self.in_features = in_features
    self.basis = build_module(basis, BASES, in_features)
    self.combiner = build_module(combiner, COMBINERS, self.basis.num_components)
    self.out_features = self.combiner.out_features

FUTON

FUTON(in_features: int, out_features: int, basis: ModuleSpec, combiner: ModuleSpec, head: ModuleSpec, output_activation: Activation | None = None)

Fourier Tensor Network for implicit neural representations.

Computes output_activation(head(combiner(basis(x)))). Each stage is a spec: a module instance, a registry key, a module class, or a (key_or_class, kwargs) pair. Registry keys:

  • Bases: "cosine", "sinc", "legendre", "chebyshev", "triangle", "lanczos".
  • Combiners: "cp", "tr", "tucker".
  • Heads: "linear", "mlp", "identity".

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • basis (ModuleSpec) –

    Basis spec; see FUTONEncoding.

  • combiner (ModuleSpec) –

    Combiner spec; see FUTONEncoding.

  • head (ModuleSpec) –

    Head spec, built as head(encoding.out_features, out_features, **kwargs).

  • output_activation (Activation | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
Note

The cosine basis, CP combiner and linear head implement Eq. (9) of the paper: encoding.combiner.factors[c] stores \(\boldsymbol{U}^{(c)}\) and head.weight stores \(\boldsymbol{V}\). The other bases and combiners and the MLP head are extensions.

Examples:

>>> model = FUTON(
...     in_features=2,
...     out_features=3,
...     basis=("cosine", {"num_components": 256}),
...     combiner=("cp", {"rank": 256}),
...     head="linear",
...     output_activation=torch.tanh,
... )
>>> model(torch.rand(64, 64, 2) * 2 - 1).shape
torch.Size([64, 64, 3])
Source code in inro/models/futon.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    basis: ModuleSpec,
    combiner: ModuleSpec,
    head: ModuleSpec,
    output_activation: Activation | None = None,
):
    super().__init__()
    self.in_features, self.out_features = in_features, out_features
    self.encoding = FUTONEncoding(in_features, basis, combiner)
    width = self.encoding.out_features
    self.head = build_module(head, HEADS, width, out_features)
    self.output_activation = output_activation or nn.Identity()

Starts

futon_inits

Starts of FUTON fitted to a signal, and the CP algebra its bounds and solvers share.

On a grid, a CP FUTON, a CP combiner with a linear head, is a CP tensor: its factor of axis c is the axis's basis times the factor matrix, and its factor of the channels is the head's weight. A start sets these blocks of a built FUTON from the signal, in place, before training:

  • SVDInit gives the R components the R heaviest rank-one terms of the signal's low frequencies, found by nested SVDs;
  • ALSInit refines those terms by alternating least squares first.

SVDInit dataclass

SVDInit(*, cut: int = 4, scale: float = 1.0)

Start a FUTON from the heaviest rank-one terms of a signal's low frequencies.

Called on a built FUTON and a signal (N_1, ..., N_C, D) on the grid linspace(-1, 1, N_c) of each axis, such as a coordinate dataset's target, it sets the model's parameters in place:

  1. Project the signal onto the lowest K_c // cut cosines of each axis c, normalized per point as in CosineBasis.
  2. Expand the projection, a tensor of the C axes and the channels, into orthonormal rank-one terms by nested SVDs: the first splits off the shortest mode, and each next one the shortest mode left, which leaves the most terms.
  3. Give the R heaviest terms, as refine returns them, to the R components: a term's vector of axis c, fitted to the model's basis on the grid, becomes a column of factor c, and its vector of the channels a column of the output factor, the head's weight.

Each term's weight, times scale, is shared evenly by its component's factors and output column. Components beyond the terms keep their draw behind a zero output column, silent until trained, and a head bias starts at zero. The start runs on the model's device, in its precision, and raises a ValueError for a model that is not a CP FUTON.

Parameters:

  • cut (int, default: 4 ) –

    Divisor of each axis's component count K_c that leaves the frequencies kept.

  • scale (float, default: 1.0 ) –

    Factor of the output before the activation. With tanh, 2 starts a signal of -1 and 1 near -0.96 and 0.96.

Examples:

>>> model = FUTON(
...     in_features=2,
...     out_features=3,
...     basis=("sinc", {"num_components": 32}),
...     combiner=("cp", {"rank": 16}),
...     head="linear",
... )
>>> SVDInit(scale=2.0)(model, torch.rand(64, 64, 3) * 2 - 1)

refine

refine(coefficients: Tensor, factors: list[Tensor]) -> list[Tensor]

Return the terms' CP factors as the start loads them: here, as they are.

Parameters:

  • coefficients (Tensor) –

    The projection (K_1 // cut, ..., K_C // cut, D).

  • factors (list[Tensor]) –

    CP factors (I_m, T) of its T heaviest terms, a factor per axis and then the channels', the terms' weights in the first.

Source code in inro/models/futon_inits.py
def refine(self, coefficients: Tensor, factors: list[Tensor]) -> list[Tensor]:
    """Return the terms' CP factors as the start loads them: here, as they are.

    Args:
        coefficients: The projection ``(K_1 // cut, ..., K_C // cut, D)``.
        factors: CP factors ``(I_m, T)`` of its ``T`` heaviest terms, a factor per
            axis and then the channels', the terms' weights in the first.
    """
    return factors

ALSInit dataclass

ALSInit(*, cut: int = 4, scale: float = 1.0, sweeps: int = 10)

Start a FUTON as SVDInit does, its terms refined by ALS first.

Each sweep of alternating least squares solves every factor's least-squares problem on the projection exactly, the channels' included, the others held, as futon_bound refines its upper bound, so the misfit never rises. A sweep works on the projection, not on the samples.

Parameters:

  • sweeps (int, default: 10 ) –

    Most sweeps; they stop sooner once one lowers the misfit by less than 3e-5 of it.

  • cut (int, default: 4 ) –

    See SVDInit.

  • scale (float, default: 1.0 ) –

    See SVDInit.

refine

refine(coefficients: Tensor, factors: list[Tensor]) -> list[Tensor]

Return the factors after at most sweeps sweeps of ALS on the projection.

Source code in inro/models/futon_inits.py
def refine(self, coefficients: Tensor, factors: list[Tensor]) -> list[Tensor]:
    """Return the factors after at most ``sweeps`` sweeps of ALS on the projection."""
    return _als(coefficients, factors, self.sweeps)

Solvers

A FUTON with a CP combiner and a linear head, both without bias, is multilinear: with the other blocks held, its output is linear in each factor matrix and in the head, so fitting one block is a least-squares problem. inro.LeastSquares and inro.MultiplicativeUpdate sweep the blocks once per batch, and inro.train takes them in place of Adam, on the same datasets:

model = inro.FUTON(
    2, 3,
    basis=("cosine", {"num_components": [H // 2, W // 2], "grid_size": (H, W)}),
    combiner=("cp", {"rank": 280}),
    head=("linear", {"bias": False}),
)
solver = inro.LeastSquares(model, momentum=0)
result = inro.train(model, dataset, dataset, optimizer=solver, num_steps=500, metrics={"psnr": inro.psnr})

On a grid with \(M_c\) coordinates along axis \(c\), the model is a CP tensor with one factor per axis, the response \(\boldsymbol{T}_c \boldsymbol{U}^{(c)}\) of the basis's table \(\boldsymbol{T}_c \in \mathbb{R}^{M_c \times K_c}\), and one for the channels, the head \(\boldsymbol{V}\). With the other blocks held, the normal equations of the mean squared error for \(\boldsymbol{U}^{(c)}\) are small, the size of the block itself:

\[ \boldsymbol{P}_c \boldsymbol{U}^{(c)} \boldsymbol{Q}_c = \boldsymbol{B}_c , \qquad \boldsymbol{P}_c = \tfrac{1}{M_c} \boldsymbol{T}_c^\top \boldsymbol{T}_c , \qquad \boldsymbol{Q}_c = \boldsymbol{V}^\top \boldsymbol{V} \odot \bigodot_{j \ne c} \tfrac{1}{M_j} \big(\boldsymbol{T}_j \boldsymbol{U}^{(j)}\big)^\top \big(\boldsymbol{T}_j \boldsymbol{U}^{(j)}\big) , \]

where \(\boldsymbol{P}_c\) is the basis's Gram matrix and \(\boldsymbol{B}_c\) the targets' correlation with the rest of the model. The head's are \(\boldsymbol{V} \boldsymbol{G} = \boldsymbol{C}\), with \(\boldsymbol{G}\) the elementwise product of every axis's Gram matrix and \(\boldsymbol{C}\) the targets' correlation with the whole CP product.

Alternating least squares: every block takes its problem's solution, one side at a time, each with a ridge against collinear functions or components.

  • From the left, by the inverse of \(\boldsymbol{P}_c\), which depends on the basis and the grid alone and is made once on a declared grid; from the right, by the Cholesky factorization of \(\boldsymbol{Q}_c\).
  • Any basis will do, signed or not, and the parameters start from normal draws.
  • With an activation at the output the targets are Gauss-Newton ones, refitted every sweep, so the sweeps descend the activated model's error rather than minimize a fixed one.

Lee and Seung's updates, for a nonnegative model. With everything nonnegative, the two parts of a block's gradient are too, and the update multiplies the block elementwise by their ratio:

\[ \boldsymbol{U}^{(c)} \leftarrow \boldsymbol{U}^{(c)} \odot \boldsymbol{B}_c \oslash \big(\boldsymbol{P}_c \boldsymbol{U}^{(c)} \boldsymbol{Q}_c\big). \]
  • That is the minimizer of the auxiliary function Lee and Seung build for a quadratic with a nonnegative Hessian: the block's problem is never worse after the update, the block stays nonnegative, and one at a fixed point stays there. An entry is set to zero where a minibatch's estimate of \(\boldsymbol{B}_c\) is negative.
  • A block's problem is fixed for the sweep and an update costs two small products, so each block takes inner of them in a row, as Gillis and Glineur propose.
  • The model needs a nonnegative basis, the triangle with normalize=False so that its functions sum to one, nonnegative targets and nonnegative parameters. A signal in \([-1, 1]\) is fitted through the affine output activation lambda z: 2 * z - 1.

A few facts hold for both:

  • Batches. A batch is taken as points of the grid the basis declares, or else of the grid of its distinct coordinates. The Gram matrices are the grid's, and a right-hand side is what the model explains on the grid plus what it misses at the batch's points, \(\boldsymbol{B}_c = \boldsymbol{P}_c \boldsymbol{U}^{(c)} \boldsymbol{Q}_c + \boldsymbol{T}_c^\top \boldsymbol{E}_c / N\) for the residual's correlation \(\boldsymbol{E}_c\) over the \(N\) points. That is exact on a whole grid and, on a random share of it, the whole grid's in expectation, with the noise of the residual rather than of the signal. A sweep thus costs a few products of matrices over the grid whatever share of it the batch holds.
  • Monotony. On a whole grid at lr=1, no update of MultiplicativeUpdate, or of LeastSquares without its ridge, raises the batch's error.
  • Scales. Every axis's response is held at unit mean square, the head's columns taking the scales, which fixes the freedom of a CP model to trade scale between its blocks.
  • Momentum. momentum averages each block's statistics over the batches, as online matrix factorization does (Mairal et al., 2010), with weights that sum to one from the first batch on. The newest batch's weight is the solver's lr: it starts at 1 - momentum, and inro.train anneals it to a hundredth. Without momentum it stays one, and every sweep solves the batch's own problem.

The notebook notebooks/futon_solvers.ipynb compares these solvers with Adam, and the error bound page uses LeastSquares on the whole grid.

futon_solvers

Block coordinate descent for FUTON: solvers in place of a gradient optimizer.

A FUTON with a CP combiner and a linear head, both without bias, is multilinear, so fitting one block with the others held is a least-squares problem. BCD sweeps the blocks once per batch and leaves each block's update to a subclass: LeastSquares solves the block's problem, for any basis, and MultiplicativeUpdate takes Lee and Seung's step, which keeps a nonnegative model nonnegative. inro.train takes both in place of Adam, on the same datasets.

BCD

BCD(model: FUTON, momentum: float = 0.9, initialize: bool = True, **defaults: Any)

Block coordinate descent over a FUTON's factor matrices and head.

On a grid the model is a CP tensor, so with the other blocks held each block has small normal equations, which a subclass's solve handles. step updates the factors in order, then the head, and holds every axis's response at unit mean square, the head's columns taking the scales.

A batch is taken as points of the grid the basis declares, or else of the grid of its distinct coordinates. Only the residual comes from its points, so a sweep costs about the same on any share of the grid, and a whole grid with momentum=0 is the fast setting. On a whole grid at lr=1, no update of MultiplicativeUpdate, or of LeastSquares without its ridge, raises the batch's error.

Parameters:

  • model (FUTON) –

    FUTON with a CP combiner and a linear head, both without bias, whose basis evaluates one axis at a time through axis_features, as the built-in ones do. Updated in place.

  • momentum (float, default: 0.9 ) –

    Weight of the past batches in a block's statistics, in [0, 1), as in online matrix factorization; 0 takes each batch on its own. The solver's lr, the newest batch's weight, starts at 1 - momentum. Above 0, inro.train anneals it to a hundredth, as it would a gradient step's; at 0 it stays one.

  • initialize (bool, default: True ) –

    Whether to draw the parameters anew with reset_parameters; False keeps the model's.

  • **defaults (Any, default: {} ) –

    A subclass's settings, kept in self.defaults.

Raises:

  • ValueError –

    If model is not such a FUTON or momentum is out of range.

References

Mairal et al., "Online learning for matrix factorization and sparse coding", JMLR 2010.

Source code in inro/models/futon_solvers.py
def __init__(
    self,
    model: FUTON,
    momentum: float = 0.9,
    initialize: bool = True,
    **defaults: Any,
) -> None:
    factors, output, _ = _cp_parts(model)
    if not 0 <= momentum < 1:
        raise ValueError("momentum must be in [0, 1)")
    # The blocks, the factors then the output factor, which the solver reads here
    # rather than from param_groups, whose lr changes at every step and would make
    # a compiled sweep recompile.
    self._blocks = [*factors, output]
    settings = {"lr": 1 - momentum, "initialize": initialize}
    super().__init__(self._blocks, settings | defaults)
    self.model = model
    self._sweep = self._sweep_eagerly
    if initialize:
        self.reset_parameters()

reset_parameters

reset_parameters() -> None

Draw the parameters as the solver wants them.

A subclass draws the entries. Each factor is then divided by the sum of its basis's magnitudes in the middle of the domain, so that each axis's response there is at most about one, and the head by 2 * R, so that its components sum to about one half: the first sweep starts from an output of the signal's order.

Source code in inro/models/futon_solvers.py
@torch.no_grad()
def reset_parameters(self) -> None:
    """Draw the parameters as the solver wants them.

    A subclass draws the entries. Each factor is then divided by the sum of its
    basis's magnitudes in the middle of the domain, so that each axis's response
    there is at most about one, and the head by ``2 * R``, so that its components
    sum to about one half: the first sweep starts from an output of the signal's
    order.
    """
    basis: Any = self.model.encoding.basis  # any basis with axis_features
    *factors, weight = self._blocks
    for c, factor in enumerate(factors):
        # Magnitudes, since a signed basis's values can cancel: the cosines' sum to
        # zero when K_c is 0 or 3 modulo 4.
        middle = basis.axis_features(factor.new_zeros(1), c)
        factor.copy_(self._draw(factor) / middle.abs().sum())
    weight.copy_(self._draw(weight) / (2 * weight.shape[1]))

solve

solve(value: Tensor, right: Tensor, gram: Tensor, left: Tensor | None = None) -> Tensor

Return a block's new value, from its problem left @ value @ gram = right.

Parameters:

  • value (Tensor) –

    Current block, (K_c, R) for a factor and (D, R) for the head.

  • right (Tensor) –

    Right-hand side, of the block's shape.

  • gram (Tensor) –

    Elementwise product of the other blocks' Gram matrices, (R, R).

  • left (Tensor | None, default: None ) –

    Gram matrix of the basis, in the form the subclass prepares it; None is the identity, for the head.

Source code in inro/models/futon_solvers.py
def solve(
    self, value: Tensor, right: Tensor, gram: Tensor, left: Tensor | None = None
) -> Tensor:
    """Return a block's new value, from its problem ``left @ value @ gram = right``.

    Args:
        value: Current block, ``(K_c, R)`` for a factor and ``(D, R)`` for the
            head.
        right: Right-hand side, of the block's shape.
        gram: Elementwise product of the other blocks' Gram matrices, ``(R, R)``.
        left: Gram matrix of the basis, in the form the subclass prepares it;
            ``None`` is the identity, for the head.
    """
    raise NotImplementedError

compile

compile(batch: Mapping[str, Tensor] | None = None, **options: Any) -> None

Compile the sweep with torch.compile.

Parameters:

  • batch (Mapping[str, Tensor] | None, default: None ) –

    Batch to compile on now, at its shapes, leaving the model and the solver as they were, so that no timed step pays for it; None compiles at the first step.

  • **options (Any, default: {} ) –

    Keyword arguments of torch.compile.

Source code in inro/models/futon_solvers.py
def compile(
    self, batch: Mapping[str, Tensor] | None = None, **options: Any
) -> None:
    """Compile the sweep with ``torch.compile``.

    Args:
        batch: Batch to compile on now, at its shapes, leaving the model and the
            solver as they were, so that no timed step pays for it; ``None``
            compiles at the first step.
        **options: Keyword arguments of ``torch.compile``.
    """
    self._sweep = torch.compile(self._sweep_eagerly, **options)
    if batch is not None:
        saved = [block.detach().clone() for block in self._blocks]
        self.step(batch)
        with torch.no_grad():
            for block, value in zip(self._blocks, saved):
                block.copy_(value)
        self.state.clear()

step

step(batch: Mapping[str, Tensor]) -> Tensor

Run one sweep over the factor matrices and the head on a batch.

Parameters:

  • batch (Mapping[str, Tensor]) –

    input, coordinates (..., C) in [-1, 1], and target, their values (..., D), as a dataset yields them. The points must lie on a grid, as a coordinate dataset's do.

Returns:

  • Tensor –

    The batch's mean squared error before the sweep.

Source code in inro/models/futon_solvers.py
@torch.no_grad()
def step(self, batch: Mapping[str, Tensor]) -> Tensor:  # type: ignore[override]
    """Run one sweep over the factor matrices and the head on a batch.

    Args:
        batch: ``input``, coordinates ``(..., C)`` in ``[-1, 1]``, and ``target``,
            their values ``(..., D)``, as a dataset yields them. The points must lie
            on a grid, as a coordinate dataset's do.

    Returns:
        The batch's mean squared error before the sweep.
    """
    x = batch["input"].reshape(-1, self.model.in_features)
    y = batch["target"].reshape(-1, self.model.out_features)
    if not self.state:
        self._start(x)
        self._check(x, y)
    elif not hasattr(self, "_lr"):  # the state came from load_state_dict
        self._lr = x.new_zeros(())
    self._lr.fill_(self.param_groups[0]["lr"])
    return self._sweep(x, y)

LeastSquares

LeastSquares(model: FUTON, ridge: float = 0.001, momentum: float = 0.9, initialize: bool = True)

Alternating least squares: every block takes its problem's solution.

A block's problem is solved one side at a time, each with a ridge against collinear functions or components, for any basis, signed or not. With an output activation the targets are Gauss-Newton ones, refitted every sweep, so the sweeps descend the activated model's error rather than minimize a fixed one. The parameters start from normal draws.

Parameters:

  • model (FUTON) –

    See BCD.

  • ridge (float, default: 0.001 ) –

    Added to the diagonal of each Gram matrix, relative to its mean diagonal.

  • momentum (float, default: 0.9 ) –

    See BCD.

  • initialize (bool, default: True ) –

    See BCD.

Examples:

>>> model = FUTON(
...     in_features=2,
...     out_features=3,
...     basis=("sinc", {"num_components": 16}),
...     combiner=("cp", {"rank": 8}),
...     head=("linear", {"bias": False}),
... )
>>> solver = LeastSquares(model, momentum=0)
>>> axis = torch.linspace(-1, 1, 32)
>>> x = torch.cartesian_prod(axis, axis)
>>> loss = solver.step({"input": x, "target": torch.rand(len(x), 3)})
>>> loss.shape
torch.Size([])
Source code in inro/models/futon_solvers.py
def __init__(
    self,
    model: FUTON,
    ridge: float = 1e-3,
    momentum: float = 0.9,
    initialize: bool = True,
) -> None:
    super().__init__(model, momentum, initialize, ridge=ridge)

solve

solve(value: Tensor, right: Tensor, gram: Tensor, left: Tensor | None = None) -> Tensor

Return the block's ridged solution; see BCD.solve.

Source code in inro/models/futon_solvers.py
def solve(
    self, value: Tensor, right: Tensor, gram: Tensor, left: Tensor | None = None
) -> Tensor:
    """Return the block's ridged solution; see :meth:`BCD.solve`."""
    if left is not None:
        right = left @ right
    # The Gram matrix is symmetric, so the right solve is a left one transposed.
    factor = torch.linalg.cholesky_ex(_ridged(gram, self.defaults["ridge"])).L
    return torch.cholesky_solve(right.T, factor).T

MultiplicativeUpdate

MultiplicativeUpdate(model: FUTON, inner: int = 16, floor: float = 1e-12, momentum: float = 0.9, initialize: bool = True)

Lee and Seung's multiplicative updates, for a nonnegative model.

The update multiplies a block elementwise by the ratio of its gradient's two parts, value * right / (left @ value @ gram). It never worsens the block's problem and keeps the block nonnegative; each block takes inner of them in a row, as Gillis and Glineur propose.

The first step checks that the basis, the targets and the parameters are nonnegative. Use the triangle basis with normalize=False, whose functions sum to one, and fit a signal in [-1, 1] through the affine output activation lambda z: 2 * z - 1, under which its targets are in [0, 1]. The parameters start uniformly in [0.5, 1.5], before scaling.

Parameters:

  • model (FUTON) –

    See BCD.

  • inner (int, default: 16 ) –

    Number of updates of each block per sweep.

  • floor (float, default: 1e-12 ) –

    Lower bound on a denominator, against division by zero where a basis function sees no sample.

  • momentum (float, default: 0.9 ) –

    See BCD.

  • initialize (bool, default: True ) –

    See BCD.

References

Lee and Seung, "Algorithms for non-negative matrix factorization", NeurIPS 2001. Gillis and Glineur, "Accelerated multiplicative updates and hierarchical ALS algorithms for nonnegative matrix factorization", Neural Computation 2012.

Source code in inro/models/futon_solvers.py
def __init__(
    self,
    model: FUTON,
    inner: int = 16,
    floor: float = 1e-12,
    momentum: float = 0.9,
    initialize: bool = True,
) -> None:
    if inner < 1:
        raise ValueError("inner must be at least 1")
    super().__init__(model, momentum, initialize, inner=inner, floor=floor)

solve

solve(value: Tensor, right: Tensor, gram: Tensor, left: Tensor | None = None) -> Tensor

Return the block after inner updates; see BCD.solve.

Source code in inro/models/futon_solvers.py
def solve(
    self, value: Tensor, right: Tensor, gram: Tensor, left: Tensor | None = None
) -> Tensor:
    """Return the block after ``inner`` updates; see :meth:`BCD.solve`."""
    floor = self.defaults["floor"]
    for _ in range(self.defaults["inner"]):
        denominator = value @ gram if left is None else left @ value @ gram
        # A coefficient no sample of the batch reaches has both parts at
        # zero and keeps its value; one whose right-hand side is negative,
        # which a minibatch's estimate can be, goes to zero, the
        # constrained minimizer of its auxiliary function.
        ratio = (right / denominator.clamp_min(floor)).clamp_min(0)
        value = value * torch.where(denominator > floor, ratio, 1.0)
    return value

Error bounds

The error bound page states and proves the theorem these functions compute.

futon_bounds

Bounds on how well a FUTON with a linear head can fit a sampled signal.

On a signal's sampling grid, the error of a FUTON with a CP combiner and a linear head, both without bias, splits exactly in two: the error of the best approximation in the basis's span, which no rank removes, and the CP misfit of the signal's coefficients in an orthonormal basis of that span. futon_bound computes the first exactly, brackets the best second, and returns the factors of a FUTON that attains the upper bound. The documentation's error bound page states and proves the theorem.

FUTONBound dataclass

FUTONBound(ranks: list[int], truncation: float, lower: Tensor, upper: Tensor, factors: list[tuple[list[Tensor], Tensor]])

Bounds on the best mean squared error of a FUTON with a linear head.

Errors are float64 means over the grid's points and the signal's channels, as inro.train reports them.

Attributes:

  • ranks (list[int]) –

    CP ranks, ascending.

  • truncation (float) –

    Error of the best approximation in the basis's span, which no rank goes below; exact.

  • lower (Tensor) –

    Per rank, a lower bound on the error of every FUTON of that rank; exact for a matrix of coefficients, that is, one axis, or two and one channel.

  • upper (Tensor) –

    Per rank, the error of the FUTON in factors, and so an upper bound on the best one's.

  • factors (list[tuple[list[Tensor], Tensor]]) –

    Per rank, the factors \(\boldsymbol{U}^{(c)}\), (K_c, R), and the head \(\boldsymbol{V}\), (D, R), that attain upper.

initialize

initialize(model: FUTON) -> None

Give model the factors and the head that attain upper at its rank.

The model's error on the bound's grid is then the bound, up to its precision.

Parameters:

  • model (FUTON) –

    FUTON with the bound's basis, a CP combiner and a linear head without bias.

Raises:

  • ValueError –

    If model has another combiner or head, or a rank without a bound.

Source code in inro/models/futon_bounds.py
@torch.no_grad()
def initialize(self, model: FUTON) -> None:
    """Give ``model`` the factors and the head that attain ``upper`` at its rank.

    The model's error on the bound's grid is then the bound, up to its precision.

    Args:
        model: FUTON with the bound's basis, a CP combiner and a linear head
            without bias.

    Raises:
        ValueError: If ``model`` has another combiner or head, or a rank without
            a bound.
    """
    factors, output, _ = _cp_parts(model)
    bound_factors, bound_output = self.factors[self.ranks.index(output.shape[1])]
    for parameter, factor in zip(
        [*factors, output], [*bound_factors, bound_output]
    ):
        parameter.copy_(factor)

futon_bound

futon_bound(signal: Tensor, features: Sequence[Tensor], ranks: int | Sequence[int], sweeps: int = 500, tol: float = 3e-05) -> FUTONBound

Bound the best grid error of a FUTON with a linear head on a signal.

Every FUTON with a CP combiner and a linear head, both without bias, with these per-axis features and rank R has a mean squared error on the grid of at least the bound's lower at R, and the one given by its factors has exactly its upper; no FUTON goes below truncation. The bounds depend on the basis only through each axis's span.

Parameters:

  • signal (Tensor) –

    Samples (N_1, ..., N_C, D) on a grid.

  • features (Sequence[Tensor]) –

    Per axis c, the basis (N_c, K_c) at that axis's grid coordinates, orthonormal or not. For a FUTON's basis on a coordinate dataset's grid, it is basis.axis_features(torch.linspace(-1, 1, N_c), c).

  • ranks (int | Sequence[int]) –

    CP ranks R to bound.

  • sweeps (int, default: 500 ) –

    Maximum number of alternating least-squares sweeps that refine each upper bound, in float32, from its SVD form, the R heaviest terms of the nested SVDs; 0 keeps the SVD form.

  • tol (float, default: 3e-05 ) –

    Fraction of the bound below which a sweep's decrease stops the sweeps.

Returns:

  • FUTONBound –

    The bounds and the factors that attain the upper ones, on the CPU.

Raises:

  • ValueError –

    If features is not one matrix of N_c rows per axis.

Examples:

>>> from inro import SincBasis
>>> signal = torch.rand(32, 32, 3) * 2 - 1
>>> basis = SincBasis(2, num_components=16)
>>> grid = torch.linspace(-1, 1, 32)
>>> features = [basis.axis_features(grid, c) for c in range(2)]
>>> bound = futon_bound(signal, features, ranks=[4, 8], sweeps=10)
>>> psnr = 10 * torch.log10(4 / bound.upper)  # guaranteed, in dB
>>> psnr.shape
torch.Size([2])
Source code in inro/models/futon_bounds.py
@torch.no_grad()
def futon_bound(
    signal: Tensor,
    features: Sequence[Tensor],
    ranks: int | Sequence[int],
    sweeps: int = 500,
    tol: float = 3e-5,
) -> FUTONBound:
    """Bound the best grid error of a FUTON with a linear head on a signal.

    Every FUTON with a CP combiner and a linear head, both without bias, with these
    per-axis features and rank ``R`` has a mean squared error on the grid of at least
    the bound's ``lower`` at ``R``, and the one given by its ``factors`` has exactly its
    ``upper``; no FUTON goes below ``truncation``. The bounds depend on the basis only
    through each axis's span.

    Args:
        signal: Samples ``(N_1, ..., N_C, D)`` on a grid.
        features: Per axis ``c``, the basis ``(N_c, K_c)`` at that axis's grid
            coordinates, orthonormal or not. For a FUTON's basis on a coordinate
            dataset's grid, it is
            ``basis.axis_features(torch.linspace(-1, 1, N_c), c)``.
        ranks: CP ranks ``R`` to bound.
        sweeps: Maximum number of alternating least-squares sweeps that refine each
            upper bound, in float32, from its SVD form, the ``R`` heaviest terms of the
            nested SVDs; ``0`` keeps the SVD form.
        tol: Fraction of the bound below which a sweep's decrease stops the sweeps.

    Returns:
        The bounds and the factors that attain the upper ones, on the CPU.

    Raises:
        ValueError: If ``features`` is not one matrix of ``N_c`` rows per axis.

    Examples:
        >>> from inro import SincBasis
        >>> signal = torch.rand(32, 32, 3) * 2 - 1
        >>> basis = SincBasis(2, num_components=16)
        >>> grid = torch.linspace(-1, 1, 32)
        >>> features = [basis.axis_features(grid, c) for c in range(2)]
        >>> bound = futon_bound(signal, features, ranks=[4, 8], sweeps=10)
        >>> psnr = 10 * torch.log10(4 / bound.upper)  # guaranteed, in dB
        >>> psnr.shape
        torch.Size([2])
    """
    ranks = sorted({ranks} if isinstance(ranks, int) else set(ranks))
    signal = signal.double()
    *sizes, channels = signal.shape
    if len(features) != len(sizes) or any(
        len(phi) != size for phi, size in zip(features, sizes)
    ):
        raise ValueError(
            f"expected one feature matrix of N_c rows per axis of {tuple(sizes)}"
        )

    # Phi_c / sqrt(N_c) = Q_c B_c with Q_c an orthonormal basis of its span;
    # directions the grid cannot tell from zero are left out. The small
    # decompositions here and below run on the CPU, where they are fastest.
    bases, inverses = [], []
    for phi in features:
        phi = phi.double().cpu() / math.sqrt(len(phi))
        q, s, vh = torch.linalg.svd(phi, full_matrices=False)
        keep = s > s[0] * max(phi.shape) * torch.finfo(s.dtype).eps
        bases.append(q[:, keep].to(signal))
        inverses.append(vh[keep].T / s[keep])  # B_c^+, (K_c, K_c')

    # The signal's coefficients in the orthonormal basis of the span.
    coefficients = _project(signal, bases) / math.sqrt(math.prod(sizes))
    energy = float(coefficients.square().sum())
    truncation = max(float(signal.square().mean()) * channels - energy, 0.0)
    small = coefficients.cpu()
    lower = _tails(small, ranks)

    # The upper bounds' SVD form: per rank, the R heaviest terms of the nested SVDs
    # along the mode order that keeps the most energy in them. Modes of size one are
    # peeled first, at no cost, and the last two of an order share one SVD.
    single = [mode for mode in range(small.ndim) if small.shape[mode] == 1]
    rest = [mode for mode in range(small.ndim) if small.shape[mode] > 1]
    orders = [
        single + list(order)
        for order in permutations(rest)
        if len(order) < 2 or order[-2] < order[-1]
    ]
    expansions = []
    for order in orders:
        weights, vectors = _expansion(small, order)
        weights, index = weights.sort(descending=True)
        expansions.append((weights, [vector[index] for vector in vectors]))

    # The search runs in float32 on the signal's device; each bound is the
    # exact float64 error of the factors it finds.
    working = coefficients.float()
    upper, factors = [], []
    for rank in ranks:
        kept = [float(weights[:rank].square().sum()) for weights, _ in expansions]
        weights, vectors = expansions[kept.index(max(kept))]
        start = [vector[:rank].T.to(working) for vector in vectors]
        start[0] = start[0] * weights[:rank].to(working)
        chosen = [f.double() for f in _als(working, start, sweeps, tol, truncation)]
        residual = _residual(coefficients, chosen)
        # A rank's FUTON is one of every larger rank, with zero components added.
        if upper and upper[-1] < residual:
            residual, chosen = upper[-1], factors[-1]
        upper.append(residual)
        factors.append([F.pad(f, (0, rank - f.shape[1])) for f in chosen])

    # Back to the basis: U^(c) = B_c^+ U~^(c), with the head unchanged.
    mapped = [
        ([inverse @ f.cpu() for inverse, f in zip(inverses, chosen)], chosen[-1].cpu())
        for chosen in factors
    ]
    to_mse = 1 / channels
    return FUTONBound(
        ranks=ranks,
        truncation=truncation * to_mse,
        lower=(truncation + lower) * to_mse,
        upper=(truncation + torch.tensor(upper, dtype=torch.float64)) * to_mse,
        factors=mapped,
    )