21 · Content revision 3

Count only the tokens that are really in the batch

Shuffle complete causal windows into mini-batches of fixed-length rows, keep the smaller final batch, and average loss and gradients over its actual target tokens.

Predict the smaller final batch

Chapter 5 already turns each document into shifted windows without crossing a boundary. Use context length $T=2$ on two separate training documents:

train-a = [0, 10, 11, 12, 1]  ->  train-a@0, train-a@1, train-a@2
train-b = [0, 20, 21, 1]      ->  train-b@0, train-b@1

Seed 7 orders those five window identities as:

train-b@1, train-a@1, train-b@0, train-a@0, train-a@2

With requested capacity $3$ , predict the two stacked shapes. The first three windows form $\left[3,2\right]$ ; the remaining two form $\left[2,2\right]$ . The final batch stays smaller. It is neither padded to width $3$ nor dropped.

Now count targets. The first batch has $3\cdot2=6$ target occurrences. The final batch has $2\cdot2=4$ . The requested capacity would permit a third row, but the final batch does not create one. Therefore no third-row token, loss, or gradient exists to include in the mean.

The worked example constructs that epoch from separate documents and preserves the same provenance shown above:

Build the same five-window shuffled epoch reproducibly rust/demos/ch21-mini-batches/src/lib.rs#chapter-fixture

pub fn learner_evidence() -> LearnerEvidence {
    let epoch = build_epoch(SHUFFLE_SEED);
    let replay = build_epoch(SHUFFLE_SEED);
    let changed = build_epoch(SHUFFLE_SEED + 1);
    let batches = epoch
        .batches()
        .iter()
        .enumerate()
        .map(|(index, batch)| batch_evidence(index, batch))
        .collect();

    let ordered_origins = origins(&epoch);
    let mut unique_origins = ordered_origins.clone();
    unique_origins.sort_unstable();
    unique_origins.dedup();
    let expected_origins = [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1)];
    let covered_windows = unique_origins
        .iter()
        .filter(|origin| expected_origins.contains(origin))
        .count();
    let duplicate_windows = ordered_origins.len() - unique_origins.len();
    let padding_ids_added = epoch
        .batches()
        .iter()
        .map(|batch| {
            let expected_stack_len = batch.batch_width() * batch.context_length();
            batch.inputs().len().saturating_sub(expected_stack_len)
                + batch.targets().len().saturating_sub(expected_stack_len)
        })
        .sum();
    let cross_partition_windows = epoch
        .batches()
        .iter()
        .flat_map(MiniBatch::provenance)
        .filter(|origin| origin.partition() != Partition::Train)
        .count();
    LearnerEvidence {
        replay_matches: epoch == replay,
        different_seed_changes_order: origins_in_order(&epoch) != origins_in_order(&changed),
        complete_coverage: unique_origins == expected_origins && duplicate_windows == 0,
        covered_windows,
        expected_windows: expected_origins.len(),
        duplicate_windows,
        padding_ids_added,
        cross_partition_windows,
        epoch,
        batches,
    }
}

fn build_epoch(seed: u64) -> MiniBatchEpoch {
    let documents = [
        BatchDocument::new("train-a", Partition::Train, TRAIN_A).expect("fixture ID is valid"),
        BatchDocument::new("train-b", Partition::Train, TRAIN_B).expect("fixture ID is valid"),
    ];
    let windows = CausalWindowConfig::new(CONTEXT_LENGTH, 1).expect("positive window sizes");
    let batches = MiniBatchConfig::new(REQUESTED_BATCH_SIZE, BatchOrder::Shuffled { seed })
        .expect("positive batch size");
    MiniBatchEpoch::build(Partition::Train, &documents, windows, batches)
        .expect("separate training documents make complete batches")
}

Divide by actual target tokens

For current batch $B$ , average one scalar loss from every target position:

\mathcal{L}_B=\frac{1}{|B|T}\sum_{b\in B}\sum_{t=1}^{T}\mathcal{L}_{b,t}

The parameter gradient has the same denominator:

\nabla_{\theta}\mathcal{L}_B= \frac{1}{|B|T} \sum_{b\in B}\sum_{t=1}^{T} \nabla_{\theta}\mathcal{L}_{b,t}

The assigned losses in the final batch are $\left[0.125,0.25,0.625,0.75\right]$ . Their sum is $1.75$ , so

\mathcal{L}_{B_1}=\frac{1.75}{2\cdot2}=0.4375

Dividing by capacity would use $3\cdot2=6$ and produce about $0.291667$ . That is not a harmless reporting change: it multiplies every final-batch gradient by $4/6$ as well.

If contributions arrive in pieces, keep raw sums $S_j$ and counts $N_j$ :

\bar g= \frac{\sum_j S_j}{\sum_j N_j}, \qquad S_j=\sum_{i=1}^{N_j}g_{j,i}.

Here $j$ identifies one accumulated piece, $i$ identifies one target contribution in that piece, $g_{j,i}$ is its parameter-gradient vector, $S_j$ is the raw vector sum, $N_j$ is the actual contribution count, and $\bar g$ is the token-weighted mean after all pieces are merged.

Do not average piece means without weighting by their token counts. A short piece must not receive the same weight as a full piece.

Keep batch and sequence axes distinct

$B$ is the set of complete windows actually present now.
$|B|$ is current width, not maximum capacity.
$T$ is the fixed input and target length of every admitted window.
$b$ identifies one window in $B$ .
$t$ identifies one target position from $1$ through $T$ .
$\mathcal{L}_{b,t}$ is one target occurrence’s negative log-likelihood.
$\mathcal{L}_B$ is the mean after every admitted occurrence is counted once.
$|B|T$ is the actual token denominator; absent rows and padding are excluded because they do not exist in the batch.
$j$ identifies one independently accumulated piece, and $i$ identifies one target contribution within it.
$g_{j,i}$ is target occurrence $i$ ‘s gradient vector, $S_j$ is the piece’s raw vector sum, and $N_j$ is its actual target-occurrence count.
$\bar g$ is the token-weighted mean gradient after every piece is merged.

An overlapping corpus token may appear as a target in more than one causal window. Those are separate training-example occurrences, and each occurrence contributes once. “Count each target once” does not mean deduplicating equal token IDs.

From one word update to token-sized LLM batches

Bengio et al.’s early neural language model describes a stochastic parameter update after each training-corpus word-context example. That online endpoint exposes every example directly but offers no shared update across several examples; full-batch training moves to the opposite endpoint by waiting for the entire training set.

Bengio et al., A Neural Probabilistic Language Model: Bengio et al. define a stochastic update after presenting one training-corpus word and later discuss grouping $K$ examples before communication as a mini-batch.

Their one-example update is:

\theta\leftarrow\theta+ \varepsilon\frac{\partial\log\widehat P(w_t\mid w_{t-1},\ldots,w_{t-n+1})}{\partial\theta}.

The same paper discusses communicating every $K$ language-model examples as a mini-batch. Transformer training later grouped sentence pairs by approximate length and reported about $25{,}000$ source plus $25{,}000$ target tokens per batch, making token volume an explicit batching quantity.

Vaswani et al., Attention Is All You Need: Vaswani et al. batch sentence pairs by approximate length and report about $25{,}000$ source tokens and $25{,}000$ target tokens in each Transformer training batch.

Those translation batches are not evidence for this course’s fixed-window, no-padding policy. They are evidence that Transformer training made token volume an explicit unit of work.

GPT-3 reports batch size directly in tokens, from $0.5$ million to $3.2$ million across its model scales, with a $2{,}048$ -token context. This course’s decoder training preserves a smaller implementation invariant: each admitted target token contributes once, and loss plus gradients use the actual token count.

Brown et al., Language Models are Few-Shot Learners: Brown et al. label GPT-3 batch size in tokens, report $0.5$ million through $3.2$ million tokens across model scales, and use a $2{,}048$ -token context.

This is the road from stochastic neural-language-model examples to modern LLM token batches, not a programming-language history. The exact seed, shuffle, widths, no-padding rule, losses, gradients, and trace remain course choices.

The historical executable contrast groups the same five example identities as online widths $\left[1,1,1,1,1\right]$ , mini-batch widths $\left[3,2\right]$ , or one full width $\left[5\right]$ . Rust is only the medium for running that language-independent contrast:

Contrast online, mini-batch, and full-set update widths rust/demos/ch21-mini-batches/src/lib.rs#historical-update-grouping

/// Contrasts one-example, three-example, and full-set update widths.
pub fn historical_update_widths(example_count: usize) -> [Vec<usize>; 3] {
    let online = group_widths(example_count, NonZeroUsize::MIN);
    let mini_batch = group_widths(
        example_count,
        NonZeroUsize::new(3).expect("three is positive"),
    );
    let full_capacity = NonZeroUsize::new(example_count).unwrap_or(NonZeroUsize::MIN);
    let full_batch = group_widths(example_count, full_capacity);
    [online, mini_batch, full_batch]
}

fn group_widths(example_count: usize, capacity: NonZeroUsize) -> Vec<usize> {
    let mut widths = Vec::new();
    let mut remaining = example_count;
    while remaining > 0 {
        let width = remaining.min(capacity.get());
        widths.push(width);
        remaining -= width;
    }
    widths
}

Shuffle windows, then merge raw sums

The public configuration distinguishes document provenance, requested capacity, and sequential or seeded order:

Keep each document and epoch-order policy explicit rust/crates/llm-from-scratch/src/training/batch.rs#batch-configuration

/// The stable order used for one materialized epoch.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BatchOrder {
    /// Retain document order, then increasing window start within each document.
    Sequential,
    /// Apply a deterministic Fisher-Yates permutation using the supplied seed.
    Shuffled { seed: u64 },
}

/// A positive requested batch width plus its epoch-order policy.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MiniBatchConfig {
    batch_size: usize,
    order: BatchOrder,
}

impl MiniBatchConfig {
    pub const fn new(batch_size: usize, order: BatchOrder) -> Result<Self, BatchError> {
        if batch_size == 0 {
            return Err(BatchError::ZeroBatchSize);
        }
        Ok(Self { batch_size, order })
    }

    pub const fn batch_size(self) -> usize {
        self.batch_size
    }

    pub const fn order(self) -> BatchOrder {
        self.order
    }
}

/// One separately owned document exposed to the batch builder by reference.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BatchDocument<'a> {
    id: &'a str,
    partition: Partition,
    token_ids: &'a [u32],
}

impl<'a> BatchDocument<'a> {
    pub fn new(
        id: &'a str,
        partition: Partition,
        token_ids: &'a [u32],
    ) -> Result<Self, BatchError> {
        if id.is_empty() {
            return Err(BatchError::EmptyDocumentId);
        }
        Ok(Self {
            id,
            partition,
            token_ids,
        })
    }

    /// Borrows the already-validated provenance and token IDs of one encoded document.
    pub fn from_encoded(document: &'a EncodedDocument) -> Self {
        Self {
            id: document.id(),
            partition: document.partition(),
            token_ids: document.token_ids(),
        }
    }

    pub const fn id(self) -> &'a str {
        self.id
    }

    pub const fn partition(self) -> Partition {
        self.partition
    }

    pub const fn token_ids(self) -> &'a [u32] {
        self.token_ids
    }
}

The epoch builder validates one partition, then records each complete window as a WindowDescriptor. It contains only the source document index and start offset, stored in document_index and start. Because the source documents stay borrowed during construction, those two values are enough to find the window’s $T+1$ source tokens later. Fisher-Yates shuffles only these lightweight descriptors. For each final batch, the builder creates input and target buffers for that batch’s actual logical shape, then copies the selected input and target occurrences directly from the borrowed document into their final row-major destinations. No token-owning staging window is created:

Shuffle window descriptors and write final batch rows directly rust/crates/llm-from-scratch/src/training/batch.rs#mini-batch-epoch

/// The immutable origin of one complete causal window.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WindowProvenance {
    partition: Partition,
    document_index: usize,
    document_id: String,
    start: usize,
}

impl WindowProvenance {
    pub const fn partition(&self) -> Partition {
        self.partition
    }

    pub const fn document_index(&self) -> usize {
        self.document_index
    }

    pub fn document_id(&self) -> &str {
        &self.document_id
    }

    pub const fn start(&self) -> usize {
        self.start
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct WindowDescriptor {
    document_index: usize,
    start: usize,
}

/// One row-major `[batch, sequence]` stack with no padding rows or token IDs.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MiniBatch {
    partition: Partition,
    context_length: usize,
    inputs: Vec<u32>,
    targets: Vec<u32>,
    provenance: Vec<WindowProvenance>,
}

impl MiniBatch {
    pub const fn partition(&self) -> Partition {
        self.partition
    }

    pub const fn context_length(&self) -> usize {
        self.context_length
    }

    pub fn batch_width(&self) -> usize {
        self.provenance.len()
    }

    pub fn shape(&self) -> [usize; 2] {
        [self.batch_width(), self.context_length]
    }

    pub fn token_count(&self) -> usize {
        self.targets.len()
    }

    pub fn inputs(&self) -> &[u32] {
        &self.inputs
    }

    pub fn targets(&self) -> &[u32] {
        &self.targets
    }

    pub fn provenance(&self) -> &[WindowProvenance] {
        &self.provenance
    }

    pub fn input_row(&self, row: usize) -> Option<&[u32]> {
        let start = row.checked_mul(self.context_length)?;
        let end = start.checked_add(self.context_length)?;
        self.inputs.get(start..end)
    }

    pub fn target_row(&self, row: usize) -> Option<&[u32]> {
        let start = row.checked_mul(self.context_length)?;
        let end = start.checked_add(self.context_length)?;
        self.targets.get(start..end)
    }

    /// Averages one checked loss and parameter-gradient vector per target token.
    pub fn average_token_contributions(
        &self,
        contributions: &[TokenContribution],
    ) -> Result<TokenMean, BatchError> {
        let expected = self.token_count();
        if contributions.len() != expected {
            return Err(BatchError::ContributionCountMismatch {
                expected,
                actual: contributions.len(),
            });
        }
        let gradient_width = contributions
            .first()
            .map(TokenContribution::gradient_width)
            .ok_or(BatchError::EmptyAccumulator)?;
        let mut accumulator = TokenMeanAccumulator::new(gradient_width)?;
        for contribution in contributions {
            accumulator.add_token(contribution)?;
        }
        accumulator.finish()
    }
}

/// Every mini-batch in one reproducible traversal of complete windows.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MiniBatchEpoch {
    partition: Partition,
    context_length: usize,
    config: MiniBatchConfig,
    window_count: usize,
    shuffle_state_after: Option<u64>,
    batches: Vec<MiniBatch>,
}

impl MiniBatchEpoch {
    pub fn build(
        partition: Partition,
        documents: &[BatchDocument<'_>],
        window_config: CausalWindowConfig,
        config: MiniBatchConfig,
    ) -> Result<Self, BatchError> {
        validate_documents(partition, documents)?;

        let mut window_count = 0_usize;
        for document in documents {
            window_count = window_count
                .checked_add(window_config.window_count(document.token_ids().len()))
                .ok_or(BatchError::WindowCountOverflow)?;
        }

        let mut descriptors = Vec::new();
        descriptors
            .try_reserve_exact(window_count)
            .map_err(|_| BatchError::AllocationFailed {
                elements: window_count,
            })?;
        for (document_index, document) in documents.iter().copied().enumerate() {
            for window in window_config.windows(document.token_ids()) {
                descriptors.push(WindowDescriptor {
                    document_index,
                    start: window.start(),
                });
            }
        }
        debug_assert_eq!(descriptors.len(), window_count);

        let shuffle_state_after = match config.order() {
            BatchOrder::Sequential => None,
            BatchOrder::Shuffled { seed } => {
                let mut rng = SplitMix64::from_seed(seed);
                fisher_yates(&mut descriptors, &mut rng);
                Some(rng.state())
            }
        };

        let batch_count = if window_count == 0 {
            0
        } else {
            (window_count - 1) / config.batch_size() + 1
        };
        let mut batches = Vec::new();
        batches
            .try_reserve_exact(batch_count)
            .map_err(|_| BatchError::AllocationFailed {
                elements: batch_count,
            })?;

        let context_length = window_config.context_length();
        let required_source_tokens = window_config.required_source_tokens();
        let mut descriptors = descriptors.into_iter();
        let mut remaining = window_count;
        while remaining > 0 {
            let width = remaining.min(config.batch_size());
            let token_count = width
                .checked_mul(context_length)
                .ok_or(BatchError::TokenCountOverflow)?;
            let mut inputs = Vec::new();
            let mut targets = Vec::new();
            let mut provenance = Vec::new();
            inputs
                .try_reserve_exact(token_count)
                .map_err(|_| BatchError::AllocationFailed {
                    elements: token_count,
                })?;
            targets
                .try_reserve_exact(token_count)
                .map_err(|_| BatchError::AllocationFailed {
                    elements: token_count,
                })?;
            provenance
                .try_reserve_exact(width)
                .map_err(|_| BatchError::AllocationFailed { elements: width })?;

            for _ in 0..width {
                let descriptor = descriptors
                    .next()
                    .expect("pre-counted descriptor must exist while batching");
                let document = documents[descriptor.document_index];
                let source_end = descriptor.start + required_source_tokens;
                let source = document
                    .token_ids()
                    .get(descriptor.start..source_end)
                    .expect("descriptor must name one complete causal window");

                inputs.extend_from_slice(&source[..context_length]);
                targets.extend_from_slice(&source[1..]);
                provenance.push(WindowProvenance {
                    partition,
                    document_index: descriptor.document_index,
                    document_id: document.id().to_owned(),
                    start: descriptor.start,
                });
            }
            batches.push(MiniBatch {
                partition,
                context_length,
                inputs,
                targets,
                provenance,
            });
            remaining -= width;
        }

        Ok(Self {
            partition,
            context_length,
            config,
            window_count,
            shuffle_state_after,
            batches,
        })
    }

    pub const fn partition(&self) -> Partition {
        self.partition
    }

    pub const fn context_length(&self) -> usize {
        self.context_length
    }

    pub const fn config(&self) -> MiniBatchConfig {
        self.config
    }

    pub const fn window_count(&self) -> usize {
        self.window_count
    }

    pub fn batch_count(&self) -> usize {
        self.batches.len()
    }

    pub const fn shuffle_state_after(&self) -> Option<u64> {
        self.shuffle_state_after
    }

    pub fn batches(&self) -> &[MiniBatch] {
        &self.batches
    }
}

fn validate_documents(
    partition: Partition,
    documents: &[BatchDocument<'_>],
) -> Result<(), BatchError> {
    for (document_index, document) in documents.iter().enumerate() {
        if document.partition() != partition {
            return Err(BatchError::PartitionMismatch {
                document_index,
                expected: partition,
                actual: document.partition(),
            });
        }
        if let Some(first) = documents[..document_index]
            .iter()
            .position(|candidate| candidate.id() == document.id())
        {
            return Err(BatchError::DuplicateDocumentId {
                id: document.id().to_owned(),
                first,
                repeated: document_index,
            });
        }
    }
    Ok(())
}

fn fisher_yates<T>(values: &mut [T], rng: &mut SplitMix64) {
    for upper_index in (1..values.len()).rev() {
        let selected = sample_below(rng, upper_index + 1);
        values.swap(upper_index, selected);
    }
}

fn sample_below(rng: &mut SplitMix64, exclusive_upper: usize) -> usize {
    debug_assert!(exclusive_upper > 0);
    let bound = exclusive_upper as u64;
    let rejection_threshold = bound.wrapping_neg() % bound;
    loop {
        let draw = rng.next_u64();
        if draw >= rejection_threshold {
            return (draw % bound) as usize;
        }
    }
}

Loss and gradient accumulation use an all-or-nothing preflight. add_token and merge first check the prospective token count, loss sum, and every gradient coordinate. Only after every check succeeds do they update the existing gradient-sum vector in place, retaining its allocation. A failure at a later coordinate therefore cannot leave an earlier coordinate changed:

Preflight raw totals before updating gradient storage in place rust/crates/llm-from-scratch/src/training/batch.rs#token-gradient-averaging

/// One target token's scalar loss and parameter-gradient coordinates.
#[derive(Clone, Debug, PartialEq)]
pub struct TokenContribution {
    loss: f64,
    gradient: Vec<f64>,
}

impl TokenContribution {
    pub fn new(loss: f64, gradient: Vec<f64>) -> Result<Self, BatchError> {
        if !loss.is_finite() {
            return Err(BatchError::NonFiniteLoss { value: loss });
        }
        if gradient.is_empty() {
            return Err(BatchError::ZeroGradientWidth);
        }
        if let Some((coordinate, &value)) = gradient
            .iter()
            .enumerate()
            .find(|(_, value)| !value.is_finite())
        {
            return Err(BatchError::NonFiniteGradient { coordinate, value });
        }
        Ok(Self { loss, gradient })
    }

    pub const fn loss(&self) -> f64 {
        self.loss
    }

    pub fn gradient(&self) -> &[f64] {
        &self.gradient
    }

    pub fn gradient_width(&self) -> usize {
        self.gradient.len()
    }
}

/// Raw sums that can be merged before one final division by token count.
#[derive(Clone, Debug, PartialEq)]
pub struct TokenMeanAccumulator {
    loss_sum: f64,
    gradient_sums: Vec<f64>,
    token_count: usize,
}

impl TokenMeanAccumulator {
    pub fn new(gradient_width: usize) -> Result<Self, BatchError> {
        if gradient_width == 0 {
            return Err(BatchError::ZeroGradientWidth);
        }
        let mut gradient_sums = Vec::new();
        gradient_sums
            .try_reserve_exact(gradient_width)
            .map_err(|_| BatchError::AllocationFailed {
                elements: gradient_width,
            })?;
        gradient_sums.resize(gradient_width, 0.0);
        Ok(Self {
            loss_sum: 0.0,
            gradient_sums,
            token_count: 0,
        })
    }

    pub const fn token_count(&self) -> usize {
        self.token_count
    }

    pub const fn loss_sum(&self) -> f64 {
        self.loss_sum
    }

    pub fn gradient_sums(&self) -> &[f64] {
        &self.gradient_sums
    }

    pub fn add_token(&mut self, contribution: &TokenContribution) -> Result<(), BatchError> {
        if contribution.gradient_width() != self.gradient_sums.len() {
            return Err(BatchError::GradientWidthMismatch {
                expected: self.gradient_sums.len(),
                actual: contribution.gradient_width(),
            });
        }
        let next_count = self
            .token_count
            .checked_add(1)
            .ok_or(BatchError::TokenCountOverflow)?;
        let next_loss = self.loss_sum + contribution.loss();
        if !next_loss.is_finite() {
            return Err(BatchError::NonFiniteAccumulation {
                quantity: "loss",
                coordinate: None,
                value: next_loss,
            });
        }
        validate_gradient_sum(&self.gradient_sums, contribution.gradient())?;

        self.loss_sum = next_loss;
        for (sum, &value) in self.gradient_sums.iter_mut().zip(contribution.gradient()) {
            *sum += value;
        }
        self.token_count = next_count;
        Ok(())
    }

    /// Merges raw sums without averaging either side first.
    pub fn merge(&mut self, other: &Self) -> Result<(), BatchError> {
        if other.gradient_sums.len() != self.gradient_sums.len() {
            return Err(BatchError::GradientWidthMismatch {
                expected: self.gradient_sums.len(),
                actual: other.gradient_sums.len(),
            });
        }
        let next_count = self
            .token_count
            .checked_add(other.token_count)
            .ok_or(BatchError::TokenCountOverflow)?;
        let next_loss = self.loss_sum + other.loss_sum;
        if !next_loss.is_finite() {
            return Err(BatchError::NonFiniteAccumulation {
                quantity: "loss",
                coordinate: None,
                value: next_loss,
            });
        }
        validate_gradient_sum(&self.gradient_sums, &other.gradient_sums)?;

        self.loss_sum = next_loss;
        for (sum, &value) in self.gradient_sums.iter_mut().zip(&other.gradient_sums) {
            *sum += value;
        }
        self.token_count = next_count;
        Ok(())
    }

    pub fn finish(self) -> Result<TokenMean, BatchError> {
        if self.token_count == 0 {
            return Err(BatchError::EmptyAccumulator);
        }
        let denominator = self.token_count as f64;
        let mean_loss = self.loss_sum / denominator;
        let mut mean_gradient = self.gradient_sums;
        for value in &mut mean_gradient {
            *value /= denominator;
        }
        Ok(TokenMean {
            token_count: self.token_count,
            mean_loss,
            mean_gradient,
        })
    }
}

fn validate_gradient_sum(left: &[f64], right: &[f64]) -> Result<(), BatchError> {
    debug_assert_eq!(left.len(), right.len());
    for (coordinate, (&left, &right)) in left.iter().zip(right).enumerate() {
        let value = left + right;
        if !value.is_finite() {
            return Err(BatchError::NonFiniteAccumulation {
                quantity: "gradient",
                coordinate: Some(coordinate),
                value,
            });
        }
    }
    Ok(())
}

/// One scalar mean loss and one equally normalized parameter-gradient vector.
#[derive(Clone, Debug, PartialEq)]
pub struct TokenMean {
    token_count: usize,
    mean_loss: f64,
    mean_gradient: Vec<f64>,
}

impl TokenMean {
    pub const fn token_count(&self) -> usize {
        self.token_count
    }

    pub const fn mean_loss(&self) -> f64 {
        self.mean_loss
    }

    pub fn mean_gradient(&self) -> &[f64] {
        &self.mean_gradient
    }
}

The worked example assigns one deterministic finite loss and two gradient coordinates to each flattened target position, then proves that merging two raw accumulators equals one-pass averaging. These values are stand-ins chosen to isolate token averaging; this example does not recompute logits or negative log-likelihood:

Average one finite contribution per actual target token rust/demos/ch21-mini-batches/src/lib.rs#token-contributions

fn batch_contributions(batch: &MiniBatch) -> Vec<TokenContribution> {
    batch
        .provenance()
        .iter()
        .flat_map(|origin| {
            (0..batch.context_length()).map(move |token| {
                let numerator = origin.document_index() * 8 + origin.start() * 2 + token + 1;
                let loss = numerator as f64 / 8.0;
                TokenContribution::new(loss, vec![2.0 * loss, 2.0 - loss])
                    .expect("binary-fraction fixture remains finite")
            })
        })
        .collect()
}

fn batch_evidence(index: usize, batch: &MiniBatch) -> BatchEvidence {
    let contributions = batch_contributions(batch);
    let batch_mean = batch
        .average_token_contributions(&contributions)
        .expect("one contribution exists for every target token");
    let direct = accumulate(&contributions);
    let loss_sum = direct.loss_sum();
    let mean = direct.finish().expect("batch has target tokens");
    assert_eq!(mean, batch_mean);

    let split = contributions.len() / 2;
    let mut left = accumulate(&contributions[..split]);
    let right = accumulate(&contributions[split..]);
    left.merge(&right).expect("gradient widths match");
    let accumulated = left.finish().expect("batch has target tokens");

    BatchEvidence {
        index,
        loss_sum,
        losses: contributions.iter().map(TokenContribution::loss).collect(),
        accumulation_matches: accumulated == mean,
        mean,
    }
}

fn accumulate(contributions: &[TokenContribution]) -> TokenMeanAccumulator {
    let mut accumulator = TokenMeanAccumulator::new(2).expect("two fixture coordinates");
    for contribution in contributions {
        accumulator
            .add_token(contribution)
            .expect("finite fixture contribution");
    }
    accumulator
}

The learner executable prints the batch rows, means, historical grouping widths, replay result, coverage, padding count, and partition evidence:

Print the deterministic Chapter 21 learner report rust/demos/ch21-mini-batches/src/main.rs#learner-mini-batch-output

fn main() {
    print!("{}", ch21_mini_batches::learner_report());
}

Run cargo run --quiet --locked -p ch21-mini-batches to follow the worked epoch from the shuffle seed through both batch means. The smaller final batch makes the contrast between requested capacity and actual target count visible in the report.

Trace every shuffled window and denominator

The figure follows five shuffled windows into two batch summaries, then compares the final batch’s requested capacity with its actual denominator:

Emit exact batch provenance and token-normalization evidence rust/demos/ch21-mini-batches/src/diagram_trace.rs#mini-batches-trace

pub fn diagram_trace() -> String {
    let evidence = learner_evidence();
    let epoch = &evidence.epoch;
    let mut lines = vec![format!(
        "META|context={}|capacity={REQUESTED_BATCH_SIZE}|seed={SHUFFLE_SEED}|windows={}|batches={}",
        epoch.context_length(),
        epoch.window_count(),
        epoch.batch_count(),
    )];

    let mut slot = 0;
    for (batch, batch_evidence) in epoch.batches().iter().zip(&evidence.batches) {
        for (row, origin) in batch.provenance().iter().enumerate() {
            let loss_start = row * batch.context_length();
            let loss_end = loss_start + batch.context_length();
            lines.push(format!(
                "WINDOW|slot={slot}|batch={}|row={row}|document={}|document_index={}|start={}|input={}|target={}|losses={}",
                batch_evidence.index,
                origin.document_id(),
                origin.document_index(),
                origin.start(),
                format_ids(batch.input_row(row).expect("row exists")),
                format_ids(batch.target_row(row).expect("row exists")),
                format_values(&batch_evidence.losses[loss_start..loss_end]),
            ));
            slot += 1;
        }
        lines.push(format!(
            "BATCH|index={}|width={}|shape=[{}, {}]|tokens={}|loss_sum={:.6}|mean_loss={:.6}|mean_gradient={}|accumulation={}",
            batch_evidence.index,
            batch.batch_width(),
            batch.batch_width(),
            batch.context_length(),
            batch.token_count(),
            batch_evidence.loss_sum,
            batch_evidence.mean.mean_loss(),
            format_values(batch_evidence.mean.mean_gradient()),
            if batch_evidence.accumulation_matches { "equal" } else { "different" },
        ));
    }

    let final_batch = epoch.batches().last().expect("fixture has batches");
    lines.push(format!(
        "FINAL|width={}|tokens={}|capacity_tokens={}|actual_denominator={}",
        final_batch.batch_width(),
        final_batch.token_count(),
        REQUESTED_BATCH_SIZE * epoch.context_length(),
        final_batch.token_count(),
    ));
    lines.push(format!(
        "PROOF|coverage={}/{}|duplicates={}|padding={}|cross_partition={}|replay={}|different_seed={}|accumulation={}",
        evidence.covered_windows,
        evidence.expected_windows,
        evidence.duplicate_windows,
        evidence.padding_ids_added,
        evidence.cross_partition_windows,
        if evidence.replay_matches { "same" } else { "different" },
        if evidence.different_seed_changes_order { "changed" } else { "same" },
        if evidence.batches.iter().all(|batch| batch.accumulation_matches) {
            "equal"
        } else {
            "different"
        },
    ));
    lines.join("\n") + "\n"
}

Follow five complete windows into two token-normalized batches

Read the worked shuffle order, row-major input and target IDs, per-token losses, actual denominators, mean gradients, and coverage proofs for one reproducible epoch.

Context length: $T=2$
Requested capacity: $|B|_{\mathrm{max}}=3$
Shuffle seed: 7
Complete windows: $5$
Batches emitted: $2$

Shuffle complete window identities

The seed permutes complete-window identities, each named by a document and start offset. It never shuffles a flattened stream of document tokens.

#0 train-b@1 $B_{0}$
#1 train-a@1 $B_{0}$
#2 train-b@0 $B_{0}$
#3 train-a@0 $B_{1}$
#4 train-a@2 $B_{1}$

Stack actual rows and count target tokens

Each row contributes exactly two target losses and two gradient vectors. The example supplies every value shown below.

Shuffled windows and their exact token contributions: Batch $B_{0}$ $|B_0|=3$
Window origin	Input IDs	Target IDs	One loss per target position
Window 0 `train-b@1`	`[20, 21]`	`[21, 1]`	$\mathcal{L}_{0,1}=1.375000$ $\mathcal{L}_{0,2}=1.500000$
Window 1 `train-a@1`	`[10, 11]`	`[11, 12]`	$\mathcal{L}_{1,1}=0.375000$ $\mathcal{L}_{1,2}=0.500000$
Window 2 `train-b@0`	`[0, 20]`	`[20, 21]`	$\mathcal{L}_{2,1}=1.125000$ $\mathcal{L}_{2,2}=1.250000$

Shuffled windows and their exact token contributions: Batch $B_{1}$ $|B_1|=2$
Window origin	Input IDs	Target IDs	One loss per target position
Window 3 `train-a@0`	`[0, 10]`	`[10, 11]`	$\mathcal{L}_{3,1}=0.125000$ $\mathcal{L}_{3,2}=0.250000$
Window 4 `train-a@2`	`[11, 12]`	`[12, 1]`	$\mathcal{L}_{4,1}=0.625000$ $\mathcal{L}_{4,2}=0.750000$

Possible third row Not created — contributes nothing

Stacked shape: $\operatorname{shape}(B_{0})=\left[3,2\right]$; $\operatorname{shape}(B_{1})=\left[2,2\right]$
Target tokens: $N_{B_{0}}=6$; $N_{B_{1}}=4$
Loss sum: $B_{0}:\;6.125$; $B_{1}:\;1.75$
Actual denominator: $|B_{0}|T=6$; $|B_{1}|T=4$
Mean loss: $\mathcal{L}_{B_{0}}=\frac{6.125}{6}=1.020833$; $\mathcal{L}_{B_{1}}=\frac{1.75}{4}=0.4375$
Mean gradient: $\bar g_{B_{0}}=\left[2.041667,0.979167\right]$; $\bar g_{B_{1}}=\left[0.875000,1.562500\right]$
Raw accumulation: $B_{0}:\;=$ Equal; $B_{1}:\;=$ Equal

Exclude unused capacity from the final mean

The requested capacity permits three rows, but the final batch creates only two. Its denominator is four target tokens, not six possible positions.

Requested capacity

3\times2=6

Not created — contributes nothing

Actual final width

2\times2=4

Actual denominator

Check coverage, boundaries, replay, and accumulation

Changing the seed may change order, never membership. Merging raw sums before division reproduces one-pass averaging.

Complete windows: $5/5$
Duplicate windows: $0$
Padding IDs: $0$
Cross-partition windows: $0$
Same-seed replay: Same
Different-seed order: Changed
Raw accumulation: Equal

Read the shuffled origins in order, then inspect the input IDs, target IDs, and two loss contributions in every admitted row. The first batch contains three rows and divides by six targets. The second contains two rows and divides by four; a dashed hypothetical third row shows what the requested capacity would permit, but that row is not stored and contributes nothing. The final proof records complete coverage, no duplicates, no padding, no partition crossing, reproducible same-seed order, changed order for the other seed, and equivalence with merging raw accumulators before division.

Predict before checking the exact epoch

List every complete window origin before shuffling.
With requested capacity $3$ , predict both batch widths and shapes.
Count the first batch’s actual target-token denominator.
Count the final batch’s actual denominator and compare it with capacity.
Predict the scale error caused by dividing the final sum by $6$ instead of $4$ .
Explain why raw accumulator merging is valid but an unweighted mean of means is not.
State what same and different seeds may change, and what they must preserve.
Match Bengio, Vaswani, and Brown to one-example updates, Transformer token batches, and large language-model batch sizes in tokens.

Check the predictions

The origins are train-a@0, train-a@1, train-a@2, train-b@0, and train-b@1.
Widths are $3$ and $2$ ; shapes are $\left[3,2\right]$ and $\left[2,2\right]$ .
The first denominator is $3\cdot2=6$ target occurrences.
The final denominator is $2\cdot2=4$ , not capacity $3\cdot2=6$ .
The wrong denominator multiplies the correct final loss and gradient by $4/6$ .
Raw sums retain their token counts; unweighted piece means erase those weights.
The same seed replays order. Another seed may change order, but coverage, shapes by chunk width, boundaries, and one contribution per target remain.
Bengio supplies the online and $K$ -example language-model context, Vaswani the Transformer token batches, and Brown the later LLM batch sizes in tokens.

Hand token-mean gradients to AdamW

The cumulative training path can now turn separate causal windows into reproducible mini-batches of fixed-length rows and produce token-mean loss plus gradient coordinates. Chapter 22 maps those averaged coordinates to stable named parameters and applies AdamW.

Batching changes which examples share an update; it does not yet change a parameter. Chapter 22 gives each stable named parameter first- and second-moment state, applies bias correction, and keeps weight decay separate from the token-mean gradient learned here.