Network depth and gradient propagation

Depth means the number of hidden layers; the output is always linear. Each hidden layer uses the selected width and activation. The zero-hidden-layer model is a straight-line regression model with two parameters.

For hidden width m and depth d ≥ 1, the parameter count includes input weights, connections between hidden layers, output weights, and all biases:

All models see the same 32 training points in each epoch and use the same learning rate. One epoch computes the average gradient over all 32 examples, then updates every parameter once. Deeper networks have more parameters and require more computation: equal epochs do not mean equal time or equal capacity.

The clean-curve MSE uses 101 points from the known target function without noise. It checks interpolation against that function; it is not evidence of performance on a different problem. Training noise is uniform between minus and plus the selected amplitude.

More hidden layers compose more nonlinear transformations, but they also lengthen the chain of derivatives. Gradients can shrink or grow along that chain (Goodfellow et al., 2016, §8.2.5). Sigmoid derivatives are at most ¼; tanh saturates at large magnitudes; ReLU has zero derivative on its negative side. In this implementation, the derivative at ReLU’s exact zero is also taken as zero.

The bars show the root mean square of the last batch’s weight gradients in each layer. Because a batch gradient averages across examples, cancellation can reduce it too. The bars are a diagnostic, not a proof of vanishing gradients.

Tanh and sigmoid use Xavier uniform initialization, with weights drawn from ±√(6/(n_in + n_out)) (Glorot & Bengio, 2010, §4.2.1). ReLU hidden layers use He uniform initialization, ±√(6/n_in): a uniform distribution with the variance 2/n_in that He and colleagues derived for rectifiers, who used a Gaussian (He et al., 2015, §2.2). Every linear output uses Xavier. Biases start at zero. Initialization and activation both affect training, as explored in Glorot and Bengio’s study of deep feedforward networks (Glorot & Bengio, 2010).

Seeded initialization makes resets reproducible. Different architectures still have different shapes and parameter values. Try multiple seeds before drawing a conclusion. There is no dropout, batch normalization, momentum, or adaptive optimizer here.

A model with P parameters lives in a P-dimensional parameter space. A full update moves through that space. The graph chooses two deterministic, random orthonormal directions u and v, then projects the parameter displacement from initialization onto them:

The directions are fixed for the run and independent of its outcomes; they are not fitted to make the trajectory look smooth. The loss height is measured at the actual full parameter vector. We draw no surface because, for P > 2, a projected coordinate does not uniquely determine a model or its loss.

Crossing paths in this view need not mean the same weights were revisited. Two networks with different widths or depths use different parameter spaces, so comparing the absolute lengths of their projected paths is not meaningful. The two-parameter linear model is the exception: its two projection coordinates preserve the entire parameter displacement.

Checkpoints begin at epoch zero. Long runs are progressively downsampled, keeping the beginning and the full time span. The latest state is added to the display even between stored checkpoints. Lines connect sampled states; they do not represent every intermediate update. Axis ranges expand to include the complete displayed path.

This is the implementation running in the experiment. A model is stopped if an update becomes nonfinite or a parameter magnitude exceeds one million; other models can continue. Reset with a smaller learning rate to try again.

import { randomSource } from './star-training';

export type Activation = 'tanh' | 'relu' | 'sigmoid';
export type Task = 'wave' | 'line';
export type Point = { x: number; y: number };
export const targetValue = (x: number, task: Task) => task === 'line' ? 0.7 * x + 0.1 : 0.6 * Math.sin(Math.PI * x) + 0.3 * Math.sin(3 * Math.PI * x);
export function makeData(task: Task, noise: number, seed: number): Point[] {
  const random = randomSource(seed);
  return Array.from({ length: 32 }, (_, i) => {
    const x = -1 + i * 2 / 31;
    return { x, y: targetValue(x, task) + (random() * 2 - 1) * noise };
  });
}
const activate = (z: number, kind: Activation) => kind === 'tanh' ? Math.tanh(z) : kind === 'relu' ? Math.max(0, z) : 1 / (1 + Math.exp(-z));
const derivative = (a: number, kind: Activation) => kind === 'tanh' ? 1 - a * a : kind === 'relu' ? Number(a > 0) : a * (1 - a);
type Layer = { weights: number[][]; biases: number[] };

export class DepthNetwork {
  layers: Layer[];
  gradientRms: number[];
  epochs = 0;
  failed = false;
  constructor(public depth: number, public width: number, public activation: Activation, seed: number) {
    const random = randomSource(seed);
    const sizes = [1, ...Array(depth).fill(width), 1];
    this.layers = sizes.slice(1).map((outputs, i) => {
      const inputs = sizes[i];
      const isHidden = i < depth;
      // He uniform for ReLU hidden layers; Xavier uniform otherwise.
      const bound = activation === 'relu' && isHidden ? Math.sqrt(6 / inputs) : Math.sqrt(6 / (inputs + outputs));
      return { weights: Array.from({ length: outputs }, () => Array.from({ length: inputs }, () => (2 * random() - 1) * bound)), biases: Array(outputs).fill(0) };
    });
    this.gradientRms = Array(this.layers.length).fill(0);
  }
  get parameterCount() {
    return this.layers.reduce((total, layer) => total + layer.weights.length * layer.weights[0].length + layer.biases.length, 0);
  }
  parameters() { return this.layers.flatMap(layer => [...layer.weights.flat(), ...layer.biases]); }
  forward(x: number) {
    const activations = [[x]];
    for (let l = 0; l < this.layers.length; l++) {
      const layer = this.layers[l];
      const values = layer.weights.map((row, j) => {
        const z = row.reduce((sum, w, i) => sum + w * activations[l][i], layer.biases[j]);
        return l === this.depth ? z : activate(z, this.activation);
      });
      activations.push(values);
    }
    return activations;
  }
  predict(x: number) { return this.forward(x).at(-1)![0]; }
  loss(data: Point[]) { return data.reduce((sum, point) => sum + (this.predict(point.x) - point.y) ** 2, 0) / data.length; }
  gradients(data: Point[]) {
    const gradients = this.layers.map(layer => ({ weights: layer.weights.map(row => row.map(() => 0)), biases: layer.biases.map(() => 0) }));
    for (const point of data) {
      const a = this.forward(point.x);
      // MSE derivative, averaged over the same complete dataset for every model.
      let delta = [2 * (a.at(-1)![0] - point.y) / data.length];
      for (let l = this.layers.length - 1; l >= 0; l--) {
        const layer = this.layers[l];
        for (let j = 0; j < delta.length; j++) {
          gradients[l].biases[j] += delta[j];
          for (let i = 0; i < a[l].length; i++) gradients[l].weights[j][i] += delta[j] * a[l][i];
        }
        if (l > 0) {
          delta = a[l].map((value, i) => layer.weights.reduce((sum, row, j) => sum + row[i] * delta[j], 0) * derivative(value, this.activation));
        }
      }
    }
    return gradients;
  }
  train(data: Point[], rate: number) {
    if (this.failed) return;
    const gradients = this.gradients(data);
    const next = this.layers.map((layer, l) => ({ weights: layer.weights.map((row, j) => row.map((w, i) => w - rate * gradients[l].weights[j][i])), biases: layer.biases.map((b, j) => b - rate * gradients[l].biases[j]) }));
    if (next.some(layer => [...layer.weights.flat(), ...layer.biases].some(value => !Number.isFinite(value) || Math.abs(value) > 1e6))) {
      this.failed = true; return;
    }
    this.gradientRms = gradients.map(layer => {
      const weights = layer.weights.flat();
      return Math.sqrt(weights.reduce((sum, value) => sum + value ** 2, 0) / weights.length);
    });
    this.layers = next;
    this.epochs++;
  }
}

export type DepthConfig = { depths: number[]; width: number; activation: Activation; task: Task; noise: number; seed: number };
export const DEFAULT_DEPTH_CONFIG: DepthConfig = { depths: [0, 1, 3, 6], width: 8, activation: 'tanh', task: 'wave', noise: 0, seed: 42 };

// A fixed orthonormal pair in parameter space, chosen independently of training.
// Using two directions leaves the third display axis available for actual loss.
export function projectionBasis(size: number, seed = 918) {
  const random = randomSource(seed);
  const normalize = (vector: number[]) => { const length = Math.hypot(...vector); return vector.map(value => value / length); };
  const u = normalize(Array.from({ length: size }, () => random() * 2 - 1));
  const raw = Array.from({ length: size }, () => random() * 2 - 1);
  const dot = raw.reduce((sum, value, i) => sum + value * u[i], 0);
  const v = normalize(raw.map((value, i) => value - dot * u[i]));
  return { u, v };
}
export function projectParameters(values: number[], initial: number[], basis: ReturnType<typeof projectionBasis>) {
  return values.reduce((point, value, i) => ({ x: point.x + (value - initial[i]) * basis.u[i], y: point.y + (value - initial[i]) * basis.v[i] }), { x: 0, y: 0 });
}
export type DepthCheckpoint = { epoch: number; losses: number[]; positions: { x: number; y: number }[] };
export class DepthExperiment {
  models: DepthNetwork[];
  data: Point[];
  clean: Point[];
  epoch = 0;
  history: DepthCheckpoint[] = [];
  stride = 10;
  initial: number[][];
  bases: ReturnType<typeof projectionBasis>[];
  constructor(public config: DepthConfig) {
    this.data = makeData(config.task, config.noise, config.seed);
    this.clean = Array.from({ length: 101 }, (_, i) => { const x = -1 + i * 2 / 100; return { x, y: targetValue(x, config.task) }; });
    this.models = config.depths.map(depth => new DepthNetwork(depth, config.width, config.activation, config.seed));
    this.initial = this.models.map(model => model.parameters());
    this.bases = this.initial.map(values => projectionBasis(values.length));
    this.record();
  }
  step(rate: number) {
    this.models.forEach(model => model.train(this.data, rate));
    this.epoch++;
    if (this.epoch % this.stride === 0) this.record();
  }
  checkpoint(): DepthCheckpoint {
    return { epoch: this.epoch, losses: this.models.map(model => model.failed ? NaN : model.loss(this.data)), positions: this.models.map((model, i) => projectParameters(model.parameters(), this.initial[i], this.bases[i])) };
  }
  record() {
    this.history.push(this.checkpoint());
    // Keep the origin and the whole time span with bounded memory.
    if (this.history.length > 201) { this.history = this.history.filter((_, i) => i % 2 === 0); this.stride *= 2; }
  }
}