Accelerate the rep-reading

[Y-L-LIU]

I found the calculation of the reading pipeline is super slow. I send the projection and recenter to GPU and make it run 20x faster. May I open a pull request?

def project_onto_direction(H, direction):
    """Project matrix H (n, d_1) onto direction vector (d_2,)"""
    # Calculate the magnitude of the direction vector
     # Ensure H and direction are on the same device (CPU or GPU)
    device = H.device
    if type(direction) != torch.Tensor:
        direction = torch.Tensor(direction)
    direction = direction.to(device)
    mag = torch.norm(direction)
    assert not torch.isinf(mag).any()
    # Calculate the projection
    projection = H.matmul(direction) / mag
    return projection

def recenter(x, mean=None):
    if mean is None:
        # mean = x.mean(axis=0, keepdims=True)
        mean = torch.mean(x,axis=0,keepdims=True)
    else:
        mean = torch.Tensor(mean).cuda()
    # print(type(x), type(mean))
    return x - mean

[justinphan3110cais]

hi @Y-L-LIU , yeah that would be great. I can help verify the results of your implementation on repe_eval

[Y-L-LIU]

Thanks for your kind reply. I have opened a pull request.

[caesar-jojo]

@Y-L-LIU @justinphan3110cais Author, I've encountered the same issue. Could you please tell me how to speed up the process? I'm not very experienced with coding, and after spending half a day today, I believe there should be a way to accelerate it. However, it still took 20 hours to load the reading pipeline. Thank you~

[Y-L-LIU]

@caesar-jojo Hi, I have opened a pull request for the acceleration, you can check the pull request for further deatails.

[caesar-jojo]

@Y-L-LIU Thanks for the reply, I downloaded the code on your github (URL: https://github.com/Y-L-LIU/representation-engineering) yesterday and then I ran the file "./representation-engineering-main/examples/honesty/honesty.ipynb" directly. Is it a function (rep_reading_pipeline.get_directions) that will automatically run on the GPU without me having to change the code

[Y-L-LIU]

@caesar-jojo Actually, the bottleneck of the code lies in the scoring part, i.e. the RepReader.transform function, which requires lots of computation. I use torch to accelerate this part. And for the rep_reading_pipeline.get_directions, it calls a PCA.fit() and can be computed on cpu within 10 minutes or so.

[sorcererrandy]

@Y-L-LIU Sorry, it seems that the updated code is not running properly for me. I have a dataset of 20,000 sentences, and in the original code, it took about half an hour to execute rep_reading_pipeline.get_directions. However, when using the updated code, it has been running for over an hour and still hasn't finished. I'm not sure what the issue might be. It seems that it stuck at 'get_signs'