Clorch Parameter Initialization: Xavier and Kaiming

Initialize weights with clorch.nn.init — constant, uniform, normal, Xavier/Glorot, and Kaiming/He initializers — with defmodel integration patterns.

The initial values assigned to a network’s weights have a profound effect on whether training converges, how quickly it does so, and how deep the network can be. Poorly initialised weights lead to either vanishing gradients (weights shrink toward zero) or exploding gradients (signals amplify out of control). The clorch.nn.init namespace provides a set of in-place initializers that modify a tensor’s values directly — wrapping the corresponding LibTorch C++ routines.

(require '[clorch.torch :as t]
         '[clorch.nn.init :as init])

All init/ functions modify the tensor in place and return the same tensor. They run inside a no-grad scope so the initialization operation itself is not recorded in the autograd graph.

---

Basic Initializers

These functions are useful for quick experiments, zeroing bias tensors, or setting a known baseline before applying a more sophisticated scheme.

(def x (t/zeros [3 3]))

;; Fill every element with a fixed scalar
(init/constant! x 3.14)

;; Fill with ones
(init/ones! x)

;; Fill with zeros
(init/zeros! x)

;; Draw uniformly from [from, to)
(init/uniform! x -1.0 1.0)

;; Draw from N(mean, std)
(init/normal! x 0.0 0.01)

---

Xavier (Glorot) Initialization

Xavier initialization scales the random weights based on the number of input and output connections (fan-in and fan-out) so that the variance of activations and gradients stays roughly constant across layers. It is designed for layers followed by symmetric activations such as tanh or sigmoid.

Xavier Uniform

Draws weights from a uniform distribution bounded by ± sqrt(6 / (fan_in + fan_out)) scaled by an optional gain.

(init/xavier-uniform! weight)

;; With explicit gain (e.g. sqrt(2) for tanh)
(init/xavier-uniform! weight (Math/sqrt 2.0))

Xavier Normal

Draws weights from a normal distribution with standard deviation sqrt(2 / (fan_in + fan_out)) scaled by gain.

(init/xavier-normal! weight)

(init/xavier-normal! weight 1.0)

Use Xavier initialization when your network uses tanh, sigmoid, or softmax activations. Avoid it with ReLU — use Kaiming instead, because ReLU zeroes half the activations and changes the effective fan.

---

Kaiming (He) Initialization

Kaiming initialization accounts for the fact that ReLU-like activations zero out half their inputs. It scales weights so that the variance is preserved through the forward pass (:fan-in mode) or the backward pass (:fan-out mode).

Kaiming Uniform

;; For ReLU — the most common default
(init/kaiming-uniform! weight :non-linearity :relu)

;; For LeakyReLU with slope 0.1
(init/kaiming-uniform! weight
                        :non-linearity :leaky-relu
                        :a             0.1
                        :mode          :fan-in)

Kaiming Normal

(init/kaiming-normal! weight :non-linearity :relu)

(init/kaiming-normal! weight
                       :non-linearity :leaky-relu
                       :a             0.01
                       :mode          :fan-out)

Kaiming Options

Option	Default	Description
:non-linearity	:leaky-relu	The activation following this layer. Supported: :relu, :leaky-relu, :tanh, :sigmoid, :linear
:a	0	Negative slope for :leaky-relu. Ignored for other non-linearities
:mode	:fan-in	:fan-in preserves variance in the forward pass; :fan-out preserves variance in the backward pass

The default :non-linearity in clorch.nn.init is :leaky-relu (matching PyTorch’s default). Always pass :non-linearity :relu explicitly when initializing layers before a standard ReLU activation.

---

Usage Pattern Inside defmodel

The recommended place to initialize custom weights is immediately after constructing the model, operating directly on the tensor fields exposed by the record.

(nn/defmodel MyModel [in out]
  [l1 (nn/linear in out)]
  (forward [x]
    (nn/forward l1 x)))

(def model (MyModel 128 64))

;; Apply Kaiming normal to the weight matrix, zero the bias
(init/kaiming-normal! (.weight (:l1 model)) :non-linearity :relu)
(init/zeros!          (.bias   (:l1 model)))

For a model with multiple layers, initialize each layer in sequence:

(nn/defmodel Encoder [in hidden out]
  [l1 (nn/linear in hidden)
   l2 (nn/linear hidden out)]
  (forward [x]
    (nn/forward l2 (F/relu (nn/forward l1 x)))))

(def enc (Encoder 256 128 64))

(init/kaiming-uniform! (.weight (:l1 enc)) :non-linearity :relu)
(init/zeros!           (.bias   (:l1 enc)))

(init/xavier-uniform!  (.weight (:l2 enc)))
(init/zeros!           (.bias   (:l2 enc)))

Bias tensors are almost always initialised to zero regardless of the weight scheme. Kaiming and Xavier initializers operate only on the weight matrix and have no meaningful interpretation for a 1-D bias vector.

---

Choosing the Right Initializer

Scenario	Recommended Initializer
Layer followed by tanh or sigmoid	xavier-uniform! or xavier-normal!
Layer followed by relu	kaiming-uniform! with :non-linearity :relu
Layer followed by leaky-relu	kaiming-normal! with :non-linearity :leaky-relu, :a slope
Embedding or output layer	normal! with a small std (e.g. 0.02)
Bias terms	zeros!
Debug / ablation baseline	constant! or ones!