Kahibaro
Discord Login Register

11.6 Managing Tokenizers and Batching

What Tokenizers Produce and Why Batching Is Tricky

Transformer models do not consume raw text, they consume integer token IDs. A tokenizer turns a string into a sequence of IDs from a fixed vocabulary and usually also produces helper fields such as an attention mask that marks which positions are real tokens and which positions are padding. The core batching problem is that different texts produce different sequence lengths, but PyTorch tensors in a batch must have a single rectangular shape. The standard solution is padding shorter sequences up to a chosen length and providing an attention mask so the model ignores padded positions.

In most transformer pipelines for beginners, you should think in terms of a batch having shapes like $(B, L)$ for token IDs and masks, where $B$ is batch size and $L$ is the padded sequence length for that batch.

A batch of variable length sequences cannot be stacked into a tensor without padding. If you forget to pad, your DataLoader will fail when it tries to stack examples of different lengths.

Special Tokens, Attention Masks, and Truncation

Tokenizers typically insert special tokens such as a start of sequence and end of sequence marker, or a classification token, depending on the model family. These tokens count toward sequence length, so any maximum length you set must include them. When you pad, you usually pad with a dedicated pad token ID. The attention mask is commonly a tensor of 1s and 0s with the same shape as the input IDs, where 1 means keep and 0 means padding. Many transformer implementations rely on this mask to prevent attention from reading padded positions.

Truncation is equally important. If a text tokenizes to more than your maximum length $L_{\max}$, you either need to truncate or handle longer sequences some other way. In beginner workflows you almost always truncate to a fixed maximum. The model has a maximum context length, and you cannot exceed it.

Always apply the same truncation and padding policy in training and evaluation. Changing it between splits can change what information the model sees and make comparisons misleading.

Padding Strategies: Fixed Length vs Dynamic Padding

There are two common padding strategies. Fixed length padding pads every example to the same global length, like 128 or 256 tokens. This makes shapes stable and sometimes simplifies compilation or static graph assumptions, but it wastes compute on short texts.

Dynamic padding pads to the longest sequence in each batch. This reduces wasted computation and usually improves throughput. With dynamic padding, different batches can have different $L$, which is fine as long as you handle it consistently.

A practical beginner guideline is to start with dynamic padding for efficiency, then consider fixed length padding if you need strictly consistent shapes for deployment or certain performance experiments.

If you use dynamic padding, do not assume a constant sequence length inside your training loop. Always write code that reads shapes from tensors instead of hardcoding $L$.

Using a Tokenizer to Produce Batch Tensors

A typical tokenizer call can return PyTorch tensors directly. Conceptually, you want to request input IDs and attention masks, enable truncation, and choose a padding strategy. When you tokenize a list of strings at once, the tokenizer can handle padding across the list, producing a single batch tensor. This is convenient for inference. In training, tokenization is often done per example in a Dataset and then padded in a collate function, or done in the collate function itself.

When producing tensors, the main fields you will see are input_ids and attention_mask. Some models also require token_type_ids for sentence pair tasks, but many modern architectures do not use them.

Collate Functions: Where Batching Actually Happens

A PyTorch DataLoader builds batches by collecting $B$ dataset items and then collating them into tensors. For text, default collation often fails because sequences are different lengths. A custom collate function solves this by taking a list of examples, tokenizing or gathering already tokenized outputs, and padding them into batch tensors.

One clean beginner pattern is for the Dataset to return raw text and labels, and for the collate function to run the tokenizer on the list of texts. This centralizes padding and truncation behavior and avoids storing variable length tensors in the Dataset.

Another pattern is for the Dataset to tokenize each item and return token IDs as lists, then the collate function pads those lists. This can be faster if tokenization is expensive and you can cache results, but it is more work to implement cleanly.

Do not pad inside __getitem__ to a batch dependent length, because __getitem__ does not know the other items in the batch. Padding belongs in the collate function when using dynamic padding.

Labels and Shapes: Making the Batch Trainable

For classification, labels are usually a tensor of shape $(B,)$ with integer class indices. For regression, labels are often $(B,)$ or $(B, 1)$ depending on how your model outputs are shaped. The collate function should ensure labels are stacked into a tensor with a consistent dtype, commonly torch.long for class indices and torch.float32 for regression targets.

If your model returns logits of shape $(B, C)$ for $C$ classes, then your labels should typically be shape $(B,)$ with values in $\{0, \dots, C-1\}$. If your labels are one hot vectors, you may need a different loss setup, but that belongs in the loss chapter, not here.

A frequent bug is a mismatch between label dtype and loss expectations, for example float labels passed to a loss that expects integer class indices. Check both shape and dtype at the batch level.

Choosing a Maximum Length and Avoiding Wasted Compute

The maximum length you choose is a tradeoff. Larger $L_{\max}$ can capture more context but increases memory and compute roughly proportional to $L$ and, for attention, often proportional to $L^2$ inside the model. For many beginner text classification tasks, 128 or 256 tokens is a reasonable starting point. If your texts are very short, 64 may be sufficient. If texts are long, consider analyzing tokenized length statistics and choosing a length that covers most examples without excessive padding.

A simple approach is to tokenize a sample of your dataset, measure the distribution of lengths, and pick a length that covers, for example, 90 to 95 percent of cases, then accept truncation for the rest.

Batch Size, Memory, and Gradient Stability

Batch size interacts strongly with sequence length. If you double $L$, you may have to reduce batch size to fit in GPU memory. When you see out of memory errors, the first levers are reducing batch size and reducing maximum length. Mixed precision can also help but that is covered elsewhere.

Keep in mind that dynamic padding means your effective memory use depends on the longest sequence in each batch. If you randomly batch examples, a batch that contains one very long example will force all examples in that batch to pad to that long length.

With dynamic padding, one outlier long example can inflate the padded length for the entire batch. Consider filtering extreme lengths or using a batching strategy that groups similar lengths if this becomes a problem.

Length Aware Batching and Sorting

A common efficiency trick is to reduce padding by grouping examples of similar tokenized lengths into the same batch. This can be done by sorting the dataset by length and batching consecutive examples, or by using a sampler that creates length buckets. The idea is to keep the per batch maximum length close to the average length, which reduces wasted padding.

You should be careful to preserve randomness for training. A typical approach is to shuffle within buckets or shuffle bucket order per epoch. Exact implementations vary, but the principle is simple: less length variance inside a batch yields less padding.

Sentence Pair Inputs and Token Type IDs

Some tasks involve two pieces of text, such as premise and hypothesis. Many tokenizers support encoding text pairs and will insert separators and possibly generate token_type_ids that mark which tokens belong to the first or second segment. If your model expects token_type_ids, you must include them in the batch. If it does not, passing them is usually harmless but unnecessary.

The batching logic stays the same: you still pad and create an attention mask, and the optional token type IDs have the same $(B, L)$ shape.

Final Checklist for Tokenizer Based Batches

A good batch for transformer text models typically contains input_ids and attention_mask as integer tensors of shape $(B, L)$ and a labels tensor of shape $(B,) or $(B, 1) depending on the task. Truncation should be enabled with a chosen $L_{\max}$. Padding should be either fixed to $L_{\max}$ or dynamic to the batch maximum, but you should consciously choose one and keep it consistent. When something breaks, print the batch keys, shapes, dtypes, and the min and max token IDs to catch padding and vocabulary issues early.

If your model performs strangely, inspect a decoded version of a few tokenized examples and verify that truncation, special tokens, and text pair ordering match your intended task setup.

Views: 62

Comments

Please login to add a comment.

Don't have an account? Register now!