Implementation:Kornia Kornia ZCA Whitening

Knowledge Sources	Kornia
Domains	Vision, Image_Enhancement
Last Updated	2026-02-09 15:00 GMT

Overview

This module implements ZCA (Zero-phase Component Analysis) whitening for PyTorch tensors, providing both a stateful nn.Module class and standalone functional APIs for computing and applying whitening transforms.

Description

zca.py is part of the kornia.enhance module in the Kornia computer vision library. It provides tools for ZCA whitening, a technique that decorrelates the data while preserving the original data space structure. The module contains:

• ZCAWhitening -- an nn.Module class that fits a whitening transform to data and then applies it. Supports a two-step workflow (fit then transform) or a single-step workflow (include_fit=True). Can optionally compute the inverse transform and detach gradients from the fitting process.
• zca_mean -- computes the ZCA whitening transformation matrix and mean vector using SVD of the covariance matrix. The transform is T = U * S^(-1/2) * U^T where U and S come from the SVD of the covariance.
• zca_whiten -- a convenience function that computes the transform and applies it in one call.
• linear_transform -- applies a general linear transform (X - mean) * T to flattened data along a specified dimension, then reshapes back.

The whitening is computed as: X_zca = (X - mu)(U * S^(-1/2) * U^T)^T

Usage

Users should import from this module when they need to whiten data for feature decorrelation, preprocessing for machine learning models, or when working with techniques like PCA/ZCA whitening in computer vision pipelines.

Code Reference

Source Location

• Repository: Kornia
• File: kornia/enhance/zca.py
• Lines: 1-395

Signature

class ZCAWhitening(nn.Module):
    def __init__(
        self,
        dim: int = 0,
        eps: float = 1e-6,
        unbiased: bool = True,
        detach_transforms: bool = True,
        compute_inv: bool = False,
    ) -> None
    def fit(self, x: torch.Tensor) -> "ZCAWhitening"
    def forward(self, x: torch.Tensor, include_fit: bool = False) -> torch.Tensor
    def inverse_transform(self, x: torch.Tensor) -> torch.Tensor

def zca_mean(
    inp: torch.Tensor,
    dim: int = 0,
    unbiased: bool = True,
    eps: float = 1e-6,
    return_inverse: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]

def zca_whiten(
    inp: torch.Tensor,
    dim: int = 0,
    unbiased: bool = True,
    eps: float = 1e-6,
) -> torch.Tensor

def linear_transform(
    inp: torch.Tensor,
    transform_matrix: torch.Tensor,
    mean_vector: torch.Tensor,
    dim: int = 0,
) -> torch.Tensor

Import

from kornia.enhance import ZCAWhitening, zca_mean, zca_whiten, linear_transform

I/O Contract

Inputs (ZCAWhitening)

Name	Type	Required	Description
x	torch.Tensor	Yes	Input data tensor of shape (D_0, ..., D_dim, ..., D_N)
dim	int	No	Dimension that represents the samples axis. Default: 0
eps	float	No	Small number for numerical stability. Default: 1e-6
unbiased	bool	No	Whether to use unbiased covariance estimate. Default: True
detach_transforms	bool	No	Whether to detach gradients from fitting. Default: True
compute_inv	bool	No	Whether to compute inverse transform matrix. Default: False
include_fit	bool	No (forward)	Whether to include fitting in forward pass. Default: False

Inputs (zca_mean)

Name	Type	Required	Description
inp	torch.Tensor	Yes	Input data tensor
dim	int	No	Samples dimension. Default: 0
unbiased	bool	No	Unbiased covariance. Default: True
eps	float	No	Numerical stability constant. Default: 1e-6
return_inverse	bool	No	Whether to return inverse transform. Default: False

Outputs (zca_mean)

Name	Type	Description
transform_matrix	torch.Tensor	ZCA whitening matrix
mean_vector	torch.Tensor	Mean vector of shape (1, D)
inv_transform	Optional[torch.Tensor]	Inverse ZCA matrix (None if return_inverse=False)

Usage Examples

import torch
from kornia.enhance import ZCAWhitening, zca_mean, zca_whiten

# Basic ZCA whitening with the module
x = torch.tensor([[0, 1], [1, 0], [-1, 0], [0, -1]], dtype=torch.float32)
zca = ZCAWhitening().fit(x)
x_whiten = zca(x)

# Single-step fit and transform
zca = ZCAWhitening()
x_whiten = zca(x, include_fit=True)

# Enable backprop through the fitting process
zca = ZCAWhitening(detach_transforms=False)
x_whiten = zca(x, include_fit=True)

# Using functional API
transform_matrix, mean_vector, _ = zca_mean(x)
x_whiten = zca_whiten(x)

# Multi-dimensional data with inverse transform
x = torch.rand(3, 20, 2, 2)
zca = ZCAWhitening(dim=1, compute_inv=True).fit(x)
x_whiten = zca(x)
x_reconstructed = zca.inverse_transform(x_whiten)

# Using zca_mean with inverse
transform_matrix, mean_vector, inv_transform = zca_mean(x, dim=1, return_inverse=True)

Related Pages

• Principle:Kornia_Kornia_ZCA_Whitening