Skip to main content

qualia_core_db/solvers/
activation.rs

1//! Activation & normalization functions — the STEM definitions of the element-wise and
2//! reduction operations a transformer forward pass is built from.
3//!
4//! These are not a proprietary "AI engine"; they are standard mathematics:
5//! - **activations** (ReLU, sigmoid, tanh, SiLU/Swish, GELU) — element-wise nonlinear maps;
6//! - **softmax** — the normalized exponential, a projection onto the probability simplex;
7//! - **RMS / layer normalization** — statistical rescaling (variance / mean-and-variance).
8//!
9//! This module is the canonical, inspectable home for that math. The LLM runtime
10//! (`gguf_bridge`) holds inline `f32` hot-path versions and promoted GPU kernels; those are
11//! *backends* of these definitions and are checked against them (the same arrangement proved
12//! for GEMM in `solvers::linear_algebra::gemm`). `gguf` itself is only a weight *file format* —
13//! the mathematics lives here.
14//!
15//! All functions operate **in place on caller-owned `f64` slices** — zero allocation.
16
17/// ReLU: `max(0, x)`, element-wise.
18pub fn relu(x: &mut [f64]) {
19    for v in x.iter_mut() {
20        if *v < 0.0 {
21            *v = 0.0;
22        }
23    }
24}
25
26/// Logistic sigmoid: `σ(x) = 1 / (1 + e^{-x})`, element-wise.
27pub fn sigmoid(x: &mut [f64]) {
28    for v in x.iter_mut() {
29        *v = 1.0 / (1.0 + (-*v).exp());
30    }
31}
32
33/// Hyperbolic tangent, element-wise.
34pub fn tanh(x: &mut [f64]) {
35    for v in x.iter_mut() {
36        *v = v.tanh();
37    }
38}
39
40/// SiLU / Swish: `x · σ(x) = x / (1 + e^{-x})`, element-wise (Llama/SmolLM2 gate activation).
41pub fn silu(x: &mut [f64]) {
42    for v in x.iter_mut() {
43        *v = *v / (1.0 + (-*v).exp());
44    }
45}
46
47/// GELU (Gaussian Error Linear Unit), tanh approximation:
48/// `0.5·x·(1 + tanh(√(2/π)·(x + 0.044715·x³)))`. The standard GPT-2/transformer GELU.
49pub fn gelu(x: &mut [f64]) {
50    const C: f64 = 0.797_884_560_802_865_4; // sqrt(2/π)
51    for v in x.iter_mut() {
52        let x3 = *v * *v * *v;
53        *v = 0.5 * *v * (1.0 + (C * (*v + 0.044_715 * x3)).tanh());
54    }
55}
56
57/// Softmax in place: `softmax(x)_i = e^{x_i} / Σ_j e^{x_j}`, computed in the numerically
58/// stable shifted form `e^{x_i − max} / Σ e^{x_j − max}`. After the call `x` sums to 1
59/// (a probability distribution). A length-0 slice is left unchanged.
60pub fn softmax(x: &mut [f64]) {
61    if x.is_empty() {
62        return;
63    }
64    let mut max = f64::NEG_INFINITY;
65    for &v in x.iter() {
66        if v > max {
67            max = v;
68        }
69    }
70    let mut sum = 0.0;
71    for v in x.iter_mut() {
72        *v = (*v - max).exp();
73        sum += *v;
74    }
75    if sum > 0.0 {
76        let inv = 1.0 / sum;
77        for v in x.iter_mut() {
78            *v *= inv;
79        }
80    }
81}
82
83/// RMS normalization in place: `x_i ← (x_i / sqrt(mean(x²) + eps)) · weight_i`.
84/// No mean subtraction (the Llama/transformer RMSNorm). `weight` must match `x` in length;
85/// shorter is honoured up to the common length.
86pub fn rms_norm(x: &mut [f64], weight: &[f64], eps: f64) {
87    let n = x.len().min(weight.len());
88    if n == 0 {
89        return;
90    }
91    let mut ss = 0.0;
92    for i in 0..n {
93        ss += x[i] * x[i];
94    }
95    let inv_rms = 1.0 / (ss / n as f64 + eps).sqrt();
96    for i in 0..n {
97        x[i] = x[i] * inv_rms * weight[i];
98    }
99}
100
101/// Layer normalization in place: `x_i ← ((x_i − μ) / sqrt(σ² + eps)) · weight_i + bias_i`,
102/// with `μ`, `σ²` the mean and (population) variance over `x`. `weight`/`bias` match `x`.
103pub fn layer_norm(x: &mut [f64], weight: &[f64], bias: &[f64], eps: f64) {
104    let n = x.len().min(weight.len()).min(bias.len());
105    if n == 0 {
106        return;
107    }
108    let mut mean = 0.0;
109    for i in 0..n {
110        mean += x[i];
111    }
112    mean /= n as f64;
113    let mut var = 0.0;
114    for i in 0..n {
115        let d = x[i] - mean;
116        var += d * d;
117    }
118    var /= n as f64;
119    let inv_std = 1.0 / (var + eps).sqrt();
120    for i in 0..n {
121        x[i] = (x[i] - mean) * inv_std * weight[i] + bias[i];
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    fn approx(a: f64, b: f64, tol: f64) {
130        assert!((a - b).abs() < tol, "{a} != {b} (tol {tol})");
131    }
132
133    #[test]
134    fn relu_clamps_negatives() {
135        let mut x = [-2.0, -0.1, 0.0, 0.5, 3.0];
136        relu(&mut x);
137        assert_eq!(x, [0.0, 0.0, 0.0, 0.5, 3.0]);
138    }
139
140    #[test]
141    fn sigmoid_known_values() {
142        let mut x = [0.0, 1.0, -1.0];
143        sigmoid(&mut x);
144        approx(x[0], 0.5, 1e-12);
145        approx(x[1], 1.0 / (1.0 + (-1.0f64).exp()), 1e-12);
146        approx(x[2], 1.0 / (1.0 + 1.0f64.exp()), 1e-12);
147    }
148
149    #[test]
150    fn silu_equals_x_times_sigmoid() {
151        let mut x = [1.5, -0.7, 2.0];
152        let orig = x;
153        silu(&mut x);
154        for i in 0..3 {
155            let s = 1.0 / (1.0 + (-orig[i]).exp());
156            approx(x[i], orig[i] * s, 1e-12);
157        }
158        // SiLU(0) = 0.
159        let mut z = [0.0];
160        silu(&mut z);
161        approx(z[0], 0.0, 1e-12);
162    }
163
164    #[test]
165    fn gelu_zero_and_monotone() {
166        let mut x = [0.0];
167        gelu(&mut x);
168        approx(x[0], 0.0, 1e-12); // GELU(0) = 0
169                                  // Large positive ≈ identity, large negative ≈ 0.
170        let mut big = [10.0, -10.0];
171        gelu(&mut big);
172        approx(big[0], 10.0, 1e-3);
173        approx(big[1], 0.0, 1e-3);
174    }
175
176    #[test]
177    fn softmax_sums_to_one_and_orders() {
178        let mut x = [1.0, 2.0, 3.0];
179        softmax(&mut x);
180        approx(x.iter().sum::<f64>(), 1.0, 1e-12);
181        assert!(x[2] > x[1] && x[1] > x[0]); // monotone in the inputs
182                                             // Uniform inputs ⇒ uniform distribution.
183        let mut u = [5.0, 5.0, 5.0, 5.0];
184        softmax(&mut u);
185        for &v in &u {
186            approx(v, 0.25, 1e-12);
187        }
188    }
189
190    #[test]
191    fn softmax_is_shift_invariant_and_stable() {
192        let mut a = [1.0, 2.0, 3.0];
193        let mut b = [1.0 + 1000.0, 2.0 + 1000.0, 3.0 + 1000.0];
194        softmax(&mut a);
195        softmax(&mut b); // would overflow without the max-shift
196        for i in 0..3 {
197            approx(a[i], b[i], 1e-12);
198        }
199    }
200
201    #[test]
202    fn rms_norm_scales_to_unit_rms() {
203        // With unit weights, the output RMS is 1 (for eps→0).
204        let mut x = [3.0, -4.0, 0.0, 5.0];
205        let w = [1.0, 1.0, 1.0, 1.0];
206        rms_norm(&mut x, &w, 0.0);
207        let ms = x.iter().map(|v| v * v).sum::<f64>() / 4.0;
208        approx(ms.sqrt(), 1.0, 1e-12);
209    }
210
211    #[test]
212    fn layer_norm_zero_mean_unit_var() {
213        let mut x = [1.0, 2.0, 3.0, 4.0];
214        let w = [1.0, 1.0, 1.0, 1.0];
215        let b = [0.0, 0.0, 0.0, 0.0];
216        layer_norm(&mut x, &w, &b, 0.0);
217        let mean = x.iter().sum::<f64>() / 4.0;
218        approx(mean, 0.0, 1e-12);
219        let var = x.iter().map(|v| (v - mean) * (v - mean)).sum::<f64>() / 4.0;
220        approx(var, 1.0, 1e-12);
221    }
222}