xbtorch.decomposition.sbpca

Streaming Batch PCA (SBPCA) decomposition module.

Implements the SBPCA class, which performs an incremental, streaming principal component analysis on gradients. This allows low-rank approximations without computing the full gradient matrix, reducing memory and computational overhead.

Useful for memristive neural networks to reduce device writes and improve training efficiency.

Classes

SBPCA([rank, sub_batch_size])

Streaming Batch Principal Component Analysis (SBPCA) gradient decomposition.
class xbtorch.decomposition.sbpca.SBPCA(rank=1, sub_batch_size=32)[source]

Bases: GenericDecomposition

Streaming Batch Principal Component Analysis (SBPCA) gradient decomposition.

This method applies an incremental, streaming variant of PCA to approximate the full gradient matrix with a low-rank factorization. By operating on sub-batches of data, SBPCA avoids forming the complete gradient explicitly, thereby reducing memory and compute costs.

SBPCA is particularly useful in memristive neural networks where gradient compression reduces the number of device writes, mitigating endurance limitations and improving training efficiency.

Parameters:

  • rank (int, optional (default=1)) – Target rank for the decomposition. Smaller values yield greater compression but less accurate approximations.

  • sub_batch_size (int, optional (default=32)) – Size of the sub-batches used for incremental updates during streaming PCA.

info

Stores intermediate estimates for each parameter group: - X: right singular vectors - DELTA: left singular vectors - sigma: singular values

Type:

dict

Notes

  • SBPCA incrementally updates low-rank approximations across sub-batches using QR decompositions for numerical stability.

  • This implementation follows the methodology in: Huang et al., “Low-rank gradient descent for memory-efficient training of deep in-memory arrays”, ACM JETC 2023.

References

  • Huang et al., “Low-rank gradient descent for memory-efficient training of deep in-memory arrays”, ACM JETC, 2023.

  • Hoskins et al., “Streaming batch eigenupdates for hardware neural networks”, Frontiers in Neuroscience, 2019.

decompose(input, delta, gradient, group_param_idx)[source]

Perform SBPCA-based gradient decomposition.

Parameters:

  • input (torch.Tensor) – Input activations for the current layer, shape (batch_size, input_dim).

  • delta (torch.Tensor) – Backpropagated errors for the current layer, shape (batch_size, output_dim).

  • gradient (torch.Tensor) – Gradient tensor (used for shape and initialization), shape (output_dim, input_dim).

  • group_param_idx (int) – Identifier for parameter grouping, used to track state across multiple calls.

Returns:

Low-rank approximation of the gradient with shape (output_dim, input_dim), reconstructed from:

approx_grad = DELTA @ diag(sigma) @ X^T

scaled by the batch size.

Return type:

torch.Tensor

Notes

  • If state for group_param_idx does not exist, orthogonal initializations are created for X and DELTA.

  • QR decompositions are used to maintain orthogonality of factors during updates.

  • For rank > 1, only the leading rank components are kept.