Transformer training often starts with a simple batching rule: shuffle the examples, take the next B sequences, and pad every sequence in the batch to the length of the longest one. The rule is correct, but it can waste substantial computation when sequence lengths vary widely.

A batch containing a 900-token document and several 100-token documents must usually represent every sequence with 900 token positions. Attention masks prevent padding from acting like real input, but they do not necessarily make the padded positions free to process.

Length bucketing addresses this problem by forming batches from examples with roughly similar lengths. This article develops the mental model behind bucketing, shows how to measure padding waste, and explains the trade-off between efficient batches and sufficiently random training data.

Padding cost comes from the longest sequence in a batch

Consider four tokenized examples with these lengths:

[100, 110, 120, 900]

If they form one batch, padding to the longest sequence gives a rectangular batch with:

4 * 900 = 3,600 token positions

Only this many positions contain real tokens:

100 + 110 + 120 + 900 = 1,230 real tokens

The remaining 2,370 positions are padding. A useful batch-level metric is padding fraction:

padding_fraction = 1 - real_tokens / padded_token_positions

For this batch:

1 - 1,230 / 3,600 ≈ 0.658

About 65.8% of its token positions are padding.

An attention mask is still required so padded positions do not participate as ordinary content. But masking and computational efficiency are separate concerns. Depending on the model and implementation, operations can still be performed over rectangular tensors that include padded positions. Reducing padding therefore reduces the amount of unnecessary tensor work presented to those operations.

The core idea: batch examples of similar length

Suppose a dataset contains eight examples:

[90, 100, 110, 120, 850, 880, 900, 920]

With a batch size of four, an unlucky random shuffle could produce:

batch A: [90, 110, 850, 920]
batch B: [100, 120, 880, 900]

Those batches require:

batch A: 4 * 920 = 3,680 positions
batch B: 4 * 900 = 3,600 positions

Now group similar lengths instead:

batch A: [90, 100, 110, 120]
batch B: [850, 880, 900, 920]

They require:

batch A: 4 * 120 = 480 positions
batch B: 4 * 920 = 3,680 positions

The examples have not changed. The model has not changed. Only batch composition changed, yet the two batches now contain 4,160 padded positions instead of 7,280.

That is the reusable mental model: padding waste depends on length variation inside each batch, not only on the dataset’s overall length distribution.

Do not confuse bucketing with masking

Length bucketing and padding masks solve different problems.

A padding mask tells the model which positions are artificial padding and should not be treated as normal sequence content. Bucketing tries to create less padding in the first place.

You generally still need correct masking after introducing bucketing. A bucket contains similar lengths, not necessarily identical lengths, so most batches still require some padding.

Bucketing also does not change a model’s maximum supported sequence length. An input that is too long must still be truncated, chunked, rejected, or handled by a model and inference path that supports it.

Preserve randomness without returning to fully random batches

The simplest implementation would sort the entire training dataset by sequence length and then take consecutive batches. That minimizes local length variation, but it creates a new problem: training order becomes strongly correlated with sequence length.

If length correlates with other properties of the data, consecutive updates may become less representative of the dataset. For example, short customer messages might mostly be simple requests while long messages contain complicated troubleshooting conversations. A strict shortest-to-longest pass would then organize training by more than length.

A practical approach is to randomize at more than one level:

  1. Shuffle examples or sample a randomized working set.
  2. Group that set into a moderately sized pool.
  3. Sort or partition examples by length inside the pool.
  4. Form batches from nearby lengths.
  5. Shuffle the resulting batches before processing them.

This is often called sortish sampling or local bucketing. The exact algorithm is an implementation choice; the important property is that it reduces within-batch length spread without making the entire epoch globally length-sorted.

A simplified teaching version looks like this:

shuffle(examples)

for pool in chunks(examples, pool_size):
    sort pool by token_length
    batches = chunks(pool, batch_size)
    shuffle(batches)
    yield each batch

Production data loaders may implement bucketing differently, especially for streaming datasets or distributed training. The principle remains the same.

Measure tokenized length, not a convenient proxy

The cost relevant to a Transformer is normally related to the sequence presented to the model. Character count and word count can correlate with token count, but neither is guaranteed to match it.

If tokenization happens before batching, use the post-tokenization sequence length when practical. Include any special tokens that the actual model input adds if they affect the padded length.

This matters near bucket boundaries. Two strings with similar character counts can tokenize to noticeably different lengths, so a character-based bucket can still produce uneven tensor shapes.

For preprocessing pipelines where exact token lengths are expensive to obtain early, an approximate proxy may still be worthwhile. Treat it as an optimization heuristic and measure the resulting padding fraction rather than assuming it works well.

Choose buckets from the observed length distribution

Fixed-width buckets such as 1-128, 129-256, and 257-384 tokens are easy to understand, but they are not automatically efficient. Their usefulness depends on the dataset.

Imagine that almost every example is between 40 and 80 tokens, with a small tail around 1,000 tokens. A single 1-128 bucket may already be adequate for most data. In another dataset, examples may be spread continuously from 50 to 2,000 tokens, making finer or adaptive grouping more useful.

Start by measuring a few quantities from actual batches:

  • real token count;
  • padded token-position count;
  • padding fraction;
  • maximum sequence length per batch;
  • examples or real tokens processed per second;
  • accelerator memory use, if relevant.

Padding fraction is especially useful because it separates useful token positions from positions introduced only to make a rectangular batch.

Do not optimize padding fraction in isolation, however. A complicated sampler that saves a small amount of padding but starves workers, increases synchronization overhead, or makes the pipeline difficult to reproduce may reduce end-to-end throughput.

Batch by a token budget when lengths vary dramatically

A fixed number of examples per batch can still produce large swings in memory and compute. A batch of 32 sequences with 64 tokens each is very different from a batch of 32 sequences with 2,048 tokens each.

For highly variable data, a token budget can be more useful than a fixed example count. The batching rule can add examples while an estimate of padded positions stays under a target:

prospective_cost = prospective_batch_size * prospective_max_length

If adding the next sequence would exceed the budget, close the current batch and start another.

For example, with a budget of 4,096 padded token positions:

16 examples * 256 tokens = 4,096
 8 examples * 512 tokens = 4,096
 4 examples * 1,024 tokens = 4,096

This does not imply those batches have identical runtime or memory use. Transformer costs depend on architecture, kernels, hardware, precision, and other details; self-attention in a standard dense attention layer also has sequence-length-dependent work that is not captured by token count alone. The budget is a practical batching heuristic, not a performance guarantee.

Variable example counts also affect how training statistics should be interpreted. If the loss is averaged per token or per example, verify that gradient accumulation and logging produce the weighting you intend. A nominal “batch” is no longer a constant amount of training data.

Understand the interaction with attention cost

Padding can be particularly expensive for dense self-attention because attention operates over pairs of sequence positions. For a sequence length L, the attention score matrix has L * L entries per attention head.

This does not mean total Transformer runtime is simply proportional to L²; feed-forward layers, projections, memory traffic, optimized kernels, and hardware utilization also matter. It does mean that increasing the padded sequence dimension can increase attention work sharply in implementations that operate on the padded dense shape.

That is why mixing one very long sequence with many short ones can be more costly than the raw count of padding tokens suggests.

Some specialized kernels and packed-sequence implementations can avoid part of this waste. When such an implementation truly operates on variable-length sequences rather than a padded dense rectangle, the benefit of conventional length bucketing may be smaller. Verify the behavior of the actual training stack rather than assuming that an attention mask alone skips padded computation.

Distributed training needs balanced work too

In data-parallel training, workers commonly synchronize after a training step. If one worker receives a much more expensive batch than the others, faster workers can spend time waiting at synchronization points.

Length-aware batching can therefore help with more than single-device padding. But naive bucketing can also create imbalance if one worker receives mostly long batches while another receives short ones.

A distributed sampler should preserve the intended data partitioning and try to keep per-step work reasonably comparable across workers. The exact solution depends on the framework and whether batches use fixed example counts, token budgets, gradient accumulation, or sequence packing.

Also ensure that randomization remains reproducible when that matters. Distributed samplers often require epoch-dependent seeds or coordinated shuffling so workers do not accidentally process the same examples.

Common mistakes

Sorting the whole epoch and never shuffling batches. This maximizes length ordering and can make training order correlate with properties associated with sequence length. Prefer local sorting or shuffled buckets when training relies on stochastic batches.

Using character count as if it were token count. It may be a useful approximation, but tokenization determines the model sequence length. Measure whether the approximation actually reduces padding.

Removing padding masks after adding buckets. Similar-length examples still differ in length. Bucketing reduces padding; it does not identify padding for the model.

Comparing throughput only in examples per second. If batch sizes vary, examples per second can hide how much real text is processed. Real tokens per second and padding fraction provide useful additional context.

Assuming fewer padded tokens guarantees a proportional speedup. Runtime depends on the complete input pipeline, model, kernels, device utilization, and shape-dependent implementation details. Benchmark end-to-end training.

Making buckets too narrow. Extremely strict buckets can complicate sampling, leave partially filled batches, or reduce useful randomness. The goal is not zero padding at any cost; it is better overall training efficiency.

When length bucketing is worth using

Length bucketing is most attractive when sequence lengths vary substantially and the training path uses padded dense batches. It is easy to justify when profiling shows a high padding fraction or when occasional long examples inflate memory use for otherwise short batches.

A simpler random batcher can be preferable when sequences already have similar lengths, padding is a small part of the workload, the dataset is tiny, or the execution stack efficiently supports packed variable-length inputs. Simplicity has value: fewer sampler rules mean fewer opportunities for ordering bugs and easier reproducibility.

The decision should come from measurement. Record the baseline padding fraction and throughput, introduce a modest length-aware batching strategy, then compare end-to-end behavior under the same model and training settings.

Conclusion

Length bucketing is a batching optimization, not a change to the model. It works by controlling which sequence lengths share a rectangular tensor, reducing the padding created by large within-batch length differences.

The practical pattern is straightforward: measure tokenized lengths, group nearby lengths without globally sorting away training randomness, keep padding masks correct, and evaluate real-token throughput rather than assuming less padding automatically means proportional speed gains. When sequence lengths vary widely, this small data-loading decision can make each batch represent substantially more useful input and substantially less padding.

Artificial Intelligence Machine Learning Transformers Model Training