metrale_model_layers/layers/
ple.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: PLE, hashed n-gram injection into the hyper-connection
4//! highway: the layer (`layer`, with its per-sequence carry in
5//! `layer/aux_state`), its ids (`ids`) and debug taps (`dump`). The
6//! reference forward, `Qwen4ExpTextPLELayer.forward` in
7//! `bench/qwen4_exp/ref/modeling_qwen4_exp.py`:
8//!
9//! ```text
10//! embeddings   = ple_embedding(input_ids)                  # [T, ple_embed_dim]
11//! key_normed   = norm_key(key_proj(emb)) -> [T, hc, H]
12//! value        = value_proj(emb)                           # [T, H]
13//! query_normed = norm_query(hidden)      -> [T, hc, H]     # hidden is [T, hc * H]
14//! gate  = (key_normed * query_normed).sum(-1) / sqrt(H)    # [T, hc]
15//! gate  = sign(gate) * sqrt(max(|gate|, 1e-6))             # signed square root
16//! gated = sigmoid(gate) * value                            # [T, hc, H]
17//! out   = gated.flatten() + silu(conv1d(norm_conv(gated.flatten())))
18//! ```
19//!
20//! The reference decoder layer adds the result to its input first
21//! (`hidden_states = hidden_states + self.ple(...)`). Three details:
22//!
23//! 1. The gate takes a signed square root, `sign(g) * sqrt(max(|g|, 1e-6))`.
24//! 2. `conv1d` is depthwise (`groups` equals its channel count) and dilated
25//!    (`dilation = ngram_size`), so the carried state is
26//!    `(kernel_size - 1) * dilation` steps.
27//! 3. All three norms use the `normed * (1 + w)` form and are grouped with
28//!    `group_size = hidden_size`.
29//!
30//! Owner: model-layers (PLE).
31//! Invariants: none beyond the types.
32
33#[path = "ple/ids.rs"]
34pub mod ids;
35
36#[cfg(test)]
37#[path = "ple/tests.rs"]
38mod tests;
39
40#[path = "ple/dump.rs"]
41pub mod dump;
42
43#[path = "ple/layer.rs"]
44mod layer;
45
46pub use ids::{PleIdDims, ple_ngram_ids};
47pub use layer::{PleLayer, PleSeqState, PleWeights};
metrale_model_layers/layers/ple.rs

metrale_model_layers/layers/
ple.rs
