Memory-efficient TCA - #5
Open
quannfa wants to merge 1 commit into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR improves the memory efficiency of TCA on large datasets and fixes a normalization edge case that could produce NaN values. The changes are focused on the TCA fit implementation.
Background
In the original implementation, the primal mode explicitly constructs n×n M and H matrices. With large sample sizes (for example, around 56k samples), this leads to very high memory usage and can cause OOM failures.
Also, column normalization did not handle zero-norm columns, which could result in NaN values.
What Changed
Reworked the primal-kernel computation path to avoid explicitly building n×n M/H matrices.
Replaced matrix products with low-rank equivalent forms:
K·M·Kᵀ is computed via an outer product based on Xe.
K·H·Kᵀ is computed as XXᵀ minus a rank-1 centering update.
Added safe normalization for both input X and projected Z, normalizing only columns with norm > 0.
Kept the non-primal branches (linear/rbf) as fallback with the original logic.
Expected Impact
Significantly lower memory usage in primal mode, from O(n²) to approximately O(m·n + m²).
Better scalability and reduced OOM risk on large datasets.
Improved numerical stability by preventing NaN generation from zero columns.
Compatibility
No public API changes.
Primal branch internal computation path changed but remains mathematically equivalent.
Linear/RBF branches preserve existing behavior.
Validation Notes
The commit modifies only the TCA implementation file.
Full end-to-end regression/training was not run in this step; it is recommended to run a full pipeline check on representative datasets.