TensorFlow
variable-length sequences
batching
deep learning
machine learning

How to deal with batches with variable-length sequences in TensorFlow?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Variable-length sequences are common in text, audio, event streams, and time-series tasks, but batching them efficiently requires explicit handling. TensorFlow offers several approaches including padding with masks, RaggedTensors, and bucketed batching. The best choice depends on model type, performance goals, and serving constraints.

Why Variable Length Is a Batching Problem

GPU and TPU execution benefits from regular tensor shapes. Raw sequence data usually has uneven lengths, so direct stacking is not possible. You need a representation strategy that keeps semantic correctness while remaining efficient.

Typical options:

  • pad to fixed or batch-local max length
  • keep ragged structure
  • group similar lengths to reduce padding waste

Option 1: Padding with Masks

Padding is the most common approach and works well with many Keras layers.

python
1import tensorflow as tf
2
3sequences = [
4    [1, 5, 8, 2],
5    [3, 7],
6    [4, 9, 6]
7]
8
9padded = tf.keras.preprocessing.sequence.pad_sequences(
10    sequences,
11    padding="post",
12    value=0
13)
14
15print(padded)

Then use a masking layer so padded tokens are ignored.

python
1inputs = tf.keras.Input(shape=(None,), dtype="int32")
2x = tf.keras.layers.Embedding(input_dim=10000, output_dim=64, mask_zero=True)(inputs)
3x = tf.keras.layers.LSTM(64)(x)
4outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
5model = tf.keras.Model(inputs, outputs)

mask_zero=True propagates mask info through compatible layers.

Option 2: RaggedTensors

RaggedTensors represent variable-length sequences without explicit padding.

python
1rt = tf.ragged.constant([
2    [1, 5, 8, 2],
3    [3, 7],
4    [4, 9, 6]
5])
6
7print(rt.shape)

Some TensorFlow and Keras operations support ragged inputs directly. This can reduce wasted compute from padding, but layer support is not universal. Always verify compatibility for your full model path.

Option 3: Bucketed Batching

Bucket by sequence length so each batch has similar lengths and minimal padding overhead.

python
1def element_length(x, y):
2    return tf.shape(x)[0]
3
4
5def make_dataset(seqs, labels, batch_size=32):
6    ds = tf.data.Dataset.from_tensor_slices((seqs, labels))
7
8    return ds.apply(
9        tf.data.experimental.bucket_by_sequence_length(
10            element_length_func=element_length,
11            bucket_boundaries=[20, 50, 100, 200],
12            bucket_batch_sizes=[batch_size] * 5,
13            padded_shapes=([None], [])
14        )
15    )

Bucketed batching is often a strong compromise between simplicity and runtime efficiency.

Build Input Pipelines with tf.data

For scalable training, perform padding and batching in tf.data pipeline rather than ad hoc Python preprocessing.

python
1ds = tf.data.Dataset.from_generator(
2    lambda: iter([([1, 2, 3], 0), ([4, 5], 1), ([6], 0)]),
3    output_signature=(
4        tf.TensorSpec(shape=(None,), dtype=tf.int32),
5        tf.TensorSpec(shape=(), dtype=tf.int32),
6    ),
7)
8
9ds = ds.padded_batch(
10    batch_size=2,
11    padded_shapes=([None], []),
12    padding_values=(0, 0),
13)

This keeps data processing graph-friendly and reduces host bottlenecks.

Training and Loss Considerations

If your target is sequence-level classification, masking often handles padded tokens automatically in recurrent layers. For token-level tasks, you may need loss masking so padded positions do not affect gradients.

Example strategy:

  • create boolean mask where token is not padding value
  • compute per-token loss
  • zero out masked positions
  • normalize by valid token count

This is critical for fair optimization on uneven lengths.

Serving Considerations

At inference, ensure input contract matches training representation. If you trained with padded tensors and mask logic, keep the same preprocessing path in serving. Mismatch between training and serving sequence handling is a common source of silent accuracy drop.

Common Pitfalls

  • Padding sequences but forgetting to apply masks in compatible layers.
  • Assuming every Keras layer supports RaggedTensors.
  • Using global max sequence length and wasting large amounts of compute.
  • Ignoring loss masking in token-level prediction tasks.
  • Training with one preprocessing path and serving with another.

Summary

  • Variable-length batching in TensorFlow requires explicit representation strategy.
  • Padding plus masking is the most widely supported approach.
  • Ragged tensors reduce padding waste but need layer compatibility checks.
  • Bucketed batching improves efficiency for diverse sequence lengths.
  • Keep training and serving preprocessing consistent to preserve model quality.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.