ml-kem: Don't use FieldElement for compressed values

FieldElement should specifically represent a member of $Z_q$ (can't figure out how to get GitHub to do \mathbb). Right now, it is also used as the output of the compress() method as well, even though these are elements of $Z_{2^d}$. We should define a separate type for this latter type of number, something like CompressedValue<D: CompressionFactor>

Just started on a branch to implement this! Would love to be assigned for this if that makes sense--hoping to have a full ready-to-review PR up in the next weekend or two.

@pqcfox

Cool, thanks for picking this up @pqcfox. If you're interested in a slightly bigger project in this zone, the ml-dsa crate has a "module_lattice" module that I isolated with the idea of applying it to ML-KEM as well. You could take a stab at porting ML-KEM to that framework, though we might need help from @tarcieri to actually get the shared crate set up. If you do go that way, let me know; I have a partial starting point I could share.

Ooh, that sounds fantastic!

Would it make sense to you for me to go ahead and complete this issue as described, then move the Compress trait and new implementations over into the module_lattice module once the shared crate is ready?

Thanks!

Yeah, I think probably. You should take a look at the module_lattice module to make sure I'm not misremembering, but I think its field element representation would also be sad about representing compressed values. So this work would be complementary.

Actually, it looks like there are encoding routines in there, so maybe more thought is needed. Let me follow up once I'm back at my computer in a bit.

@pqcfox

OK, having glanced at the code and thought about this for a moment:

The module_lattice module defines:

A Field trait that lets you define a representation of a finite field
A struct Elem<F: Field> for elements of that field, and similar structs for doing linear algebra on polynomials with Elem coefficients.
An Encode trait for encoding polynomials.

So it's tempting to say that the right order of operations here is:

Move ml_kem over to use module_lattice, including using its Encode for encoding
Define a new Field to represent $Z/2^dZ$
Change the compress() methods in ml_kem so that their range is the new field, not $Z/qZ$
Verify that the Encode traits produce the expected values with this new scheme

The main thing that gives me pause there is that the Mul implementation will be wrong for fields that do not have prime order (in particular, binary fields). It's not a practical issue, since ML-KEM doesn't actually do anything with compressed values except encode them. If we wanted to be extra safe, we could define a trait PrimeField: Field trait and only enable Mul when F: PrimeField.

@pqcfox - Sorry if that sounds a little more involved than this issue might have appeared at the start. Totally understand if you're not interested in the bigger project! In any case, I went ahead and filed #112 to track the module_lattice transition.

@bifurcation

Ooh, I like this idea a lot, esp. using a trait to gate the Mul implementation. Would be very happy to give this a shot once the module_lattice repo is ready for PRs!

Would it make sense to rename Field in this case e.g. to Ring to accommodate $\mathbb{Z}/2^d\mathbb{Z}$?

Thanks so much @bifurcation!

Hm. I was going to object that $Z / 2^d Z$ is a field, just with a different multiplication algorithm. But actually $Z / 2^d Z$ is not the same as $GF(2^d)$, which is what I had in mind. But this doesn't actually matter in the context of ML-KEM, since the only thing ML-KEM does with these values is encode/decode them.

So maybe the right answer is actually to define another type that just represents "a value that fits in $d$ bits". Something like FixedWidthInt<F, D> where F: Field, D: EncodingSize. Then you would use this as the input and output of the encode/decode methods.

The trouble with this approach will be avoiding copies, since you'll want to be able to convert between Array<Elem<F>, U256> and Array<FixedWidthInt<F, D>, U256>. In principle, this should be doable, since Elem and FixedWidthInt will have the same representation in memory (an F::Int). But it might require transmute or something similar.

That is, you would want to be able to do the following:

For ML-KEM:

Array<Elem> ---(compress)--> Array<FixedWidthInt> --(encode)-> Array<u8>
Array<Elem> <-(decompress)-- Array<FixedWidthInt> <-(decode)-- Array<u8>

For ML-DSA

Array<Elem> --(transmute)-> Array<FixedWidthInt> --(encode)-> Array<u8>
Array<Elem> <-(transmute)-- Array<FixedWidthInt> <-(decode)-- Array<u8>

Of course, maybe just having the "transmute" be array.map(|x| x.into()) is cheap enough.

This approach seems perfect to me! Thanks for detailing this out :)

I'll go ahead and move forward with this FixedWidthInt strategy once the module_lattice crate is in place, excited to dive in on this!

Nit: would probably suggest naming the crate module-lattice

Not sure where it should go but perhaps it can be developed in the KEMs tree for now and extracted somewhere else

@tarcieri

Sounds like a plan! @tarcieri, would it make sense for me to go ahead and add a module-lattice crate myself to this repo as part of the draft PR for this issue, then before full review all the details of publishing, where it'll end up, etc. can be figured out?

@bifurcation

Sounds good to me, but double check with @bifurcation

@tarcieri

Yeah, that sounds fine to me. Just elaborating the plan a little bit, it seems like the following is roughly the right order of operations:

Locally:
- Clone signatures and KEMs
- Move the module-lattice code to a new crate under KEMs
- Update ml-kem to use module-lattice (incl. any changes to module-lattice needed, hopefully none!)
- Update ml-dsa to refer to the new module-lattice crate (incl. any changes)
Then make PRs:
1. PR on KEMs to add module-lattice crate and update ml-kem
2. PR on signatures to convert to module-lattice crate

Basically, since we have a cross-repo dependency, seems like a good idea to do cross-repo testing before we start making PRs.

@tarcieri - It seems like maybe we should publish the module-lattice crate in between (1) and (2) above, so that ml-dsa can refer to it via crates.io. Does that sound right to you? And if so, should we agree on its ultimate home before we do that? I think you had mentioned the utils repo before.

@bifurcation

@bifurcation that all sounds good to me. we could also consider a new repo for module-lattice rather than utils

Ooh, perfect--if I wanted to get started on this, would it make okay for me to instead then make a draft PR on utils extracting module-lattice (and then optionally we could copy it to a new module-lattice repo instead if desired)?

@tarcieri

I would be fine with it going in utils, as it's not a huge chunk of code. But @tarcieri is the authority here. Like I said before, my main concern here is just avoiding a bunch of cross-repo complexity. Maybe the thing to do is:

Make all the code changes locally
Open PRs on utils, ml-dsa, and ml-kem and review them in parallel
Merge the utils PR and publish the crate
Update the ml-dsa and ml-kem PRs to point to the new crate, then merge them

Okay, we can start with utils, I just perhaps have PTSD from a number of other crates we started there and eventually extracted

Fair point, I'll go ahead and make a separate module-lattice crate locally. For getting that available for review, should I just create a fresh repo under my username, or wait to push that until there's a module-lattice crate here?

@pqcfox

@pqcfox you can keep it in the KEMs repo for now, which should simplify development

Small update, am still working on this--most of the slowdown is coming from running into various issues with external traits and orphan rule workarounds, but progress has been steady :)

@pqcfox

@pqcfox if you're still interested in this stuff, I got a module-lattice crate extracted (and used by ml-dsa), putting it here for now: #199

I have closed #112 now that we have fully migrated to the module-lattice crate.

@pqcfox

@pqcfox Are you still working this? If not, I could have a PR for this soon.

@MavenRain

@MavenRain I don't think anyone's working on it if you'd like to

module-lattice, ml-kem: introduce FixedWidthInt for compressed values

Navigation Menu

Repository navigation

ml-kem: Don't use FieldElement for compressed values #26

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

ml-kem: Don't use FieldElement for compressed values #26

Description

Activity

pqcfox commented on Feb 25, 2025

bifurcation commented on Feb 25, 2025

pqcfox commented on Feb 25, 2025

bifurcation commented on Feb 25, 2025

bifurcation commented on Feb 25, 2025

bifurcation commented on Feb 26, 2025

pqcfox commented on Feb 26, 2025

bifurcation commented on Feb 26, 2025

bifurcation commented on Feb 26, 2025

pqcfox commented on Feb 27, 2025

tarcieri commented on Feb 27, 2025

pqcfox commented on Feb 27, 2025

tarcieri commented on Feb 27, 2025

bifurcation commented on Feb 28, 2025

tarcieri commented on Feb 28, 2025

pqcfox commented on Mar 2, 2025

bifurcation commented on Mar 2, 2025

tarcieri commented on Mar 2, 2025

pqcfox commented on Mar 8, 2025

tarcieri commented on Mar 8, 2025

pqcfox commented on Apr 24, 2025

tarcieri commented on Jan 27, 2026

tarcieri commented on Jan 29, 2026

MavenRain commented on Apr 28, 2026

tarcieri commented on Apr 29, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions