Momentum Viz

A 1-D ball that rolls down a loss curve under heavy-ball momentum dynamics: v_{t+1} = β·v_t + ∇L(x_t), then x_{t+1} = x_t − α·v_{t+1}. Set the momentum coefficient β to zero and steps stay tiny — pure SGD. Slide β toward 1 and velocity compounds: the ball rolls past the minimum, overshoots, and oscillates back. The velocity vector is rendered as an arrow on top of the ball so the about-to-take step is always legible.

Installation

npx shadcn@latest add https://craftbits.dev/r/momentum-viz.json

Usage

import { MomentumViz } from "@craft-bits/core";
 
<MomentumViz />

Drive the momentum coefficient and step from outside the component:

const [coeff, setCoeff] = useState(0.9);
const [step, setStep] = useState(0);
 
<MomentumViz
  momentumCoeff={coeff}
  onMomentumCoeffChange={setCoeff}
  currentStep={step}
  onCurrentStepChange={setStep}
/>

Provide a different 1-D loss function:

// A double-well so the ball can get stuck on the wrong side.
<MomentumViz
  loss={(x) => (x * x - 1) ** 2}
  initialX={-1.5}
  xRange={[-2, 2]}
/>

Autoplay through the trajectory:

<MomentumViz playing playSpeed={120} />

Understanding the component

1. Two coupled scalar updates. Every step the optimizer evaluates the gradient at x_t, blends it into the velocity (v_{t+1} = β·v_t + g_t), then takes a step proportional to the new velocity (x_{t+1} = x_t − α·v_{t+1}). At β = 0 the velocity collapses to the current gradient — pure SGD. As β rises, the velocity term remembers more of the past, so the ball builds momentum even when the gradient flattens.
2. The arrow shows the next step. The velocity vector drawn on top of the ball is −α · v_{t+1} — the signed displacement the optimizer is about to apply. Watch it grow as the ball accelerates downhill, flip sign when the ball climbs past the minimum, and shrink as the system damps toward zero.
3. Trajectory is precomputed, then scrubbed. The component simulates steps iterations up-front each time momentumCoeff, learningRate, loss, initialX, or xRange changes. currentStep is a cursor into the precomputed trajectory — autoplay and the slider both shift the cursor without rerunning the math. The same inputs always yield the same trajectory, so server and client renders agree.
4. Numerical gradient. The default loss f(x) = x^2 is differentiated via a central difference (f(x+h) − f(x−h)) / 2h with h = 1e-3 — accurate enough that the rendered trajectory is indistinguishable from the analytic one, and zero-config for arbitrary user-supplied loss functions.
5. Reduced-motion fallback. With prefers-reduced-motion: reduce the ball snaps directly to its new position and autoplay is suppressed.

Props

Accessibility

• The figure is role="figure" with a labelled title and an aria-live="polite" summary that names the current β, learning rate, step, position x, velocity v, and loss L(x) — every value the bar visualises is also exposed textually.
• Both sliders are native <input type="range"> instances wrapped in LabeledSlider, so keyboard arrows nudge the value and aria-valuetext is read aloud by screen readers.
• Color is never the only signal — the trail, ball, and arrow all share the same accent hue but the textual readouts below the curve repeat the same information.
• prefers-reduced-motion: reduce snaps the ball directly to its target position and suppresses autoplay; the velocity arrow updates in lockstep.

Credits

• Extracted from: craftingattention (app/src/lessons/primitives/viz/MomentumViz.tsx). Stripped the multi-optimizer (SGD / Momentum / Adam) ravine racetrack, the predict-quiz mode, the 2-D contour plot, the usePredictRounds curriculum integration, and the lesson-specific narration. Generalised to the textbook heavy-ball story on a 1-D curve: configurable loss, controlled / uncontrolled momentumCoeff and currentStep pairs, autoplay, and a velocity arrow rendered as the next-step displacement.