PyTorch
register_parameter
register_buffer
deep learning
machine learning

What is the difference between register_parameter and register_buffer in PyTorch?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

register_parameter and register_buffer both attach tensors to a torch.nn.Module, but they represent different kinds of state. Parameters are model weights or other trainable values that belong in optimization. Buffers are non-parameter tensors that should still move with the module and usually be saved with it. Choosing the right one affects training, serialization, device transfer, and optimizer behavior.

What a Registered Parameter Means

A parameter is part of the module's learnable state. Registered parameters appear in:

  • 'model.parameters()'
  • 'model.named_parameters()'
  • 'state_dict()'

Because optimizers iterate over parameters, registered parameters are what training code usually updates.

python
1import torch
2import torch.nn as nn
3
4
5class ScaleLayer(nn.Module):
6    def __init__(self):
7        super().__init__()
8        self.register_parameter("scale", nn.Parameter(torch.tensor(1.0)))
9
10    def forward(self, x):
11        return x * self.scale
12
13
14layer = ScaleLayer()
15print([name for name, _ in layer.named_parameters()])
16print(layer.state_dict())

This is the right choice when the value is conceptually a weight.

What a Registered Buffer Means

A buffer is state that belongs to the module but is not meant to be optimized. Buffers do not appear in model.parameters(), but they do move with the module when you call .to(device) and they are included in state_dict() by default.

Typical examples are:

  • running statistics in batch normalization
  • fixed masks
  • lookup tables or constants used during forward passes
python
1import torch
2import torch.nn as nn
3
4
5class OffsetLayer(nn.Module):
6    def __init__(self):
7        super().__init__()
8        self.register_buffer("offset", torch.tensor([1.0, 2.0, 3.0]))
9
10    def forward(self, x):
11        return x + self.offset
12
13
14layer = OffsetLayer()
15print(list(layer.named_parameters()))
16print(layer.state_dict())

The buffer participates in model state, but the optimizer ignores it.

Why Not Just Store a Plain Tensor Attribute

You can assign self.offset = torch.tensor(...), but plain attributes do not behave like registered state. In particular:

  • they are not included in state_dict()
  • they do not move automatically with .cuda() or .to(device)
  • they are easier to forget during checkpointing

Registration tells PyTorch that the tensor is part of the module state model.

The Key Behavioral Differences

Use a parameter when:

  • the value should be learned
  • optimizers should see it
  • gradients should flow into it

Use a buffer when:

  • the tensor is part of the model state
  • it should follow device transfers
  • it should not be optimized

That distinction is more important than whether the tensor currently happens to require gradients.

requires_grad=False Is Not the Same as a Buffer

A common mistake is to create a parameter with requires_grad=False and assume it is equivalent to a buffer. It is not.

A parameter with requires_grad=False is still a parameter:

  • it still appears in model.parameters()
  • optimizers still receive it unless filtered out
  • semantically it still looks like a weight

If the tensor is not supposed to be a trainable weight at all, a buffer is usually the better representation.

Persistent and Non-Persistent Buffers

PyTorch buffers are saved in state_dict() by default, but you can opt out for temporary state:

python
1import torch
2import torch.nn as nn
3
4
5class Example(nn.Module):
6    def __init__(self):
7        super().__init__()
8        self.register_buffer("cache", torch.zeros(3), persistent=False)

This is useful for runtime caches that should move with the module during execution but should not be checkpointed.

A Practical Mental Model

Ask two questions:

  1. Should an optimizer treat this tensor as a model weight?
  2. Should this tensor still move with the module and usually be saved?

If the answer is yes to the first, use a parameter. If the answer is no to the first but yes to the second, use a buffer.

Common Pitfalls

  • Using a parameter for fixed constants just because they are tensors.
  • Storing important tensor state as a plain attribute and then losing it during save or device transfer.
  • Assuming requires_grad=False makes a parameter behave like a buffer.
  • Forgetting that buffers are included in state_dict() unless marked persistent=False.
  • Choosing based on syntax convenience instead of the tensor’s semantic role in the model.

Summary

  • Parameters are learnable module state and are exposed to optimizers.
  • Buffers are non-parameter module state that still move with the model and are usually saved.
  • Plain tensor attributes are not automatically tracked like parameters or buffers.
  • 'requires_grad=False does not make a parameter equivalent to a buffer.'
  • Choose based on whether the tensor is a weight or just state the module must carry around.

Course illustration
Course illustration

All Rights Reserved.