[BugFix][mHC] Address numerical instability in the mHC kernel

[superAngGao]

Description

Currently, the Sinkhorn-Knopp algorithm in the mhc_pre kernel suffers from numerical instability, which may produce abnormal results under extreme conditions. This issue stems from the algorithm's inherent alternating row-wise and column-wise normalization procedure, which can involve summations and divisions of extremely small elements, thereby amplifying rounding errors.

To address this problem, an epsilon (eps) parameter can be added to the denominator during normalization, preventing it from becoming too small and thus enhancing the algorithm's numerical stability. In the corresponding unit test environment, the eps parameter is set to 0.02—an empirically determined value that effectively mitigates numerical errors in bfloat16 tensors.

[gemini-code-assist]

Summary of Changes

Hello @superAngGao, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the robustness of the mhc_pre kernel by resolving numerical instability issues inherent in its Sinkhorn-Knopp algorithm. The core change involves introducing an explicit epsilon parameter into the normalization process, which prevents division by extremely small numbers. This modification significantly improves the algorithm's reliability and accuracy, particularly when dealing with bfloat16 tensors under challenging computational conditions.

Highlights

• Numerical Stability Fix: Addressed numerical instability in the Sinkhorn-Knopp algorithm within the mhc_pre kernel by introducing an eps parameter during normalization.
• Epsilon Parameterization: The eps parameter is now passed through the ManifoldConstrainedHyperConnectionPreOp and mhc_pre kernel functions, replacing a hardcoded value.
• Unit Test Update: The unit tests for mhc_pre_op were updated to include and pass the eps parameter, setting it to 0.02 for bfloat16 tensors.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• tests/ops/test_mhc_pre.py
  • Introduced a local eps variable set to 0.02 for testing purposes.
  • Modified the call to ManifoldConstrainedHyperConnectionPreOp.forward to include the new eps parameter.
  • Removed a previously hardcoded eps value from the reference implementation within the test function.
• top/kernels/mhc/mhc_pre.py
  • Updated the _get_H_res function signature to accept sinkhorn_eps as an argument.
  • Replaced the hardcoded eps = 0.0001 with the new sinkhorn_eps parameter within _get_H_res.
  • Modified the mhc_pre function signature to accept sinkhorn_eps and passed it down to _get_H_res.
  • Adjusted the _mhc_pre_wrapped_kernel signature to include sinkhorn_eps and ensured it's passed to the underlying kernel.
  • Updated the ManifoldConstrainedHyperConnectionPreOp.forward method signature to accept sinkhorn_eps and forward it to the kernel.
• top/ops/mhc_pre.py
  • Modified the ManifoldConstrainedHyperConnectionPreOp.forward method signature to accept sinkhorn_eps.
  • Updated the internal kernel call within forward to pass the new sinkhorn_eps parameter.

Activity

• The pull request is categorized as a bug fix.
• All linting issues were addressed and fixed.
• Local unit tests have been verified and passed.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request effectively addresses a numerical instability issue in the mHC kernel by introducing a configurable sinkhorn_eps parameter for the Sinkhorn-Knopp algorithm. The changes are consistently applied through the operator and kernel layers, and the corresponding unit test has been updated. My main feedback is that adding the new parameter without a default value constitutes a breaking API change. I have provided suggestions to add a default value to ensure backward compatibility for existing users of this operator.