diff --git a/CHANGELOG.md b/CHANGELOG.md index c2d43be8..3e741012 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,8 @@ Please view [SomeB1oody/RustyML](https://github.com/SomeB1oody/RustyML) for more - **`LayerPath`, the address of 1 layer inside a model.** It holds the model position and the sublayer names below it, and its text form is `2.forward`. `Ctx::layer_path` gives the current path. `Ctx::state_left_under` lists the state values of a model position that no layer took back. `neural_network` re-exports `LayerPath`, `Sublayer`, `SublayerMut`, and `SublayerName`. - **Public tree walks for a layer that a caller drives by hand.** `layer_path::walk` and `walk_mut` visit every node of 1 layer tree in pre-order. `update_tree` runs `Optimizer::update` on every node, and `apply_state_tree` moves the state of every node. `Ctx::state_left_below` lists the state values at a path, or below it, that no layer took back. - **`layer_path::total_param_count` sums the parameters of a layer and all of its sublayers.** `Sequential::summary` and `Graph::summary` use it. `LayerBase::param_count` still counts the arrays of 1 layer alone. +- **3 new activations: GELU, SiLU, and Mish.** `Activation::GELU { approximate }` is `x * Phi(x)`, where `Phi` is the standard normal cumulative distribution function. With `approximate: true` it uses the tanh approximation `x * sigmoid(2 * sqrt(2/pi) * (x + 0.044715 * x^3))`. `Activation::SiLU`, also named Swish, is `x * sigmoid(x)`. `Activation::Mish` is `x * tanh(softplus(x))`. Each one is also a standalone layer: `GELU`, `SiLU`, and `Mish`. `GELU::new()` gives the exact form, and `with_approximate(true)` selects the tanh approximation. `Dense`, all 10 convolution layers, `SimpleRNN`, `LSTM`, and `GRU` accept the 3 new activations. +- **`ActivationInput` and `ActivationCache` name the tensor that the backward pass of an activation reads.** `Activation::saves()` returns `ActivationInput::PreActivation` for GELU, SiLU, and Mish, and `ActivationInput::Output` for every other variant. `Activation::forward_train(z)` returns the activated output and an `ActivationCache` that holds the saved tensor. `Activation::output_of(&cache)` gives the activated output that a cache stands for. A host layer caches 1 activation tensor, never both. ### Changed - **Breaking: `ParamId` is `{ layer: LayerPath, name }` in place of `{ scope, name }`, and it is no longer `Copy`.** `ParamId::new(scope, name)` still addresses the layer at a model position. `ParamId::at(path, name)` and `LayerPath::param(name)` address a sublayer. `ParamId` implements `Display` and prints its checkpoint path. @@ -21,6 +23,8 @@ Please view [SomeB1oody/RustyML](https://github.com/SomeB1oody/RustyML) for more - **Breaking: `Ctx::has_state` and `Ctx::state_slot` take `&LayerPath` in place of a model position.** Pass `&LayerPath::root(scope)` for the layer at a model position. - **Breaking: `MODEL_FORMAT_VERSION` is 4, and no version 3 checkpoint loads.** `LayerCheckpoint` gains a `sublayers` field of the new `SublayerCheckpoint` records, so a file holds the same tree as the model. A checkpoint path is `.`, such as `2.kernel` or `2.forward.kernel`. A layer with no sublayer keeps the path it had in version 3. Rebuild the model, retrain, and save again. - **Breaking: `weight_path` moves from `layers::checkpoint` to `layer_path`, and it takes a `&LayerPath` in place of a model position.** +- **Breaking: `Activation::backward` takes `&ActivationCache` in place of the activated output tensor.** Get the cache from `Activation::forward_train(z)`. The inference path still calls `Activation::forward(&z)`, which makes no cache and no copy. +- **Behavior change: `Activation::backward` returns an error for a bad argument.** It returns `Error::InvalidInput` for a cache that holds the other tensor than `Activation::saves()` names. It returns `Error::ShapeMismatch` for a gradient whose shape is not the shape of the cache. Before, a gradient of another shape caused a panic in most variants. - **Behavior change: a state value that no layer takes back stops the training step.** `Sequential::train_batch` and `Graph::train_batch` return `Error::Computation`, and the message names each value. Before, a debug assertion caught this case, and a release build ignored the value. - **Behavior change: `UnaryLayer::forward_mut` and `Layer::forward_many_mut` return `Error::Computation` for a state value that no layer takes back.** - **Behavior change: a training step refuses a sublayer call that reaches another layer than the tree holds at its path.** A training pass records the layer of each `Ctx::sublayer` call. The model compares each record against the tree after the forward pass and after the backward pass. A call of sublayer `b` under the name of sublayer `a` therefore stops the step, also when `a` and `b` have the same type and shape. diff --git a/Cargo.lock b/Cargo.lock index 3fa97a5d..91752af0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -833,6 +833,7 @@ dependencies = [ "criterion", "gemmkit-ndarray", "indicatif", + "libm", "nalgebra", "ndarray", "ndarray-rand", diff --git a/Cargo.toml b/Cargo.toml index 4540cff8..79fc48d1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,7 @@ autobenches = false # Classical machine learning algorithms for supervised and unsupervised learning machine_learning = ["math", "dep:ndarray", "dep:rayon", "dep:ndarray-rand", "dep:ahash", "dep:serde", "dep:postcard", "dep:thiserror"] # Complete neural network framework with flexible architecture design -neural_network = ["math", "dep:ndarray", "dep:rayon", "dep:ndarray-rand", "dep:indicatif", "dep:serde", "dep:postcard", "dep:thiserror"] +neural_network = ["math", "dep:ndarray", "dep:rayon", "dep:ndarray-rand", "dep:indicatif", "dep:libm", "dep:serde", "dep:postcard", "dep:thiserror"] # Data preprocessing and dimensionality reduction utilities utils = ["math", "dep:ndarray", "dep:rayon", "dep:ahash", "dep:serde", "dep:postcard", "dep:ndarray-rand", "dep:thiserror"] # Comprehensive evaluation metrics for model performance assessment @@ -44,6 +44,7 @@ serde = { version = "1.0.229", features = ["derive"], optional = true } postcard = { version = "1.1.3", features = ["use-std"], optional = true } thiserror = { version = "2.0.20", optional = true } gemmkit-ndarray = { version = "0.1.2", features = ["epilogue"], optional = true } +libm = { version = "0.2.16", optional = true } [dev-dependencies] approx = "0.5.1" diff --git a/guide/en/src/Chapter-01/1.5._The_Prelude_and_Imports.md b/guide/en/src/Chapter-01/1.5._The_Prelude_and_Imports.md index a63e8d88..b70a4eb2 100644 --- a/guide/en/src/Chapter-01/1.5._The_Prelude_and_Imports.md +++ b/guide/en/src/Chapter-01/1.5._The_Prelude_and_Imports.md @@ -66,7 +66,7 @@ The `prelude` is a hand-picked list: it re-exports only the items you reach for | Graph model | `GraphBuilder` (collects nodes and layers for a non-chain topology), `Graph` (what `build` gives back), `NodeId` (a handle to 1 node of the graph) | | Core layers | `Dense`, `Embedding`, `Flatten`, `Identity`, `Reshape`, `Reverse`, `Permute`, `RepeatVector`, `Rescaling`, `Activation` | | Merge layers | `Add`, `Average`, `Concatenate`, `Maximum`, `Minimum`, `Multiply`, `Subtract` | -| Activation layers | `Linear`, `ReLU`, `LeakyReLU`, `ELU`, `SELU`, `Sigmoid`, `HardSigmoid`, `Tanh`, `Softplus`, `Softsign`, `Exponential`, `Softmax`, `PReLU` | +| Activation layers | `Linear`, `ReLU`, `LeakyReLU`, `ELU`, `SELU`, `Sigmoid`, `HardSigmoid`, `Tanh`, `Softplus`, `Softsign`, `Exponential`, `GELU`, `SiLU`, `Mish`, `Softmax`, `PReLU` | | Convolution | `Conv1D`, `Conv2D`, `Conv3D`, `Conv1DTranspose`, `Conv2DTranspose`, `Conv3DTranspose`, `DepthwiseConv1D`, `DepthwiseConv2D`, `SeparableConv1D`, `SeparableConv2D`, `PaddingType`, `ConvPadding` | | Border | `ZeroPadding1D/2D/3D`, `Cropping1D/2D/3D`, `Border1D/2D/3D` | | Pooling | `MaxPooling1D/2D/3D`, `AveragePooling1D/2D/3D`, `GlobalMaxPooling1D/2D/3D`, `GlobalAveragePooling1D/2D/3D` | diff --git a/guide/en/src/Chapter-03/3.2._Dense_Layers_and_Activations.md b/guide/en/src/Chapter-03/3.2._Dense_Layers_and_Activations.md index 8933de1d..15912fe5 100644 --- a/guide/en/src/Chapter-03/3.2._Dense_Layers_and_Activations.md +++ b/guide/en/src/Chapter-03/3.2._Dense_Layers_and_Activations.md @@ -2,7 +2,7 @@ `Dense` is the main layer of a feedforward network. It applies a linear map, `input * W + b`, then an elementwise nonlinearity. RustyML fuses that nonlinearity into the layer itself. It does not treat the nonlinearity as a separate stage. -This page covers the constructor and its fused-activation design. It also covers the exact initialization scheme and the 12 activations, with their gradients and failure modes. It covers the GEMM-backed compute path and how to read or inject weights. It assumes you have already read [the Sequential model page](./3.1._The_Sequential_Model.md). The layers here are what you collect in a `SequentialBuilder`, and the input width of each one comes from the model build. +This page covers the constructor and its fused-activation design. It also covers the exact initialization scheme and the 15 activations, with their gradients and failure modes. It covers the GEMM-backed compute path and how to read or inject weights. It assumes you have already read [the Sequential model page](./3.1._The_Sequential_Model.md). The layers here are what you collect in a `SequentialBuilder`, and the input width of each one comes from the model build. ## 3.2.1. The fused-activation design @@ -12,13 +12,13 @@ The constructor takes the activation as its second argument, not as a separate l pub fn new(units: usize, activation: impl Into) -> Result ``` -Internally the layer stores an `Activation` *value*. This is a plain `Copy` enum with 12 variants. It is not a generic type parameter `Dense`, and it is not a `Box`. That choice is deliberate. A generic parameter would monomorphize `Dense` 12 ways. +Internally the layer stores an `Activation` *value*. This is a plain `Copy` enum with 15 variants. It is not a generic type parameter `Dense`, and it is not a `Box`. That choice is deliberate. A generic parameter would monomorphize `Dense` 15 ways. It would also force weight deserialization to probe every `Dense` pairing to find the concrete type on disk. A trait object would add an indirection to every elementwise call. The runtime enum keeps `Dense` a single concrete type. The persistence layer can then downcast every saved layer to exactly 1 struct (see [Saving and Loading Weights](./3.9._Saving_and_Loading_Weights.md)). The activation math stays a pure, stateless function, and the layer calls it inside its own forward and backward passes. Keras makes the same fusion choice, with `Dense(units, activation="relu")`. RustyML has no implicit "no activation". You must always pass one. To get a pure linear layer, pass `Activation::Linear`, the identity. This is the equivalent of Keras' `activation=None`. Every example below that ends in a regression head does exactly this. -Fusing is the right default. RustyML also provides every activation as a standalone layer (`Linear`, `ReLU`, `LeakyReLU`, `ELU`, `SELU`, `Sigmoid`, `HardSigmoid`, `Tanh`, `Softplus`, `Softsign`, `Exponential`, `Softmax`). Use these when you need something between the linear map and the nonlinearity. A normalization layer is the most common case (see [Regularization and Normalization Layers](./3.8._Regularization_and_Normalization_Layers.md)). The 2 forms are equivalent: +Fusing is the right default. RustyML also provides every activation as a standalone layer (`Linear`, `ReLU`, `LeakyReLU`, `ELU`, `SELU`, `Sigmoid`, `HardSigmoid`, `Tanh`, `Softplus`, `Softsign`, `Exponential`, `GELU`, `SiLU`, `Mish`, `Softmax`). Use these when you need something between the linear map and the nonlinearity. A normalization layer is the most common case (see [Regularization and Normalization Layers](./3.8._Regularization_and_Normalization_Layers.md)). The 2 forms are equivalent: ```rust,ignore // These two stacks compute the same thing. @@ -28,7 +28,7 @@ builder.add(Dense::new(32, Activation::Linear).unwrap()) // split: identit .add(ReLU::new()); // ... then a standalone ReLU ``` -Prefer the fused form. It caches a single activated tensor instead of 2. The backward pass then differentiates the activation in terms of that cached output (see 3.2.6). Use the split form only when a layer must sit between `W*x + b` and the nonlinearity. +Prefer the fused form. It caches 1 activation tensor in place of 2. That tensor is the activated output for most activations, and the pre-activation for `GELU`, `SiLU`, and `Mish` (see 3.2.4). Use the split form only when a layer must sit between `W*x + b` and the nonlinearity. ## 3.2.2. Constructing and sizing a layer @@ -165,9 +165,9 @@ Either way the seed leaves the global stream untouched. The full seeding model i ## 3.2.4. The activations -`Activation` has 12 variants. Each variant's backward pass is expressed in terms of the *activated output* `a = f(x)`, not the pre-activation `x`. The layer caches that output, not the input. That contract admits every activation whose derivative has a closed form in `a`, which covers all 12 below. +`Activation` has 15 variants. The backward pass of each variant reads exactly 1 saved tensor. For 12 variants that tensor is the *activated output* `a = f(x)`, because their derivative has a closed form in `a`. GELU, SiLU (also named Swish), and Mish have the shape `a = x * g(x)`. That shape has no closed-form inverse, so no derivative in `a` alone exists. These 3 variants therefore save the pre-activation `x`. -The same contract is exactly why GELU, SiLU (Swish), and Mish are absent. Their `a = x * g(x)` shape has no closed-form inverse, so no derivative in `a` alone exists. Adding them needs a wider contract that also hands the backward pass the pre-activation. +A host layer never saves both tensors. 3.2.6 explains the contract that names the saved tensor. The tables below give the backward pass in terms of the saved tensor, for an upstream gradient `g`. | `Activation` | Forward `f(x)` | Backward (given upstream `g`, output `a`) | Output range | | --- | --- | --- | --- | @@ -184,6 +184,15 @@ The same contract is exactly why GELU, SiLU (Swish), and Mish are absent. Their | `HardSigmoid` | `clip(x/6 + 0.5, 0, 1)` | `g / 6` where `0 < a < 1`, else `0` | `[0, 1]` | | `Exponential` | `e^x` | `g * a` | `(0, inf)` | +The 3 variants that save the pre-activation `x` use `sigmoid(x) = 1 / (1 + e^-x)`, `softplus(x) = ln(1 + e^x)`, and the standard normal distribution: `Phi(x)` is its cumulative distribution function, and `phi(x)` is its density. + +| `Activation` | Forward `f(x)` | Backward (given upstream `g`, pre-activation `x`) | Output range | +| --- | --- | --- | --- | +| `GELU { approximate: false }` | `x * Phi(x)` | `g * (Phi(x) + x * phi(x))` | about `[-0.17, inf)` | +| `GELU { approximate: true }` | `x * sigmoid(u)`, with `u = 2 * sqrt(2/pi) * (x + 0.044715 * x^3)` | `g * (s + x * s * (1 - s) * du/dx)`, with `s = sigmoid(u)` | about `[-0.17, inf)` | +| `SiLU` | `x * sigmoid(x)` | `g * s * (1 + x * (1 - s))`, with `s = sigmoid(x)` | about `[-0.28, inf)` | +| `Mish` | `x * tanh(softplus(x))` | `g * (t + x * (1 - t^2) * sigmoid(x))`, with `t = tanh(softplus(x))` | about `[-0.31, inf)` | + **ReLU** is the default hidden-layer choice. It is cheap, and it does not saturate on the positive side. Its failure mode is the *dead neuron*. The derivative is `0` for `x <= 0`, so a neuron whose pre-activation is negative for every example in the batch gets zero gradient. It never updates, and it stays off for good. A high learning rate makes this worse. It pushes neurons into the dead region early. Glorot init has no leaky variant. Your options are a smaller learning rate and well-scaled inputs, or, if you inject your own weights, a better initial scale. @@ -194,7 +203,7 @@ It never updates, and it stays off for good. A high learning rate makes this wor `negative_slope` and `alpha` must both be finite and greater than 0. `Activation::validate` enforces that bound. All 14 trainable layer constructors call it: `Dense`, `Conv1D`, `Conv2D`, `Conv3D`, `Conv1DTranspose`, `Conv2DTranspose`, `Conv3DTranspose`, `DepthwiseConv1D`, `DepthwiseConv2D`, `SeparableConv1D`, `SeparableConv2D`, `SimpleRNN`, `LSTM`, and `GRU`. An unusable value therefore returns `Error::InvalidParameter` where you build the model, not on the first forward pass. -The bound is strict, and the output-only contract is the reason. The backward pass reads the branch off the sign of `a`. A value of 0 collapses the whole negative side onto `a = 0`, which erases the branch. A negative value inverts the sign, so the backward pass reads the wrong branch. Use `Activation::ReLU` when you want a slope of 0. +The bound is strict, because these 2 variants save the output. The backward pass reads the branch off the sign of `a`. A value of 0 collapses the whole negative side onto `a = 0`, which erases the branch. A negative value inverts the sign, so the backward pass reads the wrong branch. Use `Activation::ReLU` when you want a slope of 0. `LeakyReLU` uses `x >= 0` for the positive branch, so its derivative at exactly 0 is 1. `ELU` and `SELU` use `x > 0`, which puts 0 itself on the negative branch. Their derivatives at exactly 0 are therefore `alpha` and `scale * alpha`, not 1. `LeakyReLU::default()` uses a slope of `0.3`, and `ELU::default()` uses an `alpha` of `1.0`. @@ -257,6 +266,31 @@ fn main() { **HardSigmoid** approximates sigmoid with a clipped straight line, and it uses no exponential. It is `clip(x/6 + 0.5, 0, 1)`, so it reaches exactly 0 at `x = -3` and exactly 1 at `x = 3`. True sigmoid only approaches both ends and never arrives. The derivative is the constant `1/6` on the linear segment, and exactly 0 outside it. That flat outer gradient is the same dead-unit risk ReLU carries. Use `HardSigmoid` as a gate where the cost of `exp` matters. +**GELU, SiLU, and Mish** are the smooth, self-gated activations of transformer and modern convolutional models. Each one multiplies `x` by a gate `g(x)` that rises from 0 to 1. Positive inputs pass almost unchanged, and large negative inputs go to 0. Between those 2 regions, each curve dips below 0 and then comes back up, so none of the 3 is monotonic. The derivative is not 0 on the negative side, so a unit whose pre-activation stays negative still receives a small gradient. That gradient goes toward 0 as `x` goes toward `-inf`. + +`GELU` gates `x` by `Phi(x)`, the probability that a standard normal value is below `x` (Hendrycks and Gimpel, 2016). `Activation::GELU { approximate: false }` computes `Phi` exactly. It uses the complementary error function, `Phi(x) = erfc(-x / sqrt(2)) / 2`, which keeps its relative precision in the negative tail. `Activation::GELU { approximate: true }` uses the tanh approximation `0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`. RustyML evaluates it as the equal form `x * sigmoid(2 * sqrt(2/pi) * (x + 0.044715 * x^3))`, which does not cancel in the negative tail. The exact form is the default, as in Keras and PyTorch. Use the approximate form to match a model that was trained with it, such as GPT-2. + +`SiLU` is `x * sigmoid(x)`, the sigmoid-weighted linear unit. Keras also names it `swish`. `Mish` is `x * tanh(softplus(x))` (Misra, 2019). It costs 1 exponential, 1 logarithm, and 1 hyperbolic tangent per element, which makes it the most expensive of the 3. + +The standalone layers are `GELU`, `SiLU`, and `Mish`. `SiLU::new()` and `Mish::new()` take no argument. `GELU::new()` gives the exact form, and the `with_approximate(true)` builder selects the tanh approximation. Each layer converts into the matching `Activation` value, so you can pass the layer itself where a host layer takes an activation: + +```rust,ignore +// Fused: the activation runs inside the host layer. +builder.add(Dense::new(64, Activation::GELU { approximate: false }).unwrap()) + .add(Dense::new(64, Activation::SiLU).unwrap()) + .add(Conv2D::new(16, (3, 3), (1, 1), Activation::Mish).unwrap()); + +// Split: an identity Dense, then the standalone layer. +builder.add(Dense::new(64, Activation::Linear).unwrap()) + .add(GELU::new().with_approximate(true)); + +// The layer converts into its variant, so the 2 lines below build the same Dense. +let a = Dense::new(64, GELU::new().with_approximate(true)).unwrap(); +let b = Dense::new(64, Activation::GELU { approximate: true }).unwrap(); +``` + +All 14 trainable host layers accept the 3 variants: `Dense`, every convolution layer, `SimpleRNN`, `LSTM`, and `GRU`. In a recurrent layer the activation applies to the hidden state or the candidate, as [Recurrent Layers](./3.7._Recurrent_Layers.md) explains. + **Exponential** is `e^x`. Its output is strictly positive, and its derivative is the output itself. Use it on a head that must emit a positive quantity, such as a rate, a variance, or a scale. It has no upper bound, so a pre-activation above about 88 overflows `f32` to `inf`. Keep the head's input scaled, or emit a log-scale value through `Linear` instead. **Softmax** turns a lane of logits into a probability distribution. `Softmax { axis }` names the axis that the lanes run along, and the default is `-1`, the last axis. The forward pass subtracts each lane's maximum before it takes the exponent. This makes the largest term `exp(0) = 1`, and the sum always `>= 1`. That shift makes softmax overflow-proof and shift-invariant: adding a constant to every logit leaves the output unchanged. @@ -413,7 +447,7 @@ The layer holds 1 array, named `alpha`. `weight("alpha")` reads it, and `set_wei ## 3.2.6. Forward, backward, and the GEMM path -The forward pass is `activation(input * W + b)`, computed as a single call into the [gemmkit](https://docs.rs/gemmkit) backend. The product, the per-column bias, and, for `ReLU`, the activation, all run in 1 pass. The bias and activation apply in the kernel's epilogue while the output tile is still in registers. `ReLU` is the only one of the 12 activations with a fused epilogue. `Linear` needs no separate step, since the biased product is already the output. Every other activation runs as a separate vectorized `Activation::forward` pass over the biased product. +The forward pass is `activation(input * W + b)`, computed as a single call into the [gemmkit](https://docs.rs/gemmkit) backend. The product, the per-column bias, and, for `ReLU`, the activation, all run in 1 pass. The bias and activation apply in the kernel's epilogue while the output tile is still in registers. `ReLU` is the only one of the 15 activations with a fused epilogue. `Linear` needs no separate step, since the biased product is already the output. Every other activation runs as a separate vectorized `Activation::forward` pass over the biased product. The fused result matches the unfused product plus scalar activation bit for bit, with 1 exception. The fused `ReLU` epilogue maps a `NaN` pre-activation to `0.0`. The standalone `ReLU` activation used elsewhere in the crate instead propagates `NaN`. A `NaN` pre-activation already means a diverged model. The cached-output `ReLU` derivative then treats that `0.0` like a dead unit. It backpropagates a zero gradient there, instead of propagating the `NaN` further back. @@ -421,17 +455,52 @@ Whether the product threads is gemmkit's decision, not the layer's. It compares Small layers in a tight loop stay serial by design, because the per-call dispatch would dominate. These products also nest safely inside an outer parallel region. The backend is described in full in [Matrix Multiplication](../Chapter-06/6.2._Matrix_Multiplication.md), and the tunable gates in [Performance Tuning and Parallelism](../Chapter-07/7.3._Performance_Tuning_and_Parallelism.md). -The elementwise activation that follows has its own, separate parallelism gate. Each activation falls into 1 of 2 cost classes. `ReLU`, `LeakyReLU`, `Softsign`, and `HardSigmoid` are memory-bound "cheap maps" whose crossover sits at 4,000,000 elements. At any practical layer size, they run serial. `Sigmoid`, `Tanh`, `Softmax`, `ELU`, `SELU`, `Softplus`, and `Exponential` are `exp`-dominated, and go parallel above 131,072 elements. `Linear` copies the tensor and has no gate. +The elementwise activation that follows has its own, separate parallelism gate. Each activation falls into 1 of 2 cost classes. `ReLU`, `LeakyReLU`, `Softsign`, and `HardSigmoid` are memory-bound "cheap maps" whose crossover sits at 4,000,000 elements. At any practical layer size, they run serial. `Sigmoid`, `Tanh`, `Softmax`, `ELU`, `SELU`, `Softplus`, `Exponential`, `GELU`, `SiLU`, and `Mish` are `exp`-dominated, and go parallel above 131,072 elements. `Linear` copies the tensor and has no gate. + +The backward pass uses the same 2 classes, but it picks the class from the *derivative*, not from the forward pass. `Sigmoid`, `Tanh`, `Softmax`, `ELU`, `SELU`, and `Exponential` each compute their derivative from the cached output alone, with no exponential call. Their backward pass therefore uses the cheap-map gate, even though their forward pass uses the `exp` gate. `Softplus` is the exception: its derivative still needs 1 exponential per element, so its backward pass keeps the `exp` gate too. `GELU`, `SiLU`, and `Mish` compute their derivative from the pre-activation, with the same transcendental calls as their forward pass, so their backward pass also keeps the `exp` gate. Moving either gate only trades serial for parallel. The results are identical, and every product is run-to-run deterministic on the same machine. + +`forward` takes `&self`, and it parks the input and 1 activation tensor in the context. `Ctx::training()` selects the mode that parks them, and `Ctx::inference()` selects the mode that parks nothing. The 2 modes compute the same values, and the inference mode saves the cost of the parking. The layer itself holds no cache at all, which is why several threads can run inference through 1 model. The backward pass takes the parked activation tensor back and differentiates the activation with it. -The backward pass uses the same 2 classes, but it picks the class from the *derivative*, not from the forward pass. `Sigmoid`, `Tanh`, `Softmax`, `ELU`, `SELU`, and `Exponential` each compute their derivative from the cached output alone, with no exponential call. Their backward pass therefore uses the cheap-map gate, even though their forward pass uses the `exp` gate. `Softplus` is the 1 exception: its derivative still needs 1 exponential per element, so its backward pass keeps the `exp` gate too. Moving either gate only trades serial for parallel. The results are identical, and every product is run-to-run deterministic on the same machine. +**The activation cache.** `Activation::saves()` returns an `ActivationInput`, which names the tensor that the backward pass reads. It is `ActivationInput::PreActivation` for `GELU`, `SiLU`, and `Mish`, and `ActivationInput::Output` for the other 12 variants. A training pass calls `Activation::forward_train(z)`. That call returns the activated output and an `ActivationCache`. The cache holds a copy of the output, or `z` itself, as `saves()` names. The fused `ReLU` epilogue has no separate `z`, so the layer wraps the output of the product as an `Output` cache. That is correct, because `ReLU` saves the output. -`forward` takes `&self`, and it parks the input and the activated output in the context. `Ctx::training()` selects the mode that parks them, and `Ctx::inference()` selects the mode that parks nothing. The 2 modes compute the same values, and the inference mode saves the cost of the parking. The layer itself holds no cache at all, which is why several threads can run inference through 1 model. The backward pass takes the parked output back and differentiates the activation with it. +`Activation::backward(&cache, &grad_output)` gives the gradient with respect to `z`. It returns `Error::InvalidInput` for a cache that holds the other tensor, and `Error::ShapeMismatch` for a gradient whose shape is not the shape of the cache. `ActivationCache::holds()`, `tensor()`, `shape()`, and `into_tensor()` read the cache. `Activation::output_of(&cache)` gives the activated output that a cache stands for. An inference pass calls `Activation::forward(&z)`, which makes no cache and no copy. + +```rust +use ndarray::Array; +use rustyml::neural_network::layers::{Activation, ActivationInput}; + +fn main() { + let z = Array::from_vec(vec![-2.0_f32, -0.5, 0.0, 1.5]).into_dyn(); + let gelu = Activation::GELU { approximate: false }; + + // GELU saves the pre-activation, and ReLU saves the output. + assert_eq!(gelu.saves(), ActivationInput::PreActivation); + assert_eq!(Activation::ReLU.saves(), ActivationInput::Output); + + // A training pass returns the output and the cache. The inference pass agrees. + let (a, cache) = gelu.forward_train(z.clone()).unwrap(); + assert_eq!(cache.holds(), ActivationInput::PreActivation); + assert_eq!(gelu.forward(&z).unwrap(), a); + assert_eq!(gelu.output_of(&cache).unwrap(), a); + + // The backward pass reads the cache, and its gradient has the shape of the cache. + let grad = gelu.backward(&cache, &Array::ones(4).into_dyn()).unwrap(); + assert_eq!(grad.shape(), cache.shape()); + + // A cache of the other kind, or a gradient of another shape, is an error. + let (_, relu_cache) = Activation::ReLU.forward_train(z.clone()).unwrap(); + assert!(gelu.backward(&relu_cache, &Array::ones(4).into_dyn()).is_err()); + assert!(gelu.backward(&cache, &Array::ones(3).into_dyn()).is_err()); + + println!("dL/dz: {grad}"); +} +``` It then computes 3 quantities. The weight gradient is `input^T * grad`, a GEMM. The bias gradient is the column-sum over the batch. The input gradient is `grad * W^T`, another GEMM. An input of rank 3 or more folds its leading axes into the row axis first. The 3 quantities therefore stay the same 3 products at every rank. `backward` returns the input gradient, and it puts the other 2 in the gradient store of the context, under the names `kernel` and `bias`. Read one back with `ctx.grads().get(&ParamId::new(scope, name))`, where `scope` is the position of the layer in the model that drives it. A caller that drives 1 layer by hand takes position 0. A `Dense` layer that is a sublayer of another layer has a longer address, such as `0.first.kernel` (see [3.1](./3.1._The_Sequential_Model.md#a-layer-that-holds-other-layers)). A bias-free layer computes no bias gradient, so the store holds none. The store sums, so a layer that 2 positions of a graph model share receives the total of the 2 backward passes. -Calling `backward` before `forward` returns `Error::NeuralNetwork(NnError::ForwardPassNotRun)`. An upstream gradient whose shape does not match the parked output returns `Error::ShapeMismatch`. Both are errors, never panics. +Calling `backward` before `forward` returns `Error::NeuralNetwork(NnError::ForwardPassNotRun)`. An upstream gradient whose shape does not match the parked activation tensor returns `Error::ShapeMismatch`. Both are errors, never panics. ## 3.2.7. Reading and setting weights diff --git a/guide/en/src/Chapter-03/3.7._Recurrent_Layers.md b/guide/en/src/Chapter-03/3.7._Recurrent_Layers.md index 22ddd43a..047222bc 100644 --- a/guide/en/src/Chapter-03/3.7._Recurrent_Layers.md +++ b/guide/en/src/Chapter-03/3.7._Recurrent_Layers.md @@ -225,7 +225,11 @@ All 3 layers are the same layer over a different cell. 1 implementation holds th Training uses backpropagation through time with a **full unroll**. There is no truncation window. A training pass parks everything the backward pass needs in the context. -Every layer parks the state that enters each timestep. For `SimpleRNN` and `GRU` that state is the hidden state, and for `LSTM` it is the hidden state and the cell state. Each cell then parks what only it reads. `SimpleRNN` parks nothing more. `LSTM` parks `activation(c_t)` and all 4 gate activations per timestep. `GRU` parks the reset and update gates, the candidate, and the `r_t * h_{t-1}` product. +Every layer parks the state that enters each timestep. For `SimpleRNN` and `GRU` that state is the hidden state, and for `LSTM` it is the hidden state and the cell state. Each cell then parks what only it reads. What it parks for the layer activation depends on [`Activation::saves`](3.2._Dense_Layers_and_Activations.md): + +- `SimpleRNN` parks nothing more when the activation saves its output, because the hidden state is that output. For GELU, SiLU, and Mish it parks the pre-activation of each timestep. +- `LSTM` parks the input, forget, and output gates per timestep. It also parks the candidate when the activation saves its output, or the candidate pre-activation for GELU, SiLU, and Mish. When the activation saves its output, it also parks `activation(c_t)`. For GELU, SiLU, and Mish the backward pass reads `c_t` from the parked state. +- `GRU` parks the reset and update gates and the `r_t * h_{t-1}` product per timestep. It also parks the candidate, or the candidate pre-activation for GELU, SiLU, and Mish. An inference pass records none of it, and it therefore costs less than a training pass. Memory scales linearly with `timesteps`. Long sequences cost RAM, not just time. The layer itself holds none of these values, whatever the mode, so a built model serves several threads at once. diff --git a/guide/en/src/Chapter-06/6.2._Matrix_Multiplication.md b/guide/en/src/Chapter-06/6.2._Matrix_Multiplication.md index a313c066..ff830a19 100644 --- a/guide/en/src/Chapter-06/6.2._Matrix_Multiplication.md +++ b/guide/en/src/Chapter-06/6.2._Matrix_Multiplication.md @@ -331,10 +331,10 @@ and elementwise gates. ## 6.2.6. When parallelism pays, and how to measure it -The gates encode where threading helps. A `Dense` forward runs 1 fused backend GEMM. The bias and -the activation run inside the kernel's epilogue, not as extra passes. A training context also parks an -input cache and an output cache for the backward pass, so an inference context times the -projection alone. +The gates encode where threading helps. A `Dense` forward runs 1 fused backend GEMM. The bias, and +a ReLU activation, run inside the kernel's epilogue, not as extra passes. Every other activation +runs as 1 more pass over the output. A training context also parks the input and 1 activation +tensor for the backward pass, so an inference context times the projection alone. The shape of the product decides which regime the backend picks. Label a shape `batch x in_features x out_features`, which is `m x k x n`. diff --git "a/guide/zh-Hans/src/Chapter-01/1.5._Prelude\344\270\216\346\250\241\345\235\227\345\257\274\345\205\245.md" "b/guide/zh-Hans/src/Chapter-01/1.5._Prelude\344\270\216\346\250\241\345\235\227\345\257\274\345\205\245.md" index 550f186a..30a9b197 100644 --- "a/guide/zh-Hans/src/Chapter-01/1.5._Prelude\344\270\216\346\250\241\345\235\227\345\257\274\345\205\245.md" +++ "b/guide/zh-Hans/src/Chapter-01/1.5._Prelude\344\270\216\346\250\241\345\235\227\345\257\274\345\205\245.md" @@ -66,7 +66,7 @@ use rustyml::metrics::r2_score; | 图模型 | `GraphBuilder`(为非链式拓扑收集节点和层)、`Graph`(`build` 交回的类型)、`NodeId`(图中 1 个节点的句柄) | | 核心层 | `Dense`, `Embedding`, `Flatten`, `Identity`, `Reshape`, `Reverse`, `Permute`, `RepeatVector`, `Rescaling`, `Activation` | | 合并层 | `Add`, `Average`, `Concatenate`, `Maximum`, `Minimum`, `Multiply`, `Subtract` | -| 激活层 | `Linear`, `ReLU`, `LeakyReLU`, `ELU`, `SELU`, `Sigmoid`, `HardSigmoid`, `Tanh`, `Softplus`, `Softsign`, `Exponential`, `Softmax`, `PReLU` | +| 激活层 | `Linear`, `ReLU`, `LeakyReLU`, `ELU`, `SELU`, `Sigmoid`, `HardSigmoid`, `Tanh`, `Softplus`, `Softsign`, `Exponential`, `GELU`, `SiLU`, `Mish`, `Softmax`, `PReLU` | | 卷积 | `Conv1D`, `Conv2D`, `Conv3D`, `Conv1DTranspose`, `Conv2DTranspose`, `Conv3DTranspose`, `DepthwiseConv1D`, `DepthwiseConv2D`, `SeparableConv1D`, `SeparableConv2D`, `PaddingType`, `ConvPadding` | | 边界 | `ZeroPadding1D/2D/3D`, `Cropping1D/2D/3D`, `Border1D/2D/3D` | | 池化 | `MaxPooling1D/2D/3D`, `AveragePooling1D/2D/3D`, `GlobalMaxPooling1D/2D/3D`, `GlobalAveragePooling1D/2D/3D` | diff --git "a/guide/zh-Hans/src/Chapter-03/3.2._\345\205\250\350\277\236\346\216\245\345\261\202\344\270\216\346\277\200\346\264\273\345\207\275\346\225\260.md" "b/guide/zh-Hans/src/Chapter-03/3.2._\345\205\250\350\277\236\346\216\245\345\261\202\344\270\216\346\277\200\346\264\273\345\207\275\346\225\260.md" index dd7d7a23..677dc2d5 100644 --- "a/guide/zh-Hans/src/Chapter-03/3.2._\345\205\250\350\277\236\346\216\245\345\261\202\344\270\216\346\277\200\346\264\273\345\207\275\346\225\260.md" +++ "b/guide/zh-Hans/src/Chapter-03/3.2._\345\205\250\350\277\236\346\216\245\345\261\202\344\270\216\346\277\200\346\264\273\345\207\275\346\225\260.md" @@ -2,7 +2,7 @@ `Dense` 是前馈网络的主力层。它先做一次线性映射 `input * W + b`,再套一个逐元素的非线性。RustyML 把这个非线性直接折叠进层本身,不会把它当成单独的一步来处理。 -本页涉及构造函数及其折叠激活的设计,也涉及精确的初始化方案,以及 12 个激活函数各自的梯度与失效模式。还涉及由 GEMM 支撑的计算路径,以及如何读取或注入权重。本页假定你已经读过 [Sequential 模型](./3.1._Sequential模型.md) 那一页。这里讲的层,就是你收集进 `SequentialBuilder` 的东西,而它们各自的输入宽度来自模型的 build。 +本页涉及构造函数及其折叠激活的设计,也涉及精确的初始化方案,以及 15 个激活函数各自的梯度与失效模式。还涉及由 GEMM 支撑的计算路径,以及如何读取或注入权重。本页假定你已经读过 [Sequential 模型](./3.1._Sequential模型.md) 那一页。这里讲的层,就是你收集进 `SequentialBuilder` 的东西,而它们各自的输入宽度来自模型的 build。 ## 3.2.1. 折叠激活的设计 @@ -12,23 +12,23 @@ pub fn new(units: usize, activation: impl Into) -> Result ``` -层内部存的是一个 `Activation` *值*。这是一个有 12 个变体、实现了 `Copy` 的普通枚举。它不是泛型参数 `Dense`,也不是 `Box`。这是有意为之的选择。泛型参数会把 `Dense` 单态化成 12 份。 +层内部存的是一个 `Activation` *值*。这是一个有 15 个变体、实现了 `Copy` 的普通枚举。它不是泛型参数 `Dense`,也不是 `Box`。这是有意为之的选择。泛型参数会把 `Dense` 单态化成 15 份。 它还会逼着权重反序列化时逐个试探 `Dense` 的每种组合,才能对上磁盘里的具体类型。trait 对象则会给每次逐元素调用都加一层间接跳转。运行时枚举让 `Dense` 始终保持为单一的具体类型。于是持久化层能把每个存下来的层向下转型到唯一一个结构体(见[权重保存与加载](./3.9._权重保存与加载.md))。激活的数学运算则始终是一个纯粹、无状态的函数,由层在自己的前向和反向传播里调用它。 Keras 用 `Dense(units, activation="relu")` 做了同样的折叠选择。RustyML 没有隐式的“不带激活”。你必须始终显式传一个。想要一个纯线性层,就传 `Activation::Linear`,也就是恒等映射。这对应的正是 Keras 的 `activation=None`。下面每个以回归输出头收尾的示例都是这么做的。 -折叠是合理的默认选择。RustyML 也把每个激活函数都提供为独立的层(`Linear`、`ReLU`、`LeakyReLU`、`ELU`、`SELU`、`Sigmoid`、`HardSigmoid`、`Tanh`、`Softplus`、`Softsign`、`Exponential`、`Softmax`)。当你需要在线性映射和非线性之间插入别的东西时,就用这些独立层。最常见的情况就是插一个归一化层(见[正则化与归一化层](./3.8._正则化与归一化层.md))。2 种写法是等价的: +折叠是合理的默认选择。RustyML 也把每个激活函数都提供为独立的层(`Linear`、`ReLU`、`LeakyReLU`、`ELU`、`SELU`、`Sigmoid`、`HardSigmoid`、`Tanh`、`Softplus`、`Softsign`、`Exponential`、`GELU`、`SiLU`、`Mish`、`Softmax`)。当你需要在线性映射和非线性之间插入别的东西时,就用这些独立层。最常见的情况就是插一个归一化层(见[正则化与归一化层](./3.8._正则化与归一化层.md))。2 种写法是等价的: ```rust,ignore // 这两种堆叠方式计算结果相同。 -builder.add(Dense::new(32, Activation::ReLU).unwrap()); // 折叠写法:只缓存一份激活输出 +builder.add(Dense::new(32, Activation::ReLU).unwrap()); // 折叠写法:只缓存一份激活张量 builder.add(Dense::new(32, Activation::Linear).unwrap()) // 拆分写法:先来一个恒等 Dense…… .add(ReLU::new()); // ……再接一个独立的 ReLU ``` -优先用折叠写法。它只缓存一份激活后的张量,而不是两份。反向传播随后就基于这份缓存输出对激活求导(见 3.2.6)。只有当某个层必须待在 `W*x + b` 和非线性之间时,才用拆分写法。 +优先用折叠写法。它只缓存一份激活张量,而不是两份。对大多数激活函数而言,这份张量是激活后的输出;对 `GELU`、`SiLU` 和 `Mish` 而言,它是激活前的值(见 3.2.4)。只有当某个层必须待在 `W*x + b` 和非线性之间时,才用拆分写法。 ## 3.2.2. 构造并确定层的尺寸 @@ -164,9 +164,9 @@ RustyML 无论激活函数是什么都套用 Glorot。它不会为 `ReLU` 层切 ## 3.2.4. 激活函数 -`Activation` 有 12 个变体。每个变体的反向传播都是用*激活后的输出* `a = f(x)` 来表达的,而不是激活前的 `x`。层缓存的正是这个输出,而不是输入。这条约定容纳了所有导数在 `a` 上有闭式表达的激活函数,也就是下面这 12 个。 +`Activation` 有 15 个变体。每个变体的反向传播都恰好读取 1 份保存下来的张量。其中 12 个变体读的是*激活后的输出* `a = f(x)`,因为它们的导数在 `a` 上有闭式表达。GELU、SiLU(又名 Swish)和 Mish 的形状是 `a = x * g(x)`。这种形状没有闭式反函数,所以不存在只用 `a` 表达的导数。因此这 3 个变体保存的是激活前的值 `x`。 -同样是这条约定,正好解释了为什么 GELU、SiLU(Swish)和 Mish 缺席。它们 `a = x * g(x)` 的形状没有闭式反函数,所以不存在只用 `a` 表达的导数。要把它们加进来,就得有一条更宽的约定,把激活前的值也一并交给反向传播。 +宿主层绝不会把两份张量都存下来。指明保存哪份张量的那条约定,见 3.2.6。下面两张表按保存的那份张量给出反向传播,上游梯度记作 `g`。 | `Activation` | 前向 `f(x)` | 反向(给定上游 `g`、输出 `a`) | 输出范围 | | --- | --- | --- | --- | @@ -183,6 +183,15 @@ RustyML 无论激活函数是什么都套用 Glorot。它不会为 `ReLU` 层切 | `HardSigmoid` | `clip(x/6 + 0.5, 0, 1)` | `0 < a < 1` 处为 `g / 6`,否则为 `0` | `[0, 1]` | | `Exponential` | `e^x` | `g * a` | `(0, inf)` | +保存激活前值 `x` 的 3 个变体用到 `sigmoid(x) = 1 / (1 + e^-x)`、`softplus(x) = ln(1 + e^x)`,以及标准正态分布:`Phi(x)` 是它的累积分布函数,`phi(x)` 是它的密度函数。 + +| `Activation` | 前向 `f(x)` | 反向(给定上游 `g`、激活前值 `x`) | 输出范围 | +| --- | --- | --- | --- | +| `GELU { approximate: false }` | `x * Phi(x)` | `g * (Phi(x) + x * phi(x))` | 约 `[-0.17, inf)` | +| `GELU { approximate: true }` | `x * sigmoid(u)`,其中 `u = 2 * sqrt(2/pi) * (x + 0.044715 * x^3)` | `g * (s + x * s * (1 - s) * du/dx)`,其中 `s = sigmoid(u)` | 约 `[-0.17, inf)` | +| `SiLU` | `x * sigmoid(x)` | `g * s * (1 + x * (1 - s))`,其中 `s = sigmoid(x)` | 约 `[-0.28, inf)` | +| `Mish` | `x * tanh(softplus(x))` | `g * (t + x * (1 - t^2) * sigmoid(x))`,其中 `t = tanh(softplus(x))` | 约 `[-0.31, inf)` | + **ReLU** 是隐藏层的默认选择。它便宜,而且在正半轴不饱和。它的失效模式是死神经元。因为 `x <= 0` 处导数为 `0`,一个在整个 batch 的每个样本上激活前值都为负的神经元,拿到的梯度就是零。 它永远不会更新,从此彻底关死。学习率过高会让情况更糟。它会早早把神经元推进死区。Glorot 初始化没有 leaky 变体可用。你能用的办法是更小的学习率和缩放得当的输入,或者,如果是你自己注入权重,用更合理的初始缩放。 @@ -193,7 +202,7 @@ RustyML 无论激活函数是什么都套用 Glorot。它不会为 `ReLU` 层切 `negative_slope` 和 `alpha` 都必须是有限值,且大于 0。`Activation::validate` 负责执行这条边界。全部 14 个带训练参数的层构造函数都会调用它:`Dense`、`Conv1D`、`Conv2D`、`Conv3D`、`Conv1DTranspose`、`Conv2DTranspose`、`Conv3DTranspose`、`DepthwiseConv1D`、`DepthwiseConv2D`、`SeparableConv1D`、`SeparableConv2D`、`SimpleRNN`、`LSTM` 和 `GRU`。所以不可用的取值会在你搭建模型的地方就返回 `Error::InvalidParameter`,而不是等到第一次前向传播。 -这条边界是严格的,原因就在于那条只认输出的约定。反向传播是靠 `a` 的符号来判分支的。取 0 会把整个负半轴塌缩到 `a = 0` 上,分支也就被抹掉了。取负值会把符号反过来,于是反向传播读到的是错误的分支。想要斜率为 0,就用 `Activation::ReLU`。 +这条边界是严格的,原因在于这 2 个变体保存的是输出。反向传播是靠 `a` 的符号来判分支的。取 0 会把整个负半轴塌缩到 `a = 0` 上,分支也就被抹掉了。取负值会把符号反过来,于是反向传播读到的是错误的分支。想要斜率为 0,就用 `Activation::ReLU`。 `LeakyReLU` 的正分支用的是 `x >= 0`,所以它在恰好 0 处的导数是 1。`ELU` 和 `SELU` 用的是 `x > 0`,这把 0 本身划到了负分支。因此它们在恰好 0 处的导数分别是 `alpha` 和 `scale * alpha`,而不是 1。`LeakyReLU::default()` 用的斜率是 `0.3`,`ELU::default()` 用的 `alpha` 是 `1.0`。 @@ -256,6 +265,31 @@ fn main() { **HardSigmoid** 用一条截断的直线去近似 sigmoid,且不用指数运算。它是 `clip(x/6 + 0.5, 0, 1)`,所以它在 `x = -3` 处恰好取到 0,在 `x = 3` 处恰好取到 1。真正的 sigmoid 只是趋近两端,永远到不了。它的导数在直线段上是常数 `1/6`,出了这一段就恰好为 0。段外这段平坦的梯度带来的正是 ReLU 那种死单元风险。当 `exp` 的开销要紧时,就把 `HardSigmoid` 当门控用。 +**GELU、SiLU 和 Mish** 是 Transformer 和现代卷积模型里常用的平滑自门控激活函数。它们各自用一个从 0 升到 1 的门控 `g(x)` 去乘 `x`。正输入几乎原样通过,很大的负输入则趋于 0。在这两段之间,每条曲线都会先降到 0 以下再回升,所以 3 个函数都不是单调的。负半轴上的导数不为 0,所以一个激活前值始终为负的单元仍能拿到一点梯度。随着 `x` 趋于 `-inf`,这点梯度也趋于 0。 + +`GELU` 用 `Phi(x)` 给 `x` 加门控,也就是一个标准正态值小于 `x` 的概率(Hendrycks 和 Gimpel,2016)。`Activation::GELU { approximate: false }` 精确计算 `Phi`。它用互补误差函数 `Phi(x) = erfc(-x / sqrt(2)) / 2`,在负尾部也能保持相对精度。`Activation::GELU { approximate: true }` 用 tanh 近似 `0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`。RustyML 按与之相等的 `x * sigmoid(2 * sqrt(2/pi) * (x + 0.044715 * x^3))` 形式求值,这种形式在负尾部不会发生抵消。精确形式是默认值,与 Keras 和 PyTorch 一致。要对齐用近似形式训练出来的模型(比如 GPT-2),就用近似形式。 + +`SiLU` 是 `x * sigmoid(x)`,即 sigmoid 加权线性单元。Keras 也把它叫作 `swish`。`Mish` 是 `x * tanh(softplus(x))`(Misra,2019)。它每个元素要算 1 次指数、1 次对数和 1 次双曲正切,是这 3 个里开销最大的。 + +对应的独立层是 `GELU`、`SiLU` 和 `Mish`。`SiLU::new()` 和 `Mish::new()` 不接收参数。`GELU::new()` 给出精确形式,`with_approximate(true)` 构建方法则选用 tanh 近似。每个层都能转换成对应的 `Activation` 值,所以凡是宿主层接收激活函数的地方,都可以直接传这个层: + +```rust,ignore +// 折叠写法:激活在宿主层内部运行。 +builder.add(Dense::new(64, Activation::GELU { approximate: false }).unwrap()) + .add(Dense::new(64, Activation::SiLU).unwrap()) + .add(Conv2D::new(16, (3, 3), (1, 1), Activation::Mish).unwrap()); + +// 拆分写法:一个恒等 Dense,再接独立层。 +builder.add(Dense::new(64, Activation::Linear).unwrap()) + .add(GELU::new().with_approximate(true)); + +// 层会转换成对应的变体,所以下面 2 行构建出的是同一个 Dense。 +let a = Dense::new(64, GELU::new().with_approximate(true)).unwrap(); +let b = Dense::new(64, Activation::GELU { approximate: true }).unwrap(); +``` + +全部 14 个带训练参数的宿主层都接受这 3 个变体:`Dense`、每一个卷积层、`SimpleRNN`、`LSTM` 和 `GRU`。在循环层里,激活作用于隐藏状态或候选值,详见[循环层](./3.7._循环层.md)。 + **Exponential** 就是 `e^x`。它的输出严格为正,导数就是输出本身。把它用在必须吐出正量的输出头上,比如速率、方差或者尺度。它没有上界,所以激活前值超过大约 88 就会让 `f32` 溢出成 `inf`。要么把这个输出头的输入缩放好,要么改用 `Linear` 输出一个对数尺度的值。 **Softmax** 把一条 logits 变成一个概率分布。`Softmax { axis }` 指定这些条沿着哪一根轴,默认值是 `-1`,也就是最后一根轴。前向传播在取指数之前,先减去每一条的最大值。这让最大的那一项变成 `exp(0) = 1`,总和也就始终 `>= 1`。这个平移让 softmax 既防溢出又平移不变:给每个 logit 加同一个常数,输出不变。 @@ -412,7 +446,7 @@ fn main() { ## 3.2.6. 前向、反向与 GEMM 路径 -前向传播是 `activation(input * W + b)`,作为对 [gemmkit](https://docs.rs/gemmkit) 后端的*一次*调用算出来的。乘积、按列的偏置,以及(对 `ReLU` 而言)激活函数,全都在一趟里跑完。偏置和激活是在输出分块还留在寄存器里的时候,由内核的尾声(epilogue)施加的。这 12 个激活函数里,只有 `ReLU` 有折叠进尾声的实现。`Linear` 不需要额外的一步,因为加过偏置的乘积本身就是输出。其他每个激活函数都是在加过偏置的乘积上另跑一趟向量化的 `Activation::forward`。 +前向传播是 `activation(input * W + b)`,作为对 [gemmkit](https://docs.rs/gemmkit) 后端的*一次*调用算出来的。乘积、按列的偏置,以及(对 `ReLU` 而言)激活函数,全都在一趟里跑完。偏置和激活是在输出分块还留在寄存器里的时候,由内核的尾声(epilogue)施加的。这 15 个激活函数里,只有 `ReLU` 有折叠进尾声的实现。`Linear` 不需要额外的一步,因为加过偏置的乘积本身就是输出。其他每个激活函数都是在加过偏置的乘积上另跑一趟向量化的 `Activation::forward`。 融合后的结果与分开做乘积再加标量激活相比,逐位相同,只有一个例外。融合的 `ReLU` 尾声会把 `NaN` 的激活前值映射成 `0.0`。crate 里其他地方用的独立 `ReLU` 激活则会传播 `NaN`。`NaN` 的激活前值本身就意味着模型已经发散。缓存输出版本的 `ReLU` 导数接下来会把那个 `0.0` 当成死神经元处理。它在那里反传回一个零梯度,而不是继续把 `NaN` 往回传。 @@ -420,17 +454,52 @@ fn main() { 紧凑循环里的小层按设计保持串行,因为每次调用的派发开销会占大头。这些乘积嵌在外层并行区域里也能安全嵌套。后端在[矩阵乘法](../Chapter-06/6.2._矩阵乘法.md)里有完整描述,可调的阈值则在[性能调优与并行](../Chapter-07/7.3._性能调优与并行.md)。 -紧随其后的逐元素激活有它自己的、独立的并行阈值。每个激活函数都落进 2 个开销类别中的一个。`ReLU`、`LeakyReLU`、`Softsign` 和 `HardSigmoid` 是受内存带宽限制的“廉价映射”,它们的临界点是 4,000,000 个元素。在任何实际的层规模下,它们都串行运行。`Sigmoid`、`Tanh`、`Softmax`、`ELU`、`SELU`、`Softplus` 和 `Exponential` 由 `exp` 主导,超过 131,072 个元素就转并行。`Linear` 只是复制张量,没有阈值。 +紧随其后的逐元素激活有它自己的、独立的并行阈值。每个激活函数都落进 2 个开销类别中的一个。`ReLU`、`LeakyReLU`、`Softsign` 和 `HardSigmoid` 是受内存带宽限制的“廉价映射”,它们的临界点是 4,000,000 个元素。在任何实际的层规模下,它们都串行运行。`Sigmoid`、`Tanh`、`Softmax`、`ELU`、`SELU`、`Softplus`、`Exponential`、`GELU`、`SiLU` 和 `Mish` 由 `exp` 主导,超过 131,072 个元素就转并行。`Linear` 只是复制张量,没有阈值。 + +反向传播用的是同样这 2 个类别,但它是按*导数*来归类的,而不是按前向传播来归类。`Sigmoid`、`Tanh`、`Softmax`、`ELU`、`SELU` 和 `Exponential` 的导数都只靠缓存输出就能算出来,不含指数运算。所以尽管它们的前向传播走 `exp` 的阈值,反向传播用的却是廉价映射的阈值。`Softplus` 是例外:它的导数每个元素仍要算一次指数,所以它的反向传播也留在 `exp` 的阈值上。`GELU`、`SiLU` 和 `Mish` 从激活前值计算导数,用到的超越函数调用与前向传播相同,所以它们的反向传播也留在 `exp` 的阈值上。挪动其中任何一个阈值,都只是在串行和并行之间做交换。结果完全相同,而且在同一台机器上,每次运行每个乘积都是确定一致的。 + +`forward` 接收 `&self`,它把输入和 1 份激活张量寄存进上下文。`Ctx::training()` 选中寄存这些值的模式,`Ctx::inference()` 选中什么都不寄存的模式。2 种模式算出同样的值,所以推理省下的恰好就是寄存这一份开销。层自身完全不持有缓存,正因为如此,多个线程可以同时用 1 个模型做推理。反向传播把寄存的激活张量取回来,再用它对激活求导。 -反向传播用的是同样这 2 个类别,但它是按*导数*来归类的,而不是按前向传播来归类。`Sigmoid`、`Tanh`、`Softmax`、`ELU`、`SELU` 和 `Exponential` 的导数都只靠缓存输出就能算出来,不含指数运算。所以尽管它们的前向传播走 `exp` 的阈值,反向传播用的却是廉价映射的阈值。`Softplus` 是唯一的例外:它的导数每个元素仍要算一次指数,所以它的反向传播也留在 `exp` 的阈值上。挪动其中任何一个阈值,都只是在串行和并行之间做交换。结果完全相同,而且在同一台机器上,每次运行每个乘积都是确定一致的。 +**激活缓存。** `Activation::saves()` 返回一个 `ActivationInput`,它指明反向传播读取哪份张量:`GELU`、`SiLU` 和 `Mish` 是 `ActivationInput::PreActivation`,其余 12 个变体是 `ActivationInput::Output`。训练时的前向传播调用 `Activation::forward_train(z)`,它返回激活后的输出和一个 `ActivationCache`。按 `saves()` 的指示,缓存里存的是输出的一份拷贝,或者 `z` 本身。融合的 `ReLU` 尾声没有单独的 `z`,所以层把乘积的输出包装成一个 `Output` 缓存。这样做是正确的,因为 `ReLU` 保存的就是输出。 -`forward` 接收 `&self`,它把输入和激活后的输出寄存进上下文。`Ctx::training()` 选中寄存这些值的模式,`Ctx::inference()` 选中什么都不寄存的模式。2 种模式算出同样的值,所以推理省下的恰好就是寄存这一份开销。层自身完全不持有缓存,正因为如此,多个线程可以同时用 1 个模型做推理。反向传播把寄存的输出取回来,再用它对激活求导。 +`Activation::backward(&cache, &grad_output)` 给出关于 `z` 的梯度。如果缓存里存的是另一份张量,它返回 `Error::InvalidInput`;如果梯度的形状与缓存的形状不同,它返回 `Error::ShapeMismatch`。`ActivationCache::holds()`、`tensor()`、`shape()` 和 `into_tensor()` 用来读取缓存。`Activation::output_of(&cache)` 给出缓存所对应的激活后输出。推理时调用 `Activation::forward(&z)`,它既不建缓存,也不做拷贝。 + +```rust +use ndarray::Array; +use rustyml::neural_network::layers::{Activation, ActivationInput}; + +fn main() { + let z = Array::from_vec(vec![-2.0_f32, -0.5, 0.0, 1.5]).into_dyn(); + let gelu = Activation::GELU { approximate: false }; + + // GELU 保存激活前值,ReLU 保存输出。 + assert_eq!(gelu.saves(), ActivationInput::PreActivation); + assert_eq!(Activation::ReLU.saves(), ActivationInput::Output); + + // 训练时的前向传播返回输出和缓存。推理时的前向传播结果与之一致。 + let (a, cache) = gelu.forward_train(z.clone()).unwrap(); + assert_eq!(cache.holds(), ActivationInput::PreActivation); + assert_eq!(gelu.forward(&z).unwrap(), a); + assert_eq!(gelu.output_of(&cache).unwrap(), a); + + // 反向传播读取缓存,梯度的形状与缓存的形状相同。 + let grad = gelu.backward(&cache, &Array::ones(4).into_dyn()).unwrap(); + assert_eq!(grad.shape(), cache.shape()); + + // 另一种缓存,或者另一种形状的梯度,都是错误。 + let (_, relu_cache) = Activation::ReLU.forward_train(z.clone()).unwrap(); + assert!(gelu.backward(&relu_cache, &Array::ones(4).into_dyn()).is_err()); + assert!(gelu.backward(&cache, &Array::ones(3).into_dyn()).is_err()); + + println!("dL/dz: {grad}"); +} +``` 接下来它算 3 个量。权重梯度是 `input^T * grad`,一次 GEMM。偏置梯度是按 batch 做的列求和。输入梯度是 `grad * W^T`,又一次 GEMM。秩为 3 或更高的输入会先把前导轴折进行轴。因此在任何秩上,这 3 个量都还是同样这 3 次乘积。 `backward` 返回的是输入梯度,另外 2 个则按 `kernel` 和 `bias` 这 2 个名字进入上下文的梯度存储。用 `ctx.grads().get(&ParamId::new(scope, name))` 把其中一个读回来,其中 `scope` 是这个层在驱动它的模型里的位置。手工驱动 1 个层的调用方取位置 0。作为另一个层的子层的 `Dense` 层,地址会更长,例如 `0.first.kernel`(见 [3.1](./3.1._Sequential模型.md#持有其他层的层))。不带偏置的层不计算偏置梯度,所以存储里没有它。存储做的是累加,所以被图模型的 2 个位置共用的层拿到的是这 2 遍反向传播之和。 -在 `forward` 之前调用 `backward` 会返回 `Error::NeuralNetwork(NnError::ForwardPassNotRun)`。形状对不上寄存输出的上游梯度会返回 `Error::ShapeMismatch`。两者都是错误,绝不会 panic。 +在 `forward` 之前调用 `backward` 会返回 `Error::NeuralNetwork(NnError::ForwardPassNotRun)`。形状对不上寄存激活张量的上游梯度会返回 `Error::ShapeMismatch`。两者都是错误,绝不会 panic。 ## 3.2.7. 读取与设置权重 diff --git "a/guide/zh-Hans/src/Chapter-03/3.7._\345\276\252\347\216\257\345\261\202.md" "b/guide/zh-Hans/src/Chapter-03/3.7._\345\276\252\347\216\257\345\261\202.md" index 4de108a8..07651bc1 100644 --- "a/guide/zh-Hans/src/Chapter-03/3.7._\345\276\252\347\216\257\345\261\202.md" +++ "b/guide/zh-Hans/src/Chapter-03/3.7._\345\276\252\347\216\257\345\261\202.md" @@ -225,7 +225,11 @@ bias [1, 32] 训练用的是随时间反向传播(BPTT),而且是**完整展开**,没有截断窗口。训练时的一遍会把反向传播需要的一切都寄存进上下文。 -每个层都寄存进入每个时间步的状态:对 `SimpleRNN` 和 `GRU` 来说这个状态就是隐藏状态,对 `LSTM` 来说则是隐藏状态加细胞状态。随后每个单元再寄存只有它自己会读的东西。`SimpleRNN` 不再多存任何东西。`LSTM` 寄存 `activation(c_t)` 以及每个时间步的 4 个门激活值。`GRU` 寄存重置门、更新门、候选值,以及 `r_t * h_{t-1}` 这个乘积。 +每个层都寄存进入每个时间步的状态:对 `SimpleRNN` 和 `GRU` 来说这个状态就是隐藏状态,对 `LSTM` 来说则是隐藏状态加细胞状态。随后每个单元再寄存只有它自己会读的东西。为层激活函数寄存什么,取决于 [`Activation::saves`](3.2._全连接层与激活函数.md): + +- `SimpleRNN`:激活函数保存输出时不再多存任何东西,因为隐藏状态就是那个输出。GELU、SiLU 和 Mish 则会寄存每个时间步的预激活值。 +- `LSTM`:每个时间步寄存输入门、遗忘门和输出门。激活函数保存输出时,还寄存候选值和 `activation(c_t)`;GELU、SiLU 和 Mish 则寄存候选值的预激活值,反向传播从已寄存的状态里读取 `c_t`。 +- `GRU`:每个时间步寄存重置门、更新门和 `r_t * h_{t-1}` 这个乘积,另外寄存候选值;GELU、SiLU 和 Mish 寄存的是候选值的预激活值。 推理时的一遍这些一概不记录,所以它比训练时的一遍更省。内存开销随 `timesteps` 线性增长:长序列吃的是内存,不只是时间。无论哪种模式,层自身都不持有这些值,所以一个构建好的模型可以同时服务多个线程。 diff --git "a/guide/zh-Hans/src/Chapter-06/6.2._\347\237\251\351\230\265\344\271\230\346\263\225.md" "b/guide/zh-Hans/src/Chapter-06/6.2._\347\237\251\351\230\265\344\271\230\346\263\225.md" index 01d45aac..1ac85297 100644 --- "a/guide/zh-Hans/src/Chapter-06/6.2._\347\237\251\351\230\265\344\271\230\346\263\225.md" +++ "b/guide/zh-Hans/src/Chapter-06/6.2._\347\237\251\351\230\265\344\271\230\346\263\225.md" @@ -153,7 +153,7 @@ gemmkit 只建一次这些分档池,然后热着复用,不会每次调用都 ## 6.2.6. 并行何时划算,以及如何测量 -这些阈值把"多线程在哪里有用"编码了进去。一次 `Dense` 前向就是 1 次融合的后端 GEMM。偏置和激活都跑在内核的尾声里,不是额外的趟数。训练上下文还会为反向传播存放一份输入缓存和一份输出缓存,所以要单独给这次投影计时,请用推理上下文。 +这些阈值把"多线程在哪里有用"编码了进去。一次 `Dense` 前向就是 1 次融合的后端 GEMM。偏置和 ReLU 激活都跑在内核的尾声里,不是额外的趟数;其他激活函数要在输出上多走 1 趟。训练上下文还会为反向传播存放输入和 1 个激活张量,所以要单独给这次投影计时,请用推理上下文。 乘积的形状决定后端选哪一种情形。形状按 `batch x in_features x out_features` 标注,也就是 `m x k x n`。 diff --git a/src/neural_network/layers/activation.rs b/src/neural_network/layers/activation.rs index 9475530c..8819ca48 100644 --- a/src/neural_network/layers/activation.rs +++ b/src/neural_network/layers/activation.rs @@ -3,7 +3,7 @@ //! Defines the [`Activation`] enum, the single source of truth for each activation's forward //! transform and derivative. The thin layer wrappers delegate their math to it: [`Linear`], //! [`ReLU`], [`LeakyReLU`], [`ELU`], [`SELU`], [`Sigmoid`], [`HardSigmoid`], [`Tanh`], -//! [`Softplus`], [`Softsign`], [`Exponential`], and [`Softmax`] +//! [`Softplus`], [`Softsign`], [`Exponential`], [`Softmax`], [`GELU`], [`SiLU`], and [`Mish`] //! //! [`PReLU`] stands apart. Its negative-side slope is a trainable array, so it holds its own //! math and its own weights instead of delegating to the enum @@ -19,12 +19,16 @@ use rayon::iter::{IntoParallelIterator, ParallelIterator}; pub mod elu; /// Exponential activation layer pub mod exponential; +/// GELU (Gaussian Error Linear Unit) activation layer +pub mod gelu; /// Hard sigmoid activation layer, a piecewise-linear approximation of the logistic sigmoid pub mod hard_sigmoid; /// Leaky ReLU activation layer pub mod leaky_relu; /// Linear (Identity) activation layer pub mod linear; +/// Mish activation layer +pub mod mish; /// PReLU activation layer, whose negative-side slope is a trainable parameter pub mod p_relu; /// ReLU (Rectified Linear Unit) activation layer @@ -33,6 +37,8 @@ pub mod relu; pub mod selu; /// Sigmoid activation layer pub mod sigmoid; +/// SiLU (Sigmoid Linear Unit) activation layer, also named Swish +pub mod silu; /// Softmax activation layer pub mod softmax; /// Softplus activation layer @@ -44,13 +50,16 @@ pub mod tanh; pub use elu::ELU; pub use exponential::Exponential; +pub use gelu::GELU; pub use hard_sigmoid::HardSigmoid; pub use leaky_relu::LeakyReLU; pub use linear::Linear; +pub use mish::Mish; pub use p_relu::PReLU; pub use relu::ReLU; pub use selu::SELU; pub use sigmoid::Sigmoid; +pub use silu::SiLU; pub use softmax::Softmax; pub use softplus::Softplus; pub use softsign::Softsign; @@ -78,6 +87,15 @@ const SELU_SCALE_ALPHA: f32 = SELU_SCALE * SELU_ALPHA; /// The slope of the hard sigmoid's linear segment, `1/6` const HARD_SIGMOID_SLOPE: f32 = 1.0 / 6.0; +/// The scale `sqrt(2 / pi)` of the tanh approximation of GELU +const GELU_TANH_SCALE: f32 = 0.797_884_6; + +/// The cubic coefficient of the tanh approximation of GELU, from Hendrycks and Gimpel (2016) +const GELU_TANH_CUBIC: f32 = 0.044_715; + +/// The density scale `1 / sqrt(2 * pi)` of the standard normal distribution +const NORMAL_DENSITY_SCALE: f32 = 0.398_942_3; + /// The default softmax axis, the last axis of the input /// /// A negative axis counts back from the end, so `-1` holds for every rank @@ -96,15 +114,13 @@ pub const DEFAULT_SOFTMAX_AXIS: i32 = -1; /// /// # Notes /// -/// [`Activation::backward`] receives the *activated output* `a`, never the pre-activation `z`. -/// Every variant's derivative is therefore expressed through `a` alone, and a host layer caches -/// 1 tensor rather than 2. This is the output-only derivative contract that every variant below -/// follows. +/// The backward pass of each variant reads 1 tensor. [`Activation::saves`] names it: the +/// activated output `a`, or the pre-activation `z`. A host layer therefore caches 1 tensor, not +/// 2. [`Activation::forward_train`] returns that tensor in an [`ActivationCache`], and +/// [`Activation::backward`] refuses a cache that holds the other tensor. /// -/// The contract admits every activation whose derivative has a closed form in `a`, which covers -/// the full family below. It excludes GELU, SiLU (Swish), and Mish, whose `a = z * g(z)` shape -/// has no closed-form inverse. Supporting them needs a wider contract that also hands the -/// backward pass the pre-activation. +/// A variant reads `a` when its derivative has a closed form in `a`. GELU, SiLU, and Mish have +/// the shape `a = z * g(z)`, which has no closed-form inverse, so these 3 variants read `z`. /// /// The 2 saturating variants pay a small accuracy cost for the contract. `ELU` recovers its /// negative branch as `a + alpha`, and `SELU` recovers its negative branch as @@ -168,6 +184,220 @@ pub enum Activation { HardSigmoid, /// Exponential, `e^x` Exponential, + /// Gaussian error linear unit, `x * Phi(x)`, where `Phi` is the standard normal + /// cumulative distribution function, from Hendrycks and Gimpel (2016) + GELU { + /// When `true`, the variant uses the tanh approximation + /// `x * sigmoid(2 * sqrt(2 / pi) * (x + 0.044715 * x^3))`. When `false`, it uses the + /// exact `Phi`. The layer form defaults to `false` + approximate: bool, + }, + /// Sigmoid linear unit, `x * sigmoid(x)`, also named Swish + SiLU, + /// Mish, `x * tanh(softplus(x))`, from Misra (2019) + Mish, +} + +/// The tensor that the backward pass of an [`Activation`] reads +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ActivationInput { + /// The activated output `a` + Output, + /// The pre-activation `z` + PreActivation, +} + +/// The tensor that a training forward pass saves for the backward pass of an [`Activation`] +/// +/// [`Activation::forward_train`] makes the cache. The cache records which tensor it holds, +/// and [`Activation::backward`] refuses a cache that holds the other tensor +#[derive(Debug, Clone, PartialEq)] +pub struct ActivationCache { + /// The tensor that the cache holds + holds: ActivationInput, + /// The saved tensor + tensor: Tensor, +} + +impl ActivationCache { + /// Wraps a tensor that a host layer saved without [`Activation::forward_train`] + /// + /// A host layer that fuses an activation into a matrix product gets the output from the + /// product. It wraps that output here as an [`ActivationInput::Output`] cache + /// + /// # Parameters + /// + /// - `holds` - The tensor that `tensor` is + /// - `tensor` - The saved tensor + /// + /// # Returns + /// + /// - `ActivationCache` - The cache + pub(crate) fn new(holds: ActivationInput, tensor: Tensor) -> Self { + Self { holds, tensor } + } + + /// The tensor that the cache holds + /// + /// # Returns + /// + /// - `ActivationInput` - [`ActivationInput::Output`] or [`ActivationInput::PreActivation`] + pub fn holds(&self) -> ActivationInput { + self.holds + } + + /// The saved tensor + /// + /// # Returns + /// + /// - `&Tensor` - The activated output or the pre-activation, as [`ActivationCache::holds`] + /// names + pub fn tensor(&self) -> &Tensor { + &self.tensor + } + + /// The shape of the saved tensor, which is also the shape of the activated output + /// + /// # Returns + /// + /// - `&[usize]` - The shape + pub fn shape(&self) -> &[usize] { + self.tensor.shape() + } + + /// Gives back the saved tensor + /// + /// # Returns + /// + /// - `Tensor` - The activated output or the pre-activation, as [`ActivationCache::holds`] + /// names + pub fn into_tensor(self) -> Tensor { + self.tensor + } +} + +/// Logistic sigmoid of 1 value +#[inline] +fn sigmoid(x: f32) -> f32 { + 1.0 / (1.0 + (-x).exp()) +} + +/// Softplus of 1 value, `ln(1 + e^x)` +/// +/// Above 0 the function factors `e^x` out of the logarithm. This keeps both tails in range. +/// The direct form overflows for a large `x`, and it loses every digit for a small `x` +#[inline] +fn softplus(x: f32) -> f32 { + if x > 0.0 { + x + (-x).exp().ln_1p() + } else { + x.exp().ln_1p() + } +} + +/// Exact GELU of 1 value, `x * Phi(x)` +/// +/// `Phi(x) = erfc(-x / sqrt(2)) / 2`. The complementary error function keeps the relative +/// precision of `Phi` in the negative tail, where `1 + erf` would cancel to 0 +#[inline] +fn gelu_exact(x: f32) -> f32 { + x * 0.5 * libm::erfcf(-x * std::f32::consts::FRAC_1_SQRT_2) +} + +/// Derivative of the exact GELU at `x`, `Phi(x) + x * phi(x)` +#[inline] +fn gelu_exact_grad(x: f32) -> f32 { + let cdf = 0.5 * libm::erfcf(-x * std::f32::consts::FRAC_1_SQRT_2); + let pdf = NORMAL_DENSITY_SCALE * (-0.5 * x * x).exp(); + cdf + x * pdf +} + +/// The sigmoid argument `2 * sqrt(2 / pi) * (x + 0.044715 * x^3)` of the tanh approximation +/// +/// `1 + tanh(u) = 2 * sigmoid(2 * u)`. The sigmoid form does not cancel in the negative tail +#[inline] +fn gelu_tanh_argument(x: f32) -> f32 { + 2.0 * GELU_TANH_SCALE * (x + GELU_TANH_CUBIC * x * x * x) +} + +/// Tanh approximation of GELU at `x` +#[inline] +fn gelu_tanh(x: f32) -> f32 { + x * sigmoid(gelu_tanh_argument(x)) +} + +/// Derivative of the tanh approximation of GELU at `x` +#[inline] +/// +/// Where the sigmoid saturates, `s * (1 - s)` is 0 and the derivative is `s`. The function +/// returns `s` there, because `x * x` overflows to infinity for `|x|` above about 1.8e19, and +/// `0 * infinity` is NaN +fn gelu_tanh_grad(x: f32) -> f32 { + let s = sigmoid(gelu_tanh_argument(x)); + let slope = s * (1.0 - s); + if slope == 0.0 { + return s; + } + let argument_grad = 2.0 * GELU_TANH_SCALE * (1.0 + 3.0 * GELU_TANH_CUBIC * x * x); + s + x * slope * argument_grad +} + +/// SiLU of 1 value, `x * sigmoid(x)` +#[inline] +fn silu(x: f32) -> f32 { + x * sigmoid(x) +} + +/// Derivative of SiLU at `x`, `sigmoid(x) * (1 + x * (1 - sigmoid(x)))` +#[inline] +fn silu_grad(x: f32) -> f32 { + let s = sigmoid(x); + s * (1.0 + x * (1.0 - s)) +} + +/// Mish of 1 value, `x * tanh(softplus(x))` +#[inline] +fn mish(x: f32) -> f32 { + x * softplus(x).tanh() +} + +/// Derivative of Mish at `x`, `t + x * (1 - t^2) * sigmoid(x)` with `t = tanh(softplus(x))` +#[inline] +fn mish_grad(x: f32) -> f32 { + let t = softplus(x).tanh(); + t + x * (1.0 - t * t) * sigmoid(x) +} + +/// Applies an element-wise map, in parallel when the tensor reaches `threshold` elements +fn map_elements(z: &Tensor, threshold: usize, f: impl Fn(f32) -> f32 + Sync + Send) -> Tensor { + if z.len() >= threshold { + Zip::from(z).par_map_collect(|&x| f(x)) + } else { + z.mapv(f) + } +} + +/// Multiplies each upstream gradient by the derivative at the matching pre-activation +/// +/// # Errors +/// +/// - `Error::ShapeMismatch` - `grad_output` and `z` differ in shape +fn scale_by_derivative( + z: &Tensor, + grad_output: &Tensor, + derivative: impl Fn(f32) -> f32 + Sync + Send, +) -> Result { + if grad_output.shape() != z.shape() { + return Err(Error::shape_mismatch(z.shape(), grad_output.shape())); + } + let mut grad = grad_output.clone(); + let scale = |g: &mut f32, &x: &f32| *g *= derivative(x); + if z.len() >= exp_map_parallel_threshold() { + Zip::from(&mut grad).and(z).par_for_each(scale); + } else { + Zip::from(&mut grad).and(z).for_each(scale); + } + Ok(grad) } impl Activation { @@ -252,23 +482,7 @@ impl Activation { }; Ok(out) } - Activation::Softplus => { - // Factoring `e^x` out of the logarithm above 0 keeps both tails in range. - // The direct form overflows for large x, and loses every digit for small x - let softplus = |x: f32| { - if x > 0.0 { - x + (-x).exp().ln_1p() - } else { - x.exp().ln_1p() - } - }; - let out = if z.len() >= exp_map_parallel_threshold() { - Zip::from(z).par_map_collect(|&x| softplus(x)) - } else { - z.mapv(softplus) - }; - Ok(out) - } + Activation::Softplus => Ok(map_elements(z, exp_map_parallel_threshold(), softplus)), Activation::Softsign => { let softsign = |x: f32| x / (1.0 + x.abs()); let out = if z.len() >= cheap_map_parallel_threshold() { @@ -296,32 +510,133 @@ impl Activation { }; Ok(out) } + Activation::GELU { approximate: false } => { + Ok(map_elements(z, exp_map_parallel_threshold(), gelu_exact)) + } + Activation::GELU { approximate: true } => { + Ok(map_elements(z, exp_map_parallel_threshold(), gelu_tanh)) + } + Activation::SiLU => Ok(map_elements(z, exp_map_parallel_threshold(), silu)), + Activation::Mish => Ok(map_elements(z, exp_map_parallel_threshold(), mish)), } } - /// Computes the gradient with respect to the pre-activation input + /// The tensor that the backward pass of this activation reads + /// + /// # Returns + /// + /// - `ActivationInput` - [`ActivationInput::PreActivation`] for GELU, SiLU, and Mish, and + /// [`ActivationInput::Output`] for every other variant + pub fn saves(&self) -> ActivationInput { + match self { + Activation::GELU { .. } | Activation::SiLU | Activation::Mish => { + ActivationInput::PreActivation + } + _ => ActivationInput::Output, + } + } + + /// Applies the activation and saves the tensor that the backward pass reads + /// + /// The cache holds a copy of the output, or `z` itself, as [`Activation::saves`] names + /// + /// # Parameters + /// + /// - `z` - Pre-activation tensor (the linear output of the host layer) + /// + /// # Returns + /// + /// - `Result<(Tensor, ActivationCache), Error>` - The activated tensor, and the cache for + /// [`Activation::backward`] + /// + /// # Errors + /// + /// - `Error::InvalidInput` - The Softmax axis resolves outside the rank of `z` + /// - `Error::Computation` - Softmax failed to reshape the input + pub fn forward_train(&self, z: Tensor) -> Result<(Tensor, ActivationCache), Error> { + let output = self.forward(&z)?; + let cache = match self.saves() { + ActivationInput::Output => { + ActivationCache::new(ActivationInput::Output, output.clone()) + } + ActivationInput::PreActivation => { + ActivationCache::new(ActivationInput::PreActivation, z) + } + }; + Ok((output, cache)) + } + + /// The activated output that a cache stands for + /// + /// An [`ActivationInput::Output`] cache gives a copy of its tensor. An + /// [`ActivationInput::PreActivation`] cache gives the forward pass of its tensor + /// + /// # Parameters + /// + /// - `cache` - A cache that [`Activation::forward_train`] of this activation made + /// + /// # Returns /// - /// Every supported activation's derivative is expressible in terms of its own - /// output, so this takes the cached activated tensor rather than the original input + /// - `Result` - The activated output + /// + /// # Errors + /// + /// - `Error::InvalidInput` - The cache holds another tensor than [`Activation::saves`] + /// names, or the Softmax axis resolves outside the rank of the tensor + /// - `Error::Computation` - Softmax failed to reshape the tensor + pub fn output_of(&self, cache: &ActivationCache) -> Result { + self.check_cache(cache)?; + match cache.holds { + ActivationInput::Output => Ok(cache.tensor.clone()), + ActivationInput::PreActivation => self.forward(&cache.tensor), + } + } + + /// Refuses a cache that holds another tensor than this activation reads + /// + /// # Errors + /// + /// - `Error::InvalidInput` - The cache holds another tensor than [`Activation::saves`] + /// names + fn check_cache(&self, cache: &ActivationCache) -> Result<(), Error> { + if cache.holds != self.saves() { + return Err(Error::invalid_input(format!( + "{self:?} reads {:?}, but the cache holds {:?}", + self.saves(), + cache.holds + ))); + } + Ok(()) + } + + /// Computes the gradient with respect to the pre-activation input /// /// This is pure math with no clamping or NaN/Inf sanitization /// /// # Parameters /// - /// - `activated` - The activated output `a` produced by [`forward`](Activation::forward) + /// - `cache` - The cache that [`Activation::forward_train`] of this activation made /// - `grad_output` - Upstream gradient `dL/da` /// /// # Returns /// - /// - `Result` - The gradient `dL/dz`, same shape as `activated` + /// - `Result` - The gradient `dL/dz`, with the shape of the cached tensor /// /// # Errors /// - /// - `Error::InvalidInput` - The Softmax axis resolves outside the rank of `activated` - /// - `Error::ShapeMismatch` - Softmax received a `grad_output` whose shape differs from - /// the shape of `activated` + /// - `Error::InvalidInput` - The cache holds another tensor than [`Activation::saves`] + /// names, or the Softmax axis resolves outside the rank of the cached tensor + /// - `Error::ShapeMismatch` - `grad_output` and the cached tensor differ in shape /// - `Error::Computation` - Softmax failed to reshape the tensors - pub fn backward(&self, activated: &Tensor, grad_output: &Tensor) -> Result { + pub fn backward(&self, cache: &ActivationCache, grad_output: &Tensor) -> Result { + self.check_cache(cache)?; + let activated = &cache.tensor; + if grad_output.shape() != activated.shape() { + return Err(Error::shape_mismatch( + activated.shape(), + grad_output.shape(), + )); + } match self { Activation::Linear => Ok(grad_output.clone()), Activation::ReLU => { @@ -492,6 +807,14 @@ impl Activation { } Ok(grad) } + Activation::GELU { approximate: false } => { + scale_by_derivative(activated, grad_output, gelu_exact_grad) + } + Activation::GELU { approximate: true } => { + scale_by_derivative(activated, grad_output, gelu_tanh_grad) + } + Activation::SiLU => scale_by_derivative(activated, grad_output, silu_grad), + Activation::Mish => scale_by_derivative(activated, grad_output, mish_grad), } } @@ -516,15 +839,13 @@ impl Activation { /// restriction is a deliberate simplification, not a limit of the math /// /// The host layers do not agree on the rank of the tensor they hand to the activation. - /// [`Dense`](crate::neural_network::layers::dense::Dense) folds every leading axis of its - /// output into 1 row axis, and it applies the activation to that rank-2 matrix. The - /// convolution layers do not fold, and they apply the activation to the full output - /// tensor. The recurrent layers apply the activation to 1 rank-2 state per timestep + /// [`Dense`](crate::neural_network::layers::dense::Dense) and the convolution layers apply + /// the activation to the full output tensor. The recurrent layers apply the activation to + /// 1 rank-2 state per timestep, and that state has no time axis /// - /// The axis `-1` names the same lane in all 3 families. Any other axis names a different - /// tensor axis in each family. In `Dense` a non-final axis would normalize the folded - /// rows, which are not the wanted lanes. A per-host axis rule is deferred, not - /// impossible. Use the standalone [`Softmax`] layer for a different axis + /// The axis `-1` names the same lane in all of these hosts. Any other axis names a + /// different tensor axis in a recurrent layer than in its output. A per-host axis rule is + /// deferred, not impossible. Use the standalone [`Softmax`] layer for a different axis /// /// # Errors /// @@ -631,6 +952,26 @@ impl From for Activation { Activation::Exponential } } +impl From for Activation { + #[inline] + fn from(layer: GELU) -> Self { + Activation::GELU { + approximate: layer.approximate, + } + } +} +impl From for Activation { + #[inline] + fn from(_: SiLU) -> Self { + Activation::SiLU + } +} +impl From for Activation { + #[inline] + fn from(_: Mish) -> Self { + Activation::Mish + } +} /// Resolves a softmax axis against the rank of the tensor it applies to /// @@ -1102,10 +1443,11 @@ mod tests { fn activation_softmax_backward_via_enum() { let output = tensor2(1, 3, vec![0.25, 0.25, 0.5]); let grad_output = tensor2(1, 3, vec![1.0, 0.0, 0.0]); + let cache = ActivationCache::new(ActivationInput::Output, output); let grad_input = Activation::Softmax { axis: DEFAULT_SOFTMAX_AXIS, } - .backward(&output, &grad_output) + .backward(&cache, &grad_output) .expect("Activation::Softmax backward failed"); let vals = grad_input.as_slice().expect("not contiguous"); assert_abs_diff_eq!(vals[0], 0.1875_f32, epsilon = 1e-6); @@ -1141,13 +1483,11 @@ mod tests { /// An all-ones upstream gradient makes the backward result the derivative itself fn check_against_reference(name: &str, activation: Activation, fwd: &[f32], grad: &[f32]) { let input = tensor2(1, PROBES.len(), PROBES.to_vec()); - let output = activation.forward(&input).expect("forward failed"); + let (output, cache) = activation.forward_train(input).expect("forward failed"); assert_pinned(&format!("{name} forward"), &output, fwd); let ones = tensor2(1, PROBES.len(), vec![1.0; PROBES.len()]); - let derivative = activation - .backward(&output, &ones) - .expect("backward failed"); + let derivative = activation.backward(&cache, &ones).expect("backward failed"); assert_pinned(&format!("{name} backward"), &derivative, grad); } @@ -1408,10 +1748,10 @@ mod tests { #[test] fn softplus_backward_keeps_the_far_negative_tail() { let input = tensor2(1, 1, vec![-40.0]); - let output = Activation::Softplus.forward(&input).expect("forward"); + let (_, cache) = Activation::Softplus.forward_train(input).expect("forward"); let ones = tensor2(1, 1, vec![1.0]); let derivative = Activation::Softplus - .backward(&output, &ones) + .backward(&cache, &ones) .expect("backward"); let got = derivative.iter().next().copied().expect("1 element"); @@ -1470,6 +1810,10 @@ mod tests { Activation::Softsign, Activation::HardSigmoid, Activation::Exponential, + Activation::GELU { approximate: false }, + Activation::GELU { approximate: true }, + Activation::SiLU, + Activation::Mish, ]; for activation in usable { assert!( @@ -1478,4 +1822,308 @@ mod tests { ); } } + + /// Exact GELU. The value at 0 is 0, and the derivative there is `Phi(0) = 0.5` + #[test] + fn gelu_matches_reference() { + check_against_reference( + "GELU", + Activation::GELU { approximate: false }, + &[ + -1.433_257_9e-6, + -0.004_049_694, + -0.045_500_264, + -0.158_655_25, + -0.154_268_77, + 0.0, + 0.345_731_23, + 0.841_344_8, + 1.954_499_7, + 2.995_950_3, + 4.999_998_6, + ], + &[ + -7.146_946e-6, + -0.011_945_647, + -0.085_231_8, + -0.083_315_47, + 0.132_504_88, + 0.5, + 0.867_495_1, + 1.083_315_5, + 1.085_231_8, + 1.011_945_6, + 1.000_007_1, + ], + ); + } + + /// The tanh approximation of GELU + #[test] + fn gelu_tanh_matches_reference() { + check_against_reference( + "GELU(approximate)", + Activation::GELU { approximate: true }, + &[ + -2.291_796_2e-7, + -0.003_637_392, + -0.045_402_306, + -0.158_808_01, + -0.154_286, + 0.0, + 0.345_714, + 0.841_192, + 1.954_597_7, + 2.996_362_6, + 5.0, + ], + &[ + -1.546_362e-6, + -0.011_584_167, + -0.086_099_26, + -0.082_964_084, + 0.132_630_1, + 0.5, + 0.867_369_9, + 1.082_964_1, + 1.086_099_3, + 1.011_584_2, + 1.000_001_5, + ], + ); + } + + /// SiLU. The derivative at 0 is `sigmoid(0) = 0.5` + #[test] + fn silu_matches_reference() { + check_against_reference( + "SiLU", + Activation::SiLU, + &[ + -0.033_464_255, + -0.142_277_62, + -0.238_405_84, + -0.268_941_42, + -0.188_770_33, + 0.0, + 0.311_229_67, + 0.731_058_6, + 1.761_594_2, + 2.857_722_4, + 4.966_535_6, + ], + &[ + -0.026_547_432, + -0.088_104_11, + -0.090_784_25, + 0.072_329_49, + 0.260_038_8, + 0.5, + 0.739_961_2, + 0.927_670_5, + 1.090_784_2, + 1.088_104_1, + 1.026_547_4, + ], + ); + } + + /// Mish. The derivative at 0 is `tanh(ln 2) = 0.6` + #[test] + fn mish_matches_reference() { + check_against_reference( + "Mish", + Activation::Mish, + &[ + -0.033_576_24, + -0.145_647_46, + -0.252_501_5, + -0.303_401_46, + -0.220_743_77, + 0.0, + 0.375_245_2, + 0.865_098_4, + 1.943_959, + 2.986_535, + 4.999_552, + ], + &[ + -0.026_747_498, + -0.093_393_12, + -0.108_355_09, + 0.059_216_756, + 0.289_510_68, + 0.6, + 0.886_424_4, + 1.049_036_2, + 1.069_317_9, + 1.021_107, + 1.000_800_2, + ], + ); + } + + /// The exact GELU keeps its relative precision in the negative tail + /// + /// At x = -12 the value is about -2.13e-32. The form `x * (1 + erf(x / sqrt(2))) / 2` + /// gives exactly 0, because `erf` rounds to -1 at f32 precision + #[test] + fn gelu_keeps_the_negative_tail() { + let input = tensor2(1, 1, vec![-12.0]); + let activation = Activation::GELU { approximate: false }; + let (output, cache) = activation.forward_train(input).expect("forward"); + let derivative = activation + .backward(&cache, &tensor2(1, 1, vec![1.0])) + .expect("backward"); + + for (name, got, want) in [ + ("value", output[[0, 0]], -2.131_778_5e-32_f32), + ("derivative", derivative[[0, 0]], -2.557_895_7e-31_f32), + ] { + assert!( + (got - want).abs() <= 1e-4 * want.abs(), + "GELU {name} at x = -12: got {got}, want {want}" + ); + } + } + + /// The derivative of the tanh approximation stays finite where the sigmoid saturates + /// + /// At `|x| = 1e30` the term `x * x` overflows to infinity. The derivative is 1 on the + /// positive side and 0 on the negative side + #[test] + fn gelu_tanh_derivative_stays_finite_at_huge_inputs() { + let activation = Activation::GELU { approximate: true }; + let input = tensor2(1, 4, vec![1e30, -1e30, f32::MAX, f32::MIN]); + let (_, cache) = activation.forward_train(input).expect("forward"); + let derivative = activation + .backward(&cache, &tensor2(1, 4, vec![1.0; 4])) + .expect("backward"); + let got: Vec = derivative.iter().copied().collect(); + assert_eq!(got, vec![1.0, 0.0, 1.0, 0.0]); + } + + /// The derivative of each pre-activation variant matches a central difference + /// + /// The difference runs in f64 over the f32 forward pass, so its error stays far below the + /// tolerance + #[test] + fn pre_activation_variants_match_a_central_difference() { + let points: Vec = (-60..=60).map(|i| i as f32 * 0.125).collect(); + let step = 1e-2_f32; + for activation in [ + Activation::GELU { approximate: false }, + Activation::GELU { approximate: true }, + Activation::SiLU, + Activation::Mish, + ] { + let input = tensor2(1, points.len(), points.clone()); + let (_, cache) = activation.forward_train(input).expect("forward"); + let ones = tensor2(1, points.len(), vec![1.0; points.len()]); + let derivative = activation.backward(&cache, &ones).expect("backward"); + + let shifted = |delta: f32| { + let moved = tensor2(1, points.len(), points.iter().map(|x| x + delta).collect()); + activation.forward(&moved).expect("forward") + }; + let (above, below) = (shifted(step), shifted(-step)); + for (i, &x) in points.iter().enumerate() { + let difference = + (above[[0, i]] as f64 - below[[0, i]] as f64) / (2.0 * step as f64); + let got = derivative[[0, i]] as f64; + assert!( + (got - difference).abs() <= 1e-3, + "{activation:?} at x = {x}: backward {got}, central difference {difference}" + ); + } + } + } + + /// Each variant names the tensor that its backward pass reads + #[test] + fn only_gelu_silu_and_mish_read_the_pre_activation() { + for activation in [ + Activation::GELU { approximate: false }, + Activation::GELU { approximate: true }, + Activation::SiLU, + Activation::Mish, + ] { + assert_eq!(activation.saves(), ActivationInput::PreActivation); + } + for activation in [ + Activation::Linear, + Activation::ReLU, + Activation::Sigmoid, + Activation::Tanh, + Activation::Softmax { + axis: DEFAULT_SOFTMAX_AXIS, + }, + Activation::LeakyReLU { + negative_slope: 0.3, + }, + Activation::ELU { alpha: 1.0 }, + Activation::SELU, + Activation::Softplus, + Activation::Softsign, + Activation::HardSigmoid, + Activation::Exponential, + ] { + assert_eq!(activation.saves(), ActivationInput::Output); + } + } + + /// The backward pass refuses a cache that holds the other tensor + /// + /// A pre-activation read as an output, or an output read as a pre-activation, gives a + /// wrong gradient with no other sign of the fault + #[test] + fn backward_refuses_a_cache_of_the_other_tensor() { + let tensor = tensor2(1, 3, vec![-1.0, 0.0, 1.0]); + let ones = tensor2(1, 3, vec![1.0; 3]); + let output_cache = ActivationCache::new(ActivationInput::Output, tensor.clone()); + let input_cache = ActivationCache::new(ActivationInput::PreActivation, tensor); + + for (activation, cache) in [ + (Activation::SiLU, &output_cache), + (Activation::Tanh, &input_cache), + ] { + assert!( + matches!( + activation.backward(cache, &ones), + Err(Error::InvalidInput(_)) + ), + "{activation:?} must refuse a {:?} cache", + cache.holds() + ); + assert!(matches!( + activation.output_of(cache), + Err(Error::InvalidInput(_)) + )); + } + } + + /// The backward pass refuses an upstream gradient of another shape, for every variant + #[test] + fn backward_refuses_a_gradient_of_another_shape() { + let ones = tensor2(1, 4, vec![1.0; 4]); + for activation in [Activation::ReLU, Activation::Mish] { + let (_, cache) = activation + .forward_train(tensor2(1, 3, vec![-1.0, 0.0, 1.0])) + .expect("forward"); + assert!(matches!( + activation.backward(&cache, &ones), + Err(Error::ShapeMismatch { .. }) + )); + } + } + + /// `output_of` gives the forward output for both kinds of cache + #[test] + fn output_of_recovers_the_forward_output() { + let input = tensor2(1, 4, vec![-2.0, -0.5, 0.5, 2.0]); + for activation in [Activation::Tanh, Activation::GELU { approximate: false }] { + let (output, cache) = activation.forward_train(input.clone()).expect("forward"); + assert_eq!(activation.output_of(&cache).expect("output_of"), output); + } + } } diff --git a/src/neural_network/layers/activation/elu.rs b/src/neural_network/layers/activation/elu.rs index 71159404..5b31567a 100644 --- a/src/neural_network/layers/activation/elu.rs +++ b/src/neural_network/layers/activation/elu.rs @@ -7,7 +7,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -128,22 +128,21 @@ impl UnaryLayer for ELU { let output = Activation::ELU { alpha: self.alpha }.forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("ELU", output.clone()); + ctx.push_cache( + "ELU", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("ELU")?; - - // ELU preserves shape, so gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("ELU")?; // ELU derivative is 1 for x > 0, and alpha * e^x, which equals `output + alpha`, below it - Activation::ELU { alpha: self.alpha }.backward(&output, grad_output) + Activation::ELU { alpha: self.alpha }.backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/activation/exponential.rs b/src/neural_network/layers/activation/exponential.rs index e58b9a0f..962e49db 100644 --- a/src/neural_network/layers/activation/exponential.rs +++ b/src/neural_network/layers/activation/exponential.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -103,22 +103,21 @@ impl UnaryLayer for Exponential { let output = Activation::Exponential.forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("Exponential", output.clone()); + ctx.push_cache( + "Exponential", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("Exponential")?; - - // Exponential preserves shape, so gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("Exponential")?; // Exponential derivative is e^x, which equals the cached output - Activation::Exponential.backward(&output, grad_output) + Activation::Exponential.backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/activation/gelu.rs b/src/neural_network/layers/activation/gelu.rs new file mode 100644 index 00000000..7e2412f1 --- /dev/null +++ b/src/neural_network/layers/activation/gelu.rs @@ -0,0 +1,169 @@ +//! GELU (Gaussian Error Linear Unit) activation layer +//! +//! The `GELU` struct holds the `approximate` flag and the shape recorded at build time. The +//! `with_approximate` method changes that flag after construction. `UnaryLayer::forward` +//! computes `Activation::GELU` and caches the input during training. `UnaryLayer::backward` +//! reads that cache to compute the gradient + +use crate::error::Error; +use crate::neural_network::layers::ParamCounts; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; +use crate::neural_network::layers::validation::start_build; +use crate::neural_network::layers::{ + built_layer_shape_functions, no_trainable_parameters_layer_functions, +}; +use crate::neural_network::traits::{LayerBase, UnaryLayer}; +use crate::neural_network::{Ctx, Shape, Tensor}; + +/// GELU (Gaussian Error Linear Unit) activation layer +/// +/// Applies `x * Phi(x)` elementwise to the input tensor, keeping the original shape. `Phi` is +/// the cumulative distribution function of the standard normal distribution. Common inputs +/// include 2D tensors for dense layers and 4D tensors for convolutional layers +/// +/// Hendrycks and Gimpel (2016) introduced this activation. The layer uses the exact `Phi` by +/// default. [`GELU::with_approximate`] selects the tanh approximation +/// `x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715 * x^3))) / 2` +/// +/// [`Activation::GELU`] provides the activation math. This layer only adds boundary +/// validation and the caching needed for backpropagation. The derivative has no closed form in +/// the output, so the layer caches the input +/// +/// # Examples +/// +/// ```rust +/// use rustyml::neural_network::Shape; +/// use rustyml::neural_network::sequential::SequentialBuilder; +/// use rustyml::neural_network::layers::activation::gelu::GELU; +/// use rustyml::neural_network::optimizers::*; +/// use rustyml::neural_network::losses::MeanSquaredError; +/// use ndarray::Array2; +/// +/// // Create a 2D input tensor +/// let x = Array2::from_shape_vec((2, 3), vec![-1.0, 2.0, -3.0, 4.0, -5.0, 6.0]) +/// .unwrap() +/// .into_dyn(); +/// +/// // Build a model with GELU activation +/// let mut model = SequentialBuilder::new() +/// .add(GELU::new()) +/// .build(&Shape::known(x.shape())) +/// .unwrap(); +/// model.compile(SGD::new(0.01, 0.0, false, 0.0).unwrap(), MeanSquaredError::new()); +/// +/// // Forward propagation +/// let output = model.predict(&x); +/// +/// // Output will be: [[-0.15865526, 1.9544997, -0.004049696], [3.9998734, -1.4332579e-6, 6.0]] +/// ``` +/// +/// Select the tanh approximation with the builder: +/// +/// ```rust +/// use rustyml::neural_network::layers::activation::gelu::GELU; +/// use rustyml::neural_network::traits::UnaryLayer; +/// use rustyml::neural_network::Ctx; +/// use ndarray::Array2; +/// +/// let x = Array2::from_shape_vec((1, 3), vec![-1.0, 0.0, 1.0]) +/// .unwrap() +/// .into_dyn(); +/// let mut layer = GELU::new().with_approximate(true); +/// let output = layer.forward_mut(&x, &mut Ctx::inference()).unwrap(); +/// +/// // Output will be: [[-0.158808, 0.0, 0.841192]] +/// ``` +#[derive(Debug)] +pub struct GELU { + /// When `true`, the layer uses the tanh approximation. When `false`, it uses the exact + /// `Phi` + pub(super) approximate: bool, + /// Shape the layer was built for, batch axis first. `None` before the build + built: Option, +} + +impl GELU { + /// Creates a new GELU activation layer that uses the exact `Phi` + /// + /// # Returns + /// + /// - `Self` - A new `GELU` layer + pub fn new() -> Self { + GELU { + approximate: false, + built: None, + } + } + + /// Selects the form of the activation + /// + /// # Parameters + /// + /// - `approximate` - `true` for the tanh approximation, and `false` for the exact `Phi` + /// + /// # Returns + /// + /// - `Self` - The layer with the selected form + pub fn with_approximate(mut self, approximate: bool) -> Self { + self.approximate = approximate; + self + } +} + +impl Default for GELU { + fn default() -> Self { + Self::new() + } +} + +impl LayerBase for GELU { + fn layer_type(&self) -> &str { + "GELU" + } + + built_layer_shape_functions!(); + + no_trainable_parameters_layer_functions!(); +} + +impl UnaryLayer for GELU { + /// Records the shape the layer serves. The layer holds no array, so nothing is allocated + fn build(&mut self, input: &Shape) -> Result<(), Error> { + let Some(built) = start_build(&self.built, "GELU", input)? else { + return Ok(()); + }; + self.compute_output_shape(&built)?; + self.built = Some(built); + Ok(()) + } + + fn forward(&self, input: &Tensor, ctx: &mut Ctx) -> Result { + if input.is_empty() { + return Err(Error::empty_input("input tensor")); + } + + let activation = Activation::GELU { + approximate: self.approximate, + }; + if !ctx.is_training() { + return activation.forward(input); + } + + // The backward pass reads the input, so the cache holds a copy of it + let (output, cache) = activation.forward_train(input.clone())?; + ctx.push_cache("GELU", cache); + + Ok(output) + } + + fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { + let cache: ActivationCache = ctx.pop_cache("GELU")?; + + // The derivative of the exact form is Phi(x) + x * phi(x), where phi is the standard + // normal density + Activation::GELU { + approximate: self.approximate, + } + .backward(&cache, grad_output) + } +} diff --git a/src/neural_network/layers/activation/hard_sigmoid.rs b/src/neural_network/layers/activation/hard_sigmoid.rs index 9fcca100..f5e98c0e 100644 --- a/src/neural_network/layers/activation/hard_sigmoid.rs +++ b/src/neural_network/layers/activation/hard_sigmoid.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -105,22 +105,21 @@ impl UnaryLayer for HardSigmoid { let output = Activation::HardSigmoid.forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("HardSigmoid", output.clone()); + ctx.push_cache( + "HardSigmoid", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("HardSigmoid")?; - - // Hard sigmoid preserves shape, so gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("HardSigmoid")?; // Hard sigmoid derivative is 1/6 on the linear segment, and 0 on both saturated ends - Activation::HardSigmoid.backward(&output, grad_output) + Activation::HardSigmoid.backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/activation/leaky_relu.rs b/src/neural_network/layers/activation/leaky_relu.rs index 23ca6b29..02f07e64 100644 --- a/src/neural_network/layers/activation/leaky_relu.rs +++ b/src/neural_network/layers/activation/leaky_relu.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -134,25 +134,24 @@ impl UnaryLayer for LeakyReLU { } .forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("LeakyReLU", output.clone()); + ctx.push_cache( + "LeakyReLU", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("LeakyReLU")?; - - // Leaky ReLU preserves shape, so gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("LeakyReLU")?; // Leaky ReLU derivative is 1 for x >= 0, and `negative_slope` below 0 Activation::LeakyReLU { negative_slope: self.negative_slope, } - .backward(&output, grad_output) + .backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/activation/mish.rs b/src/neural_network/layers/activation/mish.rs new file mode 100644 index 00000000..f804d1fc --- /dev/null +++ b/src/neural_network/layers/activation/mish.rs @@ -0,0 +1,122 @@ +//! Mish activation layer +//! +//! The `Mish` struct holds only the shape recorded at build time, because the layer takes no +//! parameter. `UnaryLayer::forward` computes `Activation::Mish` and caches the input during +//! training. `UnaryLayer::backward` reads that cache to compute the gradient + +use crate::error::Error; +use crate::neural_network::layers::ParamCounts; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; +use crate::neural_network::layers::validation::start_build; +use crate::neural_network::layers::{ + built_layer_shape_functions, no_trainable_parameters_layer_functions, +}; +use crate::neural_network::traits::{LayerBase, UnaryLayer}; +use crate::neural_network::{Ctx, Shape, Tensor}; + +/// Mish activation layer +/// +/// Applies `x * tanh(softplus(x))` elementwise to the input tensor, keeping the original shape. +/// Common inputs include 2D tensors for dense layers and 4D tensors for convolutional layers +/// +/// Misra (2019) introduced this activation. The output is smooth and not monotonic. Its +/// minimum is about -0.3088 at x = -1.1924 +/// +/// [`Activation::Mish`] provides the activation math. This layer only adds boundary +/// validation and the caching needed for backpropagation. The derivative has no closed form in +/// the output, so the layer caches the input +/// +/// # Examples +/// +/// ```rust +/// use rustyml::neural_network::Shape; +/// use rustyml::neural_network::sequential::SequentialBuilder; +/// use rustyml::neural_network::layers::activation::mish::Mish; +/// use rustyml::neural_network::optimizers::*; +/// use rustyml::neural_network::losses::MeanSquaredError; +/// use ndarray::Array2; +/// +/// // Create a 2D input tensor +/// let x = Array2::from_shape_vec((2, 3), vec![-1.0, 2.0, -3.0, 4.0, -5.0, 6.0]) +/// .unwrap() +/// .into_dyn(); +/// +/// // Build a model with Mish activation +/// let mut model = SequentialBuilder::new() +/// .add(Mish::new()) +/// .build(&Shape::known(x.shape())) +/// .unwrap(); +/// model.compile(SGD::new(0.01, 0.0, false, 0.0).unwrap(), MeanSquaredError::new()); +/// +/// // Forward propagation +/// let output = model.predict(&x); +/// +/// // Output will be: [[-0.30340147, 1.943959, -0.14564745], [3.997413, -0.033576235, 5.9999266]] +/// ``` +#[derive(Debug)] +pub struct Mish { + /// Shape the layer was built for, batch axis first. `None` before the build + built: Option, +} + +impl Mish { + /// Creates a new Mish activation layer + /// + /// # Returns + /// + /// - `Self` - A new `Mish` layer + pub fn new() -> Self { + Mish { built: None } + } +} + +impl Default for Mish { + fn default() -> Self { + Self::new() + } +} + +impl LayerBase for Mish { + fn layer_type(&self) -> &str { + "Mish" + } + + built_layer_shape_functions!(); + + no_trainable_parameters_layer_functions!(); +} + +impl UnaryLayer for Mish { + /// Records the shape the layer serves. The layer holds no array, so nothing is allocated + fn build(&mut self, input: &Shape) -> Result<(), Error> { + let Some(built) = start_build(&self.built, "Mish", input)? else { + return Ok(()); + }; + self.compute_output_shape(&built)?; + self.built = Some(built); + Ok(()) + } + + fn forward(&self, input: &Tensor, ctx: &mut Ctx) -> Result { + if input.is_empty() { + return Err(Error::empty_input("input tensor")); + } + + if !ctx.is_training() { + return Activation::Mish.forward(input); + } + + // The backward pass reads the input, so the cache holds a copy of it + let (output, cache) = Activation::Mish.forward_train(input.clone())?; + ctx.push_cache("Mish", cache); + + Ok(output) + } + + fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { + let cache: ActivationCache = ctx.pop_cache("Mish")?; + + // Mish derivative is t + x * (1 - t^2) * sigmoid(x), with t = tanh(softplus(x)) + Activation::Mish.backward(&cache, grad_output) + } +} diff --git a/src/neural_network/layers/activation/relu.rs b/src/neural_network/layers/activation/relu.rs index a3e7ad79..5cd3589a 100644 --- a/src/neural_network/layers/activation/relu.rs +++ b/src/neural_network/layers/activation/relu.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -100,22 +100,21 @@ impl UnaryLayer for ReLU { let output = Activation::ReLU.forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("ReLU", output.clone()); + ctx.push_cache( + "ReLU", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("ReLU")?; - - // ReLU preserves shape, so gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("ReLU")?; // ReLU derivative is 1 for x > 0, and 0 for x <= 0 - Activation::ReLU.backward(&output, grad_output) + Activation::ReLU.backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/activation/selu.rs b/src/neural_network/layers/activation/selu.rs index 98d77248..e01d8cd0 100644 --- a/src/neural_network/layers/activation/selu.rs +++ b/src/neural_network/layers/activation/selu.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -106,22 +106,21 @@ impl UnaryLayer for SELU { let output = Activation::SELU.forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("SELU", output.clone()); + ctx.push_cache( + "SELU", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("SELU")?; - - // SELU preserves shape, so gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("SELU")?; // SELU derivative is scale for x > 0, and scale * alpha * e^x for x <= 0 - Activation::SELU.backward(&output, grad_output) + Activation::SELU.backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/activation/sigmoid.rs b/src/neural_network/layers/activation/sigmoid.rs index 614bebca..4a0bb283 100644 --- a/src/neural_network/layers/activation/sigmoid.rs +++ b/src/neural_network/layers/activation/sigmoid.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -103,22 +103,21 @@ impl UnaryLayer for Sigmoid { // Large-magnitude inputs saturate to 0/1 by construction let output = Activation::Sigmoid.forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("Sigmoid", output.clone()); + ctx.push_cache( + "Sigmoid", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("Sigmoid")?; - - // Sigmoid preserves shape, so gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("Sigmoid")?; // Sigmoid derivative is f(x) * (1 - f(x)) - Activation::Sigmoid.backward(&output, grad_output) + Activation::Sigmoid.backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/activation/silu.rs b/src/neural_network/layers/activation/silu.rs new file mode 100644 index 00000000..694cc155 --- /dev/null +++ b/src/neural_network/layers/activation/silu.rs @@ -0,0 +1,122 @@ +//! SiLU (Sigmoid Linear Unit) activation layer, also named Swish +//! +//! The `SiLU` struct holds only the shape recorded at build time, because the layer takes no +//! parameter. `UnaryLayer::forward` computes `Activation::SiLU` and caches the input during +//! training. `UnaryLayer::backward` reads that cache to compute the gradient + +use crate::error::Error; +use crate::neural_network::layers::ParamCounts; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; +use crate::neural_network::layers::validation::start_build; +use crate::neural_network::layers::{ + built_layer_shape_functions, no_trainable_parameters_layer_functions, +}; +use crate::neural_network::traits::{LayerBase, UnaryLayer}; +use crate::neural_network::{Ctx, Shape, Tensor}; + +/// SiLU (Sigmoid Linear Unit) activation layer +/// +/// Applies `x * sigmoid(x)` elementwise to the input tensor, keeping the original shape. +/// Common inputs include 2D tensors for dense layers and 4D tensors for convolutional layers +/// +/// Elfwing et al. (2018) introduced this activation, and Ramachandran et al. (2017) named it +/// Swish. The output is smooth and not monotonic. Its minimum is about -0.2785 at x = -1.2785 +/// +/// [`Activation::SiLU`] provides the activation math. This layer only adds boundary +/// validation and the caching needed for backpropagation. The derivative has no closed form in +/// the output, so the layer caches the input +/// +/// # Examples +/// +/// ```rust +/// use rustyml::neural_network::Shape; +/// use rustyml::neural_network::sequential::SequentialBuilder; +/// use rustyml::neural_network::layers::activation::silu::SiLU; +/// use rustyml::neural_network::optimizers::*; +/// use rustyml::neural_network::losses::MeanSquaredError; +/// use ndarray::Array2; +/// +/// // Create a 2D input tensor +/// let x = Array2::from_shape_vec((2, 3), vec![-1.0, 2.0, -3.0, 4.0, -5.0, 6.0]) +/// .unwrap() +/// .into_dyn(); +/// +/// // Build a model with SiLU activation +/// let mut model = SequentialBuilder::new() +/// .add(SiLU::new()) +/// .build(&Shape::known(x.shape())) +/// .unwrap(); +/// model.compile(SGD::new(0.01, 0.0, false, 0.0).unwrap(), MeanSquaredError::new()); +/// +/// // Forward propagation +/// let output = model.predict(&x); +/// +/// // Output will be: [[-0.26894143, 1.761594, -0.14227763], [3.928055, -0.033464253, 5.9851646]] +/// ``` +#[derive(Debug)] +pub struct SiLU { + /// Shape the layer was built for, batch axis first. `None` before the build + built: Option, +} + +impl SiLU { + /// Creates a new SiLU activation layer + /// + /// # Returns + /// + /// - `Self` - A new `SiLU` layer + pub fn new() -> Self { + SiLU { built: None } + } +} + +impl Default for SiLU { + fn default() -> Self { + Self::new() + } +} + +impl LayerBase for SiLU { + fn layer_type(&self) -> &str { + "SiLU" + } + + built_layer_shape_functions!(); + + no_trainable_parameters_layer_functions!(); +} + +impl UnaryLayer for SiLU { + /// Records the shape the layer serves. The layer holds no array, so nothing is allocated + fn build(&mut self, input: &Shape) -> Result<(), Error> { + let Some(built) = start_build(&self.built, "SiLU", input)? else { + return Ok(()); + }; + self.compute_output_shape(&built)?; + self.built = Some(built); + Ok(()) + } + + fn forward(&self, input: &Tensor, ctx: &mut Ctx) -> Result { + if input.is_empty() { + return Err(Error::empty_input("input tensor")); + } + + if !ctx.is_training() { + return Activation::SiLU.forward(input); + } + + // The backward pass reads the input, so the cache holds a copy of it + let (output, cache) = Activation::SiLU.forward_train(input.clone())?; + ctx.push_cache("SiLU", cache); + + Ok(output) + } + + fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { + let cache: ActivationCache = ctx.pop_cache("SiLU")?; + + // SiLU derivative is sigmoid(x) * (1 + x * (1 - sigmoid(x))) + Activation::SiLU.backward(&cache, grad_output) + } +} diff --git a/src/neural_network/layers/activation/softmax.rs b/src/neural_network/layers/activation/softmax.rs index bf225ad2..02be7084 100644 --- a/src/neural_network/layers/activation/softmax.rs +++ b/src/neural_network/layers/activation/softmax.rs @@ -7,7 +7,9 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::{Activation, DEFAULT_SOFTMAX_AXIS}; +use crate::neural_network::layers::activation::{ + Activation, ActivationCache, ActivationInput, DEFAULT_SOFTMAX_AXIS, +}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -146,21 +148,20 @@ impl UnaryLayer for Softmax { // The axis resolves against the rank of this input, and an out-of-range axis fails here let output = Activation::Softmax { axis: self.axis }.forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("Softmax", output.clone()); + ctx.push_cache( + "Softmax", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("Softmax")?; - - // Softmax preserves shape, so the gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("Softmax")?; - Activation::Softmax { axis: self.axis }.backward(&output, grad_output) + Activation::Softmax { axis: self.axis }.backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/activation/softplus.rs b/src/neural_network/layers/activation/softplus.rs index 519bc3b1..69559eaa 100644 --- a/src/neural_network/layers/activation/softplus.rs +++ b/src/neural_network/layers/activation/softplus.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -104,22 +104,21 @@ impl UnaryLayer for Softplus { let output = Activation::Softplus.forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("Softplus", output.clone()); + ctx.push_cache( + "Softplus", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("Softplus")?; - - // Softplus preserves shape, so gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("Softplus")?; // Softplus derivative is 1 / (1 + e^(-x)) - Activation::Softplus.backward(&output, grad_output) + Activation::Softplus.backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/activation/softsign.rs b/src/neural_network/layers/activation/softsign.rs index f5053122..a7be885c 100644 --- a/src/neural_network/layers/activation/softsign.rs +++ b/src/neural_network/layers/activation/softsign.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -104,22 +104,21 @@ impl UnaryLayer for Softsign { let output = Activation::Softsign.forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("Softsign", output.clone()); + ctx.push_cache( + "Softsign", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("Softsign")?; - - // Softsign preserves shape, so gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("Softsign")?; // Softsign derivative is 1 / (1 + |x|)^2 - Activation::Softsign.backward(&output, grad_output) + Activation::Softsign.backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/activation/tanh.rs b/src/neural_network/layers/activation/tanh.rs index d080449e..730ac5e2 100644 --- a/src/neural_network/layers/activation/tanh.rs +++ b/src/neural_network/layers/activation/tanh.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::start_build; use crate::neural_network::layers::{ built_layer_shape_functions, no_trainable_parameters_layer_functions, @@ -103,22 +103,21 @@ impl UnaryLayer for Tanh { // Large-magnitude inputs saturate toward -1/+1 by construction let output = Activation::Tanh.forward(input)?; + // The backward pass reads the output, so the cache holds a copy of it if ctx.is_training() { - ctx.push_cache("Tanh", output.clone()); + ctx.push_cache( + "Tanh", + ActivationCache::new(ActivationInput::Output, output.clone()), + ); } Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { - let output: Tensor = ctx.pop_cache("Tanh")?; - - // Tanh preserves shape, so gradient must match the cached output - if grad_output.shape() != output.shape() { - return Err(Error::shape_mismatch(output.shape(), grad_output.shape())); - } + let cache: ActivationCache = ctx.pop_cache("Tanh")?; // Tanh derivative is 1 - tanh(x)^2 - Activation::Tanh.backward(&output, grad_output) + Activation::Tanh.backward(&cache, grad_output) } } diff --git a/src/neural_network/layers/convolution/conv_1d.rs b/src/neural_network/layers/convolution/conv_1d.rs index 6a2c6894..7005dfc0 100644 --- a/src/neural_network/layers/convolution/conv_1d.rs +++ b/src/neural_network/layers/convolution/conv_1d.rs @@ -5,7 +5,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; #[doc(inline)] pub use crate::neural_network::layers::convolution::convolution_engine::ConvPadding; use crate::neural_network::layers::convolution::convolution_engine::{ @@ -345,8 +345,8 @@ impl Conv1D { struct Conv1DCache { /// The input tensor the forward pass received input: Tensor, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the backward pass of the activation reads + activation: ActivationCache, } impl LayerBase for Conv1D { @@ -424,25 +424,26 @@ impl UnaryLayer for Conv1D { &[self.dilation_rate], self.padding, )?; - let activated = self.activation.forward(&output)?; - - if ctx.is_training() { - ctx.push_cache( - "Conv1D", - Conv1DCache { - input: input.clone(), - output: activated.clone(), - }, - ); + if !ctx.is_training() { + return self.activation.forward(&output); } + let (activated, activation) = self.activation.forward_train(output)?; + ctx.push_cache( + "Conv1D", + Conv1DCache { + input: input.clone(), + activation, + }, + ); + Ok(activated) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: Conv1DCache = ctx.pop_cache("Conv1D")?; - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; let grads = conv_backward( &grad_upstream, diff --git a/src/neural_network/layers/convolution/conv_1d_transpose.rs b/src/neural_network/layers/convolution/conv_1d_transpose.rs index c1fe6492..e81e3546 100644 --- a/src/neural_network/layers/convolution/conv_1d_transpose.rs +++ b/src/neural_network/layers/convolution/conv_1d_transpose.rs @@ -5,7 +5,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; use crate::neural_network::layers::convolution::PaddingType; use crate::neural_network::layers::convolution::conv_transpose_engine::{ conv_transpose_backward, conv_transpose_forward, transpose_output_length, @@ -346,8 +346,8 @@ impl Conv1DTranspose { struct Conv1DTransposeCache { /// The input the forward pass received, which the weight gradient reads input: Tensor, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the backward pass of the activation reads + activation: ActivationCache, } impl LayerBase for Conv1DTranspose { @@ -422,24 +422,25 @@ impl UnaryLayer for Conv1DTranspose { &[self.dilation_rate], self.padding, )?; - let activated = self.activation.forward(&output)?; - - if ctx.is_training() { - ctx.push_cache( - "Conv1DTranspose", - Conv1DTransposeCache { - input: input.clone(), - output: activated.clone(), - }, - ); + if !ctx.is_training() { + return self.activation.forward(&output); } + let (activated, activation) = self.activation.forward_train(output)?; + ctx.push_cache( + "Conv1DTranspose", + Conv1DTransposeCache { + input: input.clone(), + activation, + }, + ); + Ok(activated) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: Conv1DTransposeCache = ctx.pop_cache("Conv1DTranspose")?; - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; let grads = conv_transpose_backward( &grad_upstream, diff --git a/src/neural_network/layers/convolution/conv_2d.rs b/src/neural_network/layers/convolution/conv_2d.rs index bba73746..2a2f7cf7 100644 --- a/src/neural_network/layers/convolution/conv_2d.rs +++ b/src/neural_network/layers/convolution/conv_2d.rs @@ -5,7 +5,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; use crate::neural_network::layers::convolution::PaddingType; use crate::neural_network::layers::convolution::convolution_engine::{ conv_backward, conv_forward, effective_kernel, @@ -364,8 +364,8 @@ impl Conv2D { struct Conv2DCache { /// The input tensor the forward pass received input: Tensor, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the backward pass of the activation reads + activation: ActivationCache, } impl LayerBase for Conv2D { @@ -443,25 +443,26 @@ impl UnaryLayer for Conv2D { &[self.dilation_rate.0, self.dilation_rate.1], self.padding.into(), )?; - let activated = self.activation.forward(&output)?; - - if ctx.is_training() { - ctx.push_cache( - "Conv2D", - Conv2DCache { - input: input.clone(), - output: activated.clone(), - }, - ); + if !ctx.is_training() { + return self.activation.forward(&output); } + let (activated, activation) = self.activation.forward_train(output)?; + ctx.push_cache( + "Conv2D", + Conv2DCache { + input: input.clone(), + activation, + }, + ); + Ok(activated) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: Conv2DCache = ctx.pop_cache("Conv2D")?; - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; let grads = conv_backward( &grad_upstream, diff --git a/src/neural_network/layers/convolution/conv_2d_transpose.rs b/src/neural_network/layers/convolution/conv_2d_transpose.rs index 7cf5329a..5e41ce45 100644 --- a/src/neural_network/layers/convolution/conv_2d_transpose.rs +++ b/src/neural_network/layers/convolution/conv_2d_transpose.rs @@ -5,7 +5,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; use crate::neural_network::layers::convolution::PaddingType; use crate::neural_network::layers::convolution::conv_transpose_engine::{ conv_transpose_backward, conv_transpose_forward, transpose_output_length, @@ -369,8 +369,8 @@ impl Conv2DTranspose { struct Conv2DTransposeCache { /// The input the forward pass received, which the weight gradient reads input: Tensor, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the backward pass of the activation reads + activation: ActivationCache, } impl LayerBase for Conv2DTranspose { @@ -445,24 +445,25 @@ impl UnaryLayer for Conv2DTranspose { &[self.dilation_rate.0, self.dilation_rate.1], self.padding, )?; - let activated = self.activation.forward(&output)?; - - if ctx.is_training() { - ctx.push_cache( - "Conv2DTranspose", - Conv2DTransposeCache { - input: input.clone(), - output: activated.clone(), - }, - ); + if !ctx.is_training() { + return self.activation.forward(&output); } + let (activated, activation) = self.activation.forward_train(output)?; + ctx.push_cache( + "Conv2DTranspose", + Conv2DTransposeCache { + input: input.clone(), + activation, + }, + ); + Ok(activated) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: Conv2DTransposeCache = ctx.pop_cache("Conv2DTranspose")?; - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; let grads = conv_transpose_backward( &grad_upstream, diff --git a/src/neural_network/layers/convolution/conv_3d.rs b/src/neural_network/layers/convolution/conv_3d.rs index 683f1459..e48df2c2 100644 --- a/src/neural_network/layers/convolution/conv_3d.rs +++ b/src/neural_network/layers/convolution/conv_3d.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; use crate::neural_network::layers::convolution::PaddingType; use crate::neural_network::layers::convolution::convolution_engine::{ conv_backward, conv_forward, effective_kernel, @@ -370,8 +370,8 @@ impl Conv3D { struct Conv3DCache { /// The input tensor the forward pass received input: Tensor, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the backward pass of the activation reads + activation: ActivationCache, } impl LayerBase for Conv3D { @@ -453,25 +453,26 @@ impl UnaryLayer for Conv3D { ], self.padding.into(), )?; - let activated = self.activation.forward(&output)?; - - if ctx.is_training() { - ctx.push_cache( - "Conv3D", - Conv3DCache { - input: input.clone(), - output: activated.clone(), - }, - ); + if !ctx.is_training() { + return self.activation.forward(&output); } + let (activated, activation) = self.activation.forward_train(output)?; + ctx.push_cache( + "Conv3D", + Conv3DCache { + input: input.clone(), + activation, + }, + ); + Ok(activated) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: Conv3DCache = ctx.pop_cache("Conv3D")?; - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; let grads = conv_backward( &grad_upstream, diff --git a/src/neural_network/layers/convolution/conv_3d_transpose.rs b/src/neural_network/layers/convolution/conv_3d_transpose.rs index cff2e4b2..69b84b8e 100644 --- a/src/neural_network/layers/convolution/conv_3d_transpose.rs +++ b/src/neural_network/layers/convolution/conv_3d_transpose.rs @@ -5,7 +5,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; use crate::neural_network::layers::convolution::PaddingType; use crate::neural_network::layers::convolution::conv_transpose_engine::{ conv_transpose_backward, conv_transpose_forward, transpose_output_length, @@ -378,8 +378,8 @@ impl Conv3DTranspose { struct Conv3DTransposeCache { /// The input the forward pass received, which the weight gradient reads input: Tensor, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the backward pass of the activation reads + activation: ActivationCache, } impl LayerBase for Conv3DTranspose { @@ -458,24 +458,25 @@ impl UnaryLayer for Conv3DTranspose { ], self.padding, )?; - let activated = self.activation.forward(&output)?; - - if ctx.is_training() { - ctx.push_cache( - "Conv3DTranspose", - Conv3DTransposeCache { - input: input.clone(), - output: activated.clone(), - }, - ); + if !ctx.is_training() { + return self.activation.forward(&output); } + let (activated, activation) = self.activation.forward_train(output)?; + ctx.push_cache( + "Conv3DTranspose", + Conv3DTransposeCache { + input: input.clone(), + activation, + }, + ); + Ok(activated) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: Conv3DTransposeCache = ctx.pop_cache("Conv3DTranspose")?; - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; let grads = conv_transpose_backward( &grad_upstream, diff --git a/src/neural_network/layers/convolution/depthwise_conv_1d.rs b/src/neural_network/layers/convolution/depthwise_conv_1d.rs index 539bbd71..fac1c6fe 100644 --- a/src/neural_network/layers/convolution/depthwise_conv_1d.rs +++ b/src/neural_network/layers/convolution/depthwise_conv_1d.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; use crate::neural_network::layers::conv_op_helpers::{ DepthwiseGeometry, depthwise_backward, depthwise_forward, }; @@ -385,10 +385,10 @@ impl DepthwiseConv1D { }) } - /// Depthwise convolution over a channels-last sequence, followed by the activation + /// Depthwise convolution over a channels-last sequence, with no activation /// - /// The numeric body of [`UnaryLayer::forward`]. A training pass parks the input and the - /// output of this call in the context, and an inference pass parks nothing + /// The numeric body of [`UnaryLayer::forward`]. The result is the pre-activation `z`. + /// [`UnaryLayer::forward`] applies the activation to it fn convolve(&self, input: &Tensor) -> Result { validate_built_input(&self.built, "DepthwiseConv1D", input.shape())?; validate_valid_kernel_fits( @@ -419,7 +419,7 @@ impl DepthwiseConv1D { output.as_slice_mut().expect("output is contiguous"), ); - self.activation.forward(&output.into_dyn()) + Ok(output.into_dyn()) } } @@ -427,8 +427,8 @@ impl DepthwiseConv1D { struct DepthwiseConv1DCache { /// The tensor the forward pass received input: Tensor, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the activation backward pass reads, as [`Activation::saves`] names + activation: ActivationCache, } impl LayerBase for DepthwiseConv1D { @@ -496,23 +496,25 @@ impl UnaryLayer for DepthwiseConv1D { if !self.is_built() { return Err(Error::not_built("DepthwiseConv1D")); } - let activated = self.convolve(input)?; - // Park only after a successful convolution, so a rejected input leaves no partial state - if ctx.is_training() { - ctx.push_cache( - "DepthwiseConv1D", - DepthwiseConv1DCache { - input: input.clone(), - output: activated.clone(), - }, - ); + let z = self.convolve(input)?; + if !ctx.is_training() { + return self.activation.forward(&z); } + let (activated, activation) = self.activation.forward_train(z)?; + // Park only after a successful forward pass, so a rejected input leaves no partial state + ctx.push_cache( + "DepthwiseConv1D", + DepthwiseConv1DCache { + input: input.clone(), + activation, + }, + ); Ok(activated) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: DepthwiseConv1DCache = ctx.pop_cache("DepthwiseConv1D")?; - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; let input = &cache.input; diff --git a/src/neural_network/layers/convolution/depthwise_conv_2d.rs b/src/neural_network/layers/convolution/depthwise_conv_2d.rs index 4643d1d2..a809af51 100644 --- a/src/neural_network/layers/convolution/depthwise_conv_2d.rs +++ b/src/neural_network/layers/convolution/depthwise_conv_2d.rs @@ -6,7 +6,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; use crate::neural_network::layers::conv_op_helpers::{ DepthwiseGeometry, depthwise_backward, depthwise_forward, }; @@ -413,10 +413,10 @@ impl DepthwiseConv2D { ) } - /// Depthwise convolution over a channels-last tensor, followed by the activation + /// Depthwise convolution over a channels-last tensor, with no activation /// - /// The numeric body of [`UnaryLayer::forward`]. A training pass parks the input and the - /// output of this call in the context, and an inference pass parks nothing + /// The numeric body of [`UnaryLayer::forward`]. The result is the pre-activation `z`. + /// [`UnaryLayer::forward`] applies the activation to it fn convolve(&self, input: &Tensor) -> Result { validate_built_input(&self.built, "DepthwiseConv2D", input.shape())?; validate_valid_kernel_fits( @@ -448,7 +448,7 @@ impl DepthwiseConv2D { output.as_slice_mut().expect("output is contiguous"), ); - self.activation.forward(&output.into_dyn()) + Ok(output.into_dyn()) } } @@ -456,8 +456,8 @@ impl DepthwiseConv2D { struct DepthwiseConv2DCache { /// The tensor the forward pass received input: Tensor, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the activation backward pass reads, as [`Activation::saves`] names + activation: ActivationCache, } impl LayerBase for DepthwiseConv2D { @@ -525,23 +525,25 @@ impl UnaryLayer for DepthwiseConv2D { if !self.is_built() { return Err(Error::not_built("DepthwiseConv2D")); } - let activated = self.convolve(input)?; - // Park only after a successful convolution, so a rejected input leaves no partial state - if ctx.is_training() { - ctx.push_cache( - "DepthwiseConv2D", - DepthwiseConv2DCache { - input: input.clone(), - output: activated.clone(), - }, - ); + let z = self.convolve(input)?; + if !ctx.is_training() { + return self.activation.forward(&z); } + let (activated, activation) = self.activation.forward_train(z)?; + // Park only after a successful forward pass, so a rejected input leaves no partial state + ctx.push_cache( + "DepthwiseConv2D", + DepthwiseConv2DCache { + input: input.clone(), + activation, + }, + ); Ok(activated) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: DepthwiseConv2DCache = ctx.pop_cache("DepthwiseConv2D")?; - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; let input = &cache.input; diff --git a/src/neural_network/layers/convolution/separable_conv_1d.rs b/src/neural_network/layers/convolution/separable_conv_1d.rs index c55332b8..a4e4cc82 100644 --- a/src/neural_network/layers/convolution/separable_conv_1d.rs +++ b/src/neural_network/layers/convolution/separable_conv_1d.rs @@ -8,7 +8,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; use crate::neural_network::layers::conv_op_helpers::{ DepthwiseGeometry, depthwise_backward, depthwise_forward, }; @@ -488,8 +488,8 @@ struct SeparableConv1DCache { input: Tensor, /// The output of the depthwise stage, which the pointwise backward pass reads depthwise_output: Tensor, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the activation backward pass reads, as [`Activation::saves`] names + activation: ActivationCache, } impl LayerBase for SeparableConv1D { @@ -578,27 +578,28 @@ impl UnaryLayer for SeparableConv1D { let depthwise_output = self.depthwise_convolve(input)?; let output = self.pointwise_convolve(&depthwise_output); - let activated = self.activation.forward(&output)?; + if !ctx.is_training() { + return self.activation.forward(&output); + } + let (activated, activation) = self.activation.forward_train(output)?; // Park only after a successful pass, so a rejected input leaves no partial state. The - // depthwise output reaches the backward pass alone - if ctx.is_training() { - ctx.push_cache( - "SeparableConv1D", - SeparableConv1DCache { - input: input.clone(), - depthwise_output, - output: activated.clone(), - }, - ); - } + // depthwise output moves into the cache with no copy + ctx.push_cache( + "SeparableConv1D", + SeparableConv1DCache { + input: input.clone(), + depthwise_output, + activation, + }, + ); Ok(activated) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: SeparableConv1DCache = ctx.pop_cache("SeparableConv1D")?; - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; let input = &cache.input; let depthwise_output = &cache.depthwise_output; diff --git a/src/neural_network/layers/convolution/separable_conv_2d.rs b/src/neural_network/layers/convolution/separable_conv_2d.rs index a4ed40dc..2b0b523c 100644 --- a/src/neural_network/layers/convolution/separable_conv_2d.rs +++ b/src/neural_network/layers/convolution/separable_conv_2d.rs @@ -8,7 +8,7 @@ use crate::error::Error; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache}; use crate::neural_network::layers::conv_op_helpers::{ DepthwiseGeometry, depthwise_backward, depthwise_forward, }; @@ -559,8 +559,8 @@ struct SeparableConv2DCache { input: Tensor, /// The output of the depthwise stage, which the pointwise backward pass reads depthwise_output: Tensor, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the activation backward pass reads, as [`Activation::saves`] names + activation: ActivationCache, } impl LayerBase for SeparableConv2D { @@ -649,27 +649,28 @@ impl UnaryLayer for SeparableConv2D { let depthwise_output = self.depthwise_convolve(input); let output = self.pointwise_convolve(&depthwise_output); - let activated = self.activation.forward(&output.into_dyn())?; + if !ctx.is_training() { + return self.activation.forward(&output.into_dyn()); + } + let (activated, activation) = self.activation.forward_train(output.into_dyn())?; // Park only after a successful pass, so a rejected input leaves no partial state. The - // depthwise output reaches the backward pass alone - if ctx.is_training() { - ctx.push_cache( - "SeparableConv2D", - SeparableConv2DCache { - input: input.clone(), - depthwise_output, - output: activated.clone(), - }, - ); - } + // depthwise output moves into the cache with no copy + ctx.push_cache( + "SeparableConv2D", + SeparableConv2DCache { + input: input.clone(), + depthwise_output, + activation, + }, + ); Ok(activated) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: SeparableConv2DCache = ctx.pop_cache("SeparableConv2D")?; - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; let input = &cache.input; let depthwise_output = &cache.depthwise_output; diff --git a/src/neural_network/layers/dense.rs b/src/neural_network/layers/dense.rs index 7dc9c09b..8b0c0105 100644 --- a/src/neural_network/layers/dense.rs +++ b/src/neural_network/layers/dense.rs @@ -2,7 +2,7 @@ use crate::error::{Context, Error}; use crate::neural_network::layers::ParamCounts; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::validation::{ start_build, validate_optional_weight, validate_weight_shape, }; @@ -293,16 +293,15 @@ impl Dense { .context("Failed to fold the leading axes of a Dense tensor") } - /// The layer's full forward transform, `activation(input * weights + bias)`, which - /// [`UnaryLayer::forward`] runs + /// The linear transform `input * weights + bias`, with `ReLU` fused into it /// /// The bias add rides the GEMM epilogue, so the pre-activation is written exactly once, /// with no separate broadcast add of the bias. The bias is the per-column addend, so it /// lowers to [`Bias::PerCol`] /// - /// The activation stays a fused epilogue only for `ReLU` ([`FusedActivation::Relu`]). - /// `Linear` needs no separate pass, since it changes nothing. Every other activation runs - /// as a separate [`Activation::forward`] pass + /// For `ReLU` the epilogue also applies the activation ([`FusedActivation::Relu`]), and the + /// result is the activated output. For every other activation the result is the + /// pre-activation `z`, and the caller applies the activation (see [`Dense::fuses`]) /// /// A fused `f32` epilogue matches the unfused product plus scalar activation bit for bit, /// with 1 exception. gemmkit's fused `Relu` maps `NaN` to `0.0`, while this crate's @@ -310,9 +309,10 @@ impl Dense { /// and the cached-output ReLU derivative then treats the resulting `0.0` as a dead unit /// /// The input arrives folded, and the result goes back to the rank of `input_shape`. A - /// fold keeps each last-axis lane whole, so the activation gives the same result on the - /// folded matrix as on the restored tensor. This holds for `Softmax` because an embedded - /// `Softmax` must carry the last axis. [`Activation::validate`] rejects any other axis + /// fold keeps each last-axis lane whole. An activation that the caller applies to the + /// restored tensor therefore gives the same result as on the folded matrix. This holds for + /// `Softmax` because an embedded `Softmax` must carry the last axis. + /// [`Activation::validate`] rejects any other axis /// /// # Parameters /// @@ -322,8 +322,9 @@ impl Dense { /// /// # Returns /// - /// - `Result` - Activated output. It has the shape of `input_shape` with - /// the last axis replaced by `output_dim` + /// - `Result` - The activated output for `ReLU` and `Linear`, and the + /// pre-activation for every other activation. It has the shape of `input_shape` with the + /// last axis replaced by `output_dim` /// /// # Panics /// @@ -333,8 +334,7 @@ impl Dense { /// /// # Errors /// - /// - `Error::Computation` - Softmax failed to reshape the fused pre-activation, or the - /// result failed to go back to the rank of `input_shape` + /// - `Error::Computation` - The result failed to go back to the rank of `input_shape` fn project(&self, input: &ArrayView2<'_, f32>, input_shape: &[usize]) -> Result { // A bias-free layer passes no epilogue at all, so the product is exactly the product. // A zero addend would give the same value for every input except a negative zero @@ -359,21 +359,28 @@ impl Dense { Parallelism::Rayon(0), ); - let output = output.into_dyn(); - let activated = match self.activation { - Activation::Linear | Activation::ReLU => output, - _ => self.activation.forward(&output)?, - }; - - // The fused product writes C order, and so does every activation, so this call only - // relabels the axes. It never copies + // The fused product writes C order, so this call only relabels the axes. It never + // copies let mut output_shape = input_shape.to_vec(); let last_axis = output_shape.len() - 1; output_shape[last_axis] = self.output_dim; - activated + output + .into_dyn() .into_shape_with_order(output_shape) .context("Failed to restore the rank of the Dense output") } + + /// Whether [`Dense::project`] already gives the activated output + /// + /// `ReLU` runs in the GEMM epilogue, and `Linear` changes nothing. Every other activation + /// needs a separate pass + /// + /// # Returns + /// + /// - `bool` - `true` for `ReLU` and `Linear` + fn fuses(&self) -> bool { + matches!(self.activation, Activation::Linear | Activation::ReLU) + } } /// What the forward pass of [`Dense`] parks for its backward pass @@ -382,8 +389,8 @@ struct DenseCache { input: Array2, /// Shape of the tensor the forward pass received, to restore the rank of the gradient input_shape: Vec, - /// The activated output, to backpropagate through the activation - output: Tensor, + /// The tensor that the backward pass of the activation reads + activation: ActivationCache, } impl LayerBase for Dense { @@ -462,37 +469,55 @@ impl UnaryLayer for Dense { /// /// The input has rank 2 or more. The leading axes fold into 1 row axis, so every rank uses /// the same 1 matrix product for the same row count + /// + /// In training, the layer caches 1 activation tensor. GELU, SiLU, and Mish read the + /// pre-activation, and the layer moves it into the cache with no copy. Every other + /// activation reads the output, and the layer caches a copy of the output. Inference + /// caches nothing and copies nothing fn forward(&self, input: &Tensor, ctx: &mut Ctx) -> Result { if self.built.is_none() { return Err(Error::not_built("Dense")); } let input_2d = Self::fold(input, self.input_dim, "input")?; - let output = self.project(&input_2d.view(), input.shape())?; - - if ctx.is_training() { - ctx.push_cache( - "Dense", - DenseCache { - input: input_2d.into_owned(), - input_shape: input.shape().to_vec(), - output: output.clone(), - }, - ); + let projected = self.project(&input_2d.view(), input.shape())?; + + if !ctx.is_training() { + if self.fuses() { + return Ok(projected); + } + return self.activation.forward(&projected); } + // A fused activation leaves no pre-activation. `ReLU` and `Linear` both read the + // output, so a copy of the output is a valid cache for them + let (output, activation) = if self.fuses() { + let cache = ActivationCache::new(ActivationInput::Output, projected.clone()); + (projected, cache) + } else { + self.activation.forward_train(projected)? + }; + ctx.push_cache( + "Dense", + DenseCache { + input: input_2d.into_owned(), + input_shape: input.shape().to_vec(), + activation, + }, + ); + Ok(output) } fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result { let cache: DenseCache = ctx.pop_cache("Dense")?; - if grad_output.shape() != cache.output.shape() { + if grad_output.shape() != cache.activation.shape() { return Err(Error::shape_mismatch( - cache.output.shape(), + cache.activation.shape(), grad_output.shape(), )); } - let grad_upstream = self.activation.backward(&cache.output, grad_output)?; + let grad_upstream = self.activation.backward(&cache.activation, grad_output)?; // Fold the upstream gradient to [rows, output_dim]. Both operands of the weight // gradient are then 2D, and the bias gradient sums over 1 row axis that already diff --git a/src/neural_network/layers/recurrent/cell.rs b/src/neural_network/layers/recurrent/cell.rs index ad263b12..2e9b2259 100644 --- a/src/neural_network/layers/recurrent/cell.rs +++ b/src/neural_network/layers/recurrent/cell.rs @@ -13,9 +13,10 @@ //! implementations, so the shape of the protocol can change with no effect on any public name. use crate::error::Error; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput}; use crate::neural_network::layers::recurrent::gate::FusedGates; -use ndarray::{Array2, ArrayView2}; +use ndarray::{Array2, ArrayView2, Ix2}; +use std::borrow::Cow; /// 1 group of gate blocks, and the array that their recurrent projection consumed /// @@ -68,12 +69,6 @@ pub(crate) trait RnnCell: std::fmt::Debug + Send + Sync + 'static { /// carries the cell state at slot 1. const STATE_COUNT: usize; - /// Number of arrays that the cell parks per timestep for its own backward pass - /// - /// A value that the base already keeps does not belong here. The hidden state of every step - /// is 1 such value, and [`RnnCell::step_backward`] receives it directly. - const RECORD_SLOTS: usize; - /// The map from a group of gate blocks to the array that the group projected /// /// The groups must partition `0..GATE_BIASES.len()`. [`Rnn::build`](super::rnn::Rnn) asserts @@ -91,11 +86,23 @@ pub(crate) trait RnnCell: std::fmt::Debug + Send + Sync + 'static { /// - `Self` - The cell fn new(activation: Activation) -> Self; + /// Number of arrays that the cell parks per timestep for its own backward pass + /// + /// A value that the base already keeps does not belong here. The state of every step is 1 + /// such value, and [`RnnCell::step_backward`] receives it directly. The count can depend on + /// the activation, because [`Activation::saves`] names the tensor that its backward pass + /// reads. + /// + /// # Returns + /// + /// - `usize` - Number of arrays per timestep + fn record_slots(&self) -> usize; + /// Runs 1 timestep /// /// The cell reads `state`, computes the state that leaves the step, and writes it back into /// `state`. A training pass also gives a `record` buffer, and the cell appends exactly - /// [`RnnCell::RECORD_SLOTS`] arrays to it. An inference pass gives `None` and the cell parks + /// [`RnnCell::record_slots`] arrays to it. An inference pass gives `None` and the cell parks /// nothing. /// /// # Parameters @@ -132,7 +139,7 @@ pub(crate) trait RnnCell: std::fmt::Debug + Send + Sync + 'static { /// - `gates` - The fused weights of the layer /// - `state_prev` - The state entering the step /// - `state_next` - The state leaving the step - /// - `record` - The [`RnnCell::RECORD_SLOTS`] arrays that this step parked + /// - `record` - The [`RnnCell::record_slots`] arrays that this step parked /// - `grad_state` - The gradient of the state, which the cell carries back 1 step /// /// # Returns @@ -152,3 +159,142 @@ pub(crate) trait RnnCell: std::fmt::Debug + Send + Sync + 'static { grad_state: &mut [Array2], ) -> Result, Error>; } + +/// The result of 1 activation inside a cell step +/// +/// `output` is the activated tensor that the step uses. [`Activated::into_record`] gives the +/// tensor that the backward pass of the activation reads. +#[derive(Debug)] +pub(super) struct Activated { + /// The activated tensor + pub output: Array2, + /// The pre-activation, which only a training pass of a + /// [`ActivationInput::PreActivation`] activation keeps + pub pre_activation: Option>, +} + +impl Activated { + /// Gives the tensor that the backward pass of the activation reads + /// + /// # Returns + /// + /// - `Array2` - The pre-activation when the step kept it, otherwise the output + pub fn into_record(self) -> Array2 { + self.pre_activation.unwrap_or(self.output) + } +} + +/// Applies the activation of a cell to a pre-activation +/// +/// A training pass of an [`ActivationInput::PreActivation`] activation keeps `z`. Every other +/// pass drops `z`. No pass copies the output. +/// +/// # Parameters +/// +/// - `activation` - Activation of the cell +/// - `z` - The pre-activation +/// - `training` - True when the step parks a record +/// +/// # Returns +/// +/// - `Result` - The output, and the pre-activation when the step keeps it +/// +/// # Errors +/// +/// - [`Error::InvalidInput`] - If the activation refuses its input +pub(super) fn activate( + activation: &Activation, + z: Array2, + training: bool, +) -> Result { + if training && activation.saves() == ActivationInput::PreActivation { + let (output, cache) = activation.forward_train(z.into_dyn())?; + return Ok(Activated { + output: into_2d(output), + pre_activation: Some(into_2d(cache.into_tensor())), + }); + } + Ok(Activated { + output: into_2d(activation.forward(&z.into_dyn())?), + pre_activation: None, + }) +} + +/// Builds the activation cache from 1 record +/// +/// # Parameters +/// +/// - `activation` - Activation of the cell +/// - `record` - The tensor that [`Activated::into_record`] gave +/// +/// # Returns +/// +/// - `ActivationCache` - A cache that holds a copy of `record` +pub(super) fn record_cache(activation: &Activation, record: &Array2) -> ActivationCache { + ActivationCache::new(activation.saves(), record.clone().into_dyn()) +} + +/// Gives the activated output that 1 record stands for +/// +/// An [`ActivationInput::Output`] record is the output, and the function borrows it. An +/// [`ActivationInput::PreActivation`] record is `z`, and the function runs the forward pass +/// again. +/// +/// # Parameters +/// +/// - `activation` - Activation of the cell +/// - `cache` - The cache that [`record_cache`] built from `record` +/// - `record` - The tensor that [`Activated::into_record`] gave +/// +/// # Returns +/// +/// - `Result>, Error>` - The activated output +/// +/// # Errors +/// +/// - [`Error::InvalidInput`] - If the activation refuses the cache +pub(super) fn recorded_output<'a>( + activation: &Activation, + cache: &ActivationCache, + record: &'a Array2, +) -> Result>, Error> { + match activation.saves() { + ActivationInput::Output => Ok(Cow::Borrowed(record)), + ActivationInput::PreActivation => Ok(Cow::Owned(into_2d(activation.output_of(cache)?))), + } +} + +/// Runs the backward pass of the activation of a cell +/// +/// # Parameters +/// +/// - `activation` - Activation of the cell +/// - `cache` - The cache that [`record_cache`] built +/// - `grad_output` - The gradient of the activated output +/// +/// # Returns +/// +/// - `Result, Error>` - The gradient of the pre-activation +/// +/// # Errors +/// +/// - [`Error::InvalidInput`] - If the activation refuses the cache +/// - [`Error::ShapeMismatch`] - If `grad_output` and the cache differ in shape +pub(super) fn activation_backward( + activation: &Activation, + cache: &ActivationCache, + grad_output: Array2, +) -> Result, Error> { + Ok(into_2d( + activation.backward(cache, &grad_output.into_dyn())?, + )) +} + +/// Converts a tensor that an activation gave back to 2 dimensions +/// +/// An activation keeps the shape of its input, and every cell gives it a 2D input +pub(super) fn into_2d(tensor: crate::neural_network::Tensor) -> Array2 { + tensor + .into_dimensionality::() + .expect("an activation keeps the 2D shape of its input") +} diff --git a/src/neural_network/layers/recurrent/gru.rs b/src/neural_network/layers/recurrent/gru.rs index f2ff9bbf..940cbe52 100644 --- a/src/neural_network/layers/recurrent/gru.rs +++ b/src/neural_network/layers/recurrent/gru.rs @@ -9,14 +9,16 @@ use crate::error::Error; use crate::neural_network::layers::activation::Activation; use crate::neural_network::layers::recurrent::apply_sigmoid; -use crate::neural_network::layers::recurrent::cell::{RecurrentGroup, RnnCell}; +use crate::neural_network::layers::recurrent::cell::{ + RecurrentGroup, RnnCell, activate, activation_backward, record_cache, recorded_output, +}; use crate::neural_network::layers::recurrent::gate::FusedGates; use crate::neural_network::layers::recurrent::rnn::{Rnn, recurrent_layer_traits}; use crate::neural_network::layers::validation::validate_weight_shape; use crate::neural_network::traits::{LayerBase, UnaryLayer}; use gemmkit_ndarray::dot; use gemmkit_ndarray::{Bias, Parallelism}; -use ndarray::{Array2, ArrayView2, Axis, Ix2, concatenate, s}; +use ndarray::{Array2, ArrayView2, Axis, concatenate, s}; /// The arithmetic of 1 timestep of a [`GRU`] /// @@ -28,8 +30,14 @@ use ndarray::{Array2, ArrayView2, Axis, Ix2, concatenate, s}; /// # Notes /// /// The record of 1 step holds 4 arrays, in this order: the reset gate, the update gate, the -/// candidate, and `r_t .* h_prev`. The recurrent-kernel gradient of the candidate block needs -/// `r_t .* h_prev` as its operand, so the cell records it. +/// candidate tensor, and `r_t .* h_prev`. The candidate tensor is what the backward pass of the +/// activation reads, as [`Activation::saves`] names. That is the candidate for an +/// [`ActivationInput::Output`] activation, and the candidate pre-activation for an +/// [`ActivationInput::PreActivation`] activation. The recurrent-kernel gradient of the candidate +/// block needs `r_t .* h_prev` as its operand, so the cell records it. +/// +/// [`ActivationInput::Output`]: crate::neural_network::layers::activation::ActivationInput::Output +/// [`ActivationInput::PreActivation`]: crate::neural_network::layers::activation::ActivationInput::PreActivation #[derive(Debug)] pub(crate) struct GruCell { /// Activation applied to the candidate hidden state each timestep @@ -40,7 +48,8 @@ pub(crate) struct GruCell { const RESET_GATE: usize = 0; /// Record slot of the update-gate activation const UPDATE_GATE: usize = 1; -/// Record slot of the candidate hidden state +/// Record slot of the candidate hidden state, or of its pre-activation, as [`Activation::saves`] +/// names const CANDIDATE: usize = 2; /// Record slot of `r_t .* h_prev`, which the recurrent kernel's candidate block projects const RESET_HIDDEN: usize = 3; @@ -49,7 +58,6 @@ impl RnnCell for GruCell { const CELL_TYPE: &'static str = "GRU"; const GATE_BIASES: &'static [f32] = &[0.0, 0.0, 0.0]; const STATE_COUNT: usize = 1; - const RECORD_SLOTS: usize = 4; const RECURRENT_GROUPS: &'static [RecurrentGroup] = &[ RecurrentGroup { first: 0, @@ -67,6 +75,10 @@ impl RnnCell for GruCell { Self { activation } } + fn record_slots(&self) -> usize { + 4 + } + fn step( &self, gates: &FusedGates, @@ -114,17 +126,14 @@ impl RnnCell for GruCell { None, Parallelism::Rayon(0), ); - let h_candidate = act - .forward(&h_candidate.into_dyn())? - .into_dimensionality::() - .unwrap(); + let h_candidate = activate(&act, h_candidate, record.is_some())?; - let h_t = &z_t * h_prev + &(1.0 - &z_t) * &h_candidate; + let h_t = &z_t * h_prev + &(1.0 - &z_t) * &h_candidate.output; if let Some(record) = record { record.push(r_t); record.push(z_t); - record.push(h_candidate); + record.push(h_candidate.into_record()); record.push(r_h); } state[0] = h_t; @@ -144,22 +153,17 @@ impl RnnCell for GruCell { let h_prev = &state_prev[0]; let r_t = &record[RESET_GATE]; let z_t = &record[UPDATE_GATE]; - let h_candidate = &record[CANDIDATE]; + let candidate_cache = record_cache(&act, &record[CANDIDATE]); + let h_candidate = recorded_output(&act, &candidate_cache, &record[CANDIDATE])?; let batch = h_prev.shape()[0]; // Gradient through h_t = z_t .* h_{t-1} + (1 - z_t) .* h_candidate - let grad_z_t = &grad_state[0] * (h_prev - h_candidate); + let grad_z_t = &grad_state[0] * (h_prev - &*h_candidate); let grad_h_candidate = &grad_state[0] * &(1.0 - z_t); let grad_h_prev_from_update = &grad_state[0] * z_t; // Gradient through h_candidate = activation(...), via the activation backward - let grad_h_candidate_raw = act - .backward( - &h_candidate.clone().into_dyn(), - &grad_h_candidate.into_dyn(), - )? - .into_dimensionality::() - .unwrap(); + let grad_h_candidate_raw = activation_backward(&act, &candidate_cache, grad_h_candidate)?; // Gradient through r_h = r_t .* h_{t-1} (1 recurrent product shared by both terms) let grad_rh = dot( diff --git a/src/neural_network/layers/recurrent/lstm.rs b/src/neural_network/layers/recurrent/lstm.rs index 7e7ee198..6885e50b 100644 --- a/src/neural_network/layers/recurrent/lstm.rs +++ b/src/neural_network/layers/recurrent/lstm.rs @@ -1,22 +1,24 @@ //! Long Short-Term Memory (LSTM) recurrent layer with input, forget, cell, and output gates //! //! `LstmCell` implements `RnnCell` with the arithmetic of 1 LSTM timestep. The record-slot -//! constants below it name, in order, the 5 values that `step` parks for `step_backward` to read -//! back. [`LSTM`] is the public layer. It holds an `Rnn` over `LstmCell` and forwards every +//! constants below it name, in order, the 4 or 5 values that `step` parks for `step_backward` to +//! read back. [`LSTM`] is the public layer. It holds an `Rnn` over `LstmCell` and forwards every //! trait method to it, adding only the constructors and the weight setters that a caller uses //! directly. use crate::error::Error; -use crate::neural_network::layers::activation::Activation; +use crate::neural_network::layers::activation::{Activation, ActivationInput}; use crate::neural_network::layers::recurrent::apply_sigmoid; -use crate::neural_network::layers::recurrent::cell::{RecurrentGroup, RnnCell}; +use crate::neural_network::layers::recurrent::cell::{ + RecurrentGroup, RnnCell, activate, activation_backward, into_2d, record_cache, recorded_output, +}; use crate::neural_network::layers::recurrent::gate::FusedGates; use crate::neural_network::layers::recurrent::rnn::{Rnn, recurrent_layer_traits}; use crate::neural_network::layers::validation::validate_weight_shape; use crate::neural_network::traits::{LayerBase, UnaryLayer}; use gemmkit_ndarray::dot; use gemmkit_ndarray::{Bias, Parallelism}; -use ndarray::{Array2, ArrayView2, Axis, Ix2, concatenate, s}; +use ndarray::{Array2, ArrayView2, Axis, concatenate, s}; /// The arithmetic of 1 timestep of an [`LSTM`] /// @@ -29,31 +31,42 @@ use ndarray::{Array2, ArrayView2, Axis, Ix2, concatenate, s}; /// /// # Notes /// -/// The record of 1 step holds 5 arrays, in this order: the activated cell state, then the input, -/// forget, candidate and output gate activations. The cell state that enters a step is not a -/// record, because [`Rnn`] already carries it as state slot 1. +/// The record of 1 step holds these arrays, in this order: +/// +/// 1. The input-gate activation. +/// 2. The forget-gate activation. +/// 3. The tensor that the backward pass of the candidate activation reads. This is the +/// candidate for an [`ActivationInput::Output`] activation, and the candidate +/// pre-activation for an [`ActivationInput::PreActivation`] activation. +/// 4. The output-gate activation. +/// 5. The activated cell state, only for an [`ActivationInput::Output`] activation. +/// +/// The backward pass of a [`ActivationInput::PreActivation`] activation of the cell state reads +/// the cell state itself. [`Rnn`] already carries the cell state as state slot 1, so the record +/// holds no copy of it. The cell state that enters a step is not a record for the same reason. #[derive(Debug)] pub(crate) struct LstmCell { /// Activation applied to the candidate and to the cell state each timestep activation: Activation, } -/// Record slot of the activated cell state -const CELL_ACTIVATED: usize = 0; /// Record slot of the input-gate activation -const INPUT_GATE: usize = 1; +const INPUT_GATE: usize = 0; /// Record slot of the forget-gate activation -const FORGET_GATE: usize = 2; -/// Record slot of the candidate activation -const CANDIDATE: usize = 3; +const FORGET_GATE: usize = 1; +/// Record slot of the candidate, or of the candidate pre-activation, as [`Activation::saves`] +/// names +const CANDIDATE: usize = 2; /// Record slot of the output-gate activation -const OUTPUT_GATE: usize = 4; +const OUTPUT_GATE: usize = 3; +/// Record slot of the activated cell state. Only an [`ActivationInput::Output`] activation +/// fills it +const CELL_ACTIVATED: usize = 4; impl RnnCell for LstmCell { const CELL_TYPE: &'static str = "LSTM"; const GATE_BIASES: &'static [f32] = &[0.0, 1.0, 0.0, 0.0]; const STATE_COUNT: usize = 2; - const RECORD_SLOTS: usize = 5; const RECURRENT_GROUPS: &'static [RecurrentGroup] = &[RecurrentGroup { first: 0, count: 4, @@ -64,6 +77,13 @@ impl RnnCell for LstmCell { Self { activation } } + fn record_slots(&self) -> usize { + match self.activation.saves() { + ActivationInput::Output => 5, + ActivationInput::PreActivation => 4, + } + } + fn step( &self, gates: &FusedGates, @@ -94,27 +114,28 @@ impl RnnCell for LstmCell { // Gates use the recurrent activation (sigmoid). The candidate uses `act`. let i_t = apply_sigmoid(z_all.slice(s![.., 0..u]).to_owned()); let f_t = apply_sigmoid(z_all.slice(s![.., u..2 * u]).to_owned()); - let g_t = act - .forward(&z_all.slice(s![.., 2 * u..3 * u]).to_owned().into_dyn())? - .into_dimensionality::() - .unwrap(); + let g = activate( + &act, + z_all.slice(s![.., 2 * u..3 * u]).to_owned(), + record.is_some(), + )?; let o_t = apply_sigmoid(z_all.slice(s![.., 3 * u..4 * u]).to_owned()); // Update cell state, then apply the configurable activation to it - let c_t = &f_t * &state[1] + &i_t * &g_t; - let c_t_activated = act - .forward(&c_t.clone().into_dyn())? - .into_dimensionality::() - .unwrap(); + let c_t = (&f_t * &state[1] + &i_t * &g.output).into_dyn(); + let c_t_activated = into_2d(act.forward(&c_t)?); + let c_t = into_2d(c_t); let h_t = &o_t * &c_t_activated; if let Some(record) = record { - record.push(c_t_activated); record.push(i_t); record.push(f_t); - record.push(g_t); + record.push(g.into_record()); record.push(o_t); + if act.saves() == ActivationInput::Output { + record.push(c_t_activated); + } } state[0] = h_t; state[1] = c_t; @@ -125,37 +146,40 @@ impl RnnCell for LstmCell { &self, gates: &FusedGates, state_prev: &[Array2], - _state_next: &[Array2], + state_next: &[Array2], record: &[Array2], grad_state: &mut [Array2], ) -> Result, Error> { let u = gates.units(); let act = self.activation; let c_prev = &state_prev[1]; - let c_t_activated = &record[CELL_ACTIVATED]; let i_t = &record[INPUT_GATE]; let f_t = &record[FORGET_GATE]; - let g_t = &record[CANDIDATE]; let o_t = &record[OUTPUT_GATE]; let batch = o_t.shape()[0]; + // The backward pass of the cell-state activation reads the recorded activated cell + // state, or the cell state that leaves the step, as `Activation::saves` names + let cell_saved = match act.saves() { + ActivationInput::Output => &record[CELL_ACTIVATED], + ActivationInput::PreActivation => &state_next[1], + }; + let cell_cache = record_cache(&act, cell_saved); + let c_t_activated = recorded_output(&act, &cell_cache, cell_saved)?; + let candidate_cache = record_cache(&act, &record[CANDIDATE]); + let g_t = recorded_output(&act, &candidate_cache, &record[CANDIDATE])?; + // Gradient through h_t = o_t * activation(c_t) - let grad_o_t = &grad_state[0] * c_t_activated; + let grad_o_t = &grad_state[0] * &*c_t_activated; // dL/dc_t += activation'(c_t) * (grad_h * o_t), via the activation backward. The cell // state carries its own gradient across the steps, so this accumulates - let grad_cell_act = act - .backward( - &c_t_activated.clone().into_dyn(), - &(&grad_state[0] * o_t).into_dyn(), - )? - .into_dimensionality::() - .unwrap(); + let grad_cell_act = activation_backward(&act, &cell_cache, &grad_state[0] * o_t)?; grad_state[1] += &grad_cell_act; // Gradient through c_t = f_t * c_prev + i_t * g_t. Every term reads the accumulated // cell gradient, so the write of the carried gradient comes last let grad_f_t = &grad_state[1] * c_prev; - let grad_i_t = &grad_state[1] * g_t; + let grad_i_t = &grad_state[1] * &*g_t; let grad_g_t = &grad_state[1] * i_t; let grad_c_prev = &grad_state[1] * f_t; @@ -163,10 +187,7 @@ impl RnnCell for LstmCell { let grad_o_raw = &grad_o_t * o_t * &(1.0 - o_t); let grad_f_raw = &grad_f_t * f_t * &(1.0 - f_t); let grad_i_raw = &grad_i_t * i_t * &(1.0 - i_t); - let grad_g_raw = act - .backward(&g_t.clone().into_dyn(), &grad_g_t.into_dyn())? - .into_dimensionality::() - .unwrap(); + let grad_g_raw = activation_backward(&act, &candidate_cache, grad_g_t)?; // Assemble the fused dz for this timestep, blocks [i | f | g | o] let mut dz_t = Array2::::zeros((batch, 4 * u)); diff --git a/src/neural_network/layers/recurrent/rnn.rs b/src/neural_network/layers/recurrent/rnn.rs index 52388d2e..5e4cfb13 100644 --- a/src/neural_network/layers/recurrent/rnn.rs +++ b/src/neural_network/layers/recurrent/rnn.rs @@ -56,7 +56,7 @@ pub(crate) struct Rnn { /// What the forward pass parks for its backward pass /// /// The 2 record vectors are flat. The state of processing step `k` sits at -/// `k * STATE_COUNT`, and the values that step `k` parked start at `k * RECORD_SLOTS`. +/// `k * STATE_COUNT`, and the values that step `k` parked start at `k * record_slots`. /// /// The cell type is part of the cache type, so a cache of 1 recurrent layer cannot decode as the /// cache of another. The layer name that [`Ctx::push_cache`] records is the second guard. @@ -279,8 +279,9 @@ impl UnaryLayer for Rnn { } else { 0 }); + let record_slots = self.cell.record_slots(); let mut records = Vec::with_capacity(if training { - timesteps * C::RECORD_SLOTS + timesteps * record_slots } else { 0 }); @@ -295,8 +296,8 @@ impl UnaryLayer for Rnn { self.cell .step(&self.gates, xw.index_axis(Axis(1), t), &mut state, record)?; debug_assert!( - !training || records.len() == (k + 1) * C::RECORD_SLOTS, - "{} parked a number of values that RECORD_SLOTS does not name", + !training || records.len() == (k + 1) * record_slots, + "{} parked a number of values that record_slots does not name", C::CELL_TYPE ); if let Some(seq) = sequence.as_mut() { @@ -346,6 +347,7 @@ impl UnaryLayer for Rnn { let feat = x3.shape()[2]; let units = self.units; let width = C::GATE_BIASES.len() * units; + let record_slots = self.cell.record_slots(); // With `return_sequences`, every step also takes a direct contribution from `grad_seq`. // Only the hidden state takes it. Every other state is reachable through the hidden @@ -376,7 +378,7 @@ impl UnaryLayer for Rnn { &self.gates, &states[k * C::STATE_COUNT..(k + 1) * C::STATE_COUNT], &states[(k + 1) * C::STATE_COUNT..(k + 2) * C::STATE_COUNT], - &records[k * C::RECORD_SLOTS..(k + 1) * C::RECORD_SLOTS], + &records[k * record_slots..(k + 1) * record_slots], &mut grad_state, )?; // The reductions below pair this step's gate gradients with the input row they came @@ -410,7 +412,7 @@ impl UnaryLayer for Rnn { for k in 0..timesteps { let operand = match group.operand { None => &states[k * C::STATE_COUNT], - Some(slot) => &records[k * C::RECORD_SLOTS + slot], + Some(slot) => &records[k * record_slots + slot], }; operand3 .index_axis_mut(Axis(1), input_step(k, timesteps, self.go_backwards)) diff --git a/src/neural_network/layers/recurrent/simple_rnn.rs b/src/neural_network/layers/recurrent/simple_rnn.rs index c4c69434..0e6ef62b 100644 --- a/src/neural_network/layers/recurrent/simple_rnn.rs +++ b/src/neural_network/layers/recurrent/simple_rnn.rs @@ -7,14 +7,16 @@ //! weight setters that a caller uses directly. use crate::error::Error; -use crate::neural_network::layers::activation::Activation; -use crate::neural_network::layers::recurrent::cell::{RecurrentGroup, RnnCell}; +use crate::neural_network::layers::activation::{Activation, ActivationInput}; +use crate::neural_network::layers::recurrent::cell::{ + RecurrentGroup, RnnCell, activate, activation_backward, record_cache, +}; use crate::neural_network::layers::recurrent::gate::FusedGates; use crate::neural_network::layers::recurrent::rnn::{Rnn, recurrent_layer_traits}; use crate::neural_network::traits::{LayerBase, UnaryLayer}; use gemmkit_ndarray::dot; use gemmkit_ndarray::{Activation as FusedActivation, Bias, Parallelism}; -use ndarray::{Array2, ArrayView2, Ix2}; +use ndarray::{Array2, ArrayView2}; /// The arithmetic of 1 timestep of a [`SimpleRNN`] /// @@ -28,8 +30,13 @@ use ndarray::{Array2, ArrayView2, Ix2}; /// with 1 exception. The fused `Relu` maps `NaN` to `0`, while the scalar closure of this crate /// propagates `NaN` instead. /// -/// The cell parks nothing of its own. The hidden state that leaves a step is the only value that -/// its backward pass reads, and [`Rnn`] hands that state to every cell. +/// # Notes +/// +/// The record of 1 step depends on [`Activation::saves`]. For an +/// [`ActivationInput::Output`] activation the record is empty. The backward pass reads the +/// hidden state that leaves the step, which is the activated output, and [`Rnn`] hands that +/// state to every cell. For an [`ActivationInput::PreActivation`] activation the record holds 1 +/// array, the pre-activation `z` of the step. #[derive(Debug)] pub(crate) struct SimpleRnnCell { /// Activation applied at each timestep of the recurrence @@ -40,7 +47,6 @@ impl RnnCell for SimpleRnnCell { const CELL_TYPE: &'static str = "SimpleRNN"; const GATE_BIASES: &'static [f32] = &[0.0]; const STATE_COUNT: usize = 1; - const RECORD_SLOTS: usize = 0; const RECURRENT_GROUPS: &'static [RecurrentGroup] = &[RecurrentGroup { first: 0, count: 1, @@ -51,12 +57,19 @@ impl RnnCell for SimpleRnnCell { Self { activation } } + fn record_slots(&self) -> usize { + match self.activation.saves() { + ActivationInput::Output => 0, + ActivationInput::PreActivation => 1, + } + } + fn step( &self, gates: &FusedGates, xw_t: ArrayView2<'_, f32>, state: &mut [Array2], - _record: Option<&mut Vec>>, + record: Option<&mut Vec>>, ) -> Result<(), Error> { let bias = gates.bias.as_slice().expect("bias must be contiguous"); // z = x_t @ W + h_{t-1} @ U + b, with `x_t @ W` prefilled as the accumulator @@ -75,13 +88,17 @@ impl RnnCell for SimpleRnnCell { fused_act, Parallelism::Rayon(0), ); + // `Linear` and the fused `ReLU` already hold the output in `z`. Both are + // `ActivationInput::Output` activations, so the step parks nothing for them state[0] = match self.activation { Activation::Linear | Activation::ReLU => z, - _ => self - .activation - .forward(&z.into_dyn())? - .into_dimensionality::() - .unwrap(), + _ => { + let activated = activate(&self.activation, z, record.is_some())?; + if let (Some(record), Some(z)) = (record, activated.pre_activation) { + record.push(z); + } + activated.output + } }; Ok(()) } @@ -91,19 +108,18 @@ impl RnnCell for SimpleRnnCell { gates: &FusedGates, _state_prev: &[Array2], state_next: &[Array2], - _record: &[Array2], + record: &[Array2], grad_state: &mut [Array2], ) -> Result, Error> { - // The activation backward reads the state that leaves the step, which is what the - // activation produced. It also reads the total gradient of that same state. - let d_z = { - let h_t = state_next[0].clone().into_dyn(); - let grad_h = grad_state[0].clone().into_dyn(); - self.activation - .backward(&h_t, &grad_h)? - .into_dimensionality::() - .unwrap() + // The activation backward reads the total gradient of the state that leaves the step. + // It also reads the hidden state that leaves the step, or the recorded `z`, as + // `Activation::saves` names + let saved = match self.activation.saves() { + ActivationInput::Output => &state_next[0], + ActivationInput::PreActivation => &record[0], }; + let cache = record_cache(&self.activation, saved); + let d_z = activation_backward(&self.activation, &cache, grad_state[0].clone())?; grad_state[0] = dot(&d_z, &gates.recurrent_kernel.t()); Ok(d_z) } diff --git a/tests/neural_network/activation.rs b/tests/neural_network/activation.rs index fc862527..ca2fb33a 100644 --- a/tests/neural_network/activation.rs +++ b/tests/neural_network/activation.rs @@ -13,6 +13,8 @@ //! constructor and in a trainable layer that embeds the activation //! - Default() equals the layer's own default value, and a standalone layer equals the //! matching Activation variant +//! - GELU, SiLU, and Mish: the layer equals the enum, the backward pass matches the analytic +//! derivative, the GELU flag selects the form, and the training pass parks 1 cache // The pinned float32 values in this file are reference values. Keep them as written, so a // reader can compare them to the reference output character by character. @@ -26,18 +28,21 @@ use rustyml::neural_network::Tensor; use rustyml::neural_network::layers::activation::Activation; use rustyml::neural_network::layers::activation::elu::ELU; use rustyml::neural_network::layers::activation::exponential::Exponential; +use rustyml::neural_network::layers::activation::gelu::GELU; use rustyml::neural_network::layers::activation::hard_sigmoid::HardSigmoid; use rustyml::neural_network::layers::activation::leaky_relu::LeakyReLU; use rustyml::neural_network::layers::activation::linear::Linear; +use rustyml::neural_network::layers::activation::mish::Mish; use rustyml::neural_network::layers::activation::relu::ReLU; use rustyml::neural_network::layers::activation::selu::SELU; use rustyml::neural_network::layers::activation::sigmoid::Sigmoid; +use rustyml::neural_network::layers::activation::silu::SiLU; use rustyml::neural_network::layers::activation::softmax::Softmax; use rustyml::neural_network::layers::activation::softplus::Softplus; use rustyml::neural_network::layers::activation::softsign::Softsign; use rustyml::neural_network::layers::activation::tanh::Tanh; use rustyml::neural_network::layers::dense::Dense; -use rustyml::neural_network::traits::{Layer, UnaryLayer}; +use rustyml::neural_network::traits::{Layer, LayerBase, UnaryLayer}; use crate::common::{GateGuard, assert_allclose}; @@ -1465,11 +1470,13 @@ fn softmax_axis_0_of_a_rank_2_input_matches_the_reference() { /// runs before that, so a caller that pairs the wrong tensors gets `Error::ShapeMismatch` #[test] fn softmax_non_final_axis_backward_rejects_a_mismatched_grad_output() { - let activated = Array::from_elem((2, 3, 4), 0.25_f32).into_dyn(); + let activation = Activation::Softmax { axis: 1 }; + let input = Array::from_elem((2, 3, 4), 0.25_f32).into_dyn(); + let (_, cache) = activation.forward_train(input).expect("forward"); let grad_output = Array::from_elem((2, 3, 5), 1.0_f32).into_dyn(); // Axis 1 is not the last axis, so this call takes the lane path - let result = Activation::Softmax { axis: 1 }.backward(&activated, &grad_output); + let result = activation.backward(&cache, &grad_output); match result { Err(Error::ShapeMismatch { expected, found }) => { assert_eq!(expected, vec![2, 3, 4]); @@ -1478,3 +1485,229 @@ fn softmax_non_final_axis_backward_rejects_a_mismatched_grad_output() { other => panic!("expected ShapeMismatch, got {other:?}"), } } + +// GELU, SiLU, and Mish layers. +// +// The derivative of each of these layers has no closed form in the output, so each layer +// caches its input. The pinned values come from the definitions, evaluated in f64. + +/// The inputs that the GELU, SiLU, and Mish tests share +const PRE_ACTIVATION_PROBES: [f32; 7] = [-3.0, -1.0, -0.5, 0.0, 0.5, 1.0, 3.0]; + +/// An upstream gradient with a different value and sign at each probe +const PRE_ACTIVATION_GRAD: [f32; 7] = [0.5, -2.0, 1.0, 3.0, -1.0, 0.25, 2.0]; + +/// Each of the 3 layers, with the enum variant that it wraps and its derivative at each probe +fn pre_activation_layers() -> Vec<(Box, Activation, [f32; 7])> { + vec![ + ( + Box::new(GELU::new()), + Activation::GELU { approximate: false }, + [ + -0.0119456472, + -0.0833154706, + 0.132504875, + 0.5, + 0.867495125, + 1.08331547, + 1.01194565, + ], + ), + ( + Box::new(GELU::new().with_approximate(true)), + Activation::GELU { approximate: true }, + [ + -0.0115841666, + -0.0829640838, + 0.132630096, + 0.5, + 0.867369904, + 1.08296408, + 1.01158417, + ], + ), + ( + Box::new(SiLU::new()), + Activation::SiLU, + [ + -0.088104106, + 0.0723294881, + 0.260038813, + 0.5, + 0.739961187, + 0.927670512, + 1.08810411, + ], + ), + ( + Box::new(Mish::new()), + Activation::Mish, + [ + -0.0933931145, + 0.0592167559, + 0.289510678, + 0.6, + 0.886424375, + 1.04903622, + 1.02110691, + ], + ), + ] +} + +/// Each layer gives the output of its enum variant bit for bit, in both modes +#[test] +fn pre_activation_layers_match_the_activation_enum() { + let input = tensor2(1, 7, PRE_ACTIVATION_PROBES.to_vec()); + for (mut layer, activation, _) in pre_activation_layers() { + let enum_out = activation.forward(&input).expect("enum forward"); + for mut ctx in [Ctx::training(), Ctx::inference()] { + let layer_out = layer.forward_mut(&input, &mut ctx).expect("layer forward"); + assert_eq!(layer_out, enum_out, "{activation:?}"); + } + } +} + +/// The backward pass of each layer is the upstream gradient times the analytic derivative +#[test] +fn pre_activation_layers_backward_matches_the_analytic_derivative() { + let input = tensor2(1, 7, PRE_ACTIVATION_PROBES.to_vec()); + let grad_output = tensor2(1, 7, PRE_ACTIVATION_GRAD.to_vec()); + for (mut layer, activation, derivative) in pre_activation_layers() { + let mut ctx = Ctx::training(); + layer.forward_mut(&input, &mut ctx).expect("forward"); + let grad_in = layer.backward(&grad_output, &mut ctx).expect("backward"); + + let expected: Vec = derivative + .iter() + .zip(PRE_ACTIVATION_GRAD.iter()) + .map(|(d, g)| d * g) + .collect(); + for ((&got, &want), x) in grad_in.iter().zip(&expected).zip(PRE_ACTIVATION_PROBES) { + assert!( + (got - want).abs() <= 1e-6, + "{activation:?} at x = {x}: got {got}, want {want}" + ); + } + } +} + +/// The training pass of each layer parks exactly 1 cache, and the backward pass takes it back +#[test] +fn pre_activation_layers_park_1_cache() { + let input = tensor2(1, 7, PRE_ACTIVATION_PROBES.to_vec()); + let grad_output = tensor2(1, 7, PRE_ACTIVATION_GRAD.to_vec()); + for (mut layer, activation, _) in pre_activation_layers() { + let mut ctx = Ctx::training(); + layer.forward_mut(&input, &mut ctx).expect("forward"); + assert_eq!(ctx.pending_caches(), 1, "{activation:?}"); + layer.backward(&grad_output, &mut ctx).expect("backward"); + assert_eq!(ctx.pending_caches(), 0, "{activation:?}"); + } +} + +/// Each layer refuses an upstream gradient of another shape +#[test] +fn pre_activation_layers_reject_a_mismatched_grad_output() { + let input = tensor2(1, 7, PRE_ACTIVATION_PROBES.to_vec()); + let grad_output = tensor2(7, 1, PRE_ACTIVATION_GRAD.to_vec()); + for (mut layer, activation, _) in pre_activation_layers() { + let mut ctx = Ctx::training(); + layer.forward_mut(&input, &mut ctx).expect("forward"); + let result = layer.backward(&grad_output, &mut ctx); + assert!( + matches!(result, Err(Error::ShapeMismatch { .. })), + "{activation:?}: expected ShapeMismatch, got {result:?}" + ); + } +} + +/// Each layer rejects an empty input +#[test] +fn pre_activation_layers_reject_empty_input() { + let input = Array2::::zeros((0, 3)).into_dyn(); + for (mut layer, activation, _) in pre_activation_layers() { + let result = layer.forward_mut(&input, &mut Ctx::training()); + assert!( + matches!(result, Err(Error::EmptyInput(_))), + "{activation:?}: expected EmptyInput, got {result:?}" + ); + } +} + +/// The GELU flag selects the form, and the 2 forms differ on known values +/// +/// The exact form gives `Phi(x)` scaled by `x`. The tanh approximation differs from it by up +/// to about 4.1e-4 on these probes +#[test] +fn gelu_approximate_flag_selects_the_form() { + let input = tensor2(1, 7, PRE_ACTIVATION_PROBES.to_vec()); + let exact = GELU::new() + .forward_mut(&input, &mut Ctx::inference()) + .expect("exact forward"); + let approximate = GELU::new() + .with_approximate(true) + .forward_mut(&input, &mut Ctx::inference()) + .expect("approximate forward"); + + let expected_exact = [ + -0.00404969409_f32, + -0.158655254, + -0.154268769, + 0.0, + 0.345731231, + 0.841344746, + 2.99595031, + ]; + let expected_approximate = [ + -0.00363739208_f32, + -0.158808009, + -0.15428599, + 0.0, + 0.34571401, + 0.841191991, + 2.99636261, + ]; + assert_allclose(&exact, &tensor2(1, 7, expected_exact.to_vec()), 1e-6_f32); + assert_allclose( + &approximate, + &tensor2(1, 7, expected_approximate.to_vec()), + 1e-6_f32, + ); + assert_abs_diff_eq!( + exact[[0, 0]] - approximate[[0, 0]], + -4.1230201e-4, + epsilon = 1e-6 + ); + + // The builder also selects the exact form back + let reset = GELU::new() + .with_approximate(true) + .with_approximate(false) + .forward_mut(&input, &mut Ctx::inference()) + .expect("reset forward"); + assert_eq!(reset, exact); +} + +/// Default() of each layer equals new() +#[test] +fn pre_activation_layers_default_matches_new() { + let input = tensor2(1, 7, PRE_ACTIVATION_PROBES.to_vec()); + let run = |layer: &mut dyn UnaryLayer| { + layer + .forward_mut(&input, &mut Ctx::inference()) + .expect("forward") + }; + assert_eq!(run(&mut GELU::default()), run(&mut GELU::new())); + assert_eq!(run(&mut SiLU::default()), run(&mut SiLU::new())); + assert_eq!(run(&mut Mish::default()), run(&mut Mish::new())); +} + +/// The layer type names the layer, and the GELU flag does not change it +#[test] +fn pre_activation_layers_report_their_type() { + assert_eq!(GELU::new().layer_type(), "GELU"); + assert_eq!(GELU::new().with_approximate(true).layer_type(), "GELU"); + assert_eq!(SiLU::new().layer_type(), "SiLU"); + assert_eq!(Mish::new().layer_type(), "Mish"); +} diff --git a/tests/neural_network/gradient_check.rs b/tests/neural_network/gradient_check.rs index e00bed79..29d5e5b9 100644 --- a/tests/neural_network/gradient_check.rs +++ b/tests/neural_network/gradient_check.rs @@ -16,17 +16,21 @@ //! gradient back per input. use approx::assert_abs_diff_eq; -use ndarray::Array; +use ndarray::{Array, Axis}; use rustyml::neural_network::Ctx; use rustyml::neural_network::Shape; use rustyml::neural_network::Tensor; +use rustyml::neural_network::layers::Activation; use rustyml::neural_network::layers::activation::elu::ELU; use rustyml::neural_network::layers::activation::exponential::Exponential; +use rustyml::neural_network::layers::activation::gelu::GELU; use rustyml::neural_network::layers::activation::hard_sigmoid::HardSigmoid; use rustyml::neural_network::layers::activation::leaky_relu::LeakyReLU; use rustyml::neural_network::layers::activation::linear::Linear; +use rustyml::neural_network::layers::activation::mish::Mish; use rustyml::neural_network::layers::activation::p_relu::PReLU; use rustyml::neural_network::layers::activation::selu::SELU; +use rustyml::neural_network::layers::activation::silu::SiLU; use rustyml::neural_network::layers::activation::softmax::Softmax; use rustyml::neural_network::layers::activation::softplus::Softplus; use rustyml::neural_network::layers::activation::softsign::Softsign; @@ -1729,3 +1733,223 @@ fn reverse_input_gradient_matches_finite_difference() { check_input_gradient(&mut layer, &x, 1e-3, 2e-3); } } + +// GELU, SiLU and Mish in their host layers. GELU, SiLU and Mish save the pre-activation +// `z`, and the backward pass of each host reads it through the activation cache. Each test +// below runs 1 activation through a table of the 14 host layers. A host must pass the input +// check and the weight check. Each check uses the weighted loss, so a wrong scale at 1 output +// position can not cancel against another position. +// +// The derivative of these 3 activations is not a function of the output alone. Most of the +// difference is in the negative region. Thus every fixture must put pre-activations on both +// sides of 0. The parameters get a fixed signed pattern, and each check first asserts that the +// pre-activation of the host is negative at some position and positive at another. + +/// Builds 1 host layer around the activation +type HostBuilder = fn(Activation) -> Box; + +/// A host layer and the input that its checks use +struct Host { + name: &'static str, + build: HostBuilder, + input: Tensor, + /// `true` for a recurrent host. It returns the full sequence, `[batch, steps, units]` + recurrent: bool, +} + +/// A tensor of signed values in [-0.9, 0.9] that do not increase monotonically along any axis +fn signed(shape: &[usize]) -> Tensor { + let n: usize = shape.iter().product(); + let data: Vec = (0..n) + .map(|k| 0.9 * (((k * 7 + 3) % 19) as f32 / 9.0 - 1.0)) + .collect(); + Array::from_shape_vec(shape.to_vec(), data).unwrap() +} + +/// Writes a fixed signed pattern in [-0.6, 0.6] into every parameter of the layer +/// +/// The pattern replaces the random initial values. Thus the sign checks on the +/// pre-activation give the same result on each run. +fn fill_parameters(layer: &mut dyn UnaryLayer) { + for (p, param) in layer.parameters_mut().into_iter().enumerate() { + for (k, value) in param.value.iter_mut().enumerate() { + *value = 0.6 * (((k * 5 + p * 3 + 1) % 13) as f32 / 6.0 - 1.0); + } + } +} + +/// Builds the host around the activation, runs 1 inference pass to build the parameters, and +/// writes the fixed pattern into them +fn prepared_host(host: &Host, activation: Activation) -> Box { + let mut layer = (host.build)(activation); + layer + .forward_mut(&host.input, &mut Ctx::inference()) + .unwrap(); + fill_parameters(layer.as_mut()); + layer +} + +/// Asserts that the pre-activation of the host has values on both sides of 0 +/// +/// The `Linear` twin of the host has the same parameters. For a feed-forward host, the output +/// of the twin is the pre-activation of the host at every position. +/// +/// A recurrent host starts from a zero state. At the first step, the twin gives a value with the +/// same sign as the pre-activation `z` that the activation reads, and a magnitude that is not +/// larger. SimpleRNN gives `z`. LSTM gives `o * i * z`, and GRU gives 1 gate times `z`. Each +/// gate is in (0, 1). Thus the sign check on the first step is exact. The check skips the later +/// steps, because there the activation also feeds the state. +fn assert_pre_activation_straddles_zero(host: &Host) { + let twin = prepared_host(host, Activation::Linear); + let out = twin.forward(&host.input, &mut Ctx::inference()).unwrap(); + let z = if host.recurrent { + out.index_axis(Axis(1), 0).to_owned() + } else { + out + }; + let min = z.iter().cloned().fold(f32::INFINITY, f32::min); + let max = z.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + assert!( + min < -0.1 && max > 0.1, + "{}: the pre-activation must have values on both sides of 0, found [{min}, {max}]", + host.name + ); +} + +/// The 14 host layers that take an `Activation`, each with a small input +fn hosts() -> Vec { + vec![ + Host { + name: "Dense", + build: |a| Box::new(Dense::new(3, a).unwrap()), + input: signed(&[4, 3]), + recurrent: false, + }, + Host { + name: "Conv1D", + build: |a| Box::new(Conv1D::new(2, 2, 1, a).unwrap()), + input: signed(&[1, 5, 2]), + recurrent: false, + }, + Host { + name: "Conv2D", + build: |a| Box::new(Conv2D::new(2, (2, 2), (1, 1), a).unwrap()), + input: signed(&[1, 4, 4, 1]), + recurrent: false, + }, + Host { + name: "Conv3D", + build: |a| Box::new(Conv3D::new(2, (2, 2, 2), (1, 1, 1), a).unwrap()), + input: signed(&[1, 3, 3, 3, 1]), + recurrent: false, + }, + Host { + name: "Conv1DTranspose", + build: |a| Box::new(Conv1DTranspose::new(2, 3, 2, a).unwrap()), + input: signed(&[1, 4, 2]), + recurrent: false, + }, + Host { + name: "Conv2DTranspose", + build: |a| Box::new(Conv2DTranspose::new(2, (2, 2), (1, 1), a).unwrap()), + input: signed(&[1, 3, 3, 2]), + recurrent: false, + }, + Host { + name: "Conv3DTranspose", + build: |a| Box::new(Conv3DTranspose::new(2, (2, 2, 2), (1, 1, 1), a).unwrap()), + input: signed(&[1, 2, 2, 2, 1]), + recurrent: false, + }, + Host { + name: "DepthwiseConv1D", + build: |a| Box::new(DepthwiseConv1D::new(3, 1, a).unwrap()), + input: signed(&[1, 8, 2]), + recurrent: false, + }, + Host { + name: "DepthwiseConv2D", + build: |a| Box::new(DepthwiseConv2D::new((2, 2), (1, 1), a).unwrap()), + input: signed(&[1, 4, 4, 2]), + recurrent: false, + }, + Host { + name: "SeparableConv1D", + build: |a| Box::new(SeparableConv1D::new(3, 3, 1, 1, a).unwrap()), + input: signed(&[1, 8, 2]), + recurrent: false, + }, + Host { + name: "SeparableConv2D", + build: |a| Box::new(SeparableConv2D::new(2, (2, 2), (1, 1), 1, a).unwrap()), + input: signed(&[1, 4, 4, 2]), + recurrent: false, + }, + Host { + name: "SimpleRNN", + build: |a| Box::new(SimpleRNN::new(3, a).unwrap().with_return_sequences(true)), + input: signed(&[2, 3, 2]), + recurrent: true, + }, + Host { + name: "LSTM", + build: |a| Box::new(LSTM::new(3, a).unwrap().with_return_sequences(true)), + input: signed(&[2, 3, 2]), + recurrent: true, + }, + Host { + name: "GRU", + build: |a| Box::new(GRU::new(3, a).unwrap().with_return_sequences(true)), + input: signed(&[2, 3, 2]), + recurrent: true, + }, + ] +} + +/// Runs the input check and the weight check of every host with the activation, and the input +/// check of the standalone layer of the same activation +fn check_pre_activation_hosts(activation: Activation, standalone: &mut dyn UnaryLayer) { + let table = hosts(); + assert_eq!(table.len(), 14, "the table must cover every host layer"); + for host in &table { + assert_pre_activation_straddles_zero(host); + + let mut layer = prepared_host(host, activation); + check_input_gradient_weighted(layer.as_mut(), &host.input, 1e-3, 1e-2); + + let mut layer = prepared_host(host, activation); + check_weight_gradient_weighted(layer.as_mut(), &host.input, 1e-3, 1e-2); + } + + // The standalone layer has no parameters. Its input covers both sides of 0 directly. + let x = Array::from_shape_vec( + (2, 5), + vec![-3.1, -2.2, -1.5, -0.9, -0.35, 0.3, 0.8, 1.4, 2.1, 2.9], + ) + .unwrap() + .into_dyn(); + check_input_gradient_weighted(standalone, &x, 1e-3, 1e-2); +} + +#[test] +fn gelu_exact_hosts_match_finite_difference() { + check_pre_activation_hosts(Activation::GELU { approximate: false }, &mut GELU::new()); +} + +#[test] +fn gelu_tanh_approximation_hosts_match_finite_difference() { + check_pre_activation_hosts( + Activation::GELU { approximate: true }, + &mut GELU::new().with_approximate(true), + ); +} + +#[test] +fn silu_hosts_match_finite_difference() { + check_pre_activation_hosts(Activation::SiLU, &mut SiLU::new()); +} + +#[test] +fn mish_hosts_match_finite_difference() { + check_pre_activation_hosts(Activation::Mish, &mut Mish::new()); +} diff --git a/tests/neural_network/recurrent.rs b/tests/neural_network/recurrent.rs index 92bea681..b93c70ed 100644 --- a/tests/neural_network/recurrent.rs +++ b/tests/neural_network/recurrent.rs @@ -896,7 +896,16 @@ fn flag_weights(n_gates: usize) -> (Array2, Array2, Array2) { /// Builds the flag-test SimpleRNN with fixed weights fn flag_simple_rnn(return_sequences: bool, go_backwards: bool) -> SimpleRNN { - let mut layer = SimpleRNN::new(2, Tanh::new()) + flag_simple_rnn_with(Activation::Tanh, return_sequences, go_backwards) +} + +/// Builds the flag-test SimpleRNN with fixed weights and the given activation +fn flag_simple_rnn_with( + activation: Activation, + return_sequences: bool, + go_backwards: bool, +) -> SimpleRNN { + let mut layer = SimpleRNN::new(2, activation) .unwrap() .with_return_sequences(return_sequences) .with_go_backwards(go_backwards); @@ -908,7 +917,12 @@ fn flag_simple_rnn(return_sequences: bool, go_backwards: bool) -> SimpleRNN { /// Builds the flag-test LSTM with fixed weights fn flag_lstm(return_sequences: bool, go_backwards: bool) -> LSTM { - let mut layer = LSTM::new(2, Tanh::new()) + flag_lstm_with(Activation::Tanh, return_sequences, go_backwards) +} + +/// Builds the flag-test LSTM with fixed weights and the given activation +fn flag_lstm_with(activation: Activation, return_sequences: bool, go_backwards: bool) -> LSTM { + let mut layer = LSTM::new(2, activation) .unwrap() .with_return_sequences(return_sequences) .with_go_backwards(go_backwards); @@ -920,7 +934,12 @@ fn flag_lstm(return_sequences: bool, go_backwards: bool) -> LSTM { /// Builds the flag-test GRU with fixed weights fn flag_gru(return_sequences: bool, go_backwards: bool) -> GRU { - let mut layer = GRU::new(2, Tanh::new()) + flag_gru_with(Activation::Tanh, return_sequences, go_backwards) +} + +/// Builds the flag-test GRU with fixed weights and the given activation +fn flag_gru_with(activation: Activation, return_sequences: bool, go_backwards: bool) -> GRU { + let mut layer = GRU::new(2, activation) .unwrap() .with_return_sequences(return_sequences) .with_go_backwards(go_backwards); @@ -1345,6 +1364,86 @@ fn gru_gradients_match_finite_difference_for_every_flag_combination() { } } +/// The activations whose backward pass reads the pre-activation and not the output +/// +/// A cell keeps a different record for these activations. The LSTM also reads its cell state +/// from the state that leaves the step, and not from the record. +const PRE_ACTIVATION_ACTIVATIONS: [Activation; 4] = [ + Activation::GELU { approximate: false }, + Activation::GELU { approximate: true }, + Activation::SiLU, + Activation::Mish, +]; + +/// Every gradient of each cell matches a finite difference for each activation that saves its +/// pre-activation, for all 4 flag combinations +/// +/// The upstream gradient varies along the time axis, so a record of the wrong step or of the +/// wrong slot cannot hide behind a constant gradient. +#[test] +fn pre_activation_activations_match_finite_difference_for_every_flag_combination() { + let x = flag_input(); + for activation in PRE_ACTIVATION_ACTIVATIONS { + for return_sequences in [false, true] { + for go_backwards in [false, true] { + let flags = (activation, return_sequences, go_backwards); + check_flag_input_gradient( + &mut flag_simple_rnn_with(flags.0, flags.1, flags.2), + &x, + 3e-2, + ); + check_flag_weight_gradients( + &mut flag_simple_rnn_with(flags.0, flags.1, flags.2), + &x, + 3e-2, + ); + check_flag_input_gradient(&mut flag_lstm_with(flags.0, flags.1, flags.2), &x, 3e-2); + check_flag_weight_gradients( + &mut flag_lstm_with(flags.0, flags.1, flags.2), + &x, + 3e-2, + ); + check_flag_input_gradient(&mut flag_gru_with(flags.0, flags.1, flags.2), &x, 3e-2); + check_flag_weight_gradients( + &mut flag_gru_with(flags.0, flags.1, flags.2), + &x, + 3e-2, + ); + } + } + } +} + +/// A training forward pass gives the same output as an inference forward pass, for each cell +/// and each activation that saves its pre-activation +/// +/// A training pass keeps the pre-activation in the record. An inference pass keeps nothing. Both +/// passes must apply the same activation to the same pre-activation. +#[test] +fn pre_activation_training_forward_equals_inference_forward() { + let x = flag_input(); + for activation in PRE_ACTIVATION_ACTIVATIONS { + for return_sequences in [false, true] { + let layers: [Box; 3] = [ + Box::new(flag_simple_rnn_with(activation, return_sequences, false)), + Box::new(flag_lstm_with(activation, return_sequences, false)), + Box::new(flag_gru_with(activation, return_sequences, false)), + ]; + for layer in layers { + let mut ctx = Ctx::training(); + let trained = layer.forward(&x, &mut ctx).unwrap(); + let inferred = layer.forward(&x, &mut Ctx::inference()).unwrap(); + assert_eq!( + trained, + inferred, + "{activation:?}: the 2 passes of {} differ", + layer.layer_type() + ); + } + } + } +} + /// With return_sequences on, the backward pass rejects a rank-2 gradient #[test] fn return_sequences_backward_rejects_a_rank_2_gradient() { diff --git a/tests/neural_network/roster.rs b/tests/neural_network/roster.rs index 0ef1aa32..17288bca 100644 --- a/tests/neural_network/roster.rs +++ b/tests/neural_network/roster.rs @@ -17,7 +17,7 @@ use rustyml::neural_network::{Shape, Tensor}; /// /// The tests that walk the roster compare their count of distinct layer types to this number. /// A layer type that the roster skips, or holds twice, therefore fails each of these tests -pub const LAYER_TYPE_COUNT: usize = 74; +pub const LAYER_TYPE_COUNT: usize = 77; /// The values that a sample input holds #[derive(Clone, Copy, Debug)] @@ -153,6 +153,9 @@ pub fn visit_every_layer(visitor: &mut impl LayerVisitor) { visitor.visit(unary(&flat, "(None, 4)"), &Sigmoid::new); visitor.visit(unary(&flat, "(None, 4)"), &Tanh::new); visitor.visit(unary(&flat, "(None, 4)"), &Softmax::new); + visitor.visit(unary(&flat, "(None, 4)"), &GELU::new); + visitor.visit(unary(&flat, "(None, 4)"), &SiLU::new); + visitor.visit(unary(&flat, "(None, 4)"), &Mish::new); visitor.visit(unary(&flat, "(None, 4)"), &|| PReLU::new(0.25).unwrap()); // Convolution diff --git a/tests/neural_network/serialize.rs b/tests/neural_network/serialize.rs index bcd07aec..aae3a7d2 100644 --- a/tests/neural_network/serialize.rs +++ b/tests/neural_network/serialize.rs @@ -22,8 +22,11 @@ use ndarray::Array; use rustyml::error::{Error, IoError}; use rustyml::neural_network::Shape; use rustyml::neural_network::Tensor; +use rustyml::neural_network::layers::activation::gelu::GELU; use rustyml::neural_network::layers::activation::linear::Linear; +use rustyml::neural_network::layers::activation::mish::Mish; use rustyml::neural_network::layers::activation::p_relu::PReLU; +use rustyml::neural_network::layers::activation::silu::SiLU; use rustyml::neural_network::layers::checkpoint::{ LayerCheckpoint, MODEL_FORMAT_VERSION, MODEL_MAGIC, ModelCheckpoint, WeightRecord, }; @@ -531,6 +534,49 @@ fn mixed_model_trained_round_trip() { assert_allclose(&after, &before, 1e-6_f32); } +// A trained model with GELU, SiLU, and Mish between its Dense layers round-trips. +// The 3 activation layers hold no array, so the Dense weights carry the whole model. The fit +// runs the backward pass of each activation layer, so a broken cache fails here. +#[test] +fn pre_activation_layers_trained_round_trip() { + let tmp = TempFile::new("pre_activation_trained"); + + let make_arch = || { + SequentialBuilder::new() + .add(Dense::new(4, Linear::new()).unwrap()) + .add(GELU::new().with_approximate(true)) + .add(Dense::new(4, Linear::new()).unwrap()) + .add(SiLU::new()) + .add(Dense::new(4, Linear::new()).unwrap()) + .add(Mish::new()) + .add(Dense::new(2, Linear::new()).unwrap()) + .add(GELU::new()) + .build(&Shape::known(&[2, 3])) + .unwrap() + }; + + let mut model = make_arch(); + model.compile( + SGD::new(0.05, 0.0, false, 0.0).unwrap(), + MeanSquaredError::new(), + ); + + let x: Tensor = Array::from_shape_vec((2, 3), vec![0.5f32, -1.0, 1.5, -0.5, 1.0, -1.5]) + .unwrap() + .into_dyn(); + let y: Tensor = Array::from_shape_vec((2, 2), vec![1.0f32, 0.0, 0.0, 1.0]) + .unwrap() + .into_dyn(); + let untrained = model.predict(&x).unwrap(); + model.fit(&x, &y, 5).unwrap(); + + let before = model.predict(&x).unwrap(); + assert_ne!(before, untrained, "the fit did not change the model"); + let fresh = round_trip(&model, make_arch, tmp.path()); + let after = fresh.predict(&x).unwrap(); + assert_eq!(after, before); +} + // Error paths /// Nonexistent file gives Error::Io(IoError::Std)