diff --git a/src/speculators/train/noise_transforms.py b/src/speculators/train/noise_transforms.py
index d5df48737..3aef7504f 100644
--- a/src/speculators/train/noise_transforms.py
+++ b/src/speculators/train/noise_transforms.py
@@ -1,25 +1,88 @@
+"""Noise transforms for augmenting tensors (e.g. hidden states) during training."""
+
 import torch
 
+__all__ = [
+    "AddGaussianNoise",
+    "AddUniformNoise",
+    "TransformTensors",
+]
+
 
 class TransformTensors:
-    def __init__(self, std=0.05, tensors=("hidden_states",)):
+    """Base class for noise transforms applied to named tensors of a batch.
+
+    Subclasses must override :meth:`transform` to define the noise operation
+    applied to each of the configured tensors.
+    """
+
+    def __init__(
+        self, std: float = 0.05, tensors: tuple[str, ...] = ("hidden_states",)
+    ):
+        """Initialize the transform.
+
+        Args:
+            std: Scale of the noise applied to the configured tensors.
+            tensors: Keys of the batch entries to transform.
+        """
         self.tensors = tensors
         self.std = std
 
-    def __call__(self, data):
+    def __call__(self, data: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        """Apply the transform to each configured tensor of the batch.
+
+        Args:
+            data: Mapping of batch entries to tensors. Entries listed in
+                ``self.tensors`` are replaced with their transformed values.
+
+        Returns:
+            The batch mapping with the configured tensors transformed.
+        """
         for tensor in self.tensors:
             data[tensor] = self.transform(data[tensor])
         return data
 
     def transform(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Apply noise to a single tensor.
+
+        Args:
+            tensor: The tensor to transform.
+
+        Returns:
+            The transformed tensor.
+
+        Raises:
+            NotImplementedError: If the subclass does not implement this method.
+        """
         raise NotImplementedError("Subclasses must implement this method")
 
 
 class AddGaussianNoise(TransformTensors):
+    """Add zero-mean Gaussian noise to the configured tensors."""
+
     def transform(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Add Gaussian noise with standard deviation ``self.std`` to a tensor.
+
+        Args:
+            tensor: The tensor to transform.
+
+        Returns:
+            A new tensor with element-wise Gaussian noise added.
+        """
         return tensor + torch.randn_like(tensor) * self.std
 
 
 class AddUniformNoise(TransformTensors):
+    """Add uniform noise to the configured tensors."""
+
     def transform(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Add uniform noise bounded by ``self.std`` to a tensor.
+
+        Args:
+            tensor: The tensor to transform.
+
+        Returns:
+            A new tensor with element-wise uniform noise added, where each
+            noise sample lies within ``[-self.std, self.std)``.
+        """
         return tensor + 2 * (torch.rand_like(tensor) - 0.5) * self.std
diff --git a/src/speculators/train/vocab_mapping.py b/src/speculators/train/vocab_mapping.py
index cf79271d0..52ea7197f 100644
--- a/src/speculators/train/vocab_mapping.py
+++ b/src/speculators/train/vocab_mapping.py
@@ -10,6 +10,8 @@
 
 __all__ = [
     "build_vocab_mappings_from_distribution",
+    "combine_token_frequency_distributions",
+    "get_target_vocab_size",
     "save_token_frequency_distribution",
 ]
 
@@ -17,15 +19,19 @@
 def save_token_frequency_distribution(
     dataset: HFDataset,
     output_path: Path | str = "./token_freq.pt",
-):
+) -> None:
     """Save token frequency distribution from the dataset.
 
+    Only tokens where ``loss_mask`` is 1 (assistant tokens) are counted. If
+    ``output_path`` already exists, the dataset is skipped and the existing
+    file is left untouched.
+
     Args:
         dataset: HuggingFace dataset with input_ids and loss_mask
         output_path: Path where to save the token frequency distribution
 
     Returns:
-        Path to the saved frequency distribution file
+        None. The frequency distribution is written to ``output_path``.
     """
     path = Path(output_path)
     if path.exists():
@@ -49,12 +55,22 @@ def save_token_frequency_distribution(
 def combine_token_frequency_distributions(
     token_freq_paths: list[str | Path],
     output_path: str | Path,
-):
-    """Combine multiple token frequency distributions into a single file."""
-    token_freq_dicts = [
+) -> None:
+    """Combine multiple token frequency distributions into a single file.
+
+    Args:
+        token_freq_paths: Paths of token frequency files, as written by
+            :func:`save_token_frequency_distribution`.
+        output_path: Path where to save the combined frequency distribution.
+
+    Returns:
+        None. The combined frequency distribution is written to
+        ``output_path``.
+    """
+    token_freq_dicts: list[dict[int, int]] = [
         torch.load(path, weights_only=True) for path in token_freq_paths
     ]
-    combined_token_freq: Counter[str] = Counter()
+    combined_token_freq: Counter[int] = Counter()
     for token_freq_dict in token_freq_dicts:
         combined_token_freq.update(token_freq_dict)
     combined_token_freq_dict = dict(combined_token_freq)
@@ -97,22 +113,39 @@ def build_vocab_mappings_from_distribution(
 
 
 def get_target_vocab_size(
-    target_vocab_size,
-    target_model_path,
+    target_vocab_size: int | None,
+    target_model_path: str | Path | None,
     trust_remote_code: bool = False,
-):
-    has_vocab = target_vocab_size is not None
-    has_model = target_model_path is not None
+) -> int:
+    """Resolve the vocabulary size of the target (verifier) model.
 
-    if has_vocab and has_model:
-        raise ValueError("Cannot specify both target-vocab-size and target-model-path")
+    Exactly one of ``target_vocab_size`` and ``target_model_path`` must be
+    provided. When a model path is given, the vocabulary size is read from
+    the model config, unwrapping ``text_config`` for multimodal models.
 
-    if not has_vocab and not has_model:
-        raise ValueError("Must specify either target-vocab-size or target-model-path")
+    Args:
+        target_vocab_size: Explicit vocabulary size of the target model.
+        target_model_path: Path or model name of the target model to load
+            the config from.
+        trust_remote_code: Whether to trust remote code when loading the
+            model config.
+
+    Returns:
+        The target model's vocabulary size.
+
+    Raises:
+        ValueError: If both or neither of ``target_vocab_size`` and
+            ``target_model_path`` are provided.
+    """
+    if target_vocab_size is not None and target_model_path is not None:
+        raise ValueError("Cannot specify both target-vocab-size and target-model-path")
 
-    if has_vocab:
+    if target_vocab_size is not None:
         return target_vocab_size
 
+    if target_model_path is None:
+        raise ValueError("Must specify either target-vocab-size or target-model-path")
+
     config = AutoConfig.from_pretrained(
         target_model_path,
         trust_remote_code=trust_remote_code,
diff --git a/tests/unit/train/test_noise_transforms.py b/tests/unit/train/test_noise_transforms.py
new file mode 100644
index 000000000..f86e839f6
--- /dev/null
+++ b/tests/unit/train/test_noise_transforms.py
@@ -0,0 +1,120 @@
+"""Unit tests for the tensor noise transforms used during training."""
+
+import pytest
+import torch
+
+from speculators.train.noise_transforms import (
+    AddGaussianNoise,
+    AddUniformNoise,
+    TransformTensors,
+)
+
+
+def test_base_transform_raises_not_implemented_error():
+    transform = TransformTensors()
+
+    with pytest.raises(NotImplementedError, match="Subclasses must implement"):
+        transform.transform(torch.zeros(2, 2))
+
+
+def test_base_call_raises_not_implemented_error():
+    transform = TransformTensors()
+
+    with pytest.raises(NotImplementedError, match="Subclasses must implement"):
+        transform({"hidden_states": torch.zeros(2, 2)})
+
+
+def test_add_gaussian_noise_transforms_only_configured_keys(seed):
+    data = {
+        "hidden_states": torch.zeros(4, 8),
+        "labels": torch.arange(4),
+    }
+    original_hidden = data["hidden_states"].clone()
+    original_labels = data["labels"].clone()
+
+    result = AddGaussianNoise(std=0.1)(data)
+
+    assert result is data
+    assert not torch.equal(data["hidden_states"], original_hidden)
+    assert torch.equal(data["labels"], original_labels)
+
+
+def test_add_gaussian_noise_supports_custom_tensor_keys(seed):
+    data = {
+        "hidden_states": torch.zeros(2, 2),
+        "positions": torch.zeros(2, 2),
+    }
+
+    AddGaussianNoise(std=0.1, tensors=("positions",))(data)
+
+    assert torch.equal(data["hidden_states"], torch.zeros(2, 2))
+    assert not torch.equal(data["positions"], torch.zeros(2, 2))
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_add_gaussian_noise_preserves_shape_dtype_and_device(seed, dtype):
+    tensor = torch.randn(3, 5, 7, dtype=dtype)
+
+    noisy = AddGaussianNoise(std=0.05).transform(tensor)
+
+    assert noisy.shape == tensor.shape
+    assert noisy.dtype == tensor.dtype
+    assert noisy.device == tensor.device
+
+
+def test_add_gaussian_noise_with_zero_std_is_identity(seed):
+    tensor = torch.randn(6, 6)
+
+    assert torch.equal(AddGaussianNoise(std=0.0).transform(tensor), tensor)
+
+
+def test_add_gaussian_noise_matches_requested_std(seed):
+    std = 0.25
+
+    noise = AddGaussianNoise(std=std).transform(torch.zeros(100, 100))
+
+    assert noise.std().item() == pytest.approx(std, rel=0.05)
+
+
+def test_add_uniform_noise_is_strictly_bounded(seed):
+    std = 0.3
+
+    noise = AddUniformNoise(std=std).transform(torch.zeros(10, 10))
+
+    # torch.rand samples lie in [0, 1), so the noise lies in [-std, std).
+    assert (noise >= -std).all()
+    assert (noise < std).all()
+    # The noise should span most of the bounded interval.
+    assert noise.abs().max() > 0.9 * std
+
+
+def test_add_uniform_noise_transforms_only_configured_keys(seed):
+    data = {
+        "hidden_states": torch.zeros(4, 8),
+        "labels": torch.arange(4),
+    }
+    original_hidden = data["hidden_states"].clone()
+    original_labels = data["labels"].clone()
+
+    result = AddUniformNoise(std=0.1)(data)
+
+    assert result is data
+    assert not torch.equal(data["hidden_states"], original_hidden)
+    assert torch.equal(data["labels"], original_labels)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_add_uniform_noise_preserves_shape_dtype_and_device(seed, dtype):
+    tensor = torch.randn(2, 3, 4, dtype=dtype)
+
+    noisy = AddUniformNoise(std=0.1).transform(tensor)
+
+    assert noisy.shape == tensor.shape
+    assert noisy.dtype == tensor.dtype
+    assert noisy.device == tensor.device
+
+
+def test_add_uniform_noise_with_zero_std_is_identity(seed):
+    tensor = torch.randn(6, 6)
+
+    assert torch.equal(AddUniformNoise(std=0.0).transform(tensor), tensor)
diff --git a/tests/unit/train/test_vocab_mapping.py b/tests/unit/train/test_vocab_mapping.py
new file mode 100644
index 000000000..4788ef1f8
--- /dev/null
+++ b/tests/unit/train/test_vocab_mapping.py
@@ -0,0 +1,140 @@
+"""Unit tests for the vocab mapping utilities used during draft model training."""
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from datasets import Dataset as HFDataset
+
+from speculators.train import vocab_mapping
+from speculators.train.vocab_mapping import (
+    build_vocab_mappings_from_distribution,
+    combine_token_frequency_distributions,
+    get_target_vocab_size,
+    save_token_frequency_distribution,
+)
+
+
+def _frequency_dataset() -> HFDataset:
+    return HFDataset.from_dict(
+        {
+            "input_ids": [[1, 2, 2, 3], [3, 3, 4]],
+            "loss_mask": [[1, 1, 0, 1], [0, 1, 1]],
+        }
+    ).with_format("torch")
+
+
+def test_build_vocab_mappings_ranks_tokens_by_frequency():
+    token_freq = {5: 10, 3: 30, 9: 20, 1: 5, 7: 30}
+
+    # Ranked by (-frequency, token id): 3, 7 (tied at 30), 9, 5, 1; top 3 kept.
+    draft_to_target, target_to_draft = build_vocab_mappings_from_distribution(
+        token_freq, draft_vocab_size=3, target_vocab_size=16
+    )
+
+    assert draft_to_target.dtype == torch.long
+    assert draft_to_target.shape == (3,)
+    draft_idx = torch.arange(3)
+    assert (draft_idx + draft_to_target).tolist() == [3, 7, 9]
+
+    expected_target_to_draft = torch.zeros(16, dtype=torch.bool)
+    expected_target_to_draft[[3, 7, 9]] = True
+    assert torch.equal(target_to_draft, expected_target_to_draft)
+
+
+def test_build_vocab_mappings_pads_with_unused_token_ids():
+    token_freq = {2: 7, 10: 3}
+
+    # Ranked: 2, 10; padded with the smallest unused ids < draft_vocab_size.
+    draft_to_target, target_to_draft = build_vocab_mappings_from_distribution(
+        token_freq, draft_vocab_size=4, target_vocab_size=16
+    )
+
+    draft_idx = torch.arange(4)
+    assert (draft_idx + draft_to_target).tolist() == [0, 1, 2, 10]
+
+    expected_target_to_draft = torch.zeros(16, dtype=torch.bool)
+    expected_target_to_draft[[0, 1, 2, 10]] = True
+    assert torch.equal(target_to_draft, expected_target_to_draft)
+
+
+def test_build_vocab_mappings_padding_skips_already_selected_ids():
+    token_freq = {0: 5, 9: 2}
+
+    # Ranked: 0, 9; padding must skip 0 (already selected) and pick 1 instead.
+    draft_to_target, _ = build_vocab_mappings_from_distribution(
+        token_freq, draft_vocab_size=3, target_vocab_size=16
+    )
+
+    draft_idx = torch.arange(3)
+    assert (draft_idx + draft_to_target).tolist() == [0, 1, 9]
+
+
+def test_combine_token_frequency_distributions_merges_files(tmp_path: Path):
+    torch.save({1: 2, 2: 3}, tmp_path / "freq_a.pt")
+    torch.save({2: 4, 5: 1}, tmp_path / "freq_b.pt")
+    output_path = tmp_path / "combined.pt"
+
+    combine_token_frequency_distributions(
+        [tmp_path / "freq_a.pt", tmp_path / "freq_b.pt"], output_path
+    )
+
+    assert output_path.exists()
+    assert torch.load(output_path, weights_only=True) == {1: 2, 2: 7, 5: 1}
+
+
+def test_save_token_frequency_distribution_counts_masked_tokens(tmp_path: Path):
+    output_path = tmp_path / "nested" / "token_freq.pt"
+
+    save_token_frequency_distribution(_frequency_dataset(), output_path)
+
+    # Only tokens with loss_mask == 1 are counted: [1, 2, 3] and [3, 4].
+    assert output_path.exists()
+    assert torch.load(output_path, weights_only=True) == {1: 1, 2: 1, 3: 2, 4: 1}
+
+
+def test_save_token_frequency_distribution_skips_existing_file(tmp_path: Path):
+    output_path = tmp_path / "token_freq.pt"
+    output_path.write_text("sentinel")
+
+    save_token_frequency_distribution(_frequency_dataset(), output_path)
+
+    assert output_path.read_text() == "sentinel"
+
+
+def test_get_target_vocab_size_accepts_explicit_value():
+    assert get_target_vocab_size(151936, None) == 151936
+
+
+def test_get_target_vocab_size_rejects_both_options():
+    with pytest.raises(ValueError, match="both"):
+        get_target_vocab_size(100, "some/model")
+
+
+def test_get_target_vocab_size_rejects_neither_option():
+    with pytest.raises(ValueError, match="either"):
+        get_target_vocab_size(None, None)
+
+
+def test_get_target_vocab_size_loads_from_model_config(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setattr(
+        vocab_mapping.AutoConfig,
+        "from_pretrained",
+        lambda *_args, **_kwargs: SimpleNamespace(vocab_size=42),
+    )
+
+    assert get_target_vocab_size(None, "model-path") == 42
+
+
+def test_get_target_vocab_size_unwraps_text_config(monkeypatch: pytest.MonkeyPatch):
+    config = SimpleNamespace(vocab_size=0, text_config=SimpleNamespace(vocab_size=7))
+    monkeypatch.setattr(
+        vocab_mapping.AutoConfig,
+        "from_pretrained",
+        lambda *_args, **_kwargs: config,
+    )
+
+    assert get_target_vocab_size(None, "model-path", trust_remote_code=True) == 7