Skip to content

Preserve KNN initialization under world-coordinate translation - #1089

Open
Afloat16 wants to merge 1 commit into
nerfstudio-project:mainfrom
Afloat16:fix/knn-relative-distance
Open

Afloat16 wants to merge 1 commit into
nerfstudio-project:mainfrom
Afloat16:fix/knn-relative-distance

Conversation

@Afloat16

Copy link
Copy Markdown

torch.cdist switches to a matrix-multiply distance identity for larger point sets. Subtracting large squared world coordinates loses the small relative distances: translating 32 unit-spaced points by 10,000 can change initialized neighbor scales from 1 to 1e-7.

Use direct Euclidean differences in each existing chunk, mask each query's actual self index and select its k nearest other points. Chunking and the log-scale/epsilon contract remain unchanged.

Validation on Python 3.12.14 / PyTorch 2.8.0+cpu, with native full-package imports and the repository pytest configuration:

  • New regressions on unchanged main: 7 failed / 8 passed; this branch: 15 passed.
  • Existing multiframe/init tests plus regressions: 25 passed / 1 CUDA skip.
  • Independent relative-distance references cover float32/float64 translations, multiple k/chunk sizes, duplicate coordinates and coordinate gradients.
  • Project-pinned Black 22.3.0 changed-file formatting, critical Ruff checks, Python compilation and git diff --check pass.

The direct distance mode trades throughput for numerical stability. A bounded single-thread CPU microbenchmark (k=3, chunk=256, five measured repetitions) gave median 2.01→2.73 ms for 512 points and 21.28→34.13 ms for 2,048 points. These are local observations, not a GPU or complete-training performance estimate. CUDA, full reconstruction/rendering and large-scene memory benchmarks were not run.

`torch.cdist` selects a matrix-multiply distance formula when there are more than 25 points. Its subtraction of large squared world coordinates suffers catastrophic cancellation. For 32 points spaced one unit apart, translating all coordinates by 10000 changes initialized nearest-neighbour scale from 1 to 1e-7.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant