Heuristic:Kornia Kornia Numerical Stability Patterns
| Knowledge Sources | |
|---|---|
| Domains | Optimization, Numerical_Computing |
| Last Updated | 2026-02-09 15:00 GMT |
Overview
Collection of numerical stability techniques used throughout Kornia to prevent division-by-zero, overflow, and NaN propagation in differentiable operations.
Description
Kornia operates on floating-point tensors in differentiable pipelines where numerical stability is critical. The library employs several distinct patterns: (1) fixed epsilon values (e.g., `1e-6`, `1e-10`) for safe division, (2) dtype-aware epsilon via `torch.finfo(dtype).tiny` or `torch.finfo(dtype).eps`, (3) `torch.clamp(min=eps)` for denominators, (4) value scaling to prevent fp16 overflow in attention mechanisms, and (5) `torch.where` for numerically stable conditional branching in quaternion/rotation conversions.
Usage
Apply these patterns when implementing any operation involving division, square roots, logarithms, or exponentials in differentiable code paths. Particularly important for feature descriptors, geometric transformations, loss functions, and attention mechanisms that may encounter near-zero denominators or extreme values.
The Insight (Rule of Thumb)
- Action 1: Use `torch.finfo(tensor.dtype).tiny` for safe division denominators. This adapts to float16/float32/float64 automatically.
- Action 2: Use `torch.clamp(denominator, min=1e-10)` before division when the denominator could be zero.
- Action 3: Scale values before accumulation in attention to prevent fp16 overflow: `values = values / v_length`.
- Action 4: Use `torch.where` for conditional branches based on tensor values (not Python if/else) to keep gradients flowing.
- Trade-off: Epsilon values introduce small bias. Too large an epsilon distorts results; too small may not prevent instability on all dtypes.
Reasoning
In differentiable pipelines, zero denominators produce `inf` which propagates as `NaN` through gradients, corrupting the entire backward pass. Float16 has a much smaller dynamic range than float32, so operations that are stable in float32 (like attention dot products) can overflow in float16. Using dtype-aware epsilons from `torch.finfo` ensures the epsilon is appropriate for the tensor's precision. The `torch.where` pattern is needed for operations like quaternion conversion where different formulas are numerically stable in different regions of the input space.
Code Evidence
Dtype-aware safe division from `kornia/geometry/conversions.py:457-459`:
def safe_zero_division(numerator: torch.Tensor, denominator: torch.Tensor) -> torch.Tensor:
eps: float = torch.finfo(numerator.dtype).tiny
return numerator / torch.clamp(denominator, min=eps)
Clamp for numerical stability in LoFTR fine matching from `kornia/feature/loftr/utils/fine_matching.py:83`:
std = torch.sum(torch.sqrt(torch.clamp(var, min=1e-10)), -1) # clamp needed for numerical stability
FP16 overflow prevention in linear attention from `kornia/feature/loftr/loftr_module/linear_attention.py:82-86`:
v_length = values.size(1)
values = values / v_length # prevent fp16 overflow
KV = torch.einsum("nshd,nshv->nhdv", K, values)
Z = 1 / (torch.einsum("nlhd,nhd->nlh", Q, K.sum(dim=1)) + self.eps)
queried_values = torch.einsum("nlhd,nhdv,nlh->nlhv", Q, KV, Z) * v_length
Epsilon-scaled clamp for Lie group operations from `kornia/geometry/liegroup/so3.py`:
eps = torch.finfo(v.dtype).eps * 1e3
Safe log with epsilon from `kornia/feature/adalam/ransac.py:27`:
logres = torch.log(residuals + 1e-10)
Conditional branching with torch.where for quaternion stability from `kornia/geometry/conversions.py:467-502`:
def trace_positive_cond() -> torch.Tensor:
sq = torch.sqrt(trace + 1.0 + eps) * 2.0 # sq = 4 * qw
# ...
quaternion: torch.Tensor = torch.where(trace > 0.0, trace_positive_cond(), where_1)