PD-loss-balancing 0.2.1

A package for stabalizing and controllolling loss with feedback control

PD Loss Balancing

A PyTorch library for automatic loss balancing using PID control theory. Dynamically adjust the weighting between multiple loss functions during training to maintain desired target relationships.

Installation

pip install PD-loss-balancing

Quick Start

import torch
from PD_loss_balancing import PDLossWeighter, ConstantTarget

# Create a PD controller to balance two losses
target = ConstantTarget(1.0)  # Keep loss1 close to 1.0
weighter = PDLossWeighter(
    target=target,
    kp=0.01,    # Proportional gain
    kd=0.001,   # Derivative gain
    initial_balance=0.5  # Start with equal weighting
)

# In your training loop
for batch in dataloader:
    # Compute your losses
    loss1 = criterion1(outputs1, targets1)
    loss2 = criterion2(outputs2, targets2)
    
    # Get adaptive balance parameter and info
    alpha, info = weighter.get_balance_param(measured_output=loss1.item())
    
    # Combine losses: alpha * loss1 + (1-alpha) * loss2
    combined_loss = weighter.get_combined_loss(loss1, loss2, alpha)
    
    # Backpropagate
    combined_loss.backward()
    optimizer.step()

Core Components

Target Strategies

Define what value you want your controlled loss to maintain:

ConstantTarget

Maintains a fixed target value. Useful to optimize one loss subject to a constrained other loss/metric:

from PD_loss_balancing import ConstantTarget

target = ConstantTarget(2.5)  # Keep loss at 2.5

RelativeTarget

Maintains a ratio relative to another loss. Useful to optimize two losses equally, or at relative values.

from PD_loss_balancing import RelativeTarget

# Keep loss1 = 0.8 * loss2
target = RelativeTarget(ratio=0.8)

# Usage requires target_reference_values
alpha, info = weighter.get_balance_param(
    measured_output=loss1.item(),
    target_reference_values=loss2.item()
)

LinearTrajectoryTarget

Changes target linearly over time. Useful when you want to allow a loss term to increase or decrease over time.

from PD_loss_balancing import LinearTrajectoryTarget

# Start at 1.0, end at 0.1 over 1000 steps
target = LinearTrajectoryTarget(
    initial=1.0, 
    final=0.1, 
    num_steps=1000
)

Controllers

PDLossWeighter

PD controller with proportional and derivative terms. Applies correction proportional to the error in actual loss value and target loss value, and applies damping with a derivative term.

from PD_loss_balancing import PDLossWeighter, ConstantTarget

weighter = PDLossWeighter(
    target=ConstantTarget(1.0),
    kp=0.01,           # Proportional gain
    kd=0.001,          # Derivative gain
    min_alpha=0.0,     # Minimum balance parameter
    max_alpha=1.0,     # Maximum balance parameter
    arithmetic_error=True,  # Use arithmetic error (difference), if False use geometric error calculation where the error is the factor which the loss is off by rather than absolute amount. 
    update_max=0.05     # Maximum update per step
)

For kp and kd tuning tips, see Controller Gains section.

PLossWeighter

Proportional-only controller (simpler to tune, converges slower/oscillates more):

from PD_loss_balancing import PLossWeighter, ConstantTarget

weighter = PLossWeighter(
    target=ConstantTarget(1.0),
    kp=0.01,
    initial_balance=0.5
)

FixedLossWeighter

Fixed weighting (no adaptation):

from PD_loss_balancing import FixedLossWeighter

weighter = FixedLossWeighter(initial_balance=0.7)  # 70% loss1, 30% loss2

Advanced Usage

Monitoring Controller State

All controllers return diagnostic information in the info dictionary:

alpha, info = weighter.get_balance_param(measured_output=loss1.item())

# For PD/P controllers, info contains:
print(f"Error: {info['pd_controller/error']}")
print(f"Target: {info['pd_controller/target']}")
print(f"Derivative: {info['pd_controller/derivative']}")

# For FixedLossWeighter, info is an empty dict: {}

Custom Target Implementation

Create your own target strategy:

from PD_loss_balancing import Target

class CustomTarget(Target):
    def __init__(self, base_value: float):
        self.base_value = base_value
        self.step_count = 0
    
    def get_target(self, **target_input) -> float:
        self.step_count += 1
        # Custom logic here
        return self.base_value * (1 + 0.01 * self.step_count)

Error Calculation Methods

The controller supports two error calculation methods:

Arithmetic Error (Default: True)

error = measured_output - target_value

Simple difference, suitable when loss magnitudes are similar.

Geometric Error (Default: False)

if target_value > measured_output:
    error = -target_value / (measured_output + 1e-2)
else:
    error = measured_output / (target_value + 1e-2)

Scale-invariant error calculation, better for losses with different magnitudes.

Use Cases

GAN Training

Balance generator and discriminator losses:

# Keep discriminator loss around 0.7
target = ConstantTarget(0.7)
weighter = PDLossWeighter(target, kp=0.02, kd=0.005)

alpha, info = weighter.get_balance_param(measured_output=d_loss.item())
combined = alpha * g_loss + (1-alpha) * d_loss

Multi-task Learning

Balance task-specific losses:

# Keep classification loss 2x larger than regression loss
target = RelativeTarget(ratio=2.0)
weighter = PDLossWeighter(target, kp=0.01)

alpha, info = weighter.get_balance_param(
    measured_output=cls_loss.item(),
    target_reference_values=reg_loss.item()
)

Curriculum Learning

Gradually change loss emphasis:

# Reduce auxiliary loss importance over 5000 steps
target = LinearTrajectoryTarget(initial=1.0, final=0.1, num_steps=5000)
weighter = PDLossWeighter(target, kp=0.005)

Parameters Guide

Controller Gains

• kp (Proportional gain): How strongly to react to current error
  • Higher values: Faster response, possible oscillation
• kd (Derivative gain): How strongly to react to error trends
  • Higher values: Better stability, slower response
• len_errors: How many errors to track to compute the derivative.
  • Lower len_errors means faster responsivenes but more noise. Between 1 and 10 is usually reasonable. Larger len_errors will also increase derivative magnitude.
• kp + kd tuning tips:
  • scale your losses so they're similar in magnitude! e.g. if one is in the thousands and the other in the single digits, scale up the smaller on by 1000x.
  • plot/observe the error, derivative, and target values (see Monitoring Controller State)
  • The values of kp and kd to choose may depend on how responsive the system is, how quickly you need it to converge/stabalize, etc
  • to start, try having the update_max between .01 and .05, and choose a kp and kd such that the updates are usually less than this max. That is, choose kp to be ~ .01 / |usual error|, kd = ~ .01 / |usual derivative|
    • this heuristic assumes that the expected balance parameter range is 0-1. If you expect a greatly restricted range (eg .001 to .1) reduce by the appropriate factor.
  • Start with kd = 0, you shuld see oscillatoions. If there are no oscillations and it converges, no need to add kd. If it doesn't converge increase kp.
  • add in kd. If there is too much damping and doesn't converge (or is very slow to converge) reduce kd (and/or increase kp). If it still oscilates, increase kd.

Bounds

• min_alpha/max_alpha: Constrain balance parameter range
• error_max: Clip error values to prevent extreme updates, espesially when values are near zero
• update_max: Maximum change in balance parameter per step

Examples

See the examples/ directory for complete training scripts:

• dcgan.py: GAN training with loss balancing
• gradient_ascent.py: Adversarial training example

Requirements

• Python ≥ 3.8
• PyTorch ≥ 1.9.0
• NumPy
• Weights & Biases (optional, for logging)

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Citation

If you use this library in your research, please cite:

@software{pd_loss_balancing,
    title={PD Loss Balancing: Automatic Loss Weighting with PID Control},
    author={Addie Foote},
    year={2025},
    url={https://github.com/addiefoote/PD-loss-balancing}
}