diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index ff1d33a..8bf19ab 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -44,18 +44,6 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-sbt-
 
-      # TEMPORARY: dimwit 0.2-SNAPSHOT is not published to any remote yet, so CI builds it from source.
-      - name: Check out dimwit
-        uses: actions/checkout@v4
-        with:
-          repository: marcelluethi/dimwit
-          ref: version-0.2-SNAPSHOT
-          path: dimwit
-
-      - name: Publish dimwit-core to the local ivy repo
-        working-directory: dimwit
-        run: sbt core/publishLocal
-
       - name: Check formatting
         run: sbt scalafmtCheckAll
 
diff --git a/README.md b/README.md
index 896221c..c4a51d3 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,21 @@
 
 DeepWit is a deep learning library for Scala 3 built on one idea: **the code should mirror the theory.** DeepWit is build on [DimWit](https://github.com/dimwit-dev/dimwit), a statically typed tensor library with runtime performance on-par with JAX. DeepWit provides core deep learning modules with their minimal conceptual scope to both express their logic clearly and increase reusablility.
 
+## Installation
+
+DeepWit is published for Scala 3. Add both DeepWit and [DimWit](https://github.com/dimwit-dev/dimwit) —
+DeepWit's API is expressed in DimWit's tensor types, so you will import from both:
+
+```scala
+libraryDependencies ++= Seq(
+  "ch.contrafactus" %% "deepwit-core" % "0.1.0",
+  "ch.contrafactus" %% "dimwit-core" % "0.1.0"
+)
+```
+
+DimWit runs JAX through ScalaPy, so a Python environment with `jax` and `einops` is also required;
+see [DimWit's setup instructions](https://github.com/dimwit-dev/dimwit#installation).
+
 ## Why DeepWit?
 
 DeepWit strips out the framework machinery and leaves the theory standing: no module base class, no indirect loss, no hidden parameters, no hidden gradients, no hidden optimizer states. 
diff --git a/build.sbt b/build.sbt
index e2fe594..8a6ffdb 100644
--- a/build.sbt
+++ b/build.sbt
@@ -4,23 +4,58 @@ import scala.sys.process._
 run / fork := true
 Global / cancelable := true
 
-ThisBuild / version := "0.2-SNAPSHOT"
+ThisBuild / version := "0.1.0"
 ThisBuild / scalaVersion := "3.8.1"
 ThisBuild / organization := "ch.contrafactus"
+ThisBuild / versionScheme := Some("early-semver")
+
+// Publishing to Sonatype Central. The `ch.contrafactus` namespace is verified once for the whole
+// organisation, so deepwit needs no verification of its own.
+ThisBuild / sonatypeCredentialHost := "central.sonatype.com"
+ThisBuild / publishTo := {
+  if (isSnapshot.value)
+    Some("central-snapshots" at "https://central.sonatype.com/repository/maven-snapshots/")
+  else
+    sonatypePublishToBundle.value
+}
+ThisBuild / publishMavenStyle := true
+ThisBuild / homepage := Some(url("https://github.com/dimwit-dev/deepwit"))
+ThisBuild / licenses := List("Apache-2.0" -> url("https://www.apache.org/licenses/LICENSE-2.0"))
+ThisBuild / scmInfo := Some(
+  ScmInfo(
+    url("https://github.com/dimwit-dev/deepwit"),
+    "scm:git@github.com:dimwit-dev/deepwit.git"
+  )
+)
+ThisBuild / developers := List(
+  Developer(
+    id = "dimwit-dev",
+    name = "DeepWit Contributors",
+    email = "",
+    url = url("https://github.com/dimwit-dev")
+  )
+)
 
 // scalafix's RemoveUnused reads the compiler's own unused warnings on Scala 3, so it needs both
 // semanticdb and -Wunused to be on. Run it with `sbt scalafixAll`.
 ThisBuild / semanticdbEnabled := true
 ThisBuild / scalacOptions += "-Wunused:imports"
 
-// deepwit tracks the dimwit/plotwit 0.2 snapshots and is itself snapshot-only for now.
-// NOTE: dimwit-core 0.2-SNAPSHOT is not on this resolver yet (only 0.1.0-SNAPSHOT is), and
-// plotwit is not published at all, so both currently resolve from the local ivy repo via
-// `publishLocal`. CI cannot resolve either until dimwit 0.2-SNAPSHOT is published here.
+// `core` depends only on released artifacts so that it can be published and consumed without any
+// local publishing. `examples` additionally needs plotwit, which is not published anywhere yet and
+// so still resolves from the local ivy repo via `publishLocal` in the plotwit checkout.
 ThisBuild / resolvers += "Central Portal Snapshots" at "https://central.sonatype.com/repository/maven-snapshots/"
 
+// Consequence of that split: `core` asks for the released dimwit 0.1.0 while the locally published
+// plotwit asks for 0.2-SNAPSHOT, which sbt reads as a binary-incompatible conflict under
+// early-semver. Let the newer snapshot win in `examples`; `core` on its own resolves 0.1.0.
+ThisBuild / libraryDependencySchemes += "ch.contrafactus" %% "dimwit-core" % "always"
+
 addCommandAlias("testAndCoverage", "; clean; coverage; test; coverageReport")
 
+// Publishes `core` only. `examples` depends on plotwit, which is not published.
+addCommandAlias("sonaUploadCore", "; project core; sonatypeCentralUpload; project root")
+
 lazy val uvPython: String =
   sys.env.getOrElse(
     "DIMWIT_PYTHON_PATH",
@@ -47,7 +82,7 @@ lazy val core = (project in file("core"))
       "org.scalacheck" %% "scalacheck" % "1.18.0" % Test,
       "org.scalatestplus" %% "scalacheck-1-18" % "3.2.19.0" % Test,
       "dev.scalapy" %% "scalapy-core" % "0.5.3",
-      "ch.contrafactus" %% "dimwit-core" % "0.2-SNAPSHOT" changing ()
+      "ch.contrafactus" %% "dimwit-core" % "0.1.0"
     ),
     // ScalaPy drives a single embedded CPython interpreter, and two suites importing jax at the same
     // time race into a partially initialized module. Whichever suites happen to touch a tensor first
@@ -57,8 +92,12 @@ lazy val core = (project in file("core"))
     // SCALAPY_PYTHON_LIBRARY / SCALAPY_PYTHON_PROGRAMNAME being exported by the shell.
     fork := true,
     javaOptions ++= scalapyJavaOptions,
+    description := "A theory-aligned deep learning library for Scala 3, built on DimWit",
     Compile / packageSrc / publishArtifact := true,
-    Compile / packageDoc / publishArtifact := true
+    Compile / packageDoc / publishArtifact := true,
+    // Ship the library's own sources and docs, not the test ones.
+    Test / packageSrc / publishArtifact := false,
+    Test / packageDoc / publishArtifact := false
   )
 
 // Examples subproject
diff --git a/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala b/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala
index 8f4ed16..ae72e19 100644
--- a/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala
+++ b/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala
@@ -4,10 +4,13 @@ import dimwit.*
 import deepwit.base.AffineLayer
 import dimwit.Label as Λ
 
-/** TODO
-  * This implementation is conceptually clearer than [[MultiHeadAttention]] but slower. How to merge?
+/** A readable statement of what [[MultiHeadAttention]] computes, kept for reference rather than for
+  * use: each head is its own [[Attention]] in a `List`, so the per-head structure stays visible
+  * instead of being folded into batched tensors. That is also what makes it slower — the per-head
+  * parameters are separate tensors rather than one.
+  *
+  * [[MultiHeadAttention]] is tested against it, so the batched implementation is known to agree.
   */
-
 class ReferenceMultiHeadAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding: Λ, V: IsFloating](
     params: ReferenceMultiHeadAttention.Params[SourceEmbedding, TargetEmbedding, V],
     createAttentionMask: Shape2[Target, Source] => Tensor2[Target, Source, Bool],
diff --git a/core/src/main/scala/deepwit/attention/package.scala b/core/src/main/scala/deepwit/attention/package.scala
index 1efbb5d..20e503c 100644
--- a/core/src/main/scala/deepwit/attention/package.scala
+++ b/core/src/main/scala/deepwit/attention/package.scala
@@ -2,8 +2,14 @@ package deepwit.attention
 
 import dimwit.Label
 
-/** Axis labels for the per-head spaces of a [[MultiHeadAttention]]. */
+/** The multi-head attention heads. */
 trait Head derives Label
+
+/** The space a head projects queries into. */
 trait HeadQuery derives Label
+
+/** The space a head projects keys into. */
 trait HeadKey derives Label
+
+/** The space a head projects values into. */
 trait HeadValue derives Label
diff --git a/core/src/main/scala/deepwit/cnn/AffineConv2DLayer.scala b/core/src/main/scala/deepwit/cnn/AffineConv2DLayer.scala
index 0f8ed9a..707b602 100644
--- a/core/src/main/scala/deepwit/cnn/AffineConv2DLayer.scala
+++ b/core/src/main/scala/deepwit/cnn/AffineConv2DLayer.scala
@@ -3,6 +3,11 @@ package deepwit.cnn
 import dimwit.*
 import dimwit.Label as Λ
 
+/** A 2D convolution with bias. For the bias-free equivalent, see [[LinearConv2DLayer]].
+  *
+  * @param stride An `Int` stride applies to both spatial axes.
+  * @param padding `Padding.SAME` preserves the spatial extents, while `Padding.VALID` shrinks them.
+  */
 class AffineConv2DLayer[S1: Λ, S2: Λ, InChannel: Λ, OutChannel: Λ, V: IsFloating](
     params: AffineConv2DLayer.Params[S1, S2, InChannel, OutChannel, V],
     stride: Stride2[S1, S2] | Int = 1,
diff --git a/core/src/main/scala/deepwit/cnn/LinearConv2DLayer.scala b/core/src/main/scala/deepwit/cnn/LinearConv2DLayer.scala
index 0454b92..2665373 100644
--- a/core/src/main/scala/deepwit/cnn/LinearConv2DLayer.scala
+++ b/core/src/main/scala/deepwit/cnn/LinearConv2DLayer.scala
@@ -3,6 +3,11 @@ package deepwit.cnn
 import dimwit.*
 import dimwit.Label as Λ
 
+/** A 2D convolution without bias. For the bias-carrying equivalent, see [[AffineConv2DLayer]].
+  *
+  * @params stride An `Int` stride applies to both spatial axes.
+  * @params padding `Padding.SAME` preserves the spatial extents, while `Padding.VALID` shrinks them.
+  */
 class LinearConv2DLayer[S1: Λ, S2: Λ, InChannel: Λ, OutChannel: Λ, V: IsFloating](
     params: LinearConv2DLayer.Params[S1, S2, InChannel, OutChannel, V],
     stride: Stride2[S1, S2] | Int = 1,
diff --git a/core/src/main/scala/deepwit/cnn/MaxPool2DLayer.scala b/core/src/main/scala/deepwit/cnn/MaxPool2DLayer.scala
index cebad27..2ff8ad6 100644
--- a/core/src/main/scala/deepwit/cnn/MaxPool2DLayer.scala
+++ b/core/src/main/scala/deepwit/cnn/MaxPool2DLayer.scala
@@ -5,9 +5,15 @@ import dimwit.jax.Jax
 import dimwit.python.PyBridge.{liftPyTensor, toPyTensor}
 import dimwit.Label as Λ
 
+/** A sliding-window maximum over the two spatial axes.
+  *
+  * @param window An `Int` window applies to both spatial axes.
+  * @param stride An `Int` stride applies to both spatial axes.
+  * @param padding `Padding.SAME` preserves the spatial extents, while `Padding.VALID` shrinks them.
+  */
 class MaxPool2DLayer[S1: Λ, S2: Λ, V: IsFloating](
     window: Window2[S1, S2] | Int,
-    stride: Stride2[S1, S2] | Int = 1,
+    stride: Stride2[S1, S2] | Int,
     padding: Padding = Padding.SAME
 ) extends (Tensor2[S1, S2, V] => Tensor2[S1, S2, V]):
 
diff --git a/core/src/main/scala/deepwit/cnn/TransposeAffineConv2DLayer.scala b/core/src/main/scala/deepwit/cnn/TransposeAffineConv2DLayer.scala
index 23a8791..7a9b0c0 100644
--- a/core/src/main/scala/deepwit/cnn/TransposeAffineConv2DLayer.scala
+++ b/core/src/main/scala/deepwit/cnn/TransposeAffineConv2DLayer.scala
@@ -3,6 +3,14 @@ package deepwit.cnn
 import dimwit.*
 import dimwit.Label as Λ
 
+/** The adjoint of a 2D convolution, with bias. For the bias-free equivalent, see
+  * [[TransposeLinearConv2DLayer]].
+  *
+  * @tparam InChannel The forward convolution's input channels, so this layer's output.
+  * @tparam OutChannel The forward convolution's output channels, so this layer's input.
+  * @param stride An `Int` stride applies to both spatial axes. A stride above 1 grows the spatial extents, which is how these layers upsample.
+  * @param padding `Padding.SAME` preserves the spatial extents, while `Padding.VALID` grows them.
+  */
 class TransposeAffineConv2DLayer[S1: Λ, S2: Λ, InChannel: Λ, OutChannel: Λ, V: IsFloating](
     params: TransposeAffineConv2DLayer.Params[S1, S2, InChannel, OutChannel, V],
     stride: Stride2[S1, S2] | Int = 1,
diff --git a/core/src/main/scala/deepwit/cnn/TransposeLinearConv2DLayer.scala b/core/src/main/scala/deepwit/cnn/TransposeLinearConv2DLayer.scala
index 0742081..a738169 100644
--- a/core/src/main/scala/deepwit/cnn/TransposeLinearConv2DLayer.scala
+++ b/core/src/main/scala/deepwit/cnn/TransposeLinearConv2DLayer.scala
@@ -3,6 +3,14 @@ package deepwit.cnn
 import dimwit.*
 import dimwit.Label as Λ
 
+/** The adjoint of a 2D convolution, without bias. For the bias-carrying equivalent, see
+  * [[TransposeAffineConv2DLayer]].
+  *
+  * @tparam InChannel The forward convolution's input channels, so this layer's output.
+  * @tparam OutChannel The forward convolution's output channels, so this layer's input.
+  * @param stride An `Int` stride applies to both spatial axes. A stride above 1 grows the spatial extents, which is how these layers upsample.
+  * @param padding `Padding.SAME` preserves the spatial extents, while `Padding.VALID` grows them.
+  */
 class TransposeLinearConv2DLayer[S1: Λ, S2: Λ, InChannel: Λ, OutChannel: Λ, V: IsFloating](
     params: TransposeLinearConv2DLayer.Params[S1, S2, InChannel, OutChannel, V],
     stride: Stride2[S1, S2] | Int = 1,
diff --git a/core/src/main/scala/deepwit/embedder/PositionalEncoding.scala b/core/src/main/scala/deepwit/embedder/PositionalEncoding.scala
index f2a8ab1..62be2b8 100644
--- a/core/src/main/scala/deepwit/embedder/PositionalEncoding.scala
+++ b/core/src/main/scala/deepwit/embedder/PositionalEncoding.scala
@@ -6,7 +6,7 @@ import dimwit.Label as Λ
 
 object PositionalEncoding:
 
-  /** The axis the frequencies live on, half of which become sines and half cosines. */
+  /** The axis the frequencies live on, each contributing both a sine and a cosine. */
   private trait Scale derives Label
 
   /** The default ratio between the fastest and the slowest oscillation, as chosen for sequences of
@@ -18,11 +18,13 @@ object PositionalEncoding:
   def gridPositions[P: Λ, V: IsFloating](extent: AxisExtent[P], vtype: VType[V]): Tensor1[P, V] =
     Tensor1(Axis[P]).fromArray(Array.range(0, extent.size)).asFloat(vtype)
 
-  /** Encodes each position as sines and cosines of it, at geometrically spaced frequencies.
+  /** Encodes each position as sines and cosines of it, at geometrically spaced frequencies, as
+    * described in [Attention Is All You Need](https://arxiv.org/abs/1706.03762).
     *
     * The positions are given rather than derived, so they need not be a grid's indices: any
     * position the encoding is evaluated at means the same thing to a model trained on any other.
     *
+    * @param embeddingExtent Must be even, to pair each frequency's sine with its cosine.
     * @param frequencyRange The ratio between the fastest and the slowest oscillation.
     */
   def sinusoidal[P: Λ, Embedding: Λ, V: IsFloating](
@@ -41,7 +43,11 @@ object PositionalEncoding:
     val scaled = positions.vmap(Axis[P])(_ *! scales)
     concatenate(scaled.sin, scaled.cos, concatAxis = Axis[Scale]).relabel(Axis[Scale], Axis[Embedding])
 
-  /** Encodes each point of a plane, giving each of the two axes half of the embedding. */
+  /** Encodes each point of a plane, giving each of the two axes half of the embedding.
+    *
+    * @param embeddingExtent Must be divisible by four: each axis takes half, and each half pairs sines with cosines.
+    * @param frequencyRange The ratio between the fastest and the slowest oscillation.
+    */
   def sinusoidal2D[X: Λ, Y: Λ, Embedding: Λ, V: IsFloating](
       xPositions: Tensor1[X, V],
       yPositions: Tensor1[Y, V],
diff --git a/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala b/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
index e55ba16..a425a1a 100644
--- a/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
+++ b/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
@@ -3,18 +3,17 @@ package deepwit.embedder
 import dimwit.*
 import dimwit.stats.Normal
 import dimwit.stats.Uniform
-import dimwit.jax.Jax
-import dimwit.python.PyBridge.{toPyTensor, liftPyTensor}
 import dimwit.Label as Λ
 
+/** Maps a token to its embedding. [[VocabularyEmbedder.unembed]] maps back (weight typing). */
 class VocabularyEmbedder[Vocab: Λ, Embedding: Λ, V: IsFloating](params: VocabularyEmbedder.Params[Vocab, Embedding, V]) extends (Tensor0[Int32] => Tensor1[Embedding, V]):
 
   override def apply(token: Tensor0[Int32]): Tensor1[Embedding, V] =
-    // params.vocabularyEmbeddings.slice(Axis[Vocab].at(token)) // TODO
-    // params.vocabularyEmbeddings.take(Axis[Vocab])(token)
-    val rawJax = Jax.jnp.take(toPyTensor(params.vocabularyEmbeddings), toPyTensor(token), axis = 0)
-    liftPyTensor(rawJax)
+    params.vocabularyEmbeddings.slice(Axis[Vocab].at(token))
 
+  /** Scores every token by projecting onto its embedding, through the matrix the lookup uses —
+    * weight tying, as described in [Using the Output Embedding to Improve Language Models](https://arxiv.org/abs/1608.05859).
+    */
   def unembed(embedding: Tensor1[Embedding, V]): Tensor1[Vocab, V] =
     embedding.dot(Axis[Embedding])(params.vocabularyEmbeddings)
 
@@ -22,6 +21,9 @@ object VocabularyEmbedder:
 
   case class Params[Vocab, Embedding, V](vocabularyEmbeddings: Tensor2[Vocab, Embedding, V])
 
+  /** Scaled by the embedding size — a lookup has no fan-in of its own, and this leaves each row
+    * near unit norm.
+    */
   object Params:
 
     def init[Vocab: Λ, Embedding: Λ, V: IsFloating](vocabExtent: AxisExtent[Vocab], embeddingExtent: AxisExtent[Embedding], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1.0): Params[Vocab, Embedding, V] =
diff --git a/core/src/main/scala/deepwit/init/Initialization.scala b/core/src/main/scala/deepwit/init/Initialization.scala
index 73f1c08..a4054d8 100644
--- a/core/src/main/scala/deepwit/init/Initialization.scala
+++ b/core/src/main/scala/deepwit/init/Initialization.scala
@@ -5,6 +5,9 @@ import dimwit.stats.Normal
 import dimwit.stats.Uniform
 import dimwit.Label as Λ
 
+/** Xavier/Glorot initializers, drawing weights with variance `2 / (fanIn + fanOut)` as described in
+  * [Understanding the difficulty of training deep feedforward neural networks](https://proceedings.mlr.press/v9/glorot10a.html).
+  */
 object Init:
 
   def xavierNormal[FanIn: Λ, FanOut: Λ, V: IsFloating](fanIn: AxisExtent[FanIn], fanOut: AxisExtent[FanOut], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1f): Tensor2[FanIn, FanOut, V] =
@@ -13,14 +16,18 @@ object Init:
 
   def xavierUniform[FanIn: Λ, FanOut: Λ, V: IsFloating](fanIn: AxisExtent[FanIn], fanOut: AxisExtent[FanOut], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1f): Tensor2[FanIn, FanOut, V] =
     val variance = Tensor0(vtype)(2.0f / (fanIn.size + fanOut.size))
+    // A uniform on [-a, a] has variance a²/3, so a = √(3·variance) hits the target spread.
     val a = gain * (3f * variance).sqrt
     IndependentDistribution.fromUnivariate(Shape(fanIn, fanOut), Uniform(-a, a)).sample(key)
 
   def xavierNormalVector[FanIn: Λ, V: IsFloating](fanIn: AxisExtent[FanIn], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1f): Tensor1[FanIn, V] =
-    val variance = Tensor0(vtype)(2.0f / (fanIn.size + 1))
+    val fanOut = 1 // linear form has fan-out of 1
+    val variance = Tensor0(vtype)(2.0f / (fanIn.size + fanOut))
     Normal.standardIsotropic(Shape(fanIn), scale = gain * variance.sqrt).sample(key)
 
   def xavierUniformVector[FanIn: Λ, V: IsFloating](fanIn: AxisExtent[FanIn], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1f): Tensor1[FanIn, V] =
-    val variance = Tensor0(vtype)(2.0f / (fanIn.size + 1))
+    val fanOut = 1 // linear form has fan-out of 1
+    val variance = Tensor0(vtype)(2.0f / (fanIn.size + fanOut))
+    // A uniform on [-a, a] has variance a²/3, so a = √(3·variance) hits the target spread.
     val a = gain * (3f * variance).sqrt
     IndependentDistribution.fromUnivariate(Shape(fanIn), Uniform(-a, a)).sample(key)
diff --git a/core/src/main/scala/deepwit/loss/Regression.scala b/core/src/main/scala/deepwit/loss/Regression.scala
index 0b2ece4..3d1ad59 100644
--- a/core/src/main/scala/deepwit/loss/Regression.scala
+++ b/core/src/main/scala/deepwit/loss/Regression.scala
@@ -15,9 +15,13 @@ object AbsoluteError:
 
 object Huber:
 
-  def apply[V: IsFloating](target: Tensor0[V], prediction: Tensor0[V], threshold: Float): Tensor0[V] =
-    require(threshold > 0f, s"A transition point must be positive, but was $threshold.")
+  /** Quadratic like [[SquaredError]] if residual within `transitionPoint`, linear like [[AbsoluteError]] beyond it as described in
+    * [Robust Estimation of a Location Parameter](https://doi.org/10.1214/aoms/1177703732).
+    */
+  def apply[V: IsFloating](target: Tensor0[V], prediction: Tensor0[V], transitionPoint: Float): Tensor0[V] =
+    require(transitionPoint > 0f, s"A transition point must be positive, but was $transitionPoint.")
     val residual = AbsoluteError(target, prediction)
+    // Scale squared and absolute errors to meet in value and slope at the transition point.
     val squared = 0.5f * SquaredError(target, prediction)
-    val absolute = threshold * (residual - 0.5f * threshold)
-    where(residual <= threshold, squared, absolute)
+    val absolute = transitionPoint * (residual - 0.5f * transitionPoint)
+    where(residual <= transitionPoint, squared, absolute)
diff --git a/core/src/main/scala/deepwit/normalization/LayerNorm.scala b/core/src/main/scala/deepwit/normalization/LayerNorm.scala
index 0efe825..aec5ba8 100644
--- a/core/src/main/scala/deepwit/normalization/LayerNorm.scala
+++ b/core/src/main/scala/deepwit/normalization/LayerNorm.scala
@@ -6,6 +6,11 @@ import dimwit.Label as Λ
 
 import deepwit.{defaultEpsilon, unwrapEpsilon}
 
+/** Standardizes over the `L` axis, then scales and shifts by the learned parameters, as described
+  * in [Layer Normalization](https://arxiv.org/abs/1607.06450).
+  *
+  * @param epsilon Guards the division. Defaults to the machine epsilon of data type; pass a `Float` to fix it. Pass a function to derive it from the data type.
+  */
 class LayerNorm[L: Λ, V: IsFloating](
     params: LayerNorm.Params[L, V],
     epsilon: Float | (DType => Float) = defaultEpsilon
diff --git a/core/src/main/scala/deepwit/normalization/RMSNorm.scala b/core/src/main/scala/deepwit/normalization/RMSNorm.scala
index 4f2964c..78c81ec 100644
--- a/core/src/main/scala/deepwit/normalization/RMSNorm.scala
+++ b/core/src/main/scala/deepwit/normalization/RMSNorm.scala
@@ -6,6 +6,11 @@ import dimwit.Label as Λ
 
 import deepwit.{defaultEpsilon, unwrapEpsilon}
 
+/** Rescales by the root mean square over the `L` axis, then scales by the learned weight, as
+  * described in [Root Mean Square Layer Normalization](https://arxiv.org/abs/1910.07467).
+  *
+  * @param epsilon Guards the division. Defaults to the machine epsilon of data type; pass a `Float` to fix it. Pass a function to derive it from the data type.
+  */
 class RMSNorm[L: Λ, V: IsFloating](
     params: RMSNorm.Params[L, V],
     epsilon: Float | (DType => Float) = defaultEpsilon
@@ -15,7 +20,6 @@ class RMSNorm[L: Λ, V: IsFloating](
 
   def apply(x: Tensor1[L, V]): Tensor1[L, V] =
     def rescale(x: Tensor1[L, V]): Tensor1[L, V] =
-      // Unlike LayerNorm, RMSNorm does not re-center: it only divides by the root mean square.
       val meanSquare = x.pow(2).mean
       x /! (meanSquare + ε).sqrt
     rescale(x) * params.weight
diff --git a/core/src/main/scala/deepwit/optimizer/LearningRateSchedule.scala b/core/src/main/scala/deepwit/optimizer/LearningRateSchedule.scala
index b58f3cd..8fa2494 100644
--- a/core/src/main/scala/deepwit/optimizer/LearningRateSchedule.scala
+++ b/core/src/main/scala/deepwit/optimizer/LearningRateSchedule.scala
@@ -10,6 +10,10 @@ case class LearningRateSchedulerState[P, State[_]](
 )
 type LearningRateSchedulerStateFor[State[_]] = [P] =>> LearningRateSchedulerState[P, State]
 
+/** Wraps an optimizer so its learning rate follows `schedule`, counting steps from 1.
+  *
+  * @param optF Builds the optimizer at a given learning rate, e.g. `lr => Adam(lr)`.
+  */
 class LearningRateScheduler[State[_]](val optF: Tensor0[Float32] => GradientOptimizer[State], schedule: Tensor0[Int32] => Tensor0[Float32]) extends GradientOptimizer[LearningRateSchedulerStateFor[State]]:
 
   def init[P: TensorTree, V](params: P)(using TreeOf[P, V])(using IsFloating[V]): LearningRateSchedulerState[P, State] =
@@ -24,6 +28,11 @@ class LearningRateScheduler[State[_]](val optF: Tensor0[Float32] => GradientOpti
     val (newParams, newOptState) = opt.update(gradients, params, optState)
     (newParams, LearningRateSchedulerState(step + 1, newOptState))
 
+/** A learning rate as a function of the step count, which starts at 1.
+  *
+  * @param steps How many steps this schedule is defined for, after which it holds its final value.
+  *              [[LearningRateSchedule.followBy]] uses it to know when to hand over.
+  */
 trait LearningRateSchedule(val steps: Tensor0[Int32]) extends (Tensor0[Int32] => Tensor0[Float32])
 
 object LearningRateSchedule:
@@ -54,28 +63,21 @@ object LearningRateSchedule:
       *
       * For all `t < steps`, the schedule evaluates as if `t = 0`, effectively locking
       * the learning rate at its initial starting value until the delay has passed.
-      *
-      * @param delaySteps The number of iterations to delay the schedule's progression.
-      * @return A time-shifted schedule.
       */
     def delay(delaySteps: Tensor0[Int32]): LearningRateSchedule = DelayedSchedule(schedule, delaySteps)
 
+    /** Runs this schedule through its `steps`, then hands over to `second` delayed by that much,
+      * so `second` starts from its own beginning rather than mid-curve.
+      */
     def followBy(second: LearningRateSchedule): LearningRateSchedule = FollowBySchedule(schedule, second)
 
 object LearningRateSchedules:
 
-  /** Creates a schedule that maintains a constant learning rate for a specified number of steps. */
+  /** A schedule of constant `learningRate` for a specified number of `steps`. */
   class ConstantLearningRate(val learningRate: Tensor0[Float32], steps: Int = Int.MaxValue) extends LearningRateSchedule(steps = Tensor0(steps)):
     override def apply(step: Tensor0[Int32]): Tensor0[Float32] = learningRate
 
-  /** Creates a schedule that rises linearly from `minLr` to `maxLr` over a specified number of warmup steps.
-    *
-    * @param vtype The floating-point type to use for the learning rate values.
-    * @param from The initial learning rate at the start of the warmup (typically 0.0).
-    * @param to The peak learning rate reached at the end of the warmup.
-    * @param warmupSteps The number of steps over which the learning rate increases linearly.
-    * @return A linear warmup schedule.
-    */
+  /** A schedule rising linearly from `from` to `to` for a specified number of `warmupSteps`. */
   class LinearSchedule(
       val from: Tensor0[Float32],
       val to: Tensor0[Float32],
@@ -85,6 +87,8 @@ object LearningRateSchedules:
     private val vtype = from.vtype
 
     override def apply(step: Tensor0[Int32]): Tensor0[Float32] =
+      // Steps are 1-based, and the +1 keeps the first step off `from` — with LinearWarmup's
+      // `from = 0` that would be a step at zero learning rate.
       val warmupRatio = minimum((step.asFloat(vtype)) / ((warmupSteps + 1).asFloat(vtype)), 1f)
       from + warmupRatio * (to - from)
 
@@ -94,14 +98,9 @@ object LearningRateSchedules:
         warmupSteps: Tensor0[Int32]
     ): LinearSchedule = new LinearSchedule(0.0f, to, warmupSteps)
 
-  /** Creates a schedule that decays from `maxLr` down to `minLr` following a half-cosine curve.
-    *
-    * This schedule has no concept of warmup; it begins decaying immediately at `t = 0`.
-    * Once `t >= decaySteps`, the learning rate locks permanently at `minLr`.
+  /** A schedule that decays from `from` down to `to` for `decaySteps` following a half-cosine curve.
     *
-    * @param from The initial maximum learning rate at `t = 0`.
-    * @param to The final baseline learning rate to reach after decaying.
-    * @param decaySteps The number of steps over which to apply the decay curve.
+    * Past `decaySteps` the learning rate stays at `to`.
     */
   class CosineDecay(
       val from: Tensor0[Float32],
diff --git a/core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala b/core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala
new file mode 100644
index 0000000..e5910ec
--- /dev/null
+++ b/core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala
@@ -0,0 +1,56 @@
+package deepwit.attention
+
+import deepwit.*
+import dimwit.*
+import dimwit.stats.Normal
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.funspec.AnyFunSpec
+
+class ReferenceMultiHeadAttentionSuite extends AnyFunSpec with Matchers:
+
+  trait Src derives Label
+  trait SrcEmb derives Label
+  trait Tgt derives Label
+  trait TgtEmb derives Label
+
+  private val numHeads = 2
+  private val srcExtent = Axis[Src] -> 3
+  private val tgtExtent = Axis[Tgt] -> 5
+  private val srcEmbExtent = Axis[SrcEmb] -> 4
+  private val tgtEmbExtent = Axis[TgtEmb] -> 4
+
+  private val batched = MultiHeadAttention.Params.xavierUniformDepthScaled(
+    numTransformerLayers = 2,
+    numHeads = numHeads,
+    sourceEmbeddingExtent = srcEmbExtent,
+    targetEmbeddingExtent = tgtEmbExtent,
+    vtype = VType[Float32],
+    key = Random.Key(42)
+  )
+
+  /** The weights the batched implementation holds, split head by head, so both see the same model. */
+  private val perHead = ReferenceMultiHeadAttention.Params(
+    heads = (0 until numHeads).toList.map(head =>
+      Attention.Params(
+        batched.queryWeights.slice(Axis[Head].at(head)),
+        batched.keyWeights.slice(Axis[Head].at(head)),
+        batched.valueWeights.slice(Axis[Head].at(head))
+      )
+    ),
+    outputProjection = batched.outputProjection
+  )
+
+  private def source = Normal.standardNormal(Shape(srcExtent, srcEmbExtent)).sample(Random.Key(7))
+  private def target = Normal.standardNormal(Shape(tgtExtent, tgtEmbExtent)).sample(Random.Key(8))
+
+  describe("ReferenceMultiHeadAttention"):
+
+    it("computes what MultiHeadAttention computes, under a full mask"):
+      val folded = MultiHeadCustomAttention(batched, fullMask[Tgt, Src], AttentionScore.scaledDotProduct)
+      val reference = ReferenceMultiHeadAttention(perHead, fullMask[Tgt, Src], AttentionScore.scaledDotProduct)
+      reference(source, target) should approxEqual(folded(source, target), 1e-5f)
+
+    it("computes what MultiHeadAttention computes, under a causal mask"):
+      val folded = MultiHeadCustomAttention(batched, causalMask[Tgt, Src], AttentionScore.scaledDotProduct)
+      val reference = ReferenceMultiHeadAttention(perHead, causalMask[Tgt, Src], AttentionScore.scaledDotProduct)
+      reference(source, target) should approxEqual(folded(source, target), 1e-5f)
diff --git a/core/src/test/scala/deepwit/loss/RegressionSuite.scala b/core/src/test/scala/deepwit/loss/RegressionSuite.scala
index 04de571..683fcc9 100644
--- a/core/src/test/scala/deepwit/loss/RegressionSuite.scala
+++ b/core/src/test/scala/deepwit/loss/RegressionSuite.scala
@@ -41,19 +41,19 @@ class RegressionSuite extends AnyFunSpec with Matchers:
   describe("Huber"):
 
     it("matches the halved squared error within the transition point"):
-      val loss = Huber(Tensor0(0f), Tensor0(0.5f), threshold = 1f)
+      val loss = Huber(Tensor0(0f), Tensor0(0.5f), transitionPoint = 1f)
       loss.item shouldBe (0.5f * 0.25f +- 1e-6f)
 
     it("grows linearly beyond the transition point"):
-      val atFive = Huber(Tensor0(0f), Tensor0(5f), threshold = 1f).item
-      val atSix = Huber(Tensor0(0f), Tensor0(6f), threshold = 1f).item
+      val atFive = Huber(Tensor0(0f), Tensor0(5f), transitionPoint = 1f).item
+      val atSix = Huber(Tensor0(0f), Tensor0(6f), transitionPoint = 1f).item
       atSix - atFive shouldBe (1f +- 1e-4f)
 
     it("joins its two branches continuously at the transition point"):
-      val threshold = 2f
-      val below = Huber(Tensor0(0f), Tensor0(threshold - 1e-3f), threshold).item
-      val above = Huber(Tensor0(0f), Tensor0(threshold + 1e-3f), threshold).item
+      val transitionPoint = 2f
+      val below = Huber(Tensor0(0f), Tensor0(transitionPoint - 1e-3f), transitionPoint).item
+      val above = Huber(Tensor0(0f), Tensor0(transitionPoint + 1e-3f), transitionPoint).item
       below shouldBe (above +- 1e-2f)
 
     it("rejects a non-positive transition point"):
-      an[IllegalArgumentException] should be thrownBy Huber(Tensor0(0f), Tensor0(1f), threshold = 0f)
+      an[IllegalArgumentException] should be thrownBy Huber(Tensor0(0f), Tensor0(1f), transitionPoint = 0f)
diff --git a/mdocs/README.md b/mdocs/README.md
index 7e61d83..3825016 100644
--- a/mdocs/README.md
+++ b/mdocs/README.md
@@ -22,6 +22,21 @@ trait Embedding derives Label
 
 DeepWit is a deep learning library for Scala 3 built on one idea: **the code should mirror the theory.** DeepWit is build on [DimWit](https://github.com/dimwit-dev/dimwit), a statically typed tensor library with runtime performance on-par with JAX. DeepWit provides core deep learning modules with their minimal conceptual scope to both express their logic clearly and increase reusablility.
 
+## Installation
+
+DeepWit is published for Scala 3. Add both DeepWit and [DimWit](https://github.com/dimwit-dev/dimwit) —
+DeepWit's API is expressed in DimWit's tensor types, so you will import from both:
+
+```scala
+libraryDependencies ++= Seq(
+  "ch.contrafactus" %% "deepwit-core" % "@VERSION@",
+  "ch.contrafactus" %% "dimwit-core" % "0.1.0"
+)
+```
+
+DimWit runs JAX through ScalaPy, so a Python environment with `jax` and `einops` is also required;
+see [DimWit's setup instructions](https://github.com/dimwit-dev/dimwit#installation).
+
 ## Why DeepWit?
 
 DeepWit strips out the framework machinery and leaves the theory standing: no module base class, no indirect loss, no hidden parameters, no hidden gradients, no hidden optimizer states. 
diff --git a/project/plugins.sbt b/project/plugins.sbt
index d0a3ad4..3dd7a80 100644
--- a/project/plugins.sbt
+++ b/project/plugins.sbt
@@ -2,5 +2,7 @@ addSbtPlugin("org.scalameta" % "sbt-scalafmt" % "2.5.5")
 addSbtPlugin("org.scalameta" % "sbt-mdoc" % "2.8.2")
 addSbtPlugin("ch.epfl.scala" % "sbt-scalafix" % "0.14.7")
 addSbtPlugin("org.scoverage" % "sbt-scoverage" % "2.3.0")
+addSbtPlugin("org.xerial.sbt" % "sbt-sonatype" % "3.11.3")
+addSbtPlugin("com.github.sbt" % "sbt-pgp" % "2.3.1")
 
 libraryDependencies += "ai.kien" %% "python-native-libs" % "0.2.2"