heir
heir copied to clipboard
Use generalized versions of Paterson-Stockmeyer that operate in the Chebyshev basis
The current polynomial approximation solver has an issue with high-degree approximations of some functions. For example, approximating ReLU with degree > 45 produces exponentially bad results as the degree grows.
The root cause of the problem seems to be the conversion from the Chebyshev basis back to the monomial basis when writing the result of the approximation back to the IR. The Chebyshev basis polynomials have coefficients (in the monomial basis) that grow exponentially, so converting to the Monomial basis has numerical stability and overflow issues. Having very large or very small coefficients also adds complexity in FHE because one has to encode those coefficients, and this causes additional complexity in the management of fixed-point scaling parameters.
(Some papers, such as https://eprint.iacr.org/2019/688.pdf page 18, suggest simply ignoring Chebyshev terms with very small coefficients, but this wouldn't work in the monomial basis for high-degree monomials because monomial basis elements have large norm while chebyshev norms are bounded)
The larger degree rapidly impacts noise growth and incurs bootstrapping ops, and bootstrapping itself in CKKS requires a sine function evaluation which has relatively high degree (n=~70-130). So improving this is worth the effort. Here's a plan:
- Leave the polynomial approximation in the Chebyshev basis (maybe with a special case for the polynomial attribute since the parser/printer won't make sense anymore)
- Add alongside a Clenshaw algorithm for evaluating polynomials, which operates directly in the Chebyshev basis and is an analogue of Horner's method.
- Add two Paterson-Stockmeyer analogues from the literature:
- Chen-Chilloti-Song 18 (https://eprint.iacr.org/2018/1043) Section 4.3, Algorithm 2 with the "long division Chebyshev analogue"
- Han-Ki-19 (https://eprint.iacr.org/2019/688) Section 4.2 Algorithm 1 ("Baby Step Giant Step"), which is claimed to be better than the P-S generalization at lower degree.
- Vos-Conti-Erkin 24 has some more in depth analysis of the two methods above https://eprint.iacr.org/2024/1200
This issue has 1 outstanding TODOs:
- lib/Transforms/LowerPolynomialEval/LowerPolynomialEval.td:26: Add support for Chebyshev-basis methods
This comment was autogenerated by todo-backlinks
@j2kun Hello Jeremy, are u working on this issue? If not, maybe I can take it from now on?
There is another Google engineer (@dvadym ) who is doing some of this. See for example https://github.com/google/heir/pull/1772 and https://github.com/google/heir/pull/1770
I believe he is going to work on (a simplified version of) Chen-Chilloti-Song 18's method (hence the Chebyshev decomposition PR. He has a reference implementation in Python he is porting to HEIR.
@dvadym is there a part of this you think @zmbekdemir should take on? I have not personally looked into the differences between Chen-Chilloti-Song 18 and Han-Ki-19 in enough detail to say how they are different and how they compare to your simplified method.