Introduction

In federated learning, "tuning non-IIDness" usually means controlling how different each client's local data distribution is from the global distribution. In TensorFlow Federated experiments, the usual way to do this is not through a single built-in slider. It is through how you partition the dataset across simulated clients.

What Non-IIDness Means In Practice

A federated dataset becomes more non-IID when different clients see systematically different samples. That can happen in several ways:

• label skew, where each client sees only a few classes
• quantity skew, where some clients have much more data than others
• feature skew, where input distributions differ by client
• concept skew, where the label meaning or decision boundary differs across clients

For most simulation work, label skew is the easiest non-IID pattern to control deliberately.

The Most Common Knob: Dirichlet Partitioning

A standard way to tune label skew is to sample each class distribution across clients from a Dirichlet distribution.

The key parameter is usually called alpha:

• small alpha means more non-IID, because each class concentrates on fewer clients
• large alpha means more IID-like, because class data spreads more evenly

A small standalone example makes the idea concrete.

1import numpy as np
2
3
4def dirichlet_partition(labels, num_clients, alpha, seed=0):
5    rng = np.random.default_rng(seed)
6    labels = np.array(labels)
7    client_indices = [[] for _ in range(num_clients)]
8
9    for label in np.unique(labels):
10        label_idx = np.where(labels == label)[0]
11        rng.shuffle(label_idx)
12
13        proportions = rng.dirichlet([alpha] * num_clients)
14        cuts = (np.cumsum(proportions) * len(label_idx)).astype(int)[:-1]
15        splits = np.split(label_idx, cuts)
16
17        for client_id, split in enumerate(splits):
18            client_indices[client_id].extend(split.tolist())
19
20    return client_indices
21
22
23labels = [0] * 100 + [1] * 100 + [2] * 100
24for alpha in [0.1, 1.0, 10.0]:
25    parts = dirichlet_partition(labels, num_clients=5, alpha=alpha)
26    counts = [len(p) for p in parts]
27    print(alpha, counts)

This does not use TFF directly, but it shows the partitioning logic that is usually applied before building federated client datasets.

Turning Partitions Into TFF Client Data

Once you have per-client indices, you can slice the original dataset and build one dataset per client.

1import tensorflow as tf
2
3
4def make_client_dataset(features, labels, indices, batch_size=16):
5    x = tf.gather(features, indices)
6    y = tf.gather(labels, indices)
7    return tf.data.Dataset.from_tensor_slices((x, y)).batch(batch_size)

The important point is that non-IIDness is expressed through client membership, not through some special TFF optimizer flag.

Other Useful Non-IID Controls

Number Of Classes Per Client

A simple alternative is to give each client only k labels.

For example, in a 10-class task:

• 'k = 10 is close to IID'
• 'k = 2 is strongly label-skewed'
• 'k = 1 is extremely non-IID'

This is a common benchmark style because it is easy to explain and reproduce.

Client Sample Counts

You can also vary how much data each client has. Even with the same label mix, unequal client sizes change training dynamics and aggregation behavior.

Preprocessing Differences

Some experiments introduce client-specific noise, resolution changes, or feature transforms to simulate device heterogeneity. That is another form of non-IIDness, but it should be introduced deliberately because it changes the interpretation of the results.

What To Tune Alongside The Data Split

As non-IIDness increases, model training usually becomes less stable. So the partitioning strategy should be tuned together with training hyperparameters such as:

• local epochs
• client learning rate
• server learning rate
• client sampling rate
• regularization or proximal terms

If you increase skew without adjusting training, you may confuse a training-instability problem with a data-partitioning problem.

Common Pitfalls

A common mistake is looking for a TensorFlow Federated API argument called something like non_iid_level. In practice, the non-IIDness comes from how you construct client datasets.

Another issue is changing several skew factors at once and then losing track of what caused the result. Tune one dimension at a time, such as Dirichlet alpha or labels-per-client.

Developers also sometimes compare experiments with different client counts, data volumes, and skew settings all at once. That makes the results hard to interpret.

Finally, do not assume stronger non-IIDness is always more realistic. Real-world heterogeneity depends on the application, so the partitioning scheme should match the phenomenon you are trying to study.

Summary

• In TFF experiments, non-IIDness is usually tuned by how data is partitioned across clients.
• Dirichlet partitioning is a common way to control label skew with one parameter, alpha.
• Smaller alpha values produce more non-IID client datasets.
• You can also tune non-IIDness with labels-per-client, unequal client sizes, or client-specific preprocessing.
• Interpret partitioning changes together with training hyperparameters, not in isolation.