TIC-TAC (Taylor-Induced Covariance)¶

← IVON | Heteroscedastic Laplace →

TIC-TAC binds the predicted covariance matrix to the local derivatives (gradient and curvature) of the mean prediction model with respect to input features.

Mathematical Background¶

Standard deep heteroscedastic regression predicts the mean \(\hat{y} = \mu(x)\) and covariance \(\Sigma(x)\) using separate output heads. This can lead to optimization instability and poor covariance alignment.

TIC-TAC links the covariance \(\Sigma(x)\) directly to the local geometry of the mean model using a Taylor expansion representation:

In this Taylor expansion, the derivative-based terms are defined as follows:

• \(J(x) \in \mathbb{R}^{D \times P}\) is the Jacobian of the mean prediction \(f(x)\) w.r.t input features \(x \in \mathbb{R}^P\): $\(J(x)_{i, j} = \frac{\partial f_i(x)}{\partial x_j}\)$
• \(H(x) \in \mathbb{R}^{D \times D}\) is the curvature matrix, computed as the inner product of the target-dimension Hessians: $\(H(x)_{i, j} = \text{Tr}\left(\frac{\partial^2 f_i(x)}{\partial x^2} \frac{\partial^2 f_j(x)}{\partial x^2}\right)\)$
• \(k_1(x), k_2(x)\) are positive scaling coefficients (either global scalar parameters or input-dependent MLP heads).
• \(k_3(x) \in \mathbb{R}^D\) is a residual diagonal variance parameter (allowing individual target dimensions to have unique noise floors).
• \(\epsilon I\) is a stabilizer jitter (default: \(10^{-6}\)) to ensure positive-definiteness.

Vectorized batch computation of \(J(x)\) and \(H(x)\) is performed efficiently using PyTorch's functional autograd API (torch.func.jacrev and torch.func.hessian vectorized via torch.vmap).

High-Level API: TaylorInducedCovarianceHead¶

TaylorInducedCovarianceHead is a PyTorch nn.Module wrapping a base mean backbone.

from torchregress.algorithms import TaylorInducedCovarianceHead

# Wrap a mean backbone
head = TaylorInducedCovarianceHead(
    base_model=my_backbone,
    target_dim=3,
    input_dim=5,  # Enable input-dependent coefficients k1, k2, k3
    jitter=1e-6,
)

# Returns mean (B, D) and covariance (B, D, D)
mean, cov = head(x)

Parameters¶

Complete Example¶

import torch
import torch.nn as nn
from torchregress.algorithms import TaylorInducedCovarianceHead
from torchregress.metrics import task_agnostic_correlations

# Setup synthetic features and target
torch.manual_seed(42)
x = torch.randn(10, 5)
y = torch.randn(10, 3)

# Simple mean network
backbone = nn.Sequential(
    nn.Linear(5, 10),
    nn.Tanh(),  # Tanh or gelu/silu is recommended to ensure non-zero second derivatives
    nn.Linear(10, 3)
)

# TIC-TAC head
model = TaylorInducedCovarianceHead(
    base_model=backbone,
    target_dim=3,
    input_dim=5,
)

# Forward pass
mean, cov = model(x)

# Evaluate using Task-Agnostic Correlations (stateless helper)
tac_error = task_agnostic_correlations(mean, y, cov)
print(f"Task-Agnostic Correlation Error: {tac_error.item():.4f}")

When to use TIC-TAC¶

Limitations¶

1. Not suited for scalar regression (\(D=1\)): TIC-TAC is designed for multivariate targets with correlated noise. For \(D=1\), the Jacobian collapse yields only a scalar variance — standard Gaussian NLL or \(\beta\)-NLL is simpler and faster.
2. Memory scaling \(\mathcal{O}(D \cdot P^2)\): The vectorized Hessian computation grows quadratically in input dimension \(P\). For \(P > 100\) or \(D > 20\), expect Out-of-Memory (OOM) errors. Use dimensionality reduction on inputs or fall back to a low-rank Gaussian head.
3. Requires smooth activation functions: ReLU has zero second derivative everywhere, nullifying the Hessian term \(H(x)\). The model must use Tanh, GELU, or SiLU throughout the backbone — this constraint may degrade performance on tasks where ReLU is optimal.
4. Autograd overhead per forward pass: Computing \(J(x)\) and \(H(x)\) via torch.func.jacrev and torch.func.hessian adds a 3–10× overhead per forward pass compared to standard Gaussian heads. Not suitable for real-time or high-throughput inference.
5. Global coefficients may underfit: When input_dim=None, the scaling coefficients \(k_1, k_2, k_3\) are scalar parameters shared across all inputs. For heteroscedastic data, this uniform covariance structure is too rigid.

Recommendations¶

1. Start with \(D \leq 10\), \(P \leq 30\) and profile memory before scaling. Monitor GPU memory with nvidia-smi or torch.cuda.memory_summary().
2. Use Tanh or GELU activations in the backbone. SiLU is also acceptable but may produce smaller Hessian magnitudes than Tanh.
3. Enable input_dim (input-dependent \(k\) coefficients) unless you have strong reason to believe covariance is homogeneous across the feature space.
4. Evaluate with Task-Agnostic Correlations (TAC) — see ../../metrics/multivariate.md. MSE alone cannot assess whether the predicted covariance aligns with the true residual structure.
5. Pair with a full-covariance Gaussian loss (MultivariateGaussianLoss) so the training objective matches the covariance parameterisation.
6. Set jitter adaptively: start at 1e-6, increase to 1e-4 if Cholesky decomposition fails during training (numerical instability signal).

Next steps¶

• Multivariate metrics — Task-Agnostic Correlations for evaluating predicted covariance quality
• Gaussian losses — alternative covariance parameterisations (diagonal, full, low-rank)
• Effective Bayesian Laplace — post-hoc uncertainty decomposition complement
• TIC-TAC + multivariate metrics — evaluate predicted covariances with Task-Agnostic Correlations

References¶