BatchNorm và LayerNorm khi batch nhỏ: thí nghiệm và lỗi train/eval

Trong PyTorch, BatchNorm1d ở chế độ train dùng thống kê của các mẫu trong batch theo từng kênh. Thay những mẫu đi kèm có thể làm đầu ra của một mẫu cố định thay đổi. LayerNorm trên đặc trưng của từng mẫu lấy thống kê từ chính mẫu đó. Với BatchNorm mặc định, chuyển sang eval làm lớp dùng các thống kê đã tích lũy, nên đầu ra của lớp không còn phụ thuộc vào những mẫu đi kèm.

Thí nghiệm dưới đây giữ nguyên đầu vào cần quan sát và chỉ thay hai mẫu còn lại. Nó cũng kiểm tra vì sao BatchNorm1d có thể báo lỗi với một mẫu dù đã dùng torch.no_grad(), và vì sao cùng batch size 1 lại chạy được khi đầu vào có thêm trục thời gian.

Nhìn vào trục lấy trung bình và phương sai

Với tensor [N, C], N là số mẫu và C là số đặc trưng. Ta có thể nhìn tensor như một bảng, mỗi dòng là một mẫu và mỗi cột là một kênh.

Cả hai lớp lấy giá trị trừ đi trung bình rồi chia cho căn của phương sai cộng eps. Điểm khác nhau là tập phần tử dùng để tính thống kê. Ví dụ tắt hệ số nhân và độ dịch có thể học để thấy riêng tác động của bước chuẩn hóa. Đọc thêm tổng quan các kỹ thuật chuẩn hóa trong deep learning. Tài liệu BatchNorm1d · Tài liệu LayerNorm

Với [N, C, L], BatchNorm1d lấy thống kê trên hai trục N và L cho mỗi kênh. LayerNorm lấy thống kê trên các chiều cuối khớp với normalized_shape. Ví dụ LayerNorm(C) trên [N, L, C] chuẩn hóa qua C đặc trưng cho từng vị trí của từng mẫu. Nếu đang lưu tensor [N, C, L] mà muốn phép tính đó, hãy đổi chỗ hai trục cuối trước khi áp dụng lớp. Tài liệu BatchNorm1d · Tài liệu LayerNorm

Giữ mẫu [1, 2, 4] và thay các mẫu đi kèm

Hai batch đều có shape [3, 3]. Batch A ghép mẫu mục tiêu [1, 2, 4] với [0, 0, 0] và [2, 4, 8]. Batch B giữ nguyên mẫu mục tiêu, đổi hai mẫu còn lại thành [3, 6, 12] và [5, 10, 20]. Chương trình in đầu ra của dòng đầu tiên, làm tròn đến sáu chữ số thập phân.

Mẫu đầu trùng với trung bình từng cột của A, nên BatchNorm train trả về ba số 0. Trong B, mẫu đầu nhỏ hơn cả ba trung bình cột. LayerNorm luôn nhìn vào cùng ba giá trị trên dòng đầu nên đầu ra giữ nguyên. Kết quả này mô tả riêng lớp chuẩn hóa trên đầu vào cố định. Khi kiểm tra cả mạng, các lớp khác vẫn có thể làm đầu vào của lớp đó thay đổi.

BatchNorm lưu gì sau mỗi lượt train?

Thí nghiệm tạo một lớp BatchNorm với momentum=0.1 rồi đưa A và B qua lớp theo thứ tự đó. Sau hai lượt, running_mean = [0.39, 0.78, 1.56], running_var = [1.3, 2.77, 8.65] và num_batches_tracked = 2. Với momentum=0.1, thống kê mới lưu bằng 90% giá trị cũ cộng 10% giá trị tính từ batch hiện tại. Chẳng hạn cột đầu có trung bình 1 ở A và 3 ở B, nên giá trị lưu đi từ 0 qua 0,1 rồi đến 0,39. Tài liệu BatchNorm1d

PyTorch dùng phương sai chia cho số phần tử để chuẩn hóa đầu ra lúc train, nhưng dùng phương sai có hiệu chỉnh Bessel khi cập nhật running_var. Chẳng hạn cột đầu của A có các giá trị [1, 0, 2]. Phương sai dùng cho đầu ra là 2/3, còn phương sai đưa vào phép cập nhật bộ đệm là 1. Vì vậy, tự tính phương sai từ một batch rồi so thẳng với running_var có thể khiến bạn tưởng lớp tính sai. Tài liệu BatchNorm1d

Cả hai lượt train trong mã chạy bên trong torch.no_grad(), nhưng vẫn cập nhật các bộ đệm. no_grad() tắt việc ghi phép toán để tính gradient. Chuyển trạng thái của lớp là việc của eval(). Sau khi gọi bn.eval(), mã xác minh các bộ đệm giữ nguyên và đầu ra mẫu mục tiêu giống nhau trong A, B và khi đưa riêng mẫu đó qua lớp. Tài liệu Module.eval

Cấu hình track_running_stats=False là ngoại lệ cần kiểm tra khi gỡ lỗi. Lớp không tích lũy các bộ đệm mean và variance, nên eval vẫn dùng thống kê của batch hiện tại. Mã chạy thêm trường hợp này và thu lại hai đầu ra khác nhau như ở BatchNorm train trong bảng. Với mạng dùng cả Dropout và BatchNorm, xem cách bật Monte Carlo Dropout mà vẫn giữ BatchNorm ở eval. Tài liệu BatchNorm1d

Batch size 1 có luôn gây lỗi cho BatchNorm?

Trên PyTorch 2.8.0, BatchNorm1d(3) ở train nhận đầu vào [1, 3] sẽ báo lỗi vì mỗi kênh chỉ có một giá trị. LayerNorm(3) nhận cùng tensor và vẫn dùng ba đặc trưng trên dòng để tính thống kê. Trong cùng lần chạy, BatchNorm1d(3) ở train chấp nhận đầu vào [1, 3, 2] vì hai vị trí trên trục L cung cấp hai giá trị cho mỗi kênh. Tài liệu BatchNorm1d

Nếu lỗi chỉ xuất hiện ở batch cuối, hãy xem shape tại lớp báo lỗi và số mẫu còn lại. Trong suy luận, kiểm tra training và track_running_stats của lớp trước khi kết luận do kích thước batch. Trong huấn luyện có tích lũy gradient, nhiều lượt forward nhỏ có thể góp vào một lần cập nhật trọng số, nhưng mỗi lượt BatchNorm vẫn tính thống kê trên đầu vào của chính lượt đó. Cách phân biệt epoch, batch size và số lượt cập nhật

Batch nhỏ làm thống kê dao động bao nhiêu?

Lấy 4.096 batch độc lập từ phân phối chuẩn có trung bình 0, phương sai 1 cho mỗi cỡ batch. Tính trung bình từng batch rồi đo độ lệch chuẩn của 4.096 giá trị trung bình. Trong mô hình lấy mẫu độc lập này, độ lệch chuẩn lý thuyết của trung bình mẫu là $1 / \sqrt{N}$.

Khi tăng từ 2 lên 32 giá trị, độ lệch chuẩn đo được giảm khoảng bốn lần. Phép đo cho biết trung bình batch nhạy thế nào với số mẫu độc lập trong tình huống này. Nó chưa đo độ chính xác, tốc độ hội tụ hay chất lượng mô hình. Các vị trí trên trục không gian hoặc thời gian của dữ liệu thực có thể tương quan, nên số phần tử BatchNorm gộp lại cũng chưa nói hết độ ổn định của thống kê.

Mã tái lập

Lưu đoạn dưới thành experiment.py và chạy python3 experiment.py. Lần chạy ghi trong bài dùng Python 3.10.11 và PyTorch 2.8.0 trên CPU, kiểu float64 và seed 2026 cho phép lấy mẫu. Phần đầu dùng tensor cố định. Các kiểm tra đối chiếu đầu ra, bộ đệm và trường hợp một mẫu để phát hiện sai khác khi chạy lại.

 1import math
 2import platform
 3
 4import torch
 5from torch import nn
 6
 7
 8def numbers(tensor):
 9    return [round(value, 6) for value in tensor.flatten().tolist()]
10
11
12def main():
13    torch.set_num_threads(1)
14    dtype = torch.float64
15    batch_a = torch.tensor([[1, 2, 4], [0, 0, 0], [2, 4, 8]], dtype=dtype)
16    batch_b = torch.tensor([[1, 2, 4], [3, 6, 12], [5, 10, 20]], dtype=dtype)
17    bn = nn.BatchNorm1d(3, affine=False, momentum=0.1, dtype=dtype)
18    ln = nn.LayerNorm(3, elementwise_affine=False, dtype=dtype)
19
20    print(f"Python {platform.python_version()}, PyTorch {torch.__version__}, CPU float64")
21    with torch.no_grad():
22        train_a = bn(batch_a)[0]
23        train_b = bn(batch_b)[0]
24        layer_a = ln(batch_a)[0]
25        layer_b = ln(batch_b)[0]
26        print("BN train A, B:", numbers(train_a), numbers(train_b))
27        print("LN train A, B:", numbers(layer_a), numbers(layer_b))
28        print("Running mean:", numbers(bn.running_mean))
29        print("Running var:", numbers(bn.running_var))
30        print("Batches tracked after no_grad train:", bn.num_batches_tracked.item())
31        assert not torch.allclose(train_a, train_b)
32        torch.testing.assert_close(layer_a, layer_b, rtol=0, atol=1e-12)
33        torch.testing.assert_close(bn.running_mean, torch.tensor([0.39, 0.78, 1.56], dtype=dtype))
34        torch.testing.assert_close(bn.running_var, torch.tensor([1.3, 2.77, 8.65], dtype=dtype))
35        assert bn.num_batches_tracked.item() == 2
36
37        before = (bn.running_mean.clone(), bn.running_var.clone(), bn.num_batches_tracked.clone())
38        bn.eval()
39        ln.eval()
40        eval_a = bn(batch_a)[0]
41        eval_b = bn(batch_b)[0]
42        eval_single = bn(batch_a[:1])[0]
43        print("BN eval A, B, singleton:", numbers(eval_a), numbers(eval_b), numbers(eval_single))
44        after = (bn.running_mean, bn.running_var, bn.num_batches_tracked)
45        print("BN buffers unchanged in eval:", all(torch.equal(x, y) for x, y in zip(before, after)))
46        assert all(torch.equal(x, y) for x, y in zip(before, after))
47        for output in (eval_b, eval_single):
48            torch.testing.assert_close(eval_a, output, rtol=0, atol=1e-12)
49        torch.testing.assert_close(ln(batch_a[:1])[0], layer_a, rtol=0, atol=1e-12)
50
51        bn_no_stats = nn.BatchNorm1d(3, affine=False, track_running_stats=False, dtype=dtype).eval()
52        no_stats_a, no_stats_b = bn_no_stats(batch_a)[0], bn_no_stats(batch_b)[0]
53        print("BN eval without running stats A, B:", numbers(no_stats_a), numbers(no_stats_b))
54        torch.testing.assert_close(no_stats_a, train_a, rtol=0, atol=1e-12)
55        torch.testing.assert_close(no_stats_b, train_b, rtol=0, atol=1e-12)
56
57        bn.train()
58        ln.train()
59        try:
60            bn(batch_a[:1])
61        except ValueError as error:
62            print("BN train singleton:", str(error))
63        else:
64            raise AssertionError("Expected BatchNorm1d to reject shape [1,3] in train mode")
65        print("LN train singleton:", numbers(ln(batch_a[:1])))
66        sequence = torch.arange(1, 7, dtype=dtype).reshape(1, 3, 2)
67        sequence_bn = nn.BatchNorm1d(3, affine=False, dtype=dtype)
68        print("BN train shape [1,3,2] succeeds:", list(sequence_bn(sequence).shape))
69
70    generator = torch.Generator().manual_seed(2026)
71    for batch_size in (2, 32):
72        samples = torch.randn((4096, batch_size), generator=generator, dtype=dtype)
73        measured = samples.mean(dim=1).std(unbiased=True).item()
74        print(f"N={batch_size}: sd(batch mean)={measured:.6f}, theory={1 / math.sqrt(batch_size):.6f}")
75
76
77if __name__ == "__main__":
78    main()