Skip to content

← Flax step 77 of 100

Medium Primitives

Test-Time Augmentation Aggregation

Jump to the task

Why this matters

Test-time augmentation runs a trained model on several versions of an input and combines the predictions. Useful transformations should preserve the answer, or have outputs that can be mapped back to the original coordinates. For example, a segmentation prediction on a flipped image needs to be unflipped before averaging it with the original prediction.

TTA can improve predictions, but it also costs computation and can hurt when transformations change the label or move inputs away from the training data. There is no universal accuracy gain or safe noise scale. Validate the chosen transformations and aggregation rule for the task.

This exercise is a small mechanics example: initialize an untrained linear layer, perturb its inputs with Gaussian noise, and average its outputs. It does not evaluate a trained model or demonstrate an accuracy improvement.

What gets averaged

The layer produces one number per example. For K augmentations and N input rows, the predictions form a (K, N) table. This exercise averages over both axes to return one number for the entire batch.

import jax.numpy as jnp

# Rows are augmentations; columns are input examples.
predictions = jnp.array([[1.0, 3.0], [2.0, 6.0]], dtype=jnp.float32)
per_example = predictions.mean(axis=0)  # [1.5, 4.5]
reported = predictions.mean().reshape(1)  # [3.0]

Ordinary TTA often keeps a prediction for each example, like per_example above. The extra batch mean is this exercise's reporting contract.

For classification, averaging logits, averaging probabilities and majority voting are different rules. Softmax of mean logits need not equal mean softmax probabilities. Voting counts class choices; averaging numeric class IDs is not majority voting. None of these rules is always best or guarantees calibrated probabilities. This exercise averages scalar regression outputs.

Problem

Implement tta_predict(seed, x, num_augs):

  1. Create rng = jax.random.PRNGKey(int(seed)) and set K = int(num_augs).
  2. Split rng once into init_rng, aug_rng, in that order.
  3. Initialize nn.Dense(1) with default initializers using init_rng and x. Extract its params collection for model.apply({'params': params}, ...).
  4. Obtain K keys with jax.random.split(aug_rng, K).
  5. For each key, form x_aug = x + 0.01 * jax.random.normal(key, x.shape). Apply the same initialized layer, then average its (N, 1) outputs.
  6. Average those K scalar predictions and return the result reshaped to (1,).

x is a finite float32 array of shape (N, D), with positive N and D. The seed is a nonnegative integer-valued scalar; num_augs is a positive integer-valued scalar. Either may arrive as a float. Return a float32 array of shape (1,). Do not include an extra unperturbed prediction.

Use the specified split order: changing the key tree changes the seeded outputs. Keep the default Threefry PRNG with partitionable sampling enabled, as used by the fixtures. More augmentations need not preserve the exact previous set of noise samples when the split count changes.

Common pitfalls

  • Reusing the same key with the same normal-sampling call repeats the noise. Split separate keys for initialization and the individual augmentations.
  • A split key is not simply the original key. Avoid reusing keys, but do not assume the first split key reproduces the parent's random draw.
  • The scale 0.01 is part of this exercise, not a universal recommendation. Its effect depends on input units and feature scales.
  • Flattening a (N, 1) prediction before its overall mean is optional: .mean() already averages all elements. Be clear whether you intend a per-example prediction or the batch-wide scalar required here.
  • Preserve the final mean as a JAX value. Python float conversion breaks JIT tracing and differentiation. Mark seed and num_augs static if jitting this function because their Python integer values control setup.
  • For an affine layer, averaging predictions equals applying the layer to the average perturbed input, up to rounding. Zero-mean noise leaves the expected prediction unchanged; a finite sample can still shift it.