Issue with Applying Drift During Inference Phase in AIHWKit

[adnanrana88]

Dear AIHWKit Team,

I am currently working on a project where I need to apply a custom drift function during the training phase while avoiding the standard PCM noise model drift function only (while other noise type work normall). However, during the inference phase, I want to use the standard PCM noise model and allow all functions to work as intended. Despite my efforts, the drift does not seem to be applied correctly during the inference phase. Here is a brief overview of my approach and the problem I am encountering:

Custom Drift Function During Training:

I have implemented a custom drift function in a subclass of PCMLikeNoiseModel that disables drift during training while allowing other PCM noise model functionalities, such as programming and read noise, to work as standard.

The relevant code for the custom noise model and RPU configuration is as follows:

class CustomPCMLikeNoiseModel(PCMLikeNoiseModel):
    def __init__(self, g_max=25.0, t_0=20.0, drift_scale=1.0):
        super().__init__(g_max=g_max, t_0=t_0, drift_scale=drift_scale)
        self.include_drift = False  # Set to False during training

    def apply_drift_noise_to_conductance(self, g_prog, drift_noise_param, t_inference):
        print("Include Drift: ", self.include_drift)
        if self.include_drift:
            return super().apply_drift_noise_to_conductance(g_prog, drift_noise_param, t_inference)
        else:
            return g_prog

Standard PCM Noise Model During Training and Inference:
RPU Configuration:
During training, I use a separate RPU configuration that includes the custom noise model, and for inference, I re-enable the drift in the noise model using the standard PCMLikeNoiseModel.

Training RPU Configuration:

def create_rpu_config():
    rpu_config = InferenceRPUConfig()
    rpu_config.modifier.std_dev = 0.06
    rpu_config.modifier.type = WeightModifierType.ADD_NORMAL
    rpu_config.mapping.digital_bias = True
    rpu_config.mapping.weight_scaling_omega = 1.0
    rpu_config.mapping.weight_scaling_columnwise = False
    rpu_config.mapping.out_scaling_columnwise = False
    rpu_config.remap.type = WeightRemapType.LAYERWISE_SYMMETRIC
    rpu_config.clip.type = WeightClipType.LAYER_GAUSSIAN
    rpu_config.clip.sigma = 2.0
    rpu_config.forward = IOParameters()
    rpu_config.forward.is_perfect = False
    rpu_config.forward.out_noise = 0.0
    rpu_config.forward.inp_bound = 1.0
    rpu_config.forward.inp_res = 1 / (2**8 - 2)
    rpu_config.forward.out_bound = 12
    rpu_config.forward.out_res = 1 / (2**8 - 2)
    rpu_config.forward.bound_management = BoundManagementType.NONE
    rpu_config.forward.noise_management = NoiseManagementType.NONE
    rpu_config.pre_post.input_range.enable = True
    rpu_config.pre_post.input_range.decay = 0.01
    rpu_config.pre_post.input_range.init_from_data = 50
    rpu_config.pre_post.input_range.init_std_alpha = 3.0
    rpu_config.pre_post.input_range.input_min_percentage = 0.995
    rpu_config.pre_post.input_range.manage_output_clipping = False
    rpu_config.noise_model = CustomPCMLikeNoiseModel(g_max=25.0)
    return rpu_config

standard PCMLikeNoiseModel for Inference RPU Configuration:.

The RPU configuration for inference is as follows:

def create_rpu_config_for_inference():
    rpu_config = create_rpu_config()
    rpu_config.noise_model.include_drift = True  # Enable drift during inference
    return rpu_config

Problem:
Despite following the steps to apply drift during inference, the weights do not seem to change after applying drift_analog_weights, indicating that drift is not being applied. The debug output shows identical weights before and after applying drift.

Inference Phase Code:
Here is the relevant part of my code where I apply drift during the inference phase. I also tried using the program_analog_weights and drift_analog_weights functions according to the example, but drift is not working in any case.

def test_inference(model, criterion, test_loader, device, plot_dir):
    t_inference_list = [1.0, 60.0, 1800.0, 3600.0, 86400.0, 2592000.0, 15768000.0, 31536000.0, 63072000.0, 94608000.0, 157680000.0]
    drift_accuracies = []
    model.eval()
    # Program weights before applying drift
    for layer in [model.fc1, model.fc3]:
        for tile in layer.analog_tiles():
            noise_model = tile.rpu_config.noise_model
            program_analog_weights(tile, noise_model)
    for t_inference in t_inference_list:
        for idx, layer in enumerate([model.fc1, model.fc3], start=1):
            for tile in layer.analog_tiles():
                weights_before_drift, _ = layer.get_weights()
                print(f'Weights before drift for layer {idx} at t={t_inference}: {weights_before_drift}')
                drift_analog_weights(tile, t_inference)
                weights_after_drift, _ = layer.get_weights()
                print(f'Weights after drift for layer {idx} at t={t_inference}: {weights_after_drift}')
                if torch.equal(weights_before_drift, weights_after_drift):
                    print(f'No change in weights for layer {layer} after applying drift.')
        # Evaluate accuracy after applying drift
        total_loss, correct, total = 0, 0, 0
        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                loss = criterion(output, target)
                total_loss += loss.item() * data.size(0)
                _, predicted = torch.max(output.data, 1)
                total += target.size(0)
                correct += (predicted == target).sum().item()

        accuracy_post = 100 * correct / total
        drift_accuracies.append(accuracy_post)
        print(f'Drift t={t_inference: .2e}\t Accuracy after inference: {accuracy_post:.2f}%')
    plot_drift_accuracies(t_inference_list, drift_accuracies, plot_dir)

Could you please advise on how to ensure that the drift is correctly applied during the inference phase while avoiding it in the training phase? Any guidance or suggestions would be greatly appreciated.

Thank you for your assistance.

[maljoras-sony]

Note that the drift model is never applied during training, so there is no need to have a flag. It will be only used for inference time.

It is not clear from the above how the modified rpu_config is applied after training was done, before the inferfence. Is the model recreated while using the new rpu_config? There is a special flag to change the rpu_config during loading, otherwise the old rpu_config will be taken that is stored in the state_dict. You can also use the model.to(new_rpu_config) to change the model to the new rpu_config. However, it is even easier if you just set include_drift=True always, as drift will never be called during training. In this case, you do not have to change the rpu_config.

Also, there is a simpler form to apply the drift for the while model, simple use model.drift_analog_weights(), so no need to go through the analog tiles.

In this section, you should call tile.get_weights and not layer.get_weights. Also best use tile.drift_weights(t) directly if you are already looping through the tiles.

weights_before_drift, _ = layer.get_weights()
print(f'Weights before drift for layer {idx} at t={t_inference}: {weights_before_drift}')
drift_analog_weights(tile, t_inference)
weights_after_drift, _ = layer.get_weights()

[maljoras-sony]

In fact you can also use directly (using a specific new noise model not defined in the rpu_config):

for tile in model.analog_tiles():
    tile.program_weights(noise_model=CustomPCMLikeNoiseModel(g_max=25.0))
    tile.drift_weights(t)