Sienna Meridian Satterwhite
3628424f43
Add documentation comments linking Lean 4 specifications to their corresponding Rust implementations (tree_predict, mlp_predict_32, PackedNode). Reference TorchLean arXiv ID for future axiom replacement. Signed-off-by: Sienna Meridian Satterwhite <sienna@sunbeam.pt>
35 lines
1.2 KiB
Lean4
35 lines
1.2 KiB
Lean4
import Sunbeam.Model.Basic
|
|
import Sunbeam.Model.Sigmoid
|
|
import Sunbeam.Model.ReLU
|
|
|
|
namespace Sunbeam
|
|
|
|
/-- Weights for a 2-layer MLP (input → hidden → scalar output).
|
|
|
|
Corresponds to `ensemble::mlp::mlp_predict_32` in Rust, which uses const generic
|
|
`INPUT` and a fixed hidden dimension of 32. Here `hiddenDim` is a parameter so
|
|
structural properties can be proved generically. -/
|
|
structure MLPWeights (inputDim hiddenDim : Nat) where
|
|
w1 : Fin hiddenDim → FloatVec inputDim
|
|
b1 : FloatVec hiddenDim
|
|
w2 : FloatVec hiddenDim
|
|
b2 : Float
|
|
|
|
/-- Forward pass: input → linear1 → ReLU → linear2 → sigmoid. -/
|
|
def mlpForward {inputDim hiddenDim : Nat}
|
|
(weights : MLPWeights inputDim hiddenDim) (input : FloatVec inputDim) : Float :=
|
|
let hidden := vecAdd (matVecMul weights.w1 input) weights.b1
|
|
let activated := reluVec hidden
|
|
let output := dot weights.w2 activated + weights.b2
|
|
sigmoid output
|
|
|
|
/-- MLP output is bounded in (0, 1) — follows directly from sigmoid bounds. -/
|
|
theorem mlp_output_bounded {inputDim hiddenDim : Nat}
|
|
(weights : MLPWeights inputDim hiddenDim) (input : FloatVec inputDim) :
|
|
0 < mlpForward weights input ∧ mlpForward weights input < 1 := by
|
|
constructor
|
|
· exact sigmoid_pos _
|
|
· exact sigmoid_lt_one _
|
|
|
|
end Sunbeam
|