PyTorch 2.14에서 cdist·pdist 이중 역전파가 열린 이유와 확인할 점 | DAKER 커뮤니티

쌍거리 기반 손실을 다루다 보면, 계산 자체보다 이차 미분이 막히는 순간이 더 크게 다가올 때가 있습니다. 특히 torch.cdist나 torch.pdist 위에 Hessian, 기울기 페널티, Hessian-벡터 곱을 얹으려 할 때 create_graph=True 경로가 바로 실패하던 경험이 있었다면 이번 변화는 분명히 확인해 볼 만합니다.

PyTorch 2.14에서는 이 지점이 달라졌습니다. 핵심은 속도보다도, 예전에는 지원되지 않던 이차 미분 경로가 이제 통과하는지 재현 가능하게 확인할 수 있다는 점입니다.

PyTorch 2.14에서 torch.cdist와 torch.pdist에 이중 역전파가 구현되었습니다.

무엇이 바뀌었나

이전에는 쌍거리 연산을 포함한 손실에 대해 한 번 더 미분을 걸면 지원되지 않는다는 오류로 멈추는 경우가 많았습니다. 이제는 작은 배치에서 cdist로 거리를 구한 뒤, torch.autograd.grad(..., create_graph=True)로 일차 기울기를 만들고, 그 결과에 다시 미분을 적용하는 흐름을 확인할 수 있습니다.

이 변화는 Hessian 계산, 기울기 노름 페널티, Hessian-벡터 곱처럼 이차 미분이 필요한 실험에서 직접적인 의미를 가집니다. 특히 cdist/pdist를 손실의 일부로 쓰는 경우, 이제는 해당 연산이 이중 역전파 경로를 막는지부터 의심하지 않아도 되는 기반이 생긴 셈입니다.

재현해 볼 실험 조건

확인은 단순한 형태로 시작하는 것이 좋습니다. 작은 실수 텐서 두 개로 torch.cdist 거리를 만든 뒤 스칼라 손실을 정의하고, torch.autograd.grad(..., create_graph=True)로 일차 기울기를 구합니다. 이어서 그 기울기 노름에 대해 다시 미분해 보면 됩니다.

비교가 가능하다면 같은 코드를 PyTorch 2.13 이전과 나란히 두고, 실패하던 지점이 2.14에서 통과하는지 한 줄로 남기면 재현에 도움이 됩니다. pdist도 같은 방식으로 확인할 수 있습니다.

재현 기록에서는 빨라졌는지보다 이차 미분이 실제로 통과하는지가 더 중요합니다.

참가자 노트에는 PyTorch 버전, 입력 shape, create_graph 설정, 성공·실패 여부를 짧게 적어 두는 편이 좋습니다. 실험이 끝난 뒤에도 사용한 버전, shape, 손실 형태를 고정해 두면 다음 사람이 같은 조건으로 바로 따라갈 수 있습니다.

이번 확인에서 분리해서 볼 것

node_creation_hook이나 set_output_grad_dtype은 그래프 메타 정보나 기울기 dtype 계약과 관련된 주제입니다. 이번 글의 초점은 어디까지나 쌍거리 연산의 이차 미분 지원입니다.

그래서 혼합정밀 dtype 선언 실습과 섞기보다는, 먼저 cdist/pdist 위에서 create_graph=True 경로가 통과하는지만 확인하는 편이 좋습니다.

실무에서 바로 볼 포인트

실무에서는 기울기 페널티나 HvP를 cdist/pdist 위에 올릴 때, 2.14부터는 이중 역전파를 전제로 실험을 설계할 수 있습니다. 다만 실제 기록에는 create_graph=True 성공 여부와 입력 shape를 함께 남겨 두는 것이 좋습니다. 그래야 같은 코드라도 어떤 조건에서 통과했는지 빠르게 비교할 수 있습니다.

참고 자료

https://pytorch.org/blog/pytorch-2-14-release-blog/

여기서라면 cdist나 pdist 위에 어떤 형태의 이차 미분 실험을 먼저 확인해 보고 싶으신가요?