diff --git a/CMakeLists.txt b/CMakeLists.txt
index c68ab5f8..19cdc6a9 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -2,7 +2,12 @@ cmake_minimum_required(VERSION 3.21)
 project(opensplat)
 
 set(OPENSPLAT_BUILD_SIMPLE_TRAINER OFF CACHE BOOL "Build simple trainer applications")
-set(GPU_RUNTIME "CUDA" CACHE STRING "HIP or CUDA or MPS")
+if(APPLE)
+    set(GPU_RUNTIME_DEFAULT "MPS")
+else()
+    set(GPU_RUNTIME_DEFAULT "CUDA")
+endif()
+set(GPU_RUNTIME "${GPU_RUNTIME_DEFAULT}" CACHE STRING "HIP or CUDA or MPS or CPU")
 set(OPENCV_DIR "OPENCV_DIR-NOTFOUND" CACHE PATH "Path to the OPENCV installation directory")
 set(OPENSPLAT_MAX_CUDA_COMPATIBILITY OFF CACHE BOOL "Build for maximum CUDA device compatibility")
 set(OPENSPLAT_BUILD_VISUALIZER OFF CACHE BOOL "Build visualizer application")
@@ -13,6 +18,7 @@ set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
 
 # Read version
 file(READ "VERSION" APP_VERSION)
+string(STRIP "${APP_VERSION}" APP_VERSION)
 
 # Read git commit
 set(GIT_REV "")
@@ -146,18 +152,28 @@ elseif(GPU_RUNTIME STREQUAL "HIP")
     endif()
     list(APPEND CMAKE_PREFIX_PATH "${ROCM_ROOT}")
 elseif(GPU_RUNTIME STREQUAL "MPS")
-    find_library(FOUNDATION_LIBRARY Foundation REQUIRED)
-    find_library(METAL_FRAMEWORK    Metal      REQUIRED)
-    find_library(METALKIT_FRAMEWORK MetalKit   REQUIRED)
-    message(STATUS "Metal framework found")
+    execute_process(COMMAND xcrun -sdk macosx metal --version
+                    RESULT_VARIABLE METAL_COMPILER_RESULT
+                    OUTPUT_QUIET ERROR_QUIET)
+    if(NOT METAL_COMPILER_RESULT EQUAL 0)
+        message(WARNING "Metal compiler not found, building with CPU support only. "
+                        "Install Xcode and the Metal toolchain "
+                        "(xcodebuild -downloadComponent MetalToolchain), then re-run cmake.")
+        set(GPU_RUNTIME "CPU")
+    else()
+        find_library(FOUNDATION_LIBRARY Foundation REQUIRED)
+        find_library(METAL_FRAMEWORK    Metal      REQUIRED)
+        find_library(METALKIT_FRAMEWORK MetalKit   REQUIRED)
+        message(STATUS "Metal framework found")
 
-    set(XC_FLAGS -O3)
-    if(OPENSPLAT_USE_FAST_MATH)
-        message(STATUS "Fast math optimizations enabled for Metal")
-        
-        set(XC_FLAGS ${XC_FLAGS} -ffast-math)
+        set(XC_FLAGS -O3)
+        if(OPENSPLAT_USE_FAST_MATH)
+            message(STATUS "Fast math optimizations enabled for Metal")
+
+            set(XC_FLAGS ${XC_FLAGS} -ffast-math)
+        endif()
+        set(USE_MPS ON CACHE BOOL "Use MPS for GPU acceleration")
     endif()
-    set(USE_MPS ON CACHE BOOL "Use MPS for GPU acceleration")
 else()
     set(GPU_RUNTIME "CPU")
 endif()
@@ -247,9 +263,9 @@ add_library(gsplat_cpu rasterizer/gsplat-cpu/gsplat_cpu.cpp)
 target_include_directories(gsplat_cpu PRIVATE ${TORCH_INCLUDE_DIRS})
 
 set(OPENSPLAT_SRC_FILES opensplat.cpp point_io.cpp nerfstudio.cpp model.cpp
-kdtree_tensor.cpp spherical_harmonics.cpp cv_utils.cpp utils.cpp project_gaussians.cpp
-rasterize_gaussians.cpp ssim.cpp optim_scheduler.cpp colmap.cpp opensfm.cpp openmvg.cpp input_data.cpp
-tensor_math.cpp rad.cpp zip_utils.cpp)
+kdtree_tensor.cpp spherical_harmonics.cpp cv_utils.cpp project_gaussians.cpp
+rasterize_gaussians.cpp ssim.cpp colmap.cpp opensfm.cpp openmvg.cpp input_data.cpp
+tensor_math.cpp rad.cpp zip_utils.cpp undistort.cpp)
 
 if (OPENSPLAT_BUILD_VISUALIZER)
     if (Pangolin_FOUND)
diff --git a/README.md b/README.md
index 46ed6905..7bdf2534 100644
--- a/README.md
+++ b/README.md
@@ -11,8 +11,6 @@ Graphics card recommended, but not required! OpenSplat runs the fastest on NVIDI
 
 Commercial use allowed and encouraged under the terms of the [AGPLv3](https://www.tldrlegal.com/license/gnu-affero-general-public-license-v3-agpl-3-0). ✅
 
-We even have a [song](https://youtu.be/1bma7XJkoDM) 🎵
-
 ## Getting Started
 
 If you're on Windows, you can [buy](http://sites.fastspring.com/masseranolabs/product/opensplatforwindows) the pre-built program. This saves you time and helps support the project ❤️. Then jump directly to the [run](#run) section. As an alternative, check the [build](#build) section below.
@@ -125,6 +123,7 @@ You will also need to install Xcode and the Xcode command line tools to compile
 1. Install Xcode from the Apple App Store.
 2. Install the command line tools with `xcode-select --install`. This might do nothing on your machine.
 3. If `xcode-select --print-path` prints `/Library/Developer/CommandLineTools`,then run `sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer`.
+4. On recent Xcode versions the Metal toolchain is a separate download. If `xcrun -sdk macosx metal --version` fails, run `xcodebuild -downloadComponent MetalToolchain`.
 
 Then run:
 
@@ -132,11 +131,11 @@ Then run:
 git clone https://github.com/pierotofy/OpenSplat OpenSplat
 cd OpenSplat
 mkdir build && cd build
-cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch/ -DGPU_RUNTIME=MPS .. && make -j$(sysctl -n hw.logicalcpu)
+cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch/ .. && make -j$(sysctl -n hw.logicalcpu)
 ./opensplat
 ```
 
-If building CPU-only, remove `-DGPU_RUNTIME=MPS`.
+On macOS `GPU_RUNTIME` defaults to `MPS` (metal acceleration). If the Metal compiler isn't available, the build automatically falls back to CPU. To force a CPU-only build, pass `-DGPU_RUNTIME=CPU`.
 
 :warning: You will probably get a *libc10.dylib can’t be opened because Apple cannot check it for malicious software* error on first run. Open **System Settings** and go to **Privacy & Security** and find the **Allow** button. You might need to repeat this several times until all torch libraries are loaded.
 
@@ -245,6 +244,16 @@ You can resume training of a .PLY file by using the `--resume` option:
 ./opensplat /path/to/banana --resume ./splat.ply
 ```
 
+### Image Masks
+
+You can exclude parts of your images by adding 2D masks. Place them in a `masks` folder (also recognized: `mask`, `segmentation`, `dynamic_masks`) inside your project, named after each image (e.g. `images/IMG_001.JPG` → `masks/IMG_001.png`). Masks are grayscale images matching the input dimensions: white marks pixels to keep, black pixels to ignore.
+
+When masks are found they are applied automatically. Use `--no-masks` to ignore them.
+
+### Coordinate Reference System
+
+By default OpenSplat preserves the input coordinate reference system of the model. If you want to automatically center the result so that it displays nicely in most viewers, use `--center`.
+
 ### AMD GPU Notes
 
 To train a model with AMD GPU using docker container, you can use the following command as a reference:
diff --git a/VERSION b/VERSION
index ab679818..26aaba0e 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-1.1.6
\ No newline at end of file
+1.2.0
diff --git a/cv_utils.cpp b/cv_utils.cpp
index 11b0191d..5a7f80b6 100644
--- a/cv_utils.cpp
+++ b/cv_utils.cpp
@@ -13,20 +13,6 @@ cv::Mat imreadRGB(const std::string &filename){
     return cImg;
 }
 
-void imwriteRGB(const std::string &filename, const cv::Mat &image){
-    cv::Mat rgb;
-    cv::cvtColor(image, rgb, cv::COLOR_RGB2BGR);
-    cv::imwrite(filename, rgb);
-}
-
-cv::Mat floatNxNtensorToMat(const torch::Tensor &t){
-    return cv::Mat(t.size(0), t.size(1), CV_32F, t.data_ptr());
-}
-
-torch::Tensor floatNxNMatToTensor(const cv::Mat &m){
-    return torch::from_blob(m.data, { m.rows, m.cols }, torch::kFloat32).clone();
-}
-
 cv::Mat tensorToImage(const torch::Tensor &t){
     int h = t.sizes()[0];
     int w = t.sizes()[1];
diff --git a/cv_utils.hpp b/cv_utils.hpp
index 87c7c2ad..67fb95d1 100644
--- a/cv_utils.hpp
+++ b/cv_utils.hpp
@@ -7,9 +7,6 @@
 #include <opencv2/imgproc.hpp>
 
 cv::Mat imreadRGB(const std::string &filename);
-void imwriteRGB(const std::string &filename, const cv::Mat &image);
-cv::Mat floatNxNtensorToMat(const torch::Tensor &t);
-torch::Tensor floatNxNMatToTensor(const cv::Mat &m);
 cv::Mat tensorToImage(const torch::Tensor &t);
 torch::Tensor imageToTensor(const cv::Mat &image);
 
diff --git a/input_data.cpp b/input_data.cpp
index c819a776..90c91dec 100644
--- a/input_data.cpp
+++ b/input_data.cpp
@@ -1,8 +1,18 @@
 #include <filesystem>
 #include <mutex>
+#include <atomic>
+#ifdef USE_CUDA
+#include <cuda_runtime_api.h>
+#elif defined(USE_HIP)
+#include <hip/hip_runtime_api.h>
+#endif
+#ifdef __APPLE__
+#include <sys/sysctl.h>
+#endif
 #include <nlohmann/json.hpp>
 #include "input_data.hpp"
 #include "cv_utils.hpp"
+#include "undistort.hpp"
 
 namespace fs = std::filesystem;
 using namespace torch::indexing;
@@ -51,9 +61,15 @@ void Camera::loadImage(float downscaleFactor){
     }
     
     cv::Mat cImg = imreadRGB(filePath);
-    
+
+    cv::Mat cMask;
+    if (!maskPath.empty()){
+        cMask = cv::imread(maskPath, cv::IMREAD_GRAYSCALE);
+        if (cMask.empty()) throw std::runtime_error("Cannot read mask " + maskPath);
+    }
+
     float rescaleF = 1.0f;
-    // If camera intrinsics don't match the image dimensions 
+    // If camera intrinsics don't match the image dimensions
     if (cImg.rows != height || cImg.cols != width){
         rescaleF = static_cast<float>(cImg.rows) / static_cast<float>(height);
     }
@@ -71,35 +87,43 @@ void Camera::loadImage(float downscaleFactor){
         cy *= scaleFactor;
     }
 
-    K = getIntrinsicsMatrix();
-    cv::Rect roi;
+    if (!cMask.empty()){
+        cv::threshold(cMask, cMask, 127, 255, cv::THRESH_BINARY);
+        if (cMask.rows != cImg.rows || cMask.cols != cImg.cols){
+            cv::resize(cMask, cMask, cv::Size(cImg.cols, cImg.rows), 0.0, 0.0, cv::INTER_LINEAR);
+        }
+    }
 
     if (hasDistortionParameters()){
-        // Undistort
-        std::vector<float> distCoeffs = undistortionParameters();
-        cv::Mat cK = floatNxNtensorToMat(K);
-        cv::Mat newK = cv::getOptimalNewCameraMatrix(cK, distCoeffs, cv::Size(cImg.cols, cImg.rows), 0, cv::Size(), &roi);
-
-        cv::Mat undistorted = cv::Mat::zeros(cImg.rows, cImg.cols, cImg.type());
-        cv::undistort(cImg, undistorted, cK, distCoeffs, newK);
-        
+        UndistortParams p = computeUndistortParams(fx, fy, cx, cy, cImg.cols, cImg.rows,
+                                                   k1, k2, k3, k4, k5, k6, p1, p2);
+        cv::Mat mapx, mapy;
+        buildUndistortMaps(p, mapx, mapy);
+        cv::Mat undistorted;
+        cv::remap(cImg, undistorted, mapx, mapy, cv::INTER_LINEAR, cv::BORDER_CONSTANT);
         image = imageToTensor(undistorted);
-        K = floatNxNMatToTensor(newK);
+        if (!cMask.empty()){
+            cv::Mat remapped;
+            cv::remap(cMask, remapped, mapx, mapy, cv::INTER_LINEAR, cv::BORDER_CONSTANT);
+            cMask = remapped;
+        }
+        fx = p.dstFx;
+        fy = p.dstFy;
+        cx = p.dstCx;
+        cy = p.dstCy;
     }else{
-        roi = cv::Rect(0, 0, cImg.cols, cImg.rows);
         image = imageToTensor(cImg);
     }
 
-    // Crop to ROI
-    image = image.index({Slice(roi.y, roi.y + roi.height), Slice(roi.x, roi.x + roi.width), Slice()});
-
-    // Update parameters
     height = image.size(0);
     width = image.size(1);
-    fx = K[0][0].item<float>();
-    fy = K[1][1].item<float>();
-    cx = K[0][2].item<float>();
-    cy = K[1][2].item<float>();
+    K = getIntrinsicsMatrix();
+
+    if (!cMask.empty()){
+        torch::Tensor m = torch::from_blob(cMask.data, {cMask.rows, cMask.cols}, torch::kU8)
+                            .to(torch::kFloat32).div(255.0f).clone();
+        mask = (m >= 0.5f).to(torch::kFloat32);
+    }
 }
 
 torch::Tensor Camera::getImage(int downscaleFactor){
@@ -126,9 +150,112 @@ bool Camera::hasDistortionParameters(){
     return k1 != 0.0f || k2 != 0.0f || k3 != 0.0f || k4 != 0.0f || k5 != 0.0f || k6 != 0.0f || p1 != 0.0f || p2 != 0.0f;
 }
 
-std::vector<float> Camera::undistortionParameters(){
-    std::vector<float> p = { k1, k2, p1, p2, k3, k4, k5, k6 };
-    return p;
+torch::Tensor Camera::getMask(int downscaleFactor){
+    if (!hasMask()) return mask;
+    if (downscaleFactor <= 1) return mask;
+    if (maskPyramids.find(downscaleFactor) != maskPyramids.end()){
+        return maskPyramids[downscaleFactor];
+    }
+    torch::Tensor m = mask.unsqueeze(0).unsqueeze(0);
+    m = torch::nn::functional::interpolate(m,
+            torch::nn::functional::InterpolateFuncOptions()
+                .size(std::vector<int64_t>{ mask.size(0) / downscaleFactor, mask.size(1) / downscaleFactor })
+                .mode(torch::kBilinear).align_corners(false));
+    m = (m.squeeze(0).squeeze(0) >= 0.5f).to(torch::kFloat32);
+    maskPyramids[downscaleFactor] = m;
+    return m;
+}
+
+bool Camera::gpuCacheEnabled = true;
+
+// Half the free VRAM at first use (CUDA/HIP), a quarter of system RAM on
+// Apple unified memory, 1GB otherwise
+static long long gpuCacheBudget(){
+#ifdef USE_CUDA
+    size_t freeB = 0, totalB = 0;
+    if (cudaMemGetInfo(&freeB, &totalB) == cudaSuccess){
+        return static_cast<long long>(freeB / 2);
+    }
+#elif defined(USE_HIP)
+    size_t freeB = 0, totalB = 0;
+    if (hipMemGetInfo(&freeB, &totalB) == hipSuccess){
+        return static_cast<long long>(freeB / 2);
+    }
+#endif
+#ifdef __APPLE__
+    int64_t ram = 0;
+    size_t size = sizeof(ram);
+    if (sysctlbyname("hw.memsize", &ram, &size, nullptr, 0) == 0){
+        return ram / 4;
+    }
+#endif
+    return 1LL << 30;
+}
+
+// Cache device-side tensors per camera to avoid re-uploading every iteration
+static torch::Tensor gpuCached(std::unordered_map<int, torch::Tensor> &cache, int key,
+                               const torch::Tensor &src, const torch::Device &device){
+    if (device == torch::kCPU || !Camera::gpuCacheEnabled) return src.to(device);
+    auto it = cache.find(key);
+    if (it != cache.end()) return it->second;
+
+    static std::atomic<long long> gpuCacheBytes{0};
+    static const long long budget = gpuCacheBudget();
+    long long bytes = src.numel() * src.element_size();
+    if (gpuCacheBytes.load() + bytes > budget) return src.to(device);
+    gpuCacheBytes += bytes;
+    torch::Tensor t = src.to(device);
+    cache[key] = t;
+    return t;
+}
+
+torch::Tensor Camera::getImageGpu(int downscaleFactor, const torch::Device &device){
+    return gpuCached(gpuImageCache, downscaleFactor, getImage(downscaleFactor), device);
+}
+
+torch::Tensor Camera::getMaskGpu(int downscaleFactor, const torch::Device &device){
+    torch::Tensor m = getMask(downscaleFactor);
+    if (!m.defined() || m.numel() == 0) return m;
+    return gpuCached(gpuMaskCache, downscaleFactor, m, device);
+}
+
+torch::Tensor Camera::getEdgeMapGpu(int downscaleFactor, const torch::Device &device){
+    return gpuCached(gpuEdgeCache, downscaleFactor, getEdgeMap(downscaleFactor).contiguous(), device);
+}
+
+torch::Tensor Camera::getEdgeMap(int downscaleFactor){
+    if (edgePyramids.find(downscaleFactor) != edgePyramids.end()){
+        return edgePyramids[downscaleFactor];
+    }
+    cv::Mat cImg = tensorToImage(getImage(downscaleFactor));
+    cv::Mat gray, edges;
+    cv::cvtColor(cImg, gray, cv::COLOR_RGB2GRAY);
+    cv::Canny(gray, edges, 50, 150);
+    torch::Tensor e = torch::from_blob(edges.data, {edges.rows, edges.cols}, torch::kU8)
+                        .to(torch::kFloat32).div(255.0f).clone();
+    edgePyramids[downscaleFactor] = e;
+    return e;
+}
+
+std::string findMaskPath(const std::string &imagePath, const std::string &projectRoot){
+    static const char *folders[] = { "masks", "mask", "segmentation", "dynamic_masks" };
+    static const char *extensions[] = { ".png", ".jpg", ".jpeg", ".mask.png" };
+
+    fs::path img(imagePath);
+    std::string stem = img.stem().string();
+    std::string name = img.filename().string();
+
+    for (const char *folder : folders){
+        fs::path dir = fs::path(projectRoot) / folder;
+        if (!fs::exists(dir) || !fs::is_directory(dir)) continue;
+        for (const char *ext : extensions){
+            fs::path cand = dir / (stem + ext);
+            if (fs::exists(cand)) return cand.string();
+            cand = dir / (name + ext);
+            if (fs::exists(cand)) return cand.string();
+        }
+    }
+    return "";
 }
 
 std::tuple<std::vector<Camera>, Camera *> InputData::getCameras(bool validate, const std::string &valImage){
diff --git a/input_data.hpp b/input_data.hpp
index a2c36f9f..9ddefe60 100644
--- a/input_data.hpp
+++ b/input_data.hpp
@@ -27,6 +27,7 @@ struct Camera{
     float p2 = 0;
     torch::Tensor camToWorld;
     std::string filePath = "";
+    std::string maskPath = "";
     CameraType cameraType = CameraType::Perspective;
 
     Camera(){};
@@ -38,14 +39,27 @@ struct Camera{
         camToWorld(camToWorld), filePath(filePath) {}
     torch::Tensor getIntrinsicsMatrix();
     bool hasDistortionParameters();
-    std::vector<float> undistortionParameters();
     torch::Tensor getImage(int downscaleFactor);
+    torch::Tensor getMask(int downscaleFactor);
+    torch::Tensor getEdgeMap(int downscaleFactor);
+    torch::Tensor getImageGpu(int downscaleFactor, const torch::Device &device);
+    torch::Tensor getMaskGpu(int downscaleFactor, const torch::Device &device);
+    torch::Tensor getEdgeMapGpu(int downscaleFactor, const torch::Device &device);
+    bool hasMask() const { return mask.numel() > 0; }
 
     void loadImage(float downscaleFactor);
     torch::Tensor K;
     torch::Tensor image;
+    torch::Tensor mask; // [H,W] float 0/1, aligned with image
 
     std::unordered_map<int, torch::Tensor> imagePyramids;
+    std::unordered_map<int, torch::Tensor> maskPyramids;
+    std::unordered_map<int, torch::Tensor> edgePyramids;
+    std::unordered_map<int, torch::Tensor> gpuImageCache;
+    std::unordered_map<int, torch::Tensor> gpuMaskCache;
+    std::unordered_map<int, torch::Tensor> gpuEdgeCache;
+
+    static bool gpuCacheEnabled;
 };
 
 struct Points{
@@ -63,5 +77,6 @@ struct InputData{
     void saveCameras(const std::string &filename, bool keepCrs);
 };
 InputData inputDataFromX(const std::string &projectRoot);
+std::string findMaskPath(const std::string &imagePath, const std::string &projectRoot);
 
 #endif
\ No newline at end of file
diff --git a/model.cpp b/model.cpp
index 71b44ab6..2a446e38 100644
--- a/model.cpp
+++ b/model.cpp
@@ -1,4 +1,7 @@
-#include <filesystem>
+#include <filesystem>
+#include <random>
+#include <numeric>
+#include <algorithm>
 #include <load-spz.h>
 #include "model.hpp"
 #include "constants.hpp"
@@ -23,16 +26,10 @@
 
 namespace fs = std::filesystem;
 
-torch::Tensor randomQuatTensor(long long n){
-    torch::Tensor u = torch::rand(n);
-    torch::Tensor v = torch::rand(n);
-    torch::Tensor w = torch::rand(n);
-    return torch::stack({
-        torch::sqrt(1 - u) * torch::sin(2 * PI * v),
-        torch::sqrt(1 - u) * torch::cos(2 * PI * v),
-        torch::sqrt(u) * torch::sin(2 * PI * w),
-        torch::sqrt(u) * torch::cos(2 * PI * w)
-    }, -1);
+torch::Tensor identityQuatTensor(long long n){
+    torch::Tensor q = torch::zeros({n, 4});
+    q.index_put_({Slice(), 0}, 1.0f);
+    return q;
 }
 
 torch::Tensor projectionMatrix(float zNear, float zFar, float fovX, float fovY, const torch::Device &device){
@@ -49,15 +46,6 @@ torch::Tensor projectionMatrix(float zNear, float zFar, float fovX, float fovY,
     }, device);
 }
 
-torch::Tensor psnr(const torch::Tensor& rendered, const torch::Tensor& gt){
-    torch::Tensor mse = (rendered - gt).pow(2).mean();
-    return (10.f * torch::log10(1.0 / mse));
-}
-
-torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt){
-    return torch::abs(gt - rendered).mean();
-}
-
 template<typename T>
 std::vector<T> tensor_to_vector(const torch::Tensor t){
     return std::vector<T>(t.data_ptr<float>(), t.data_ptr<float>() + t.numel());
@@ -66,14 +54,13 @@ std::vector<T> tensor_to_vector(const torch::Tensor t){
 void Model::setupOptimizers(){
     releaseOptimizers();
 
-    meansOpt = new torch::optim::Adam({means}, torch::optim::AdamOptions(0.00016));
-    scalesOpt = new torch::optim::Adam({scales}, torch::optim::AdamOptions(0.005));
-    quatsOpt = new torch::optim::Adam({quats}, torch::optim::AdamOptions(0.001));
-    featuresDcOpt = new torch::optim::Adam({featuresDc}, torch::optim::AdamOptions(0.0025));
-    featuresRestOpt = new torch::optim::Adam({featuresRest}, torch::optim::AdamOptions(0.000125));
-    opacitiesOpt = new torch::optim::Adam({opacities}, torch::optim::AdamOptions(0.05));
-
-    meansOptScheduler = new OptimScheduler(meansOpt, 0.0000016f, maxSteps);
+    const double eps = 1e-15;
+    meansOpt = new torch::optim::Adam({means}, torch::optim::AdamOptions(1.6e-4 * spatialLrScale).eps(eps));
+    scalesOpt = new torch::optim::Adam({scales}, torch::optim::AdamOptions(5e-3).eps(eps));
+    quatsOpt = new torch::optim::Adam({quats}, torch::optim::AdamOptions(1e-3).eps(eps));
+    featuresDcOpt = new torch::optim::Adam({featuresDc}, torch::optim::AdamOptions(2.5e-3).eps(eps));
+    featuresRestOpt = new torch::optim::Adam({featuresRest}, torch::optim::AdamOptions(2.5e-4).eps(eps)); // highfeature_lr / 20
+    opacitiesOpt = new torch::optim::Adam({opacities}, torch::optim::AdamOptions(0.025).eps(eps));
 }
 
 void Model::releaseOptimizers(){
@@ -83,8 +70,6 @@ void Model::releaseOptimizers(){
     RELEASE_SAFELY(featuresDcOpt);
     RELEASE_SAFELY(featuresRestOpt);
     RELEASE_SAFELY(opacitiesOpt);
-
-    RELEASE_SAFELY(meansOptScheduler);
 }
 
 
@@ -178,8 +163,11 @@ torch::Tensor Model::forward(Camera& cam, int step){
     
     xys.retain_grad();
 
-    if (radii.sum().item<float>() == 0.0f)
+    if (radii.sum().item<float>() == 0.0f){
+        lastAlpha = torch::zeros({height, width}, torch::TensorOptions().dtype(torch::kFloat32).device(device));
+        errorMap = torch::Tensor();
         return backgroundColor.repeat({height, width, 1});
+    }
 
     torch::Tensor viewDirs = means.detach() - T.transpose(0, 1).to(device);
     viewDirs = viewDirs / viewDirs.norm(2, {-1}, true);
@@ -199,8 +187,19 @@ torch::Tensor Model::forward(Camera& cam, int step){
     
     rgbs = torch::clamp_min(rgbs + 0.5f, 0.0f);
 
+    auto fOpts = torch::TensorOptions().dtype(torch::kFloat32).device(device);
+    torch::Tensor camEdgeMap = torch::empty({0}, fOpts);
+    if (!scoringPass){
+        errorMap = torch::empty({0}, fOpts);
+        densificationInfo = torch::empty({0}, fOpts);
+        xyAbsGrad = step <= densifyUntilIter ? torch::zeros({means.size(0), 2}, fOpts)
+                                             : torch::empty({0}, fOpts);
+    }else if (edgeGuidance){
+        camEdgeMap = cam.getEdgeMapGpu(getDownscaleFactor(step), device);
+    }
+    
     if (device == torch::kCPU){
-        rgb = RasterizeGaussiansCPU::apply(
+        auto rast = RasterizeGaussiansCPU::apply(
                 xys,
                 radii,
                 conics,
@@ -210,10 +209,16 @@ torch::Tensor Model::forward(Camera& cam, int step){
                 camDepths,
                 height,
                 width,
-                backgroundColor);
-    }else{  
+                backgroundColor,
+                errorMap,
+                camEdgeMap,
+                densificationInfo,
+                xyAbsGrad);
+        rgb = rast[0];
+        lastAlpha = rast[1];
+    }else{
         #if defined(USE_HIP) || defined(USE_CUDA) || defined(USE_MPS)
-        rgb = RasterizeGaussians::apply(
+        auto rast = RasterizeGaussians::apply(
                 xys,
                 depths,
                 radii,
@@ -223,7 +228,13 @@ torch::Tensor Model::forward(Camera& cam, int step){
                 torch::sigmoid(opacities),
                 height,
                 width,
-                backgroundColor);
+                backgroundColor,
+                errorMap,
+                camEdgeMap,
+                densificationInfo,
+                xyAbsGrad);
+        rgb = rast[0];
+        lastAlpha = rast[1];
         #endif
     }
 
@@ -232,26 +243,40 @@ torch::Tensor Model::forward(Camera& cam, int step){
     return rgb;
 }
 
-void Model::optimizersZeroGrad(){
-  meansOpt->zero_grad();
-  scalesOpt->zero_grad();
-  quatsOpt->zero_grad();
-  featuresDcOpt->zero_grad();
-  featuresRestOpt->zero_grad();
-  opacitiesOpt->zero_grad();
+static void setOptimizerLr(torch::optim::Adam *opt, double lr){
+    static_cast<torch::optim::AdamOptions&>(opt->param_groups()[0].options()).set_lr(lr);
 }
 
-void Model::optimizersStep(){
-  meansOpt->step();
-  scalesOpt->step();
-  quatsOpt->step();
-  featuresDcOpt->step();
-  featuresRestOpt->step();
-  opacitiesOpt->step();
+void Model::schedulersStep(int step){
+    double t = std::clamp(static_cast<double>(step) / maxSteps, 0.0, 1.0);
+    double lr = std::exp(std::log(1.6e-4) * (1.0 - t) + std::log(1.6e-6) * t) * spatialLrScale;
+    setOptimizerLr(meansOpt, lr);
 }
 
-void Model::schedulersStep(int step){
-  meansOptScheduler->step(step);
+void Model::optimizerStepCadence(int step){
+    auto stepAndZero = [](torch::optim::Adam *opt){
+        opt->step();
+        opt->zero_grad(true);
+    };
+    int lateStart = (std::min)(20000, maxSteps * 2 / 3);
+    if (step <= densifyUntilIter){
+        for (torch::optim::Adam *opt : {meansOpt, scalesOpt, quatsOpt, featuresDcOpt, opacitiesOpt}){
+            stepAndZero(opt);
+        }
+        if (step % 16 == 0) stepAndZero(featuresRestOpt);
+    }else if (step <= lateStart){
+        if (step % 32 == 0){
+            for (torch::optim::Adam *opt : {meansOpt, scalesOpt, quatsOpt, featuresDcOpt, opacitiesOpt, featuresRestOpt}){
+                stepAndZero(opt);
+            }
+        }
+    }else{
+        if (step % 64 == 0){
+            for (torch::optim::Adam *opt : {meansOpt, scalesOpt, quatsOpt, featuresDcOpt, opacitiesOpt, featuresRestOpt}){
+                stepAndZero(opt);
+            }
+        }
+    }
 }
 
 int Model::getDownscaleFactor(int step){
@@ -316,191 +341,321 @@ void Model::removeFromOptimizer(torch::optim::Adam *optimizer, const torch::Tens
     optimizer->state()[newPId] = std::move(paramState);
 }
 
-void Model::afterTrain(int step){
-    torch::NoGradGuard noGrad;
+void Model::zeroOptimizerRows(torch::optim::Adam *optimizer, const torch::Tensor &idcs){
+    if (idcs.numel() == 0) return;
+    torch::Tensor param = optimizer->param_groups()[0].params()[0];
+#if TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR > 1
+    auto pId = param.unsafeGetTensorImpl();
+#else
+    auto pId = c10::guts::to_string(param.unsafeGetTensorImpl());
+#endif
+    auto it = optimizer->state().find(pId);
+    if (it == optimizer->state().end()) return;
+    auto &s = static_cast<torch::optim::AdamParamState&>(*it->second);
+    s.exp_avg().index_put_({idcs}, torch::zeros_like(s.exp_avg().index({idcs})));
+    s.exp_avg_sq().index_put_({idcs}, torch::zeros_like(s.exp_avg_sq().index({idcs})));
+}
 
-    // When radii.sum() == 0
-    if (!xys.grad().defined()) return;
 
-    if (step < stopSplitAt){
-        torch::Tensor visibleMask = (radii > 0).flatten();
-        
-        torch::Tensor grads = torch::linalg_vector_norm(xys.grad().detach(), 2, { -1 }, false, torch::kFloat32);
-        if (!xysGradNorm.numel()){
-            xysGradNorm = grads;
-            visCounts = torch::ones_like(xysGradNorm);
-        }else{
-            visCounts.index_put_({visibleMask}, visCounts.index({visibleMask}) + 1);
-            xysGradNorm.index_put_({visibleMask}, grads.index({visibleMask}) + xysGradNorm.index({visibleMask}));
+// Inspired by FastGS
+std::tuple<torch::Tensor, torch::Tensor> Model::computeMultiViewScores(int step, bool densify){
+    long long N = means.size(0);
+    auto fOpts = torch::TensorOptions().dtype(torch::kFloat32).device(device);
+    torch::Tensor fullCounts = torch::zeros({N}, fOpts);
+    torch::Tensor fullScore = torch::zeros({N}, fOpts);
+    torch::Tensor edgeScores = torch::zeros({N}, fOpts);
+
+    std::vector<size_t> indices(trainCams->size());
+    std::iota(indices.begin(), indices.end(), 0);
+    static thread_local std::mt19937 rng(42 + step);
+    std::shuffle(indices.begin(), indices.end(), rng);
+    size_t numViews = (std::min)(static_cast<size_t>(numScoreViews), indices.size());
+
+    scoringPass = true;
+    for (size_t v = 0; v < numViews; v++){
+        Camera &cam = (*trainCams)[indices[v]];
+        int ds = getDownscaleFactor(step);
+        torch::Tensor gt = cam.getImageGpu(ds, device);
+
+        errorMap = torch::zeros({gt.size(0), gt.size(1)}, fOpts);
+        densificationInfo = torch::zeros({4, N}, fOpts);
+        xyAbsGrad = torch::empty({0}, fOpts);
+
+        torch::Tensor rgb = forward(cam, step);
+
+        {
+            torch::NoGradGuard noGrad;
+            torch::Tensor l1Map = (rgb.detach() - gt).abs().mean(-1);
+            float lo = l1Map.min().item<float>();
+            float hi = l1Map.max().item<float>();
+            torch::Tensor norm = (l1Map - lo) / (std::max)(hi - lo, 1e-8f);
+            errorMap.copy_((norm > lossThresh).to(torch::kFloat32));
         }
 
-        if (!max2DSize.numel()){
-            max2DSize = torch::zeros_like(radii, torch::kFloat32);
-        }
+        // Zero-gradient backward runs the rasterizer backward, which fills
+        // densificationInfo row 3 with the per-gaussian high-error pixel count
+        (rgb.sum() * 0.0f).backward();
+
+        torch::NoGradGuard noGrad;
+        float photometric = fusedL1SsimLossValue(rgb.detach(), gt, 0.2f).item<float>();
 
-        torch::Tensor newRadii = radii.detach().index({visibleMask});
-        max2DSize.index_put_({visibleMask}, torch::maximum(
-                max2DSize.index({visibleMask}), newRadii / static_cast<float>( (std::max)(lastHeight, lastWidth) )
-            ));
+        torch::Tensor counts = densificationInfo[3] * static_cast<float>(ds * ds);
+        fullCounts += counts;
+        fullScore += photometric * counts;
+        if (edgeGuidance) edgeScores += densificationInfo[2];
     }
+    scoringPass = false;
+    errorMap = torch::empty({0}, fOpts);
+    densificationInfo = torch::empty({0}, fOpts);
 
-    if (step % refineEvery == 0 && step > warmupLength){
-        int resetInterval = resetAlphaEvery * refineEvery;
-        bool doDensification = step < stopSplitAt && step % resetInterval > numCameras + refineEvery;
-        torch::Tensor splitsMask;
-        const float cullAlphaThresh = 0.1f;
-
-        if (doDensification){
-            int numPointsBefore = means.size(0);
-            torch::Tensor avgGradNorm = (xysGradNorm / visCounts) * 0.5f * static_cast<float>( (std::max)(lastWidth, lastHeight) );
-            torch::Tensor highGrads = (avgGradNorm > densifyGradThresh).squeeze();
-
-            // Split gaussians that are too large
-            torch::Tensor splits = (std::get<0>(scales.exp().max(-1)) > densifySizeThresh).squeeze();
-            if (step < stopScreenSizeAt){
-                splits |= (max2DSize > splitScreenSize).squeeze();
+    torch::NoGradGuard noGrad;
+    torch::Tensor importance;
+    if (densify){
+        importance = (fullCounts / static_cast<float>(numViews)).floor();
+        if (edgeGuidance){
+            torch::Tensor pos = edgeScores.index({edgeScores > 0});
+            if (pos.numel() > 0){
+                importance = importance * (1.0f + 0.25f * edgeScores / pos.median().clamp_min(1e-12f));
             }
-
-            splits &= highGrads;
-            const int nSplitSamples = 2;
-            int nSplits = splits.sum().item<int>();
-
-            torch::Tensor centeredSamples = torch::randn({nSplitSamples * nSplits, 3}, device);  // Nx3 of axis-aligned scales
-            torch::Tensor scaledSamples = torch::exp(scales.index({splits}).repeat({nSplitSamples, 1})) * centeredSamples;
-            torch::Tensor qs = quats.index({splits}) / torch::linalg_vector_norm(quats.index({splits}), 2, { -1 }, true, torch::kFloat32);
-            torch::Tensor rots = quatToRotMat(qs.repeat({nSplitSamples, 1}));
-            torch::Tensor rotatedSamples = torch::bmm(rots, scaledSamples.index({"...", None})).squeeze();
-            torch::Tensor splitMeans = rotatedSamples + means.index({splits}).repeat({nSplitSamples, 1});
-            
-            torch::Tensor splitFeaturesDc = featuresDc.index({splits}).repeat({nSplitSamples, 1});
-            torch::Tensor splitFeaturesRest = featuresRest.index({splits}).repeat({nSplitSamples, 1, 1});
-            
-            torch::Tensor splitOpacities = opacities.index({splits}).repeat({nSplitSamples, 1});
-        
-            const float sizeFac = 1.6f;
-            torch::Tensor splitScales = torch::log(torch::exp(scales.index({splits})) / sizeFac).repeat({nSplitSamples, 1});
-            scales.index({splits}) = torch::log(torch::exp(scales.index({splits})) / sizeFac);
-            torch::Tensor splitQuats = quats.index({splits}).repeat({nSplitSamples, 1});
-
-            // Duplicate gaussians that are too small
-            torch::Tensor dups = (std::get<0>(scales.exp().max(-1)) <= densifySizeThresh).squeeze();
-            dups &= highGrads;
-            torch::Tensor dupMeans = means.index({dups});
-            torch::Tensor dupFeaturesDc = featuresDc.index({dups});
-            torch::Tensor dupFeaturesRest = featuresRest.index({dups});
-            torch::Tensor dupOpacities = opacities.index({dups});
-            torch::Tensor dupScales = scales.index({dups});
-            torch::Tensor dupQuats = quats.index({dups});
-
-            means = torch::cat({means.detach(), splitMeans, dupMeans}, 0).requires_grad_();
-            featuresDc = torch::cat({featuresDc.detach(), splitFeaturesDc, dupFeaturesDc}, 0).requires_grad_();
-            featuresRest = torch::cat({featuresRest.detach(), splitFeaturesRest, dupFeaturesRest}, 0).requires_grad_();
-            opacities = torch::cat({opacities.detach(), splitOpacities, dupOpacities}, 0).requires_grad_();
-            scales = torch::cat({scales.detach(), splitScales, dupScales}, 0).requires_grad_();
-            quats = torch::cat({quats.detach(), splitQuats, dupQuats}, 0).requires_grad_();
-            
-            max2DSize = torch::cat({
-                max2DSize,
-                torch::zeros_like(splitScales.index({Slice(), 0})),
-                torch::zeros_like(dupScales.index({Slice(), 0}))
-            }, 0);
-
-            torch::Tensor splitIdcs = torch::where(splits)[0];
-
-            addToOptimizer(meansOpt, means, splitIdcs, nSplitSamples);
-            addToOptimizer(scalesOpt, scales, splitIdcs, nSplitSamples);
-            addToOptimizer(quatsOpt, quats, splitIdcs, nSplitSamples);
-            addToOptimizer(featuresDcOpt, featuresDc, splitIdcs, nSplitSamples);
-            addToOptimizer(featuresRestOpt, featuresRest, splitIdcs, nSplitSamples);
-            addToOptimizer(opacitiesOpt, opacities, splitIdcs, nSplitSamples);
-            
-            torch::Tensor dupIdcs = torch::where(dups)[0];
-            addToOptimizer(meansOpt, means, dupIdcs, 1);
-            addToOptimizer(scalesOpt, scales, dupIdcs, 1);
-            addToOptimizer(quatsOpt, quats, dupIdcs, 1);
-            addToOptimizer(featuresDcOpt, featuresDc, dupIdcs, 1);
-            addToOptimizer(featuresRestOpt, featuresRest, dupIdcs, 1);
-            addToOptimizer(opacitiesOpt, opacities, dupIdcs, 1);
-
-            splitsMask = torch::cat({
-                splits,
-                torch::full({nSplitSamples * splits.sum().item<int>() + dups.sum().item<int>()}, false, torch::TensorOptions().dtype(torch::kBool).device(device))
-            }, 0);
-
-            std::cout << "Added " << (means.size(0) - numPointsBefore) << " gaussians, new count " << means.size(0) << std::endl;
         }
+    }
+    float lo = fullScore.min().item<float>();
+    float hi = fullScore.max().item<float>();
+    torch::Tensor pruningScore = (fullScore - lo) / (std::max)(hi - lo, 1e-8f);
+    return std::make_tuple(importance, pruningScore);
+}
+
+void Model::resetOpacity(float value){
+    torch::NoGradGuard noGrad;
+    float cap = torch::logit(torch::tensor(value)).item<float>();
+    opacities.clamp_max_(cap);
+    torch::Tensor allIdx = torch::arange(opacities.size(0), torch::TensorOptions().dtype(torch::kLong).device(device));
+    zeroOptimizerRows(opacitiesOpt, allIdx);
+}
 
-        if (doDensification){
-            // Cull
-            int numPointsBefore = means.size(0);
+static torch::Tensor spatialSanityMask(const torch::Tensor &means, const torch::Tensor &scales, const torch::Device &device){
+    torch::NoGradGuard noGrad;
+    torch::Tensor mc = means.detach().cpu();
+    long long n = mc.size(0);
+    if (n < 8) return torch::zeros({n}, torch::TensorOptions().dtype(torch::kBool).device(device));
+    torch::Tensor lo = std::get<0>(mc.kthvalue((std::max)(static_cast<long long>(1), static_cast<long long>(0.1 * n)), 0));
+    torch::Tensor hi = std::get<0>(mc.kthvalue((std::min)(n, static_cast<long long>(0.9 * n) + 1), 0));
+    torch::Tensor center = ((lo + hi) / 2.0f).to(device);
+    float maxExtent = (std::max)(((hi - lo) / 2.0f).max().item<float>(), 1e-6f);
+    torch::Tensor escaped = std::get<0>((means.detach() - center).abs().max(-1)) > 100.0f * maxExtent;
+    torch::Tensor exploded = std::get<0>(scales.detach().exp().max(-1)) > 100.0f * maxExtent;
+    return escaped | exploded;
+}
 
-            torch::Tensor culls = (torch::sigmoid(opacities) < cullAlphaThresh).squeeze();
-            if (splitsMask.numel()){
-                culls |= splitsMask;
+void Model::densifyAndPrune(int step, const torch::Tensor &importanceScore, const torch::Tensor &pruningScore){
+    torch::NoGradGuard noGrad;
+    long long numPointsBefore = means.size(0);
+
+    float gradScale = 0.5f * static_cast<float>((std::max)(lastWidth, lastHeight));
+    torch::Tensor grads = (xyzGradAccum / gradDenom.clamp_min(1.0f)) * gradScale;
+    torch::Tensor gradsAbs = (xyzGradAbsAccum / gradDenom.clamp_min(1.0f)) * gradScale;
+    torch::Tensor maxScale = std::get<0>(scales.exp().max(-1));
+    torch::Tensor metricMask = importanceScore > 5.0f;
+
+    torch::Tensor cloneMask = (maxScale <= denseThresh * spatialLrScale) & (grads >= gradThresh) & metricMask;
+    torch::Tensor splitMask = (maxScale > denseThresh * spatialLrScale) & (gradsAbs >= gradAbsThresh) & metricMask;
+
+    if (maxGaussians > 0){
+        long long budget = maxGaussians - numPointsBefore;
+        long long requested = cloneMask.sum().item<int64_t>() + 2 * splitMask.sum().item<int64_t>();
+        if (requested > budget){
+            cloneMask &= importanceScore > (budget > 0 ? 5.0f : 1e30f);
+            if (budget <= 0){
+                splitMask &= torch::zeros_like(splitMask);
             }
+        }
+    }
 
-            if (step > refineEvery * resetAlphaEvery){
-                const float cullScaleThresh = 0.5f; // cull huge gaussians
-                const float cullScreenSize = 0.15f; // % of screen space
-                torch::Tensor huge = std::get<0>(torch::exp(scales).max(-1)) > cullScaleThresh;
-                if (step < stopScreenSizeAt){
-                    huge |= max2DSize > cullScreenSize;
-                }
-                culls |= huge;
-            }
+    // Clone: duplicate small high-error gaussians
+    torch::Tensor cloneIdx = torch::where(cloneMask)[0];
+    if (cloneIdx.numel() > 0){
+        means = torch::cat({means.detach(), means.detach().index({cloneIdx})}, 0).requires_grad_();
+        scales = torch::cat({scales.detach(), scales.detach().index({cloneIdx})}, 0).requires_grad_();
+        quats = torch::cat({quats.detach(), quats.detach().index({cloneIdx})}, 0).requires_grad_();
+        featuresDc = torch::cat({featuresDc.detach(), featuresDc.detach().index({cloneIdx})}, 0).requires_grad_();
+        featuresRest = torch::cat({featuresRest.detach(), featuresRest.detach().index({cloneIdx})}, 0).requires_grad_();
+        opacities = torch::cat({opacities.detach(), opacities.detach().index({cloneIdx})}, 0).requires_grad_();
+
+        addToOptimizer(meansOpt, means, cloneIdx, 1);
+        addToOptimizer(scalesOpt, scales, cloneIdx, 1);
+        addToOptimizer(quatsOpt, quats, cloneIdx, 1);
+        addToOptimizer(featuresDcOpt, featuresDc, cloneIdx, 1);
+        addToOptimizer(featuresRestOpt, featuresRest, cloneIdx, 1);
+        addToOptimizer(opacitiesOpt, opacities, cloneIdx, 1);
+    }
+
+    // Split
+    long long nAfterClone = means.size(0);
+    torch::Tensor splitIdx = torch::where(splitMask)[0];
+    long long nSplits = splitIdx.numel();
+    if (nSplits > 0){
+        const int nSamples = 2;
+        torch::Tensor sampled = torch::randn({nSamples * nSplits, 3}, torch::TensorOptions().dtype(torch::kFloat32).device(device));
+        torch::Tensor stds = scales.detach().index({splitIdx}).exp().repeat({nSamples, 1});
+        torch::Tensor qs = quats.detach().index({splitIdx});
+        qs = qs / qs.norm(2, {-1}, true).clamp_min(1e-12);
+        torch::Tensor rots = quatToRotMat(qs.repeat({nSamples, 1}));
+        torch::Tensor offsets = torch::bmm(rots, (sampled * stds).unsqueeze(-1)).squeeze(-1);
+        torch::Tensor newMeans = offsets + means.detach().index({splitIdx}).repeat({nSamples, 1});
+        torch::Tensor newScales = torch::log(stds / 1.6f);
+        torch::Tensor newQuats = quats.detach().index({splitIdx}).repeat({nSamples, 1});
+        torch::Tensor newFDc = featuresDc.detach().index({splitIdx}).repeat({nSamples, 1});
+        torch::Tensor newFRest = featuresRest.detach().index({splitIdx}).repeat({nSamples, 1, 1});
+        torch::Tensor newOpac = opacities.detach().index({splitIdx}).repeat({nSamples, 1});
+
+        means = torch::cat({means.detach(), newMeans}, 0).requires_grad_();
+        scales = torch::cat({scales.detach(), newScales}, 0).requires_grad_();
+        quats = torch::cat({quats.detach(), newQuats}, 0).requires_grad_();
+        featuresDc = torch::cat({featuresDc.detach(), newFDc}, 0).requires_grad_();
+        featuresRest = torch::cat({featuresRest.detach(), newFRest}, 0).requires_grad_();
+        opacities = torch::cat({opacities.detach(), newOpac}, 0).requires_grad_();
+
+        addToOptimizer(meansOpt, means, splitIdx, nSamples);
+        addToOptimizer(scalesOpt, scales, splitIdx, nSamples);
+        addToOptimizer(quatsOpt, quats, splitIdx, nSamples);
+        addToOptimizer(featuresDcOpt, featuresDc, splitIdx, nSamples);
+        addToOptimizer(featuresRestOpt, featuresRest, splitIdx, nSamples);
+        addToOptimizer(opacitiesOpt, opacities, splitIdx, nSamples);
+    }
+
+    long long N = means.size(0);
+    auto boolOpts = torch::TensorOptions().dtype(torch::kBool).device(device);
+
+    torch::Tensor parentMask = torch::zeros({N}, boolOpts);
+    if (nSplits > 0) parentMask.index_put_({splitIdx}, true);
+    torch::Tensor sanityMask = spatialSanityMask(means, scales, device);
+
+    torch::Tensor pruneMask = (torch::sigmoid(opacities.squeeze(-1)) < 0.005f);
+    if (step > opacityResetInterval){
+        pruneMask |= (std::get<0>(scales.exp().max(-1)) > 0.1f * spatialLrScale);
+    }
+    pruneMask &= ~parentMask;
+
+    long long removeBudget = pruneMask.sum().item<int64_t>() / 2;
+    torch::Tensor finalPrune = parentMask | sanityMask;
+    if (removeBudget > 0){
+        torch::Tensor weights = torch::zeros({N}, torch::TensorOptions().dtype(torch::kFloat32).device(device));
+        long long S = (std::min)(static_cast<long long>(pruningScore.size(0)), N);
+        weights.index_put_({Slice(None, S)}, 1.0f / (1e-6f + 1.0f - pruningScore.index({Slice(None, S)})));
+        torch::Tensor sampledIdx = torch::multinomial(weights.cpu(), removeBudget, false).to(device);
+        torch::Tensor sampledMask = torch::zeros({N}, boolOpts);
+        sampledMask.index_put_({sampledIdx}, true);
+        finalPrune |= (pruneMask & sampledMask);
+    }
+
+    long long cullCount = finalPrune.sum().item<int64_t>();
+    if (cullCount > 0){
+        torch::Tensor keep = ~finalPrune;
+        means = means.index({keep}).detach().requires_grad_();
+        scales = scales.index({keep}).detach().requires_grad_();
+        quats = quats.index({keep}).detach().requires_grad_();
+        featuresDc = featuresDc.index({keep}).detach().requires_grad_();
+        featuresRest = featuresRest.index({keep}).detach().requires_grad_();
+        opacities = opacities.index({keep}).detach().requires_grad_();
+
+        removeFromOptimizer(meansOpt, means, finalPrune);
+        removeFromOptimizer(scalesOpt, scales, finalPrune);
+        removeFromOptimizer(quatsOpt, quats, finalPrune);
+        removeFromOptimizer(featuresDcOpt, featuresDc, finalPrune);
+        removeFromOptimizer(featuresRestOpt, featuresRest, finalPrune);
+        removeFromOptimizer(opacitiesOpt, opacities, finalPrune);
+    }
+
+    // Opacity cap at 0.8
+    float cap = torch::logit(torch::tensor(0.8f)).item<float>();
+    opacities.clamp_max_(cap);
+    torch::Tensor allIdx = torch::arange(opacities.size(0), torch::TensorOptions().dtype(torch::kLong).device(device));
+    zeroOptimizerRows(opacitiesOpt, allIdx);
+
+    // Reset accumulators
+    xyzGradAccum = torch::Tensor();
+    xyzGradAbsAccum = torch::Tensor();
+    gradDenom = torch::Tensor();
+    maxRadii2D = torch::Tensor();
+
+    std::cout << "Densify " << step << ": +clone " << cloneIdx.numel() << " +split " << 2 * nSplits
+              << " -prune " << cullCount << ", total " << means.size(0) << std::endl;
+}
 
-            int cullCount = torch::sum(culls).item<int>();
-            if (cullCount > 0){
-                means = means.index({~culls}).detach().requires_grad_();
-                scales = scales.index({~culls}).detach().requires_grad_();
-                quats = quats.index({~culls}).detach().requires_grad_();
-                featuresDc = featuresDc.index({~culls}).detach().requires_grad_();
-                featuresRest = featuresRest.index({~culls}).detach().requires_grad_();
-                opacities = opacities.index({~culls}).detach().requires_grad_();
-
-                removeFromOptimizer(meansOpt, means, culls);
-                removeFromOptimizer(scalesOpt, scales, culls);
-                removeFromOptimizer(quatsOpt, quats, culls);
-                removeFromOptimizer(featuresDcOpt, featuresDc, culls);
-                removeFromOptimizer(featuresRestOpt, featuresRest, culls);
-                removeFromOptimizer(opacitiesOpt, opacities, culls);
-                
-                std::cout << "Culled " << (numPointsBefore - means.size(0)) << " gaussians, remaining " << means.size(0) << std::endl;
+bool Model::afterTrain(int step){
+    bool restructured = false;
+    long long N = means.size(0);
+    auto fOpts = torch::TensorOptions().dtype(torch::kFloat32).device(device);
+
+    if (step < densifyUntilIter){
+        if (xys.grad().defined()){
+            torch::NoGradGuard noGrad;
+            torch::Tensor visible = (radii > 0).flatten();
+            if (!xyzGradAccum.defined() || xyzGradAccum.size(0) != N){
+                xyzGradAccum = torch::zeros({N}, fOpts);
+                xyzGradAbsAccum = torch::zeros({N}, fOpts);
+                gradDenom = torch::zeros({N}, fOpts);
+                maxRadii2D = torch::zeros({N}, fOpts);
             }
+            torch::Tensor g = torch::linalg_vector_norm(xys.grad().detach(), 2, { -1 }, false, torch::kFloat32);
+            xyzGradAccum += g * visible.to(torch::kFloat32);
+            if (xyAbsGrad.defined() && xyAbsGrad.numel() == 2 * N){
+                torch::Tensor ga = torch::linalg_vector_norm(xyAbsGrad, 2, { -1 }, false, torch::kFloat32);
+                xyzGradAbsAccum += ga * visible.to(torch::kFloat32);
+            }
+            gradDenom += visible.to(torch::kFloat32);
+            maxRadii2D = torch::maximum(maxRadii2D, radii.detach().to(torch::kFloat32) * visible.to(torch::kFloat32));
         }
 
-        if (step < stopSplitAt && step % resetInterval == refineEvery){
-            float resetValue = cullAlphaThresh * 2.0f;
-            opacities = torch::clamp_max(opacities, torch::logit(torch::tensor(resetValue)).item<float>());
-
-            // Reset optimizer
-            torch::Tensor param = opacitiesOpt->param_groups()[0].params()[0];
-            #if TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR > 1
-                auto pId = param.unsafeGetTensorImpl();
-            #else
-                auto pId = c10::guts::to_string(param.unsafeGetTensorImpl());
-            #endif    
-            auto paramState = std::make_unique<torch::optim::AdamParamState>(static_cast<torch::optim::AdamParamState&>(*opacitiesOpt->state()[pId]));
-            paramState->exp_avg(torch::zeros_like(paramState->exp_avg()));
-            paramState->exp_avg_sq(torch::zeros_like(paramState->exp_avg_sq()));
-            std::cout << "Alpha reset" << std::endl;
+        if (step > densifyFromIter && step % densificationInterval == 0 && trainCams != nullptr){
+            auto scores = computeMultiViewScores(step, true);
+            densifyAndPrune(step, std::get<0>(scores), std::get<1>(scores));
+            restructured = true;
         }
 
-        // Clear
-        xysGradNorm = torch::Tensor();
-        visCounts = torch::Tensor();
-        max2DSize = torch::Tensor();
-
-        if (device != torch::kCPU){
-            #ifdef USE_HIP
-                    c10::hip::HIPCachingAllocator::emptyCache();
-            #elif defined(USE_CUDA)
-                    c10::cuda::CUDACachingAllocator::emptyCache();
-            #endif
+        if (step % opacityResetInterval == 0){
+            resetOpacity(0.01f);
+            std::cout << "Opacity reset" << std::endl;
+        }
+    }else if (step % 3000 == 0 && step > densifyUntilIter && step < maxSteps && trainCams != nullptr){
+        // Final pruning
+        auto scores = computeMultiViewScores(step, false);
+        torch::NoGradGuard noGrad;
+        torch::Tensor pruningScore = std::get<1>(scores);
+        torch::Tensor pruneMask = (torch::sigmoid(opacities.squeeze(-1)) < 0.1f) | (pruningScore > 0.9f)
+                                | spatialSanityMask(means, scales, device);
+        long long cullCount = pruneMask.sum().item<int64_t>();
+        if (cullCount > 0 && cullCount < means.size(0)){
+            torch::Tensor keep = ~pruneMask;
+            means = means.index({keep}).detach().requires_grad_();
+            scales = scales.index({keep}).detach().requires_grad_();
+            quats = quats.index({keep}).detach().requires_grad_();
+            featuresDc = featuresDc.index({keep}).detach().requires_grad_();
+            featuresRest = featuresRest.index({keep}).detach().requires_grad_();
+            opacities = opacities.index({keep}).detach().requires_grad_();
+
+            removeFromOptimizer(meansOpt, means, pruneMask);
+            removeFromOptimizer(scalesOpt, scales, pruneMask);
+            removeFromOptimizer(quatsOpt, quats, pruneMask);
+            removeFromOptimizer(featuresDcOpt, featuresDc, pruneMask);
+            removeFromOptimizer(featuresRestOpt, featuresRest, pruneMask);
+            removeFromOptimizer(opacitiesOpt, opacities, pruneMask);
+            std::cout << "Final prune " << step << ": -" << cullCount << ", remaining " << means.size(0) << std::endl;
         }
+        restructured = true;
+    }
+
+    if (restructured && device != torch::kCPU){
+        #ifdef USE_HIP
+                c10::hip::HIPCachingAllocator::emptyCache();
+        #elif defined(USE_CUDA)
+                c10::cuda::CUDACachingAllocator::emptyCache();
+        #endif
     }
+    return restructured;
 }
 
+
 void Model::save(const std::string &filename, int step){
     std::string extension = fs::path(filename).extension().string(); 
     if (extension == ".splat"){
@@ -865,8 +1020,24 @@ int Model::loadPly(const std::string &filename){
     throw std::runtime_error("Invalid PLY file");
 }
 
-torch::Tensor Model::mainLoss(torch::Tensor &rgb, torch::Tensor &gt, float ssimWeight){
-    torch::Tensor ssimLoss = 1.0f - ssim.eval(rgb, gt);
-    torch::Tensor l1Loss = l1(rgb, gt);
-    return (1.0f - ssimWeight) * l1Loss + ssimWeight * ssimLoss;
-}
+torch::Tensor Model::mainLoss(torch::Tensor &rgb, torch::Tensor &gt, torch::Tensor &mask, float ssimWeight){
+    bool hasMask = mask.defined() && mask.numel() > 0;
+    torch::Tensor loss;
+
+    if (ssimWeight > 0.0f){
+        torch::Tensor m = hasMask ? mask : torch::empty({0}, rgb.options());
+        loss = fusedL1SsimLoss(rgb, gt, m, ssimWeight, !hasMask);
+    }else{
+        torch::Tensor absDiff = torch::abs(gt - rgb);
+        loss = hasMask
+            ? (mask.unsqueeze(-1) * absDiff).sum() / (mask.sum() * gt.size(2) + 1e-8f)
+            : absDiff.sum() / (static_cast<float>(gt.numel()) + 1e-8f);
+    }
+
+    // Segment-mode opacity penalty: push alpha to 0 in masked areas
+    if (hasMask && lastAlpha.defined() && lastAlpha.numel() == mask.numel()){
+        loss = loss + (lastAlpha * (1.0f - mask).pow(2.0f)).mean();
+    }
+
+    return loss;
+}
\ No newline at end of file
diff --git a/model.hpp b/model.hpp
index ecb33e2b..415d3191 100644
--- a/model.hpp
+++ b/model.hpp
@@ -9,27 +9,26 @@
 #include "spherical_harmonics.hpp"
 #include "ssim.hpp"
 #include "input_data.hpp"
-#include "optim_scheduler.hpp"
 
 using namespace torch::indexing;
 using namespace torch::autograd;
 
-torch::Tensor randomQuatTensor(long long n);
+torch::Tensor identityQuatTensor(long long n);
 torch::Tensor projectionMatrix(float zNear, float zFar, float fovX, float fovY, const torch::Device &device);
-torch::Tensor psnr(const torch::Tensor& rendered, const torch::Tensor& gt);
-torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt);
 
 struct Model{
   Model(const InputData &inputData, int numCameras,
-        int numDownscales, int resolutionSchedule, int shDegree, int shDegreeInterval, 
-        int refineEvery, int warmupLength, int resetAlphaEvery, float densifyGradThresh, float densifySizeThresh, int stopScreenSizeAt, float splitScreenSize,
+        int numDownscales, int resolutionSchedule, int shDegree, int shDegreeInterval,
+        int densificationInterval, int densifyFromIter, int densifyUntilIter, int maxGaussians,
+        float lossThresh,
         int maxSteps, bool keepCrs,
         const torch::Device &device) :
     numCameras(numCameras),
-    numDownscales(numDownscales), resolutionSchedule(resolutionSchedule), shDegree(shDegree), shDegreeInterval(shDegreeInterval), 
-    refineEvery(refineEvery), warmupLength(warmupLength), resetAlphaEvery(resetAlphaEvery), stopSplitAt(maxSteps / 2), densifyGradThresh(densifyGradThresh), densifySizeThresh(densifySizeThresh), stopScreenSizeAt(stopScreenSizeAt), splitScreenSize(splitScreenSize),
+    numDownscales(numDownscales), resolutionSchedule(resolutionSchedule), shDegree(shDegree), shDegreeInterval(shDegreeInterval),
+    densificationInterval(densificationInterval), densifyFromIter(densifyFromIter), densifyUntilIter(densifyUntilIter), maxGaussians(maxGaussians),
+    lossThresh(lossThresh),
     maxSteps(maxSteps), keepCrs(keepCrs),
-    device(device), ssim(11, 3){
+    device(device){
 
     long long numPoints = inputData.points.xyz.size(0);
     scale = inputData.scale;
@@ -39,7 +38,7 @@ struct Model{
 
     means = inputData.points.xyz.to(device).requires_grad_();
     scales = PointsTensor(inputData.points.xyz).scales().repeat({1, 3}).log().to(device).requires_grad_();
-    quats = randomQuatTensor(numPoints).to(device).requires_grad_();
+    quats = identityQuatTensor(numPoints).to(device).requires_grad_();
 
     int dimSh = numShBases(shDegree);
     torch::Tensor shs = torch::zeros({numPoints, dimSh, 3}, torch::TensorOptions().dtype(torch::kFloat32).device(device));
@@ -50,8 +49,20 @@ struct Model{
     featuresDc = shs.index({Slice(), 0, Slice()}).to(device).requires_grad_();
     featuresRest = shs.index({Slice(), Slice(1, None), Slice()}).to(device).requires_grad_();
     opacities = torch::logit(0.1f * torch::ones({numPoints, 1})).to(device).requires_grad_();
-    
-    backgroundColor = torch::tensor({0.6130f, 0.0101f, 0.3984f}, device).requires_grad_(); // Nerf Studio default
+
+    backgroundColor = torch::zeros({3}, device);
+
+    // Scene extent from camera positions (vanilla 3DGS cameras_extent)
+    spatialLrScale = 1.0f;
+    if (!inputData.cameras.empty()){
+        torch::Tensor centers = torch::zeros({static_cast<long long>(inputData.cameras.size()), 3});
+        for (size_t i = 0; i < inputData.cameras.size(); i++){
+            centers[i] = inputData.cameras[i].camToWorld.index({Slice(None, 3), 3});
+        }
+        torch::Tensor avg = centers.mean(0, true);
+        spatialLrScale = (centers - avg).norm(2, 1).max().item<float>() * 1.1f;
+        if (spatialLrScale <= 0.0f) spatialLrScale = 1.0f;
+    }
 
     setupOptimizers();
   }
@@ -64,11 +75,14 @@ struct Model{
   void releaseOptimizers();
 
   torch::Tensor forward(Camera& cam, int step);
-  void optimizersZeroGrad();
-  void optimizersStep();
+  void optimizerStepCadence(int step); // FastGS stepping schedule with gradient accumulation
   void schedulersStep(int step);
   int getDownscaleFactor(int step);
-  void afterTrain(int step);
+  bool afterTrain(int step); // returns true if parameters were restructured
+  std::tuple<torch::Tensor, torch::Tensor> computeMultiViewScores(int step, bool densify);
+  void densifyAndPrune(int step, const torch::Tensor &importanceScore, const torch::Tensor &pruningScore);
+  void resetOpacity(float value);
+  void zeroOptimizerRows(torch::optim::Adam *optimizer, const torch::Tensor &idcs);
   void save(const std::string &filename, int step);
   void savePly(const std::string &filename, int step);
   void saveSplat(const std::string &filename);
@@ -76,7 +90,7 @@ struct Model{
   bool saveRad(const std::string &filename);
   void saveDebugPly(const std::string &filename, int step);
   int loadPly(const std::string &filename);
-  torch::Tensor mainLoss(torch::Tensor &rgb, torch::Tensor &gt, float ssimWeight);
+  torch::Tensor mainLoss(torch::Tensor &rgb, torch::Tensor &gt, torch::Tensor &mask, float ssimWeight);
 
   void addToOptimizer(torch::optim::Adam *optimizer, const torch::Tensor &newParam, const torch::Tensor &idcs, int nSamples);
   void removeFromOptimizer(torch::optim::Adam *optimizer, const torch::Tensor &newParam, const torch::Tensor &deletedMask);
@@ -94,38 +108,49 @@ struct Model{
   torch::optim::Adam *featuresRestOpt = nullptr;
   torch::optim::Adam *opacitiesOpt = nullptr;
 
-  OptimScheduler *meansOptScheduler = nullptr;
+  float spatialLrScale = 1.0f;
+  std::vector<Camera> *trainCams = nullptr; // set by the trainer, used for multi-view scoring
 
   torch::Tensor radii; // set in forward()
   torch::Tensor xys; // set in forward()
+  torch::Tensor lastAlpha; // set in forward()
+  torch::Tensor errorMap; // [H,W] binary metric map for scoring passes, read by rasterize backward
+  torch::Tensor densificationInfo; // [4,N] accumulated by rasterize backward
+  torch::Tensor xyAbsGrad; // [N,2] Abs-GS screen-gradient accumulation, filled by rasterize backward
   int lastHeight; // set in forward()
   int lastWidth; // set in forward()
 
-  torch::Tensor xysGradNorm; // set in afterTrain()
-  torch::Tensor visCounts; // set in afterTrain()  
-  torch::Tensor max2DSize; // set in afterTrain()
+  bool scoringPass = false; // true while computeMultiViewScores drives forward/backward
+  torch::Tensor xyzGradAccum; // [N] accumulated ||d mean2d||
+  torch::Tensor xyzGradAbsAccum; // [N] accumulated ||d mean2d|| (absolute, Abs-GS)
+  torch::Tensor gradDenom; // [N] visibility counts
+  torch::Tensor maxRadii2D; // [N] max screen radius in px
 
 
   torch::Tensor backgroundColor;
   torch::Device device;
-  SSIM ssim;
 
   int numCameras;
   int numDownscales;
   int resolutionSchedule;
   int shDegree;
   int shDegreeInterval;
-  int refineEvery;
-  int warmupLength;
-  int resetAlphaEvery;
-  int stopSplitAt;
-  float densifyGradThresh;
-  float densifySizeThresh;
-  int stopScreenSizeAt;
-  float splitScreenSize;
+  int densificationInterval;
+  int densifyFromIter;
+  int densifyUntilIter;
+  int maxGaussians;
+  float lossThresh;
   int maxSteps;
   bool keepCrs;
 
+  // FastGS hyperparameters (paper defaults)
+  float denseThresh = 0.001f;
+  float gradThresh = 0.0002f;
+  float gradAbsThresh = 0.0012f;
+  int opacityResetInterval = 3000;
+  int numScoreViews = 10;
+  bool edgeGuidance = true; // Canny-edge weighting of the densification importance
+
   float scale;
   torch::Tensor translation;
 };
diff --git a/openmvg.hpp b/openmvg.hpp
index db0afde5..94caf356 100644
--- a/openmvg.hpp
+++ b/openmvg.hpp
@@ -43,7 +43,7 @@ namespace omvg{
       // image size
       uint32_t ui_width, ui_height;
     };
-    bool read_views(const json& data, std::unordered_map<uint32_t, Intrinsic> &views);
+    bool read_views(const json& data, std::unordered_map<uint32_t, View> &views);
 
     
     struct Pose{
@@ -52,32 +52,6 @@ namespace omvg{
     };
     bool read_poses(const json& data, std::unordered_map<uint32_t, Pose> &poses);
 
-    /*
-    struct Landmark{
-        uint32_t id_view;
-        Eigen::Vector3d location_3d;
-    }
-
-    struct Observation{
-        uint32_t id_feat;
-        // two location in the image
-        Eigen::Vector2d x;
-    }
-    */
-
-
-    //std::unordered_map<Key, Value>;
-
-
-/*
-    struct Pose{
-        std::vector<float> rotation = {0.0f, 0.0f, 0.0f};
-        std::vector<float> translation = {0.0f, 0.0f, 0.0f};
-        std::string camera = "";
-    };
-    void from_json(const json& j, Shot &s);
-*/
-
     InputData inputDataFromOpenMVG(const std::string &projectRoot);
 }   
 
diff --git a/opensplat.cpp b/opensplat.cpp
index e56f44c9..5a80aa77 100644
--- a/opensplat.cpp
+++ b/opensplat.cpp
@@ -25,23 +25,24 @@ int main(int argc, char *argv[]){
         ("val", "Withhold a camera shot for validating the scene loss")
         ("val-image", "Filename of the image to withhold for validating scene loss", cxxopts::value<std::string>()->default_value("random"))
         ("val-render", "Path of the directory where to render validation images", cxxopts::value<std::string>()->default_value(""))
-        ("keep-crs", "Retain the project input's coordinate reference system")
+        ("center", "Center the model at the origin")
         ("cpu", "Force CPU execution")
         
         ("n,num-iters", "Number of iterations to run", cxxopts::value<int>()->default_value("30000"))
         ("d,downscale-factor", "Scale input images by this factor.", cxxopts::value<float>()->default_value("1"))
-        ("num-downscales", "Number of images downscales to use. After being scaled by [downscale-factor], images are initially scaled by a further (2^[num-downscales]) and the scale is increased every [resolution-schedule]", cxxopts::value<int>()->default_value("2"))
+        ("num-downscales", "Number of images downscales to use. After being scaled by [downscale-factor], images are initially scaled by a further (2^[num-downscales]) and the scale is increased every [resolution-schedule]", cxxopts::value<int>()->default_value("0"))
         ("resolution-schedule", "Double the image resolution every these many steps", cxxopts::value<int>()->default_value("3000"))
         ("sh-degree", "Maximum spherical harmonics degree (must be > 0)", cxxopts::value<int>()->default_value("3"))
         ("sh-degree-interval", "Increase the number of spherical harmonics degree after these many steps (will not exceed [sh-degree])", cxxopts::value<int>()->default_value("1000"))
         ("ssim-weight", "Weight to apply to the structural similarity loss. Set to zero to use least absolute deviation (L1) loss only", cxxopts::value<float>()->default_value("0.2"))
-        ("refine-every", "Split/duplicate/prune gaussians every these many steps", cxxopts::value<int>()->default_value("100"))
-        ("warmup-length", "Split/duplicate/prune gaussians only after these many steps", cxxopts::value<int>()->default_value("500"))
-        ("reset-alpha-every", "Reset the opacity values of gaussians after these many refinements (not steps)", cxxopts::value<int>()->default_value("30"))
-        ("densify-grad-thresh", "Threshold of the positional gradient norm (magnitude of the loss function) which when exceeded leads to a gaussian split/duplication", cxxopts::value<float>()->default_value("0.0002"))
-        ("densify-size-thresh", "Gaussians' scales below this threshold are duplicated, otherwise split", cxxopts::value<float>()->default_value("0.01"))
-        ("stop-screen-size-at", "Stop splitting gaussians that are larger than [split-screen-size] after these many steps", cxxopts::value<int>()->default_value("4000"))
-        ("split-screen-size", "Split gaussians that are larger than this percentage of screen space", cxxopts::value<float>()->default_value("0.05"))
+        ("refine-every", "Densify/prune gaussians every these many steps", cxxopts::value<int>()->default_value("500"))
+        ("densify-from", "Start densifying gaussians after these many steps", cxxopts::value<int>()->default_value("500"))
+        ("densify-until", "Stop densifying gaussians after these many steps (-1 = min(15000, half of num-iters))", cxxopts::value<int>()->default_value("-1"))
+        ("loss-thresh", "High-error pixel threshold on the normalized L1 map for multi-view scoring", cxxopts::value<float>()->default_value("0.1"))
+        ("no-edge-guidance", "Disable Canny edge weighting of the densification importance", cxxopts::value<bool>()->default_value("false"))
+        ("max-gaussians", "Maximum number of gaussians (0 = unlimited)", cxxopts::value<int>()->default_value("5000000"))
+        ("no-masks", "Ignore image masks even when present", cxxopts::value<bool>()->default_value("false"))
+        ("no-gpu-cache", "Do not cache images/masks on the GPU (reduces VRAM usage, slower)", cxxopts::value<bool>()->default_value("false"))
 #ifdef USE_VISUALIZATION
         ("has-visualization", "Show the visualization steps of training", cxxopts::value<bool>()->default_value("0"))
 #endif
@@ -78,7 +79,7 @@ int main(int argc, char *argv[]){
     const std::string valImage = result["val-image"].as<std::string>();
     const std::string valRender = result["val-render"].as<std::string>();
     if (!valRender.empty() && !fs::exists(valRender)) fs::create_directories(valRender);
-    const bool keepCrs = result.count("keep-crs") > 0;
+    const bool keepCrs = result.count("center") == 0;
     const float downScaleFactor = (std::max)(result["downscale-factor"].as<float>(), 1.0f);
     const int numIters = result["num-iters"].as<int>();
     const int numDownscales = result["num-downscales"].as<int>();
@@ -87,12 +88,12 @@ int main(int argc, char *argv[]){
     const int shDegreeInterval = result["sh-degree-interval"].as<int>();
     const float ssimWeight = result["ssim-weight"].as<float>();
     const int refineEvery = result["refine-every"].as<int>();
-    const int warmupLength = result["warmup-length"].as<int>();
-    const int resetAlphaEvery = result["reset-alpha-every"].as<int>();
-    const float densifyGradThresh = result["densify-grad-thresh"].as<float>();
-    const float densifySizeThresh = result["densify-size-thresh"].as<float>();
-    const int stopScreenSizeAt = result["stop-screen-size-at"].as<int>();
-    const float splitScreenSize = result["split-screen-size"].as<float>();
+    const int densifyFrom = result["densify-from"].as<int>();
+    int densifyUntil = result["densify-until"].as<int>();
+    if (densifyUntil < 0) densifyUntil = (std::min)(15000, result["num-iters"].as<int>() / 2);
+    const float lossThresh = result["loss-thresh"].as<float>();
+    const int maxGaussians = result["max-gaussians"].as<int>();
+    const bool noMasks = result["no-masks"].as<bool>();
     #ifdef USE_VISUALIZATION
         const bool hasVisualization = result["has-visualization"].as<bool>();
     #endif
@@ -122,6 +123,15 @@ int main(int argc, char *argv[]){
         if (isZipArchive(projectRoot)) projectPath = extractZipToCache(projectRoot);
         InputData inputData = inputDataFromX(projectPath);
 
+        int numMasks = 0;
+        if (!noMasks){
+            for (Camera &cam : inputData.cameras){
+                cam.maskPath = findMaskPath(cam.filePath, projectPath);
+                if (!cam.maskPath.empty()) numMasks++;
+            }
+        }
+        if (numMasks > 0) std::cout << "Found " << numMasks << " masks" << std::endl;
+
         parallel_for(inputData.cameras.begin(), inputData.cameras.end(), [&downScaleFactor](Camera &cam){
             cam.loadImage(downScaleFactor);
         });
@@ -133,10 +143,14 @@ int main(int argc, char *argv[]){
 
         Model model(inputData,
                     cams.size(),
-                    numDownscales, resolutionSchedule, shDegree, shDegreeInterval, 
-                    refineEvery, warmupLength, resetAlphaEvery, densifyGradThresh, densifySizeThresh, stopScreenSizeAt, splitScreenSize,
+                    numDownscales, resolutionSchedule, shDegree, shDegreeInterval,
+                    refineEvery, densifyFrom, densifyUntil, maxGaussians,
+                    lossThresh,
                     numIters, keepCrs,
                     device);
+        model.trainCams = &cams;
+        model.edgeGuidance = !result["no-edge-guidance"].as<bool>();
+        Camera::gpuCacheEnabled = !result["no-gpu-cache"].as<bool>();
 
         std::vector< size_t > camIndices( cams.size() );
         std::iota( camIndices.begin(), camIndices.end(), 0 );
@@ -152,23 +166,21 @@ int main(int argc, char *argv[]){
         for (; step <= numIters; step++){
             Camera& cam = cams[ camsIter.next() ];
 
-            model.optimizersZeroGrad();
-
             torch::Tensor rgb = model.forward(cam, step);
-            torch::Tensor gt = cam.getImage(model.getDownscaleFactor(step));
-            gt = gt.to(device);
+            torch::Tensor gt = cam.getImageGpu(model.getDownscaleFactor(step), device);
+            torch::Tensor mask = cam.getMaskGpu(model.getDownscaleFactor(step), device);
 
-            torch::Tensor mainLoss = model.mainLoss(rgb, gt, ssimWeight);
+            torch::Tensor mainLoss = model.mainLoss(rgb, gt, mask, ssimWeight);
             mainLoss.backward();
-            
+
             if (step % displayStep == 0) {
                 const float percentage = static_cast<float>(step) / numIters;
                 std::cout << "Step " << step << ": " << mainLoss.item<float>() << " [" << floor(percentage * 100) << "%]" <<  std::endl;
             }
 
-            model.optimizersStep();
-            model.schedulersStep(step);
             model.afterTrain(step);
+            model.optimizerStepCadence(step);
+            model.schedulersStep(step);
 
             if (saveEvery > 0 && step % saveEvery == 0){
                 fs::path p(outputScene);
@@ -203,8 +215,17 @@ int main(int argc, char *argv[]){
         // Validate
         if (valCam != nullptr){
             torch::Tensor rgb = model.forward(*valCam, numIters);
-            torch::Tensor gt = valCam->getImage(model.getDownscaleFactor(numIters)).to(device);
-            std::cout << valCam->filePath << " validation loss: " << model.mainLoss(rgb, gt, ssimWeight).item<float>() << std::endl; 
+            torch::Tensor gt = valCam->getImageGpu(model.getDownscaleFactor(numIters), device);
+            torch::Tensor valMask = valCam->getMaskGpu(model.getDownscaleFactor(numIters), device);
+            std::cout << valCam->filePath << " validation loss: " << model.mainLoss(rgb, gt, valMask, ssimWeight).item<float>() << std::endl;
+
+            torch::Tensor mse;
+            if (valMask.defined() && valMask.numel() > 0){
+                mse = (valMask.unsqueeze(-1) * (rgb - gt).pow(2)).sum() / (valMask.sum() * gt.size(2) + 1e-8f);
+            }else{
+                mse = (rgb - gt).pow(2).mean();
+            }
+            std::cout << valCam->filePath << " validation PSNR: " << (10.0f * torch::log10(1.0f / mse)).item<float>() << std::endl;
         }
     }catch(const std::exception &e){
         std::cerr << e.what() << std::endl;
diff --git a/optim_scheduler.cpp b/optim_scheduler.cpp
deleted file mode 100644
index 59850cef..00000000
--- a/optim_scheduler.cpp
+++ /dev/null
@@ -1,12 +0,0 @@
-#include "optim_scheduler.hpp"
-
-
-float OptimScheduler::getLearningRate(int step){
-    float t = (std::max)((std::min)(static_cast<float>(step) / static_cast<float>(maxSteps), 1.0f), 0.0f);
-    return std::exp(std::log(lrInit) * (1.0f - t) + std::log(lrFinal) * t);
-}
-
-void OptimScheduler::step(int step){
-    float lr = getLearningRate(step);
-    static_cast<torch::optim::AdamOptions&>(opt->param_groups()[0].options()).set_lr(lr);
-}
\ No newline at end of file
diff --git a/optim_scheduler.hpp b/optim_scheduler.hpp
deleted file mode 100644
index 076981b9..00000000
--- a/optim_scheduler.hpp
+++ /dev/null
@@ -1,23 +0,0 @@
-#ifndef OPTIM_SCHEDULER
-#define OPTIM_SCHEDULER
-
-#include <iostream>
-#include <torch/torch.h>
-
-class OptimScheduler{
-public:
-    OptimScheduler(torch::optim::Adam *opt, float lrFinal, int maxSteps) :
-        opt(opt), lrInit(
-            static_cast<torch::optim::AdamOptions&>(opt->param_groups()[0].options()).get_lr()
-        ), lrFinal(lrFinal), maxSteps(maxSteps) {};
-    void step(int step);
-    float getLearningRate(int step);
-
-private:
-    torch::optim::Adam *opt;
-    float lrInit;
-    float lrFinal;
-    int maxSteps;
-};
-
-#endif
\ No newline at end of file
diff --git a/point_io.cpp b/point_io.cpp
index bf918249..49d49ba6 100644
--- a/point_io.cpp
+++ b/point_io.cpp
@@ -1,65 +1,9 @@
-#include <random>
 #include <filesystem>
 
 #include "point_io.hpp"
-#include "model.hpp"
 
 namespace fs = std::filesystem;
 
-double PointSet::spacing(int kNeighbors) {
-    if (m_spacing != -1) return m_spacing;
-
-    const auto index = getIndex<KdTree>();
-
-    const size_t np = count();
-    const size_t SAMPLES = std::min<size_t>(np, 10000);
-    const int count = kNeighbors + 1;
-
-    std::unordered_map<size_t, size_t> dist_map;
-
-    std::random_device rd;
-    std::mt19937_64 gen(rd());
-    std::uniform_int_distribution<size_t> randomDis(
-        std::numeric_limits<size_t>::min(),
-        np - 1
-    );
-
-    std::vector<size_t> indices(count);
-    std::vector<float> sqr_dists(count);
-
-    for (size_t i = 0; i < SAMPLES; ++i) {
-        const size_t idx = randomDis(gen);
-        index->knnSearch(points[idx].data(), count, indices.data(), sqr_dists.data());
-
-        float sum = 0.0;
-        for (size_t j = 1; j < kNeighbors; ++j) {
-            sum += std::sqrt(sqr_dists[j]);
-        }
-        sum /= static_cast<float>(kNeighbors);
-
-        auto k = static_cast<size_t>(std::ceil(sum * 100));
-
-        if (dist_map.find(k) == dist_map.end()) {
-            dist_map[k] = 1;
-        }
-        else {
-            dist_map[k] += 1;
-        }
-    }
-
-    size_t max_val = std::numeric_limits<size_t>::min();
-    size_t d = 0;
-    for (const auto it : dist_map) {
-        if (it.second > max_val) {
-            d = it.first;
-            max_val = it.second;
-        }
-    }
-
-    m_spacing = std::max(0.01, static_cast<double>(d) / 100.0);
-    return m_spacing;
-}
-
 std::string getVertexLine(std::ifstream &reader) {
     std::string line;
 
@@ -407,102 +351,3 @@ bool hasHeader(const std::string &line, const std::string &prop) {
     //std::cout << line << " -> " << prop << " : " << line.substr(line.length() - prop.length(), prop.length()) << std::endl;
     return line.substr(0, 8) == "property" && line.substr(line.length() - prop.length(), prop.length()) == prop;
 }
-
-void savePointSet(PointSet &pSet, const std::string &filename) {
-    const fs::path p(filename);
-    if (p.extension().string() == ".ply") fastPlySavePointSet(pSet, filename);
-    else pdalSavePointSet(pSet, filename);
-}
-
-void pdalSavePointSet(PointSet &pSet, const std::string &filename) {
-    #ifdef WITH_PDAL
-    pdal::StageFactory factory;
-    const std::string driver = pdal::StageFactory::inferWriterDriver(filename);
-    if (driver.empty()) {
-        throw std::runtime_error("Can't infer point cloud writer from " + filename);
-    }
-
-    // Sync position, color data
-    if (pSet.pointView == nullptr) throw std::runtime_error("pointView is null (should not have happened)");
-    const pdal::PointViewPtr pView = pSet.pointView;
-
-    for (pdal::PointId i = 0; i < pSet.count(); i++) {
-        if (pSet.hasColors()) {
-            pView->setField(pdal::Dimension::Id::Red, i, pSet.colors[i][0]);
-            pView->setField(pdal::Dimension::Id::Green, i, pSet.colors[i][1]);
-            pView->setField(pdal::Dimension::Id::Blue, i, pSet.colors[i][2]);
-        }
-    }
-
-    pdal::PointTable table;
-    pdal::BufferReader reader;
-    reader.addView(pView);
-
-    for (const auto d : pView->dims()) {
-        table.layout()->registerOrAssignDim(pView->dimName(d), pView->dimType(d));
-    }
-
-    pdal::Stage *s = factory.createStage(driver);
-    pdal::Options opts;
-    opts.add("filename", filename);
-    s->setOptions(opts);
-    s->setInput(reader);
-
-    s->prepare(table);
-    s->execute(table);
-
-    std::cout << "Wrote " << filename << std::endl;
-    #else
-    fs::path p(filename);
-    throw std::runtime_error("Unsupported file extension " + p.extension().string() + ", build program with PDAL support for additional file types support.");
-    #endif
-}
-
-void fastPlySavePointSet(PointSet &pSet, const std::string &filename) {
-    std::ofstream o(filename, std::ios::binary);
-
-    o << "ply" << std::endl;
-    o << "format binary_little_endian 1.0" << std::endl;
-    o << "comment Generated by OpenSplat" << std::endl;
-    o << "element vertex " << pSet.count() << std::endl;
-    o << "property float x" << std::endl;
-    o << "property float y" << std::endl;
-    o << "property float z" << std::endl;
-
-    const bool hasNormals = pSet.hasNormals();
-    const bool hasColors = pSet.hasColors();
-    const bool hasViews = pSet.hasViews();
-
-    if (hasNormals) {
-        o << "property float nx" << std::endl;
-        o << "property float ny" << std::endl;
-        o << "property float nz" << std::endl;
-    }
-    if (hasColors) {
-        o << "property uchar red" << std::endl;
-        o << "property uchar green" << std::endl;
-        o << "property uchar blue" << std::endl;
-    }
-    if (hasViews) {
-        o << "property uchar views" << std::endl;
-    }
-
-    o << "end_header" << std::endl;
-
-    for (size_t i = 0; i < pSet.count(); i++) {
-        o.write(reinterpret_cast<const char *>(pSet.points[i].data()), sizeof(float) * 3);
-        if (hasNormals) o.write(reinterpret_cast<const char *>(pSet.normals[i].data()), sizeof(float) * 3);
-        if (hasColors) o.write(reinterpret_cast<const char *>(pSet.colors[i].data()), sizeof(uint8_t) * 3);
-        if (hasViews) o.write(reinterpret_cast<const char *>(&pSet.views[i]), sizeof(uint8_t));
-    }
-
-    o.close();
-    std::cout << "Wrote " << filename << std::endl;
-}
-
-bool fileExists(const std::string &path) {
-    std::ifstream fin(path);
-    const bool e = fin.good();
-    fin.close();
-    return e;
-}
diff --git a/point_io.hpp b/point_io.hpp
index 3982a006..d5e89707 100644
--- a/point_io.hpp
+++ b/point_io.hpp
@@ -14,12 +14,6 @@
 
 #include <nanoflann.hpp>
 
-struct XYZ {
-    float x;
-    float y;
-    float z;
-};
-
 #define KDTREE_MAX_LEAF 10
 
 #define RELEASE_POINTSET(__POINTER) { if (__POINTER != nullptr) { __POINTER->freeIndex<KdTree>(); delete __POINTER; __POINTER = nullptr; } }
@@ -59,17 +53,6 @@ struct PointSet {
         return false;
     }
 
-    void appendPoint(PointSet &src, size_t idx) {
-        points.push_back(src.points[idx]);
-        colors.push_back(src.colors[idx]);
-    }
-
-    bool hasNormals() const { return normals.size() > 0; }
-    bool hasColors() const { return colors.size() > 0; }
-    bool hasViews() const { return views.size() > 0; }
-
-    double spacing(int kNeighbors = 3);
-
     template <typename T>
     void freeIndex() {
         if (kdTree != nullptr) {
@@ -89,8 +72,6 @@ struct PointSet {
 
     ~PointSet() {
     }
-private:
-    double m_spacing = -1.0;
 };
 
 using KdTree = nanoflann::KDTreeSingleIndexAdaptor<
@@ -115,11 +96,5 @@ PointSet *pdalReadPointSet(const std::string &filename);
 PointSet *colmapReadPointSet(const std::string &filename);
 PointSet *readPointSet(const std::string &filename);
 
-void fastPlySavePointSet(PointSet &pSet, const std::string &filename);
-void pdalSavePointSet(PointSet &pSet, const std::string &filename);
-void savePointSet(PointSet &pSet, const std::string &filename);
-
-bool fileExists(const std::string &path);
-
 
 #endif
diff --git a/rad.cpp b/rad.cpp
index 0d7d6f71..a94d518b 100644
--- a/rad.cpp
+++ b/rad.cpp
@@ -132,7 +132,6 @@ struct F16 {
 
     F16() = default;
     static F16 fromF32(float v){ F16 h; h.bits = f16FromF32(v); return h; }
-    static F16 fromBits(uint16_t b){ F16 h; h.bits = b; return h; }
     float toF32() const { return f16ToF32(bits); }
     bool isNan() const { return (bits & 0x7C00u) == 0x7C00u && (bits & 0x03FFu) != 0; }
 };
@@ -218,7 +217,6 @@ struct Vec3 {
     float dot(const Vec3 &o) const { return (x * o.x + y * o.y) + z * o.z; }
     float lengthSquared() const { return dot(*this); }
     float length() const { return std::sqrt(dot(*this)); }
-    float distance(const Vec3 &o) const { return (*this - o).length(); }
 
     bool isFinite() const { return std::isfinite(x) && std::isfinite(y) && std::isfinite(z); }
 };
@@ -231,13 +229,6 @@ struct I64Vec3 {
 
     int64_t operator[](int i) const { return i == 0 ? x : (i == 1 ? y : z); }
     bool operator==(const I64Vec3 &o) const { return x == o.x && y == o.y && z == o.z; }
-
-    I64Vec3 minv(const I64Vec3 &o) const {
-        return I64Vec3(std::min(x, o.x), std::min(y, o.y), std::min(z, o.z));
-    }
-    I64Vec3 maxv(const I64Vec3 &o) const {
-        return I64Vec3(std::max(x, o.x), std::max(y, o.y), std::max(z, o.z));
-    }
 };
 
 struct I64Vec3Hash {
@@ -731,7 +722,6 @@ struct GsplatArray {
     size_t len() const { return splats.size(); }
 
     void prepareChildren(){ children.resize(len()); }
-    bool hasChildren() const { return !children.empty(); }
     bool hasLodTree() const { return !children.empty(); }
 
     // gsplat.rs new_merged (step is always 0.0 from bhatt_lod)
@@ -1506,23 +1496,6 @@ std::vector<uint8_t> compressToVec(const std::vector<uint8_t> &data){
 // Property encoders (port of rad.rs encode_*). All planar (dimension-major)
 // except oct88r8, which is 3 bytes per splat interleaved.
 
-std::vector<uint8_t> encodeF32(const std::vector<float> &data, size_t dims, size_t count){
-    std::vector<uint8_t> result;
-    result.reserve(4 * dims * count);
-    for (size_t d = 0; d < dims; d++){
-        size_t index = d;
-        for (size_t i = 0; i < count; i++){
-            uint32_t bits = f32Bits(data[index]);
-            result.push_back(static_cast<uint8_t>(bits));
-            result.push_back(static_cast<uint8_t>(bits >> 8));
-            result.push_back(static_cast<uint8_t>(bits >> 16));
-            result.push_back(static_cast<uint8_t>(bits >> 24));
-            index += dims;
-        }
-    }
-    return result;
-}
-
 std::vector<uint8_t> encodeF16Prop(const std::vector<float> &data, size_t dims, size_t count){
     std::vector<uint8_t> result;
     result.reserve(2 * dims * count);
diff --git a/rasterize_gaussians.cpp b/rasterize_gaussians.cpp
index 736d4c02..cb728dfa 100644
--- a/rasterize_gaussians.cpp
+++ b/rasterize_gaussians.cpp
@@ -36,7 +36,7 @@ std::tuple<torch::Tensor,
     return std::make_tuple(isectIds, gaussianIds, isectIdsSorted, gaussianIdsSorted, tileBins);
 }
 
-torch::Tensor RasterizeGaussians::forward(AutogradContext *ctx, 
+variable_list RasterizeGaussians::forward(AutogradContext *ctx,
             torch::Tensor xys,
             torch::Tensor depths,
             torch::Tensor radii,
@@ -46,7 +46,11 @@ torch::Tensor RasterizeGaussians::forward(AutogradContext *ctx,
             torch::Tensor opacity,
             int imgHeight,
             int imgWidth,
-            torch::Tensor background
+            torch::Tensor background,
+            torch::Tensor errorMap,
+            torch::Tensor edgeMap,
+            torch::Tensor densificationInfo,
+            torch::Tensor xyAbsGrad
         ){
     
     int numPoints = xys.size(0);
@@ -82,17 +86,28 @@ torch::Tensor RasterizeGaussians::forward(AutogradContext *ctx,
     // Map of tile bin IDs
     torch::Tensor finalIdx = std::get<2>(t);
 
+    torch::Tensor outAlpha = 1.0f - finalTs;
+
     ctx->saved_data["imgWidth"] = imgWidth;
     ctx->saved_data["imgHeight"] = imgHeight;
+    ctx->saved_data["errorMap"] = errorMap;
+    ctx->saved_data["edgeMap"] = edgeMap;
+    ctx->saved_data["densificationInfo"] = densificationInfo;
+    ctx->saved_data["xyAbsGrad"] = xyAbsGrad;
     ctx->save_for_backward({ gaussianIdsSorted, tileBins, xys, conics, colors, opacity, background, finalTs, finalIdx });
-    
-    return outImg;
+
+    return { outImg, outAlpha };
 }
 
 tensor_list RasterizeGaussians::backward(AutogradContext *ctx, tensor_list grad_outputs) {
     torch::Tensor v_outImg = grad_outputs[0];
+    torch::Tensor v_outAlpha = grad_outputs[1];
     int imgHeight = ctx->saved_data["imgHeight"].toInt();
     int imgWidth = ctx->saved_data["imgWidth"].toInt();
+    torch::Tensor errorMap = ctx->saved_data["errorMap"].toTensor();
+    torch::Tensor edgeMap = ctx->saved_data["edgeMap"].toTensor();
+    torch::Tensor densificationInfo = ctx->saved_data["densificationInfo"].toTensor();
+    torch::Tensor xyAbsGrad = ctx->saved_data["xyAbsGrad"].toTensor();
 
     variable_list saved = ctx->get_saved_variables();
     torch::Tensor gaussianIdsSorted = saved[0];
@@ -105,9 +120,14 @@ tensor_list RasterizeGaussians::backward(AutogradContext *ctx, tensor_list grad_
     torch::Tensor finalTs = saved[7];
     torch::Tensor finalIdx = saved[8];
 
-    torch::Tensor v_outAlpha = torch::zeros_like(v_outImg.index({"...", 0}));
-    
-    auto t = rasterize_backward_tensor(imgHeight, imgWidth, 
+    if (!v_outAlpha.defined()) v_outAlpha = torch::zeros_like(finalTs);
+    v_outAlpha = v_outAlpha.contiguous();
+    if (!errorMap.defined()) errorMap = torch::empty({0}, xys.options());
+    if (!edgeMap.defined()) edgeMap = torch::empty({0}, xys.options());
+    if (!densificationInfo.defined()) densificationInfo = torch::empty({0}, xys.options());
+    if (!xyAbsGrad.defined()) xyAbsGrad = torch::empty({0}, xys.options());
+
+    auto t = rasterize_backward_tensor(imgHeight, imgWidth,
                             gaussianIdsSorted,
                             tileBins,
                             xys,
@@ -118,7 +138,11 @@ tensor_list RasterizeGaussians::backward(AutogradContext *ctx, tensor_list grad_
                             finalTs,
                             finalIdx,
                             v_outImg,
-                            v_outAlpha);
+                            v_outAlpha,
+                            errorMap,
+                            edgeMap,
+                            densificationInfo,
+                            xyAbsGrad);
 
     torch::Tensor v_xy = std::get<0>(t);
     torch::Tensor v_conic = std::get<1>(t);
@@ -135,13 +159,17 @@ tensor_list RasterizeGaussians::backward(AutogradContext *ctx, tensor_list grad_
             v_opacity,
             none, // imgHeight
             none, // imgWidth
-            none // background
+            none, // background
+            none, // errorMap
+            none, // edgeMap
+            none, // densificationInfo
+            none // xyAbsGrad
     };
 }
 
 #endif
 
-torch::Tensor RasterizeGaussiansCPU::forward(AutogradContext *ctx, 
+variable_list RasterizeGaussiansCPU::forward(AutogradContext *ctx,
             torch::Tensor xys,
             torch::Tensor radii,
             torch::Tensor conics,
@@ -151,7 +179,11 @@ torch::Tensor RasterizeGaussiansCPU::forward(AutogradContext *ctx,
             torch::Tensor camDepths,
             int imgHeight,
             int imgWidth,
-            torch::Tensor background
+            torch::Tensor background,
+            torch::Tensor errorMap,
+            torch::Tensor edgeMap,
+            torch::Tensor densificationInfo,
+            torch::Tensor xyAbsGrad
         ){
     
     int numPoints = xys.size(0);
@@ -171,18 +203,29 @@ torch::Tensor RasterizeGaussiansCPU::forward(AutogradContext *ctx,
     torch::Tensor finalTs = std::get<1>(t);
     std::vector<int32_t> *px2gid = std::get<2>(t);
 
+    torch::Tensor outAlpha = 1.0f - finalTs;
+
     ctx->saved_data["px2gid"] = reinterpret_cast<int64_t>(px2gid);
     ctx->saved_data["imgWidth"] = imgWidth;
     ctx->saved_data["imgHeight"] = imgHeight;
+    ctx->saved_data["errorMap"] = errorMap;
+    ctx->saved_data["edgeMap"] = edgeMap;
+    ctx->saved_data["densificationInfo"] = densificationInfo;
+    ctx->saved_data["xyAbsGrad"] = xyAbsGrad;
     ctx->save_for_backward({ xys, conics, colors, opacity, background, cov2d, camDepths, finalTs });
-    
-    return outImg;
+
+    return { outImg, outAlpha };
 }
 
 tensor_list RasterizeGaussiansCPU::backward(AutogradContext *ctx, tensor_list grad_outputs) {
     torch::Tensor v_outImg = grad_outputs[0];
+    torch::Tensor v_outAlpha = grad_outputs[1];
     int imgHeight = ctx->saved_data["imgHeight"].toInt();
     int imgWidth = ctx->saved_data["imgWidth"].toInt();
+    torch::Tensor errorMap = ctx->saved_data["errorMap"].toTensor();
+    torch::Tensor edgeMap = ctx->saved_data["edgeMap"].toTensor();
+    torch::Tensor densificationInfo = ctx->saved_data["densificationInfo"].toTensor();
+    torch::Tensor xyAbsGrad = ctx->saved_data["xyAbsGrad"].toTensor();
     const std::vector<int32_t> *px2gid = reinterpret_cast<const std::vector<int32_t> *>(ctx->saved_data["px2gid"].toInt());
 
     variable_list saved = ctx->get_saved_variables();
@@ -195,9 +238,14 @@ tensor_list RasterizeGaussiansCPU::backward(AutogradContext *ctx, tensor_list gr
     torch::Tensor camDepths = saved[6];
     torch::Tensor finalTs = saved[7];
 
-    torch::Tensor v_outAlpha = torch::zeros_like(v_outImg.index({"...", 0}));
-    
-    auto t = rasterize_backward_tensor_cpu(imgHeight, imgWidth, 
+    if (!v_outAlpha.defined()) v_outAlpha = torch::zeros_like(finalTs);
+    v_outAlpha = v_outAlpha.contiguous();
+    if (!errorMap.defined()) errorMap = torch::empty({0}, xys.options());
+    if (!edgeMap.defined()) edgeMap = torch::empty({0}, xys.options());
+    if (!densificationInfo.defined()) densificationInfo = torch::empty({0}, xys.options());
+    if (!xyAbsGrad.defined()) xyAbsGrad = torch::empty({0}, xys.options());
+
+    auto t = rasterize_backward_tensor_cpu(imgHeight, imgWidth,
                             xys,
                             conics,
                             colors,
@@ -208,7 +256,11 @@ tensor_list RasterizeGaussiansCPU::backward(AutogradContext *ctx, tensor_list gr
                             finalTs,
                             px2gid,
                             v_outImg,
-                            v_outAlpha);
+                            v_outAlpha,
+                            errorMap,
+                            edgeMap,
+                            densificationInfo,
+                            xyAbsGrad);
 
     delete[] px2gid;
 
@@ -228,7 +280,11 @@ tensor_list RasterizeGaussiansCPU::backward(AutogradContext *ctx, tensor_list gr
             none, // camDepths
             none, // imgHeight
             none, // imgWidth
-            none // background
+            none, // background
+            none, // errorMap
+            none, // edgeMap
+            none, // densificationInfo
+            none // xyAbsGrad
     };
 }
 
diff --git a/rasterize_gaussians.hpp b/rasterize_gaussians.hpp
index a860cf2e..790c6cc4 100644
--- a/rasterize_gaussians.hpp
+++ b/rasterize_gaussians.hpp
@@ -1,4 +1,4 @@
-#ifndef RASTERIZE_GAUSSIANS_H
+#ifndef RASTERIZE_GAUSSIANS_H
 #define RASTERIZE_GAUSSIANS_H
 
 #include <torch/torch.h>
@@ -22,7 +22,7 @@ std::tuple<torch::Tensor,
 
 class RasterizeGaussians : public Function<RasterizeGaussians>{
 public:
-    static torch::Tensor forward(AutogradContext *ctx, 
+    static variable_list forward(AutogradContext *ctx,
             torch::Tensor xys,
             torch::Tensor depths,
             torch::Tensor radii,
@@ -32,7 +32,11 @@ class RasterizeGaussians : public Function<RasterizeGaussians>{
             torch::Tensor opacity,
             int imgHeight,
             int imgWidth,
-            torch::Tensor background);
+            torch::Tensor background,
+            torch::Tensor errorMap = torch::Tensor(),
+            torch::Tensor edgeMap = torch::Tensor(),
+            torch::Tensor densificationInfo = torch::Tensor(),
+            torch::Tensor xyAbsGrad = torch::Tensor());
     static tensor_list backward(AutogradContext *ctx, tensor_list grad_outputs);
 };
 
@@ -40,7 +44,7 @@ class RasterizeGaussians : public Function<RasterizeGaussians>{
 
 class RasterizeGaussiansCPU : public Function<RasterizeGaussiansCPU>{
 public:
-    static torch::Tensor forward(AutogradContext *ctx, 
+    static variable_list forward(AutogradContext *ctx,
             torch::Tensor xys,
             torch::Tensor radii,
             torch::Tensor conics,
@@ -50,7 +54,11 @@ class RasterizeGaussiansCPU : public Function<RasterizeGaussiansCPU>{
             torch::Tensor camDepths,
             int imgHeight,
             int imgWidth,
-            torch::Tensor background);
+            torch::Tensor background,
+            torch::Tensor errorMap = torch::Tensor(),
+            torch::Tensor edgeMap = torch::Tensor(),
+            torch::Tensor densificationInfo = torch::Tensor(),
+            torch::Tensor xyAbsGrad = torch::Tensor());
     static tensor_list backward(AutogradContext *ctx, tensor_list grad_outputs);
 };
 
diff --git a/rasterizer/gsplat-cpu/bindings.h b/rasterizer/gsplat-cpu/bindings.h
index bfcce484..fb8c077c 100644
--- a/rasterizer/gsplat-cpu/bindings.h
+++ b/rasterizer/gsplat-cpu/bindings.h
@@ -1,4 +1,4 @@
-// Originally based on https://github.com/nerfstudio-project/gsplat
+// Originally based on https://github.com/nerfstudio-project/gsplat
 // This implementation has been substantially changed and optimized 
 // Licensed under the AGPLv3
 // Piero Toffanin - 2024
@@ -69,9 +69,34 @@ std::
         const torch::Tensor &final_Ts,
         const std::vector<int32_t> *px2gid,
         const torch::Tensor &v_output, // dL_dout_color
-        const torch::Tensor &v_output_alpha
+        const torch::Tensor &v_output_alpha,
+        const torch::Tensor &error_map,
+        const torch::Tensor &edge_map,
+        const torch::Tensor &densification_info,
+        const torch::Tensor &v_xy_abs
     );
 
+// Fused L1 + DSSIM loss over [H,W,C] float images (same contract as
+// the GPU backends): stats[0] = loss, stats[1] = normalization denominator;
+// partials holds the SSIM derivative maps (empty when want_grad is false).
+std::tuple<torch::Tensor, torch::Tensor> fused_loss_forward_tensor_cpu(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask, // [H,W] float or empty
+    const float ssim_weight,
+    const bool valid_padding,
+    const bool want_grad);
+
+torch::Tensor fused_loss_backward_tensor_cpu(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask,
+    const torch::Tensor &partials,
+    const torch::Tensor &stats,
+    const torch::Tensor &v_loss,
+    const float ssim_weight,
+    const bool valid_padding);
+
 int numShBases(int degree);
 
 torch::Tensor compute_sh_forward_tensor_cpu(
diff --git a/rasterizer/gsplat-cpu/gsplat_cpu.cpp b/rasterizer/gsplat-cpu/gsplat_cpu.cpp
index 63102883..368bed2d 100644
--- a/rasterizer/gsplat-cpu/gsplat_cpu.cpp
+++ b/rasterizer/gsplat-cpu/gsplat_cpu.cpp
@@ -1,4 +1,4 @@
-// Originally started from https://github.com/nerfstudio-project/gsplat
+// Originally started from https://github.com/nerfstudio-project/gsplat
 // This implementation has been substantially changed and optimized 
 // Licensed under the AGPLv3
 // Piero Toffanin - 2024
@@ -10,9 +10,49 @@
 #include <iostream>
 #include <cmath>
 #include <tuple>
+#include <thread>
+#include <atomic>
 
 using namespace torch::indexing;
 
+namespace {
+
+int rasterWorkers(int height, size_t floatsPerWorker = 0){
+    int workers = static_cast<int>((std::max)(1u, std::thread::hardware_concurrency()));
+    workers = (std::min)(workers, (std::max)(1, height));
+    if (floatsPerWorker > 0){
+        const size_t budget = 64ull << 20;
+        const int maxWorkers = (std::max)(1, static_cast<int>(budget / (floatsPerWorker * sizeof(float) + 1)));
+        workers = (std::min)(workers, maxWorkers);
+    }
+    return workers;
+}
+
+// Runs fn(chunk, slot) for every chunk in [0, numChunks)
+template <typename F>
+void parallelChunks(int numChunks, int numWorkers, const F &fn){
+    if (numWorkers <= 1){
+        for (int c = 0; c < numChunks; c++) fn(c, 0);
+        return;
+    }
+    std::atomic<int> next(0);
+    auto worker = [&](int slot){
+        int c;
+        while ((c = next.fetch_add(1, std::memory_order_relaxed)) < numChunks) fn(c, slot);
+    };
+    std::vector<std::thread> threads;
+    threads.reserve(numWorkers - 1);
+    for (int t = 1; t < numWorkers; t++){
+        threads.emplace_back([&worker, t](){ worker(t); });
+    }
+    worker(0);
+    for (std::thread &t : threads) t.join();
+}
+
+const int CHUNKS_PER_WORKER = 4;
+
+}
+
 torch::Tensor quatToRot(const torch::Tensor &quat){
     auto u = torch::unbind(torch::nn::functional::normalize(quat, torch::nn::functional::NormalizeFuncOptions().dim(-1)), -1);
     torch::Tensor w = u[0];
@@ -125,7 +165,7 @@ project_gaussians_forward_tensor_cpu(
     torch::Tensor xys = torch::stack({u, v}, -1); // center
 
     torch::Tensor radii = radius.to(torch::kInt32);
-    torch::Tensor camDepths = pProj.index({"...", 2});
+    torch::Tensor camDepths = pView.index({"...", 2}).contiguous();
 
     return std::make_tuple(xys, radii, conic, cov2d, camDepths);
 }
@@ -185,73 +225,90 @@ std::tuple<
 
     const float alphaThresh = 1.0f / 255.0f;
 
-    for (int idx = 0; idx < numPoints; idx++){
-        int32_t gaussianId = gIndices[idx];
-
-        float A = pConics[gaussianId * 3 + 0];
-        float B = pConics[gaussianId * 3 + 1];
-        float C = pConics[gaussianId * 3 + 2];
-
-        float gX = pCenters[gaussianId * 2 + 0];
-        float gY = pCenters[gaussianId * 2 + 1];
-
-        float sqx = pSqCov2dX[gaussianId];
-        float sqy = pSqCov2dY[gaussianId];
-        
-        int minx = (std::max)(0, static_cast<int>(std::floor(gY - sqy)) - 2);
-        int maxx = (std::min)(height, static_cast<int>(std::ceil(gY + sqy)) + 2);
-        int miny = (std::max)(0, static_cast<int>(std::floor(gX - sqx)) - 2);
-        int maxy = (std::min)(width, static_cast<int>(std::ceil(gX + sqx)) + 2);
-        
-        for (int i = minx; i < maxx; i++){
-            for (int j = miny; j < maxy; j++){
-                size_t pixIdx = (i * width + j);
-                if (pDone[pixIdx]) continue;
+    const int numWorkers = rasterWorkers(height);
+    const int numBands = (std::min)(height, numWorkers * CHUNKS_PER_WORKER);
 
-                float xCam = gX - j;
-                float yCam = gY - i;
-                float sigma = (
-                    0.5f
-                    * (A * xCam * xCam + C * yCam * yCam)
-                    + B * xCam * yCam
-                );
+    parallelChunks(numBands, numWorkers, [&](int band, int){
+            for (int idx = 0; idx < numPoints; idx++){
+                int32_t gaussianId = gIndices[idx];
 
-                if (sigma < 0.0f) continue;
-                float alpha = (std::min)(0.999f, (pOpacities[gaussianId] * std::exp(-sigma)));
-                if (alpha < alphaThresh) continue;
+                float sqy = pSqCov2dY[gaussianId];
+                float gY = pCenters[gaussianId * 2 + 1];
 
-                float T = pFinalTs[pixIdx];
-                float nextT = T * (1.0f - alpha);
-                if (nextT <= 1e-4f) { // this pixel is done
-                    pDone[pixIdx] = true;
-                    continue;
-                }
+                int minx = (std::max)(0, static_cast<int>(std::floor(gY - sqy)) - 2);
+                int maxx = (std::min)(height, static_cast<int>(std::ceil(gY + sqy)) + 2);
+                // first row >= minx that belongs to this band
+                minx += ((band - minx) % numBands + numBands) % numBands;
+                if (minx >= maxx) continue;
 
-                float vis = alpha * T;
+                float sqx = pSqCov2dX[gaussianId];
+                float gX = pCenters[gaussianId * 2 + 0];
 
-                pOutImg[pixIdx * 3 + 0] += vis * pColors[gaussianId * 3 + 0];
-                pOutImg[pixIdx * 3 + 1] += vis * pColors[gaussianId * 3 + 1];
-                pOutImg[pixIdx * 3 + 2] += vis * pColors[gaussianId * 3 + 2];
-                
-                pFinalTs[pixIdx] = nextT;
-                px2gid[pixIdx].push_back(gaussianId);
-            }
-        }
-    }
+                int miny = (std::max)(0, static_cast<int>(std::floor(gX - sqx)) - 2);
+                int maxy = (std::min)(width, static_cast<int>(std::ceil(gX + sqx)) + 2);
+                if (miny >= maxy) continue;
 
-    // Background
-    for (int i = 0; i < height; i++){
-        for (int j = 0; j < width; j++){
-            size_t pixIdx = (i * width + j);
-            float T = pFinalTs[pixIdx];
+                float A = pConics[gaussianId * 3 + 0];
+                float B = pConics[gaussianId * 3 + 1];
+                float C = pConics[gaussianId * 3 + 2];
 
-            pOutImg[pixIdx * 3 + 0] += T * bgX;
-            pOutImg[pixIdx * 3 + 1] += T * bgY;
-            pOutImg[pixIdx * 3 + 2] += T * bgZ;
+                const float opacity = pOpacities[gaussianId];
+                const float r = pColors[gaussianId * 3 + 0];
+                const float g = pColors[gaussianId * 3 + 1];
+                const float b = pColors[gaussianId * 3 + 2];
+                const float sigmaCut = opacity > 0.0f
+                    ? std::log(255.0f * opacity) + 1e-3f
+                    : -1.0f;
+
+                for (int i = minx; i < maxx; i += numBands){
+                    for (int j = miny; j < maxy; j++){
+                        size_t pixIdx = (i * width + j);
+                        if (pDone[pixIdx]) continue;
+
+                        float xCam = gX - j;
+                        float yCam = gY - i;
+                        float sigma = (
+                            0.5f
+                            * (A * xCam * xCam + C * yCam * yCam)
+                            + B * xCam * yCam
+                        );
+
+                        if (sigma < 0.0f) continue;
+                        if (sigma > sigmaCut) continue;
+                        float alpha = (std::min)(0.999f, (opacity * std::exp(-sigma)));
+                        if (alpha < alphaThresh) continue;
+
+                        float T = pFinalTs[pixIdx];
+                        float nextT = T * (1.0f - alpha);
+                        if (nextT <= 1e-4f) { // this pixel is done
+                            pDone[pixIdx] = true;
+                            continue;
+                        }
+
+                        float vis = alpha * T;
+
+                        pOutImg[pixIdx * 3 + 0] += vis * r;
+                        pOutImg[pixIdx * 3 + 1] += vis * g;
+                        pOutImg[pixIdx * 3 + 2] += vis * b;
+
+                        pFinalTs[pixIdx] = nextT;
+                        px2gid[pixIdx].push_back(gaussianId);
+                    }
+                }
+            }
 
-            std::reverse(px2gid[pixIdx].begin(), px2gid[pixIdx].end());
-        }
-    }
+            // Background
+            for (int i = band; i < height; i += numBands){
+                for (int j = 0; j < width; j++){
+                    size_t pixIdx = (i * width + j);
+                    float T = pFinalTs[pixIdx];
+
+                    pOutImg[pixIdx * 3 + 0] += T * bgX;
+                    pOutImg[pixIdx * 3 + 1] += T * bgY;
+                    pOutImg[pixIdx * 3 + 2] += T * bgZ;
+                }
+            }
+    });
 
     return std::make_tuple(outImg, finalTs, px2gid);
 }
@@ -277,7 +334,11 @@ std::
         const torch::Tensor &final_Ts,
         const std::vector<int32_t> *px2gid,
         const torch::Tensor &v_output, // dL_dout_color
-        const torch::Tensor &v_output_alpha
+        const torch::Tensor &v_output_alpha,
+        const torch::Tensor &error_map,
+        const torch::Tensor &edge_map,
+        const torch::Tensor &densification_info,
+        const torch::Tensor &v_xy_abs
     ){
     torch::NoGradGuard noGrad;
 
@@ -285,16 +346,6 @@ std::
     int channels = colors.size(1);
     torch::Device device = xys.device();
 
-    torch::Tensor v_xy = torch::zeros({numPoints, 2}, torch::TensorOptions().dtype(torch::kFloat32).device(device));   
-    torch::Tensor v_conic = torch::zeros({numPoints, 3}, torch::TensorOptions().dtype(torch::kFloat32).device(device));
-    torch::Tensor v_colors = torch::zeros({numPoints, channels}, torch::TensorOptions().dtype(torch::kFloat32).device(device));
-    torch::Tensor v_opacity = torch::zeros({numPoints, 1}, torch::TensorOptions().dtype(torch::kFloat32).device(device));
-
-    float *pv_xy = static_cast<float *>(v_xy.data_ptr());
-    float *pv_conic = static_cast<float *>(v_conic.data_ptr());
-    float *pv_colors = static_cast<float *>(v_colors.data_ptr());
-    float *pv_opacity = static_cast<float *>(v_opacity.data_ptr());
-    
     float *pColors = static_cast<float *>(colors.data_ptr());
     float *pv_output = static_cast<float *>(v_output.data_ptr());
     float *pv_outputAlpha = static_cast<float *>(v_output_alpha.data_ptr());
@@ -307,17 +358,51 @@ std::
     float bgZ = background[2].item<float>();
 
     float *pFinalTs = static_cast<float *>(final_Ts.data_ptr());
+    const bool hasDinfo = densification_info.numel() > 0;
+    float *pErr = error_map.numel() > 0 ? static_cast<float *>(error_map.data_ptr()) : nullptr;
+    float *pEdge = edge_map.numel() > 0 ? static_cast<float *>(edge_map.data_ptr()) : nullptr;
+    const bool hasXyAbs = v_xy_abs.numel() > 0;
 
     const float alphaThresh = 1.0f / 255.0f;
 
-    for (int i = 0; i < height; i++){
+    auto fOpts = torch::TensorOptions().dtype(torch::kFloat32).device(device);
+    const size_t floatsPerWorker = static_cast<size_t>(numPoints) *
+        (2 + 3 + channels + 1 + (hasDinfo ? 4 : 0) + (hasXyAbs ? 2 : 0));
+    const int numWorkers = rasterWorkers(height, floatsPerWorker);
+    const int numBands = (std::min)(height, numWorkers * CHUNKS_PER_WORKER);
+
+    torch::Tensor v_xy_b = torch::zeros({numWorkers, numPoints, 2}, fOpts);
+    torch::Tensor v_conic_b = torch::zeros({numWorkers, numPoints, 3}, fOpts);
+    torch::Tensor v_colors_b = torch::zeros({numWorkers, numPoints, channels}, fOpts);
+    torch::Tensor v_opacity_b = torch::zeros({numWorkers, numPoints, 1}, fOpts);
+    torch::Tensor dinfo_b = torch::zeros({numWorkers, hasDinfo ? 4 : 0, numPoints}, fOpts);
+    torch::Tensor xyAbs_b = torch::zeros({numWorkers, hasXyAbs ? numPoints : 0, 2}, fOpts);
+
+    float *pv_xy_b = static_cast<float *>(v_xy_b.data_ptr());
+    float *pv_conic_b = static_cast<float *>(v_conic_b.data_ptr());
+    float *pv_colors_b = static_cast<float *>(v_colors_b.data_ptr());
+    float *pv_opacity_b = static_cast<float *>(v_opacity_b.data_ptr());
+    float *pDinfo_b = hasDinfo ? static_cast<float *>(dinfo_b.data_ptr()) : nullptr;
+    float *pXyAbs_b = hasXyAbs ? static_cast<float *>(xyAbs_b.data_ptr()) : nullptr;
+
+    parallelChunks(numBands, numWorkers, [&](int band, int slot){
+        float *pv_xy = pv_xy_b + static_cast<size_t>(slot) * numPoints * 2;
+        float *pv_conic = pv_conic_b + static_cast<size_t>(slot) * numPoints * 3;
+        float *pv_colors = pv_colors_b + static_cast<size_t>(slot) * numPoints * channels;
+        float *pv_opacity = pv_opacity_b + static_cast<size_t>(slot) * numPoints;
+        float *pDinfo = pDinfo_b ? pDinfo_b + static_cast<size_t>(slot) * numPoints * 4 : nullptr;
+        float *pXyAbs = pXyAbs_b ? pXyAbs_b + static_cast<size_t>(slot) * numPoints * 2 : nullptr;
+
+        for (int i = static_cast<int>(band); i < height; i += numBands){
         for (int j = 0; j < width; j++){
             size_t pixIdx = (i * width + j);
             float Tfinal = pFinalTs[pixIdx];
             float T = Tfinal;
             float buffer[3] = {0.0f, 0.0f, 0.0f};
 
-            for (const int32_t &gaussianId : px2gid[pixIdx]){
+            const std::vector<int32_t> &gids = px2gid[pixIdx];
+            for (auto it = gids.rbegin(); it != gids.rend(); ++it){
+                const int32_t gaussianId = *it;
                 float A = pConics[gaussianId * 3 + 0];
                 float B = pConics[gaussianId * 3 + 1];
                 float C = pConics[gaussianId * 3 + 2];
@@ -335,13 +420,21 @@ std::
 
                 if (sigma < 0.0f) continue;
                 float vis = std::exp(-sigma);
-                float alpha = (std::min)(0.99f, pOpacities[gaussianId] * vis);
+                float alpha = (std::min)(0.999f, pOpacities[gaussianId] * vis);
                 if (alpha < alphaThresh) continue;
 
                 float ra = 1.0f / (1.0f - alpha);
                 T *= ra;
                 float fac = alpha * T;
 
+                if (pDinfo){
+                    float err = pErr ? pErr[pixIdx] : 1.0f;
+                    pDinfo[gaussianId] += fac;
+                    pDinfo[numPoints + gaussianId] += fac * err;
+                    if (pEdge) pDinfo[2 * numPoints + gaussianId] += fac * pEdge[pixIdx];
+                    if (err > 0.5f) pDinfo[3 * numPoints + gaussianId] += 1.0f;
+                }
+
                 pv_colors[gaussianId * 3 + 0] += fac * pv_output[pixIdx * 3 + 0];
                 pv_colors[gaussianId * 3 + 1] += fac * pv_output[pixIdx * 3 + 1];
                 pv_colors[gaussianId * 3 + 2] += fac * pv_output[pixIdx * 3 + 2];
@@ -366,11 +459,23 @@ std::
 
                 pv_xy[gaussianId * 2 + 0] += v_sigma * (A * xCam + B * yCam);
                 pv_xy[gaussianId * 2 + 1] += v_sigma * (B * xCam + C * yCam);
+                if (pXyAbs){
+                    pXyAbs[gaussianId * 2 + 0] += std::fabs(v_sigma * (A * xCam + B * yCam));
+                    pXyAbs[gaussianId * 2 + 1] += std::fabs(v_sigma * (B * xCam + C * yCam));
+                }
 
                 pv_opacity[gaussianId] += vis * v_alpha;
             }
         }
-    }
+        }
+    });
+
+    torch::Tensor v_xy = v_xy_b.sum(0);
+    torch::Tensor v_conic = v_conic_b.sum(0);
+    torch::Tensor v_colors = v_colors_b.sum(0);
+    torch::Tensor v_opacity = v_opacity_b.sum(0);
+    if (hasDinfo) densification_info.add_(dinfo_b.sum(0));
+    if (hasXyAbs) v_xy_abs.add_(xyAbs_b.sum(0));
 
     return std::make_tuple(v_xy, v_conic, v_colors, v_opacity);
 }
@@ -483,4 +588,267 @@ torch::Tensor compute_sh_forward_tensor_cpu(
     }
     
     return (result.index({"...", None}) * coeffs).sum(-2);
+}
+
+// Fused L1 + DSSIM loss over [H,W,C] images: same method as the GPU
+// backends (two-pass separable 11-tap blur computing all five moments in one
+// sweep, closed-form SSIM partials saved for a direct backward), parallelized
+// over row bands.
+
+namespace {
+
+const int LOSS_HALO = 5;
+
+// 11-tap gaussian (sigma 1.5), matches the SSIM reference window
+const float lossGauss[11] = {
+    0.001028380123898387f, 0.0075987582094967365f, 0.036000773310661316f,
+    0.10936068743467331f, 0.21300552785396576f, 0.26601171493530273f,
+    0.21300552785396576f, 0.10936068743467331f, 0.036000773310661316f,
+    0.0075987582094967365f, 0.001028380123898387f};
+
+inline bool lossValid(int y, int x, int H, int W, bool validPad){
+    if (!validPad || H <= 10 || W <= 10) return true;
+    return x >= LOSS_HALO && x < W - LOSS_HALO && y >= LOSS_HALO && y < H - LOSS_HALO;
+}
+
+}
+
+std::tuple<torch::Tensor, torch::Tensor> fused_loss_forward_tensor_cpu(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask,
+    const float ssim_weight,
+    const bool valid_padding,
+    const bool want_grad
+){
+    torch::Tensor r = rendered.contiguous();
+    torch::Tensor g = gt.contiguous();
+    const int H = r.size(0);
+    const int W = r.size(1);
+    const int C = r.size(2);
+    const bool hasMask = mask.defined() && mask.numel() > 0;
+    torch::Tensor m = hasMask ? mask.contiguous() : torch::Tensor();
+
+    const float *rp = r.data_ptr<float>();
+    const float *gp = g.data_ptr<float>();
+    const float *mp = hasMask ? m.data_ptr<float>() : nullptr;
+
+    auto fOpts = r.options();
+    torch::Tensor ssimMap = torch::empty({H, W}, fOpts);
+    torch::Tensor partials = want_grad
+        ? torch::empty({3, static_cast<long long>(H) * W * C}, fOpts.dtype(torch::kHalf))
+        : torch::empty({0}, fOpts.dtype(torch::kHalf));
+    float *sp = ssimMap.data_ptr<float>();
+    at::Half *pBase = want_grad ? partials.data_ptr<at::Half>() : nullptr;
+    const long long planeSize = static_cast<long long>(H) * W * C;
+    at::Half *pMu = pBase;
+    at::Half *pS1 = pBase ? pBase + planeSize : nullptr;
+    at::Half *pS12 = pBase ? pBase + 2 * planeSize : nullptr;
+
+    // Horizontal-moment scratch, reused across channels
+    torch::Tensor hbufT = torch::empty({static_cast<long long>(H) * W * 5}, fOpts);
+    float *hbuf = hbufT.data_ptr<float>();
+
+    const int numWorkers = rasterWorkers(H);
+    const int numBands = (std::min)(H, numWorkers * CHUNKS_PER_WORKER);
+
+    for (int c = 0; c < C; c++){
+        parallelChunks(numBands, numWorkers, [&](int band, int){
+            const int yStart = band * H / numBands;
+            const int yEnd = (band + 1) * H / numBands;
+            for (int y = yStart; y < yEnd; y++){
+                for (int x = 0; x < W; x++){
+                    float sX = 0.f, sX2 = 0.f, sY = 0.f, sY2 = 0.f, sXY = 0.f;
+                    for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                        const int xx = x + d;
+                        if (xx < 0 || xx >= W) continue;
+                        const float w = lossGauss[LOSS_HALO + d];
+                        const float X = rp[(y * W + xx) * C + c];
+                        const float Y = gp[(y * W + xx) * C + c];
+                        sX += X * w;
+                        sX2 += X * X * w;
+                        sY += Y * w;
+                        sY2 += Y * Y * w;
+                        sXY += X * Y * w;
+                    }
+                    float *hb = &hbuf[(static_cast<long long>(y) * W + x) * 5];
+                    hb[0] = sX; hb[1] = sX2; hb[2] = sY; hb[3] = sY2; hb[4] = sXY;
+                }
+            }
+        });
+
+        parallelChunks(numBands, numWorkers, [&](int band, int){
+            const int yStart = band * H / numBands;
+            const int yEnd = (band + 1) * H / numBands;
+            for (int y = yStart; y < yEnd; y++){
+                for (int x = 0; x < W; x++){
+                    float m0 = 0.f, m1 = 0.f, m2 = 0.f, m3 = 0.f, m4 = 0.f;
+                    for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                        const int yy = y + d;
+                        if (yy < 0 || yy >= H) continue;
+                        const float w = lossGauss[LOSS_HALO + d];
+                        const float *hb = &hbuf[(static_cast<long long>(yy) * W + x) * 5];
+                        m0 += hb[0] * w;
+                        m1 += hb[1] * w;
+                        m2 += hb[2] * w;
+                        m3 += hb[3] * w;
+                        m4 += hb[4] * w;
+                    }
+                    const float muX = m0;
+                    const float muY = m2;
+                    const float sigmaX = m1 - muX * muX;
+                    const float sigmaY = m3 - muY * muY;
+                    const float sigmaXY = m4 - muX * muY;
+
+                    const float A = muX * muX + muY * muY + 0.0001f;
+                    const float B = sigmaX + sigmaY + 0.0009f;
+                    const float Cc = 2.f * muX * muY + 0.0001f;
+                    const float Dc = 2.f * sigmaXY + 0.0009f;
+                    const float s = (Cc * Dc) / (A * B);
+
+                    const int p = y * W + x;
+                    if (c == 0) sp[p] = s / static_cast<float>(C);
+                    else sp[p] += s / static_cast<float>(C);
+
+                    if (pMu){
+                        const long long idx = static_cast<long long>(p) * C + c;
+                        const float dMu = (muY * 2.f * Dc) / (A * B) - (muY * 2.f * Cc) / (A * B)
+                                        - (muX * 2.f * Cc * Dc) / (A * A * B) + (muX * 2.f * Cc * Dc) / (A * B * B);
+                        pMu[idx] = static_cast<at::Half>(dMu);
+                        pS1[idx] = static_cast<at::Half>((-Cc * Dc) / (A * B * B));
+                        pS12[idx] = static_cast<at::Half>((2.f * Cc) / (A * B));
+                    }
+                }
+            }
+        });
+    }
+
+    // Reduction to the scalar loss + normalization denominator
+    std::vector<double> lossSums(numBands, 0.0);
+    std::vector<double> gateSums(numBands, 0.0);
+    parallelChunks(numBands, numWorkers, [&](int band, int){
+        const int yStart = band * H / numBands;
+        const int yEnd = (band + 1) * H / numBands;
+        double lossSum = 0.0, gateSum = 0.0;
+        for (int y = yStart; y < yEnd; y++){
+            for (int x = 0; x < W; x++){
+                const int p = y * W + x;
+                const float gate = mp ? mp[p] : (lossValid(y, x, H, W, valid_padding) ? 1.0f : 0.0f);
+                if (gate == 0.0f) continue;
+                float l1 = 0.0f;
+                for (int c = 0; c < C; c++){
+                    l1 += std::fabs(rp[p * C + c] - gp[p * C + c]);
+                }
+                lossSum += gate * ((1.0f - ssim_weight) * l1
+                                 + static_cast<float>(C) * ssim_weight * (1.0f - sp[p]));
+                gateSum += gate;
+            }
+        }
+        lossSums[band] = lossSum;
+        gateSums[band] = gateSum;
+    });
+    double lossSum = 0.0, gateSum = 0.0;
+    for (int b = 0; b < numBands; b++){ lossSum += lossSums[b]; gateSum += gateSums[b]; }
+    const double denom = gateSum * C + 1e-8;
+
+    torch::Tensor stats = torch::empty({2}, fOpts);
+    stats.data_ptr<float>()[0] = static_cast<float>(lossSum / denom);
+    stats.data_ptr<float>()[1] = static_cast<float>(denom);
+    return std::make_tuple(stats, partials);
+}
+
+torch::Tensor fused_loss_backward_tensor_cpu(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask,
+    const torch::Tensor &partials,
+    const torch::Tensor &stats,
+    const torch::Tensor &v_loss,
+    const float ssim_weight,
+    const bool valid_padding
+){
+    torch::Tensor r = rendered.contiguous();
+    torch::Tensor g = gt.contiguous();
+    const int H = r.size(0);
+    const int W = r.size(1);
+    const int C = r.size(2);
+    const bool hasMask = mask.defined() && mask.numel() > 0;
+    torch::Tensor m = hasMask ? mask.contiguous() : torch::Tensor();
+
+    const float *rp = r.data_ptr<float>();
+    const float *gp = g.data_ptr<float>();
+    const float *mp = hasMask ? m.data_ptr<float>() : nullptr;
+    const at::Half *pBase = partials.data_ptr<at::Half>();
+    const long long planeSize = static_cast<long long>(H) * W * C;
+    const at::Half *pMu = pBase;
+    const at::Half *pS1 = pBase + planeSize;
+    const at::Half *pS12 = pBase + 2 * planeSize;
+    const float chainScale = v_loss.contiguous().data_ptr<float>()[0] / stats.data_ptr<float>()[1];
+
+    torch::Tensor vRendered = torch::empty_like(r);
+    float *vp = vRendered.data_ptr<float>();
+
+    torch::Tensor hbufT = torch::empty({static_cast<long long>(H) * W * 3}, r.options());
+    float *hbuf = hbufT.data_ptr<float>();
+
+    const int numWorkers = rasterWorkers(H);
+    const int numBands = (std::min)(H, numWorkers * CHUNKS_PER_WORKER);
+
+    auto gateAt = [&](int y, int x) -> float {
+        return mp ? mp[y * W + x] : (lossValid(y, x, H, W, valid_padding) ? 1.0f : 0.0f);
+    };
+
+    for (int c = 0; c < C; c++){
+        parallelChunks(numBands, numWorkers, [&](int band, int){
+            const int yStart = band * H / numBands;
+            const int yEnd = (band + 1) * H / numBands;
+            for (int y = yStart; y < yEnd; y++){
+                for (int x = 0; x < W; x++){
+                    float a0 = 0.f, a1 = 0.f, a2 = 0.f;
+                    for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                        const int xx = x + d;
+                        if (xx < 0 || xx >= W) continue;
+                        const float chain = -ssim_weight * gateAt(y, xx) * chainScale;
+                        if (chain == 0.0f) continue;
+                        const float w = lossGauss[LOSS_HALO + d];
+                        const long long idx = (static_cast<long long>(y) * W + xx) * C + c;
+                        a0 += static_cast<float>(pMu[idx]) * chain * w;
+                        a1 += static_cast<float>(pS1[idx]) * chain * w;
+                        a2 += static_cast<float>(pS12[idx]) * chain * w;
+                    }
+                    float *hb = &hbuf[(static_cast<long long>(y) * W + x) * 3];
+                    hb[0] = a0; hb[1] = a1; hb[2] = a2;
+                }
+            }
+        });
+
+        parallelChunks(numBands, numWorkers, [&](int band, int){
+            const int yStart = band * H / numBands;
+            const int yEnd = (band + 1) * H / numBands;
+            for (int y = yStart; y < yEnd; y++){
+                for (int x = 0; x < W; x++){
+                    float s0 = 0.f, s1 = 0.f, s2 = 0.f;
+                    for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                        const int yy = y + d;
+                        if (yy < 0 || yy >= H) continue;
+                        const float w = lossGauss[LOSS_HALO + d];
+                        const float *hb = &hbuf[(static_cast<long long>(yy) * W + x) * 3];
+                        s0 += hb[0] * w;
+                        s1 += hb[1] * w;
+                        s2 += hb[2] * w;
+                    }
+                    const long long p = static_cast<long long>(y) * W + x;
+                    const float p1 = rp[p * C + c];
+                    const float p2 = gp[p * C + c];
+                    const float gradSsim = s0 + 2.f * p1 * s1 + p2 * s2;
+
+                    const float gate = gateAt(y, x);
+                    const float sign = (p1 == p2) ? 0.0f : std::copysign(1.0f, p1 - p2);
+                    vp[p * C + c] = gradSsim + (1.0f - ssim_weight) * sign * gate * chainScale;
+                }
+            }
+        });
+    }
+
+    return vRendered;
 }
\ No newline at end of file
diff --git a/rasterizer/gsplat-metal/bindings.h b/rasterizer/gsplat-metal/bindings.h
index 4809cbd1..7180142d 100644
--- a/rasterizer/gsplat-metal/bindings.h
+++ b/rasterizer/gsplat-metal/bindings.h
@@ -1,4 +1,4 @@
-#include <cstdio>
+#include <cstdio>
 #include <iostream>
 #include <math.h>
 #include <tuple>
@@ -16,6 +16,28 @@ std::tuple<
     torch::Tensor> // output radii
 compute_cov2d_bounds_tensor(const int num_pts, torch::Tensor &A);
 
+// Fused L1 + DSSIM loss over [H,W,C] float images.
+// Returns {stats, partials}: stats[0] = loss, stats[1] = normalization
+// denominator (both on device); partials holds the SSIM derivative maps
+// needed by the backward pass (empty when want_grad is false).
+std::tuple<torch::Tensor, torch::Tensor> fused_loss_forward_tensor(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask, // [H,W] float or empty
+    const float ssim_weight,
+    const bool valid_padding,
+    const bool want_grad);
+
+torch::Tensor fused_loss_backward_tensor(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask,
+    const torch::Tensor &partials,
+    const torch::Tensor &stats,
+    const torch::Tensor &v_loss,
+    const float ssim_weight,
+    const bool valid_padding);
+
 torch::Tensor compute_sh_forward_tensor(
     unsigned num_points,
     unsigned degree,
@@ -179,5 +201,9 @@ std::
         const torch::Tensor &final_Ts,
         const torch::Tensor &final_idx,
         const torch::Tensor &v_output, // dL_dout_color
-        const torch::Tensor &v_output_alpha
+        const torch::Tensor &v_output_alpha,
+        const torch::Tensor &error_map, // [H,W] or empty
+        const torch::Tensor &edge_map, // [H,W] or empty
+        const torch::Tensor &densification_info, // [4,N] accumulated in place, or empty
+        const torch::Tensor &v_xy_abs // [N,2] accumulated in place, or empty
     );
\ No newline at end of file
diff --git a/rasterizer/gsplat-metal/gsplat_metal.metal b/rasterizer/gsplat-metal/gsplat_metal.metal
index 2313a239..e1bd563e 100644
--- a/rasterizer/gsplat-metal/gsplat_metal.metal
+++ b/rasterizer/gsplat-metal/gsplat_metal.metal
@@ -1,4 +1,4 @@
-#include <metal_stdlib>
+#include <metal_stdlib>
 
 using namespace metal;
 
@@ -296,6 +296,10 @@ inline int2 read_packed_int2(constant int* arr, int idx) {
     return int2(arr[2*idx], arr[2*idx+1]);
 }
 
+inline int2 read_packed_int2(device const int* arr, int idx) {
+    return int2(arr[2*idx], arr[2*idx+1]);
+}
+
 inline void write_packed_int2(device int* arr, int idx, int2 val) {
     arr[2*idx] = val.x;
     arr[2*idx+1] = val.y;
@@ -317,6 +321,10 @@ inline float2 read_packed_float2(device float* arr, int idx) {
     return float2(arr[2*idx], arr[2*idx+1]);
 }
 
+inline float2 read_packed_float2(device const float* arr, int idx) {
+    return float2(arr[2*idx], arr[2*idx+1]);
+}
+
 inline void write_packed_float2(device float* arr, int idx, float2 val) {
     arr[2*idx] = val.x;
     arr[2*idx+1] = val.y;
@@ -340,6 +348,10 @@ inline float3 read_packed_float3(device float* arr, int idx) {
     return float3(arr[3*idx], arr[3*idx+1], arr[3*idx+2]);
 }
 
+inline float3 read_packed_float3(device const float* arr, int idx) {
+    return float3(arr[3*idx], arr[3*idx+1], arr[3*idx+2]);
+}
+
 inline void write_packed_float3(device float* arr, int idx, float3 val) {
     arr[3*idx] = val.x;
     arr[3*idx+1] = val.y;
@@ -432,6 +444,110 @@ kernel void project_gaussians_forward_kernel(
     write_packed_float2(xys, idx, center);
 }
 
+// 3-channel tile rasterizer. Unlike nd_rasterize_forward_kernel it stages the
+// tile's gaussians in threadgroup memory, accumulates the pixel colour in
+// registers instead of read-modify-writing out_img once per gaussian, and lets
+// the whole tile stop early once every pixel has saturated.
+kernel void rasterize_forward_kernel(
+    constant uint3& tile_bounds,
+    constant uint3& img_size,
+    device const int32_t* gaussian_ids_sorted,
+    device const int* tile_bins, // int2
+    device const float* xys, // float2
+    device const float* conics, // float3
+    device const float* colors, // float3
+    device const float* opacities,
+    device float* final_Ts,
+    device int* final_index,
+    device float* out_img, // float3
+    constant float* background, // single float3
+    uint3 gp [[thread_position_in_grid]],
+    uint3 blockIdx [[threadgroup_position_in_grid]],
+    uint tr [[thread_index_in_threadgroup]]
+) {
+    const int32_t tile_id = blockIdx.y * tile_bounds.x + blockIdx.x;
+    const uint i = gp.y;
+    const uint j = gp.x;
+    const float px = (float)j;
+    const float py = (float)i;
+    const int32_t pix_id = (int32_t)(i * img_size.x + j);
+
+    // Out-of-bounds threads stay resident so that the threadgroup barriers
+    // below remain uniform; they simply never write anything.
+    const bool inside = (i < img_size.y && j < img_size.x);
+
+    const int2 range = read_packed_int2(tile_bins, tile_id);
+    const int num_batches = (range.y - range.x + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+    threadgroup packed_float3 xy_opacity_batch[BLOCK_SIZE];
+    threadgroup packed_float3 conic_batch[BLOCK_SIZE];
+    threadgroup packed_float3 rgbs_batch[BLOCK_SIZE];
+    threadgroup atomic_int num_done;
+
+    bool done = !inside;
+    float T = 1.f;
+    float3 pix_out = {0.f, 0.f, 0.f};
+    // Mirrors the index the equivalent scalar loop would end on
+    int idx = range.x;
+
+    for (int b = 0; b < num_batches; ++b) {
+        // Also fences the previous batch's reads before we overwrite the arrays
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (tr == 0) atomic_store_explicit(&num_done, 0, memory_order_relaxed);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (done) atomic_fetch_add_explicit(&num_done, 1, memory_order_relaxed);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (atomic_load_explicit(&num_done, memory_order_relaxed) == BLOCK_SIZE) break;
+
+        const int batch_start = range.x + BLOCK_SIZE * b;
+        const int load_idx = batch_start + (int)tr;
+        if (load_idx < range.y) {
+            const int32_t g_id = gaussian_ids_sorted[load_idx];
+            const float2 xy = read_packed_float2(xys, g_id);
+            xy_opacity_batch[tr] = packed_float3(xy.x, xy.y, opacities[g_id]);
+            conic_batch[tr] = packed_float3(read_packed_float3(conics, g_id));
+            rgbs_batch[tr] = packed_float3(read_packed_float3(colors, g_id));
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        const int batch_size = min((int)BLOCK_SIZE, range.y - batch_start);
+        for (int t = 0; t < batch_size && !done; ++t) {
+            const int g_idx = batch_start + t;
+            const float3 conic = conic_batch[t];
+            const float3 xy_opac = xy_opacity_batch[t];
+            const float2 delta = {xy_opac.x - px, xy_opac.y - py};
+            const float sigma =
+                0.5f * (conic.x * delta.x * delta.x + conic.z * delta.y * delta.y) +
+                conic.y * delta.x * delta.y;
+            idx = g_idx + 1;
+            if (sigma < 0.f) {
+                continue;
+            }
+            const float alpha = min(0.999f, xy_opac.z * exp(-sigma));
+            if (alpha < 1.f / 255.f) {
+                continue;
+            }
+            const float next_T = T * (1.f - alpha);
+            if (next_T <= 1e-4f) {
+                // render the last gaussian that contributed
+                idx = g_idx - 1;
+                done = true;
+                break;
+            }
+            const float vis = alpha * T;
+            pix_out += float3(rgbs_batch[t]) * vis;
+            T = next_T;
+        }
+    }
+
+    if (!inside) return;
+
+    final_Ts[pix_id] = T; // transmittance at last gaussian in this pixel
+    final_index[pix_id] = (idx == range.y) ? idx - 1 : idx;
+    write_packed_float3(out_img, pix_id,
+        pix_out + T * float3(background[0], background[1], background[2]));
+}
+
 kernel void nd_rasterize_forward_kernel(
     constant uint3& tile_bounds,
     constant uint3& img_size,
@@ -797,66 +913,55 @@ kernel void get_tile_bin_edges_kernel(
 }
 
 inline int warp_reduce_all_max(int val, const int warp_size) {
-    // This uses an xor so that all threads in a warp get the same result
-    for ( int mask = warp_size / 2; mask > 0; mask /= 2 )
-        val = max(val, simd_shuffle_xor(val, mask));
-
-    return val;
+    return simd_max(val);
 }
 
 inline int warp_reduce_all_or(int val, const int warp_size) {
-    // This uses an xor so that all threads in a warp get the same result
-    for ( int mask = warp_size / 2; mask > 0; mask /= 2 )
-        val = val | simd_shuffle_xor(val, mask);
-
-    return val;
+    return simd_any(val != 0) ? 1 : 0;
 }
 
 inline float warp_reduce_sum(float val, const int warp_size, const uint lane) {
-    for (int offset = warp_size / 2; offset > 0; offset /= 2) {
-        float other = 0.0f;
-        if (lane + offset <= warp_size)
-            other = simd_shuffle_xor(val, offset);
-        val += other;
-    }
-    return val;
+    return simd_sum(val);
 }
 
 inline float3 warpSum3(float3 val, const int warp_size, const uint lane) {
-    val.x = warp_reduce_sum(val.x, warp_size, lane);
-    val.y = warp_reduce_sum(val.y, warp_size, lane);
-    val.z = warp_reduce_sum(val.z, warp_size, lane);
-    return val;
+    return float3(simd_sum(val.x), simd_sum(val.y), simd_sum(val.z));
 }
 
 inline float2 warpSum2(float2 val, const int warp_size, const uint lane) {
-    val.x = warp_reduce_sum(val.x, warp_size, lane);
-    val.y = warp_reduce_sum(val.y, warp_size, lane);
-    return val;
+    return float2(simd_sum(val.x), simd_sum(val.y));
 }
 
 inline float warpSum(float val, const int warp_size, const uint lane) {
-    return warp_reduce_sum(val, warp_size, lane);
+    return simd_sum(val);
 }
 
 kernel void rasterize_backward_kernel(
     constant uint3& tile_bounds,
     constant uint2& img_size,
-    constant int32_t* gaussian_ids_sorted,
-    constant int* tile_bins, // int2
-    constant float* xys, // float2
-    constant float* conics, // float3
-    constant float* rgbs, // float3
-    constant float* opacities,
+    device const int32_t* gaussian_ids_sorted,
+    device const int* tile_bins, // int2
+    device const float* xys, // float2
+    device const float* conics, // float3
+    device const float* rgbs, // float3
+    device const float* opacities,
     constant float* background, // single float3
-    constant float* final_Ts,
-    constant int* final_index,
-    constant float* v_output, // float3
-    constant float* v_output_alpha,
+    device const float* final_Ts,
+    device const int* final_index,
+    device const float* v_output, // float3
+    device const float* v_output_alpha,
     device atomic_float* v_xy, // float2
     device atomic_float* v_conic, // float3
     device atomic_float* v_rgb, // float3
     device atomic_float* v_opacity,
+    constant int& num_points,
+    constant int& has_densification,
+    constant int& has_edge,
+    constant int& has_xy_abs,
+    device const float* error_map,
+    device const float* edge_map,
+    device atomic_float* densification_info, // [4, num_points]
+    device atomic_float* v_xy_abs, // [num_points, 2]
     uint3 gp [[thread_position_in_grid]],
     uint3 blockIdx [[threadgroup_position_in_grid]],
     uint tr [[thread_index_in_threadgroup]],
@@ -891,13 +996,15 @@ kernel void rasterize_backward_kernel(
     const int num_batches = (range.y - range.x + BLOCK_SIZE - 1) / BLOCK_SIZE;
 
     threadgroup int32_t id_batch[BLOCK_SIZE];
-    threadgroup float3 xy_opacity_batch[BLOCK_SIZE];
-    threadgroup float3 conic_batch[BLOCK_SIZE];
-    threadgroup float3 rgbs_batch[BLOCK_SIZE];
+    threadgroup packed_float3 xy_opacity_batch[BLOCK_SIZE];
+    threadgroup packed_float3 conic_batch[BLOCK_SIZE];
+    threadgroup packed_float3 rgbs_batch[BLOCK_SIZE];
 
     // df/d_out for this pixel
     const float3 v_out = read_packed_float3(v_output, pix_id);
     const float v_out_alpha = v_output_alpha[pix_id];
+    const float pix_err = (has_densification != 0 && inside) ? error_map[pix_id] : 1.0f;
+    const float pix_edge = (has_densification != 0 && has_edge != 0 && inside) ? edge_map[pix_id] : 0.0f;
 
     // collect and process batches of gaussians
     // each thread loads one gaussian at a time before rasterizing
@@ -918,9 +1025,9 @@ kernel void rasterize_backward_kernel(
             id_batch[tr] = g_id;
             const float2 xy = read_packed_float2(xys, g_id);
             const float opac = opacities[g_id];
-            xy_opacity_batch[tr] = {xy.x, xy.y, opac};
-            conic_batch[tr] = read_packed_float3(conics, g_id);
-            rgbs_batch[tr] = read_packed_float3(rgbs, g_id);
+            xy_opacity_batch[tr] = packed_float3(xy.x, xy.y, opac);
+            conic_batch[tr] = packed_float3(read_packed_float3(conics, g_id));
+            rgbs_batch[tr] = packed_float3(read_packed_float3(rgbs, g_id));
         }
         // wait for other threads to collect the gaussians in batch
         threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -946,7 +1053,7 @@ kernel void rasterize_backward_kernel(
                                             conic.z * delta.y * delta.y) +
                                     conic.y * delta.x * delta.y;
                 vis = exp(-sigma);
-                alpha = min(0.99f, opac * vis);
+                alpha = min(0.999f, opac * vis);
                 if (sigma < 0.f || alpha < 1.f / 255.f) {
                     valid = 0;
                 }
@@ -960,13 +1067,24 @@ kernel void rasterize_backward_kernel(
             float3 v_conic_local = {0.f, 0.f, 0.f};
             float2 v_xy_local = {0.f, 0.f};
             float v_opacity_local = 0.f;
+            float dens_w_local = 0.f;
+            float dens_e_local = 0.f;
+            float dens_g_local = 0.f;
+            float dens_c_local = 0.f;
+            float2 v_xy_abs_local = {0.f, 0.f};
             //initialize everything to 0, only set if the lane is valid
-            if(valid && alpha<0.99f){
+            if(valid){
                 // compute the current T for this gaussian
                 float ra = 1.f / (1.f - alpha);
                 T *= ra;
-                // update v_rgb for this gaussian
                 const float fac = alpha * T;
+                if (has_densification != 0){
+                    dens_w_local = fac;
+                    dens_e_local = fac * pix_err;
+                    dens_g_local = fac * pix_edge;
+                    dens_c_local = pix_err > 0.5f ? 1.0f : 0.0f;
+                }
+                // update v_rgb for this gaussian
                 float v_alpha = 0.f;
                 v_rgb_local = {fac * v_out.x, fac * v_out.y, fac * v_out.z};
 
@@ -987,11 +1105,14 @@ kernel void rasterize_backward_kernel(
                 buffer.z += rgb.z * fac;
 
                 const float v_sigma = -opac * vis * v_alpha;
-                v_conic_local = {0.5f * v_sigma * delta.x * delta.x, 
-                                        0.5f * v_sigma * delta.x * delta.y, 
+                v_conic_local = {0.5f * v_sigma * delta.x * delta.x,
+                                        0.5f * v_sigma * delta.x * delta.y,
                                         0.5f * v_sigma * delta.y * delta.y};
-                v_xy_local = {v_sigma * (conic.x * delta.x + conic.y * delta.y), 
+                v_xy_local = {v_sigma * (conic.x * delta.x + conic.y * delta.y),
                                     v_sigma * (conic.y * delta.x + conic.z * delta.y)};
+                if (has_xy_abs != 0){
+                    v_xy_abs_local = {fabs(v_xy_local.x), fabs(v_xy_local.y)};
+                }
                 v_opacity_local = vis * v_alpha;
             }
 
@@ -999,22 +1120,45 @@ kernel void rasterize_backward_kernel(
             v_conic_local = warpSum3(v_conic_local, warp_size, wr);
             v_xy_local = warpSum2(v_xy_local, warp_size, wr);
             v_opacity_local = warpSum(v_opacity_local, warp_size, wr);
+            if (has_densification != 0){
+                dens_w_local = warpSum(dens_w_local, warp_size, wr);
+                dens_e_local = warpSum(dens_e_local, warp_size, wr);
+                dens_g_local = warpSum(dens_g_local, warp_size, wr);
+                dens_c_local = warpSum(dens_c_local, warp_size, wr);
+            }
+            if (has_xy_abs != 0){
+                v_xy_abs_local = warpSum2(v_xy_abs_local, warp_size, wr);
+            }
 
-            if (wr == 0) {
-                int32_t g = id_batch[t];
-
-                atomic_fetch_add_explicit(v_rgb + 3*g + 0, v_rgb_local.x, memory_order_relaxed);
-                atomic_fetch_add_explicit(v_rgb + 3*g + 1, v_rgb_local.y, memory_order_relaxed);
-                atomic_fetch_add_explicit(v_rgb + 3*g + 2, v_rgb_local.z, memory_order_relaxed);
-                
-                atomic_fetch_add_explicit(v_conic + 3*g + 0, v_conic_local.x, memory_order_relaxed);
-                atomic_fetch_add_explicit(v_conic + 3*g + 1, v_conic_local.y, memory_order_relaxed);
-                atomic_fetch_add_explicit(v_conic + 3*g + 2, v_conic_local.z, memory_order_relaxed);
-                
-                atomic_fetch_add_explicit(v_xy + 2*g + 0, v_xy_local.x, memory_order_relaxed);
-                atomic_fetch_add_explicit(v_xy + 2*g + 1, v_xy_local.y, memory_order_relaxed);
-                
-                atomic_fetch_add_explicit(v_opacity + g, v_opacity_local, memory_order_relaxed);
+            const int dens_base = 9;
+            const int abs_base = dens_base + (has_densification != 0 ? 4 : 0);
+            const int n_slots = abs_base + (has_xy_abs != 0 ? 2 : 0);
+            if ((int)wr < n_slots) {
+                const int32_t g = id_batch[t];
+                device atomic_float* p;
+                float v;
+                if (wr < 3) {
+                    p = v_rgb + 3*g + wr;
+                    v = wr == 0 ? v_rgb_local.x : (wr == 1 ? v_rgb_local.y : v_rgb_local.z);
+                } else if (wr < 6) {
+                    p = v_conic + 3*g + (wr - 3);
+                    v = wr == 3 ? v_conic_local.x : (wr == 4 ? v_conic_local.y : v_conic_local.z);
+                } else if (wr < 8) {
+                    p = v_xy + 2*g + (wr - 6);
+                    v = wr == 6 ? v_xy_local.x : v_xy_local.y;
+                } else if (wr == 8) {
+                    p = v_opacity + g;
+                    v = v_opacity_local;
+                } else if ((int)wr < abs_base) {
+                    const int k = (int)wr - dens_base;
+                    p = densification_info + k * num_points + g;
+                    v = k == 0 ? dens_w_local : (k == 1 ? dens_e_local : (k == 2 ? dens_g_local : dens_c_local));
+                } else {
+                    const int k = (int)wr - abs_base;
+                    p = v_xy_abs + 2*g + k;
+                    v = k == 0 ? v_xy_abs_local.x : v_xy_abs_local.y;
+                }
+                atomic_fetch_add_explicit(p, v, memory_order_relaxed);
             }
         }
     }
@@ -1090,7 +1234,7 @@ kernel void nd_rasterize_backward_kernel(
         }
         const float opac = opacities[g];
         const float vis = exp(-sigma);
-        const float alpha = min(0.99f, opac * vis);
+        const float alpha = min(0.999f, opac * vis);
         if (alpha < 1.f / 255.f) {
             continue;
         }
@@ -1433,3 +1577,315 @@ kernel void compute_cov2d_bounds_kernel(
     conics[index + 2] = conic.z;
     radii[row] = radius;
 }
+
+// Fused L1 + DSSIM loss over [H,W,C] images. One kernel computes the
+// SSIM map and the closed-form partials in a threadgroup-memory tile
+// (two-pass separable 11-tap blur), one reduces to the scalar loss
+// on-device, and one produces dL/dimage directly.
+
+#define LOSS_BX 16
+#define LOSS_BY 16
+#define LOSS_HALO 5
+#define LOSS_SX (LOSS_BX + 2 * LOSS_HALO)
+#define LOSS_SY (LOSS_BY + 2 * LOSS_HALO)
+#define LOSS_C1 0.0001f
+#define LOSS_C2 0.0009f
+
+// 11-tap gaussian (sigma 1.5), matches the SSIM reference window
+constant float lossGauss[11] = {
+    0.001028380123898387f, 0.0075987582094967365f, 0.036000773310661316f,
+    0.10936068743467331f, 0.21300552785396576f, 0.26601171493530273f,
+    0.21300552785396576f, 0.10936068743467331f, 0.036000773310661316f,
+    0.0075987582094967365f, 0.001028380123898387f};
+
+inline float loss_pix(device const float* img, int y, int x, int c, int H, int W, int C){
+    if (x < 0 || x >= W || y < 0 || y >= H) return 0.0f;
+    return img[(y * W + x) * C + c];
+}
+
+inline bool loss_valid(int y, int x, int H, int W, bool validPad){
+    if (!validPad || H <= 10 || W <= 10) return true;
+    return x >= LOSS_HALO && x < W - LOSS_HALO && y >= LOSS_HALO && y < H - LOSS_HALO;
+}
+
+kernel void fused_loss_fwd_kernel(
+    constant int& H,
+    constant int& W,
+    constant int& C,
+    constant int& wantGrad,
+    device const float* rendered,
+    device const float* gt,
+    device float* ssimMap, // [H,W] channel mean
+    device half* pMu,      // [H,W,C] each (dummy when !wantGrad)
+    device half* pS1,
+    device half* pS12,
+    uint2 tpos [[thread_position_in_threadgroup]],
+    uint2 gpos [[threadgroup_position_in_grid]]
+) {
+    const int px = (int)(gpos.x * LOSS_BX + tpos.x);
+    const int py = (int)(gpos.y * LOSS_BY + tpos.y);
+    const int tileX = (int)(gpos.x * LOSS_BX);
+    const int tileY = (int)(gpos.y * LOSS_BY);
+
+    threadgroup float sTile[LOSS_SY][LOSS_SX][2];
+    threadgroup float sConv[LOSS_SY][LOSS_BX][5];
+
+    float ssimSum = 0.0f;
+    for (int c = 0; c < C; c++){
+        // Load tile + halo
+        {
+            const int tileSize = LOSS_SY * LOSS_SX;
+            const int threads = LOSS_BX * LOSS_BY;
+            const int tRank = (int)(tpos.y * LOSS_BX + tpos.x);
+            for (int tid = tRank; tid < tileSize; tid += threads){
+                const int ly = tid / LOSS_SX;
+                const int lx = tid % LOSS_SX;
+                const int gy = tileY + ly - LOSS_HALO;
+                const int gx = tileX + lx - LOSS_HALO;
+                sTile[ly][lx][0] = loss_pix(rendered, gy, gx, c, H, W, C);
+                sTile[ly][lx][1] = loss_pix(gt, gy, gx, c, H, W, C);
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Horizontal pass: accumulate moments; each thread covers two rows
+        {
+            const int lx = (int)tpos.x + LOSS_HALO;
+            for (int pass = 0; pass < 2; pass++){
+                const int ly = (int)tpos.y + pass * LOSS_BY;
+                if (ly >= LOSS_SY) break;
+                float sX = 0.f, sX2 = 0.f, sY = 0.f, sY2 = 0.f, sXY = 0.f;
+                for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                    const float w = lossGauss[LOSS_HALO + d];
+                    const float x = sTile[ly][lx + d][0];
+                    const float y = sTile[ly][lx + d][1];
+                    sX += x * w;
+                    sX2 += x * x * w;
+                    sY += y * w;
+                    sY2 += y * y * w;
+                    sXY += x * y * w;
+                }
+                sConv[ly][tpos.x][0] = sX;
+                sConv[ly][tpos.x][1] = sX2;
+                sConv[ly][tpos.x][2] = sY;
+                sConv[ly][tpos.x][3] = sY2;
+                sConv[ly][tpos.x][4] = sXY;
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Vertical pass + SSIM + partials
+        if (px < W && py < H){
+            const int ly = (int)tpos.y + LOSS_HALO;
+            const int lx = (int)tpos.x;
+            float m0 = 0.f, m1 = 0.f, m2 = 0.f, m3 = 0.f, m4 = 0.f;
+            for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                const float w = lossGauss[LOSS_HALO + d];
+                m0 += sConv[ly + d][lx][0] * w;
+                m1 += sConv[ly + d][lx][1] * w;
+                m2 += sConv[ly + d][lx][2] * w;
+                m3 += sConv[ly + d][lx][3] * w;
+                m4 += sConv[ly + d][lx][4] * w;
+            }
+            const float muX = m0;
+            const float muY = m2;
+            const float sigmaX = m1 - muX * muX;
+            const float sigmaY = m3 - muY * muY;
+            const float sigmaXY = m4 - muX * muY;
+
+            const float A = muX * muX + muY * muY + LOSS_C1;
+            const float B = sigmaX + sigmaY + LOSS_C2;
+            const float Cc = 2.f * muX * muY + LOSS_C1;
+            const float Dc = 2.f * sigmaXY + LOSS_C2;
+            const float s = (Cc * Dc) / (A * B);
+            ssimSum += s;
+
+            if (wantGrad){
+                const int idx = (py * W + px) * C + c;
+                const float dMu = (muY * 2.f * Dc) / (A * B) - (muY * 2.f * Cc) / (A * B)
+                                - (muX * 2.f * Cc * Dc) / (A * A * B) + (muX * 2.f * Cc * Dc) / (A * B * B);
+                pMu[idx] = half(dMu);
+                pS1[idx] = half((-Cc * Dc) / (A * B * B));
+                pS12[idx] = half((2.f * Cc) / (A * B));
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    if (px < W && py < H){
+        ssimMap[py * W + px] = ssimSum / (float)C;
+    }
+}
+
+kernel void fused_loss_reduce_kernel(
+    constant int& H,
+    constant int& W,
+    constant int& C,
+    constant int& hasMask,
+    constant float& ssimWeight,
+    constant int& validPad,
+    device const float* rendered,
+    device const float* gt,
+    device const float* ssimMap,
+    device const float* mask,
+    device atomic_float* out,
+    uint tid [[thread_position_in_threadgroup]],
+    uint gid [[thread_position_in_grid]],
+    uint gridSize [[threads_per_grid]]
+) {
+    const int numPix = H * W;
+    float lossSum = 0.0f;
+    float gateSum = 0.0f;
+    for (int p = (int)gid; p < numPix; p += (int)gridSize){
+        const int y = p / W;
+        const int x = p % W;
+        float gate;
+        if (hasMask){
+            gate = mask[p];
+        }else{
+            gate = loss_valid(y, x, H, W, validPad != 0) ? 1.0f : 0.0f;
+        }
+        if (gate != 0.0f){
+            float l1 = 0.0f;
+            for (int c = 0; c < C; c++){
+                l1 += fabs(rendered[p * C + c] - gt[p * C + c]);
+            }
+            const float contrib = (1.0f - ssimWeight) * l1
+                                + (float)C * ssimWeight * (1.0f - ssimMap[p]);
+            lossSum += gate * contrib;
+            gateSum += gate;
+        }
+    }
+
+    threadgroup float sLoss[256];
+    threadgroup float sGate[256];
+    sLoss[tid] = lossSum;
+    sGate[tid] = gateSum;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    for (uint stride = 128; stride > 0; stride >>= 1){
+        if (tid < stride){
+            sLoss[tid] += sLoss[tid + stride];
+            sGate[tid] += sGate[tid + stride];
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+    if (tid == 0){
+        atomic_fetch_add_explicit(&out[0], sLoss[0], memory_order_relaxed);
+        atomic_fetch_add_explicit(&out[1], sGate[0], memory_order_relaxed);
+    }
+}
+
+kernel void fused_loss_finalize_kernel(
+    constant int& C,
+    device float* out,
+    uint i [[thread_position_in_grid]]
+) {
+    if (i > 0) return;
+    const float denom = out[1] * (float)C + 1e-8f;
+    out[0] = out[0] / denom;
+    out[1] = denom;
+}
+
+kernel void fused_loss_bwd_kernel(
+    constant int& H,
+    constant int& W,
+    constant int& C,
+    constant int& hasMask,
+    constant float& ssimWeight,
+    constant int& validPad,
+    device const float* rendered,
+    device const float* gt,
+    device const float* mask,
+    device const half* pMu,
+    device const half* pS1,
+    device const half* pS12,
+    device const float* stats, // stats[1] = denominator
+    device const float* vLoss,
+    device float* vRendered,
+    uint2 tpos [[thread_position_in_threadgroup]],
+    uint2 gpos [[threadgroup_position_in_grid]]
+) {
+    const int px = (int)(gpos.x * LOSS_BX + tpos.x);
+    const int py = (int)(gpos.y * LOSS_BY + tpos.y);
+    const int tileX = (int)(gpos.x * LOSS_BX);
+    const int tileY = (int)(gpos.y * LOSS_BY);
+    const float chainScale = vLoss[0] / stats[1];
+
+    threadgroup float sData[LOSS_SY][LOSS_SX][3];
+    threadgroup float sConv[LOSS_SY][LOSS_BX][3];
+
+    for (int c = 0; c < C; c++){
+        float p1 = 0.f, p2 = 0.f;
+        if (px < W && py < H){
+            p1 = rendered[(py * W + px) * C + c];
+            p2 = gt[(py * W + px) * C + c];
+        }
+
+        // Load the chain-weighted partials for the tile + halo
+        {
+            const int tileSize = LOSS_SY * LOSS_SX;
+            const int threads = LOSS_BX * LOSS_BY;
+            const int tRank = (int)(tpos.y * LOSS_BX + tpos.x);
+            for (int tid = tRank; tid < tileSize; tid += threads){
+                const int ly = tid / LOSS_SX;
+                const int lx = tid % LOSS_SX;
+                const int gy = tileY + ly - LOSS_HALO;
+                const int gx = tileX + lx - LOSS_HALO;
+                const bool inside = gx >= 0 && gx < W && gy >= 0 && gy < H;
+                float chain = 0.0f;
+                if (inside){
+                    const float gate = hasMask ? mask[gy * W + gx]
+                                               : (loss_valid(gy, gx, H, W, validPad != 0) ? 1.0f : 0.0f);
+                    chain = -ssimWeight * gate * chainScale;
+                }
+                const int idx = (gy * W + gx) * C + c;
+                sData[ly][lx][0] = inside ? (float)pMu[idx] * chain : 0.0f;
+                sData[ly][lx][1] = inside ? (float)pS1[idx] * chain : 0.0f;
+                sData[ly][lx][2] = inside ? (float)pS12[idx] * chain : 0.0f;
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Horizontal pass
+        {
+            const int lx = (int)tpos.x + LOSS_HALO;
+            for (int pass = 0; pass < 2; pass++){
+                const int ly = (int)tpos.y + pass * LOSS_BY;
+                if (ly >= LOSS_SY) break;
+                float a0 = 0.f, a1 = 0.f, a2 = 0.f;
+                for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                    const float w = lossGauss[LOSS_HALO + d];
+                    a0 += sData[ly][lx + d][0] * w;
+                    a1 += sData[ly][lx + d][1] * w;
+                    a2 += sData[ly][lx + d][2] * w;
+                }
+                sConv[ly][tpos.x][0] = a0;
+                sConv[ly][tpos.x][1] = a1;
+                sConv[ly][tpos.x][2] = a2;
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Vertical pass + L1 term
+        if (px < W && py < H){
+            const int ly = (int)tpos.y + LOSS_HALO;
+            const int lx = (int)tpos.x;
+            float s0 = 0.f, s1 = 0.f, s2 = 0.f;
+            for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                const float w = lossGauss[LOSS_HALO + d];
+                s0 += sConv[ly + d][lx][0] * w;
+                s1 += sConv[ly + d][lx][1] * w;
+                s2 += sConv[ly + d][lx][2] * w;
+            }
+            const float gradSsim = s0 + 2.f * p1 * s1 + p2 * s2;
+
+            const float gate = hasMask ? mask[py * W + px]
+                                       : (loss_valid(py, px, H, W, validPad != 0) ? 1.0f : 0.0f);
+            const float sign = (p1 == p2) ? 0.0f : copysign(1.0f, p1 - p2);
+            const float gradL1 = (1.0f - ssimWeight) * sign * gate * chainScale;
+
+            vRendered[(py * W + px) * C + c] = gradSsim + gradL1;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+}
diff --git a/rasterizer/gsplat-metal/gsplat_metal.mm b/rasterizer/gsplat-metal/gsplat_metal.mm
index 29849e78..bd3cd858 100644
--- a/rasterizer/gsplat-metal/gsplat_metal.mm
+++ b/rasterizer/gsplat-metal/gsplat_metal.mm
@@ -12,6 +12,7 @@
 
     id<MTLComputePipelineState> nd_rasterize_backward_kernel_cpso;
     id<MTLComputePipelineState> nd_rasterize_forward_kernel_cpso;
+    id<MTLComputePipelineState> rasterize_forward_kernel_cpso;
     id<MTLComputePipelineState> rasterize_backward_kernel_cpso;
     id<MTLComputePipelineState> project_gaussians_forward_kernel_cpso;
     id<MTLComputePipelineState> project_gaussians_backward_kernel_cpso;
@@ -20,6 +21,10 @@
     id<MTLComputePipelineState> compute_cov2d_bounds_kernel_cpso;
     id<MTLComputePipelineState> map_gaussian_to_intersects_kernel_cpso;
     id<MTLComputePipelineState> get_tile_bin_edges_kernel_cpso;
+    id<MTLComputePipelineState> fused_loss_fwd_kernel_cpso;
+    id<MTLComputePipelineState> fused_loss_reduce_kernel_cpso;
+    id<MTLComputePipelineState> fused_loss_finalize_kernel_cpso;
+    id<MTLComputePipelineState> fused_loss_bwd_kernel_cpso;
 };
 
 unsigned num_sh_bases(const unsigned degree) {
@@ -109,6 +114,7 @@ @implementation DummyClassForPathHack
 
     GSPLAT_METAL_ADD_KERNEL(nd_rasterize_backward_kernel);
     GSPLAT_METAL_ADD_KERNEL(nd_rasterize_forward_kernel);
+    GSPLAT_METAL_ADD_KERNEL(rasterize_forward_kernel);
     GSPLAT_METAL_ADD_KERNEL(rasterize_backward_kernel);
     GSPLAT_METAL_ADD_KERNEL(project_gaussians_forward_kernel);
     GSPLAT_METAL_ADD_KERNEL(project_gaussians_backward_kernel);
@@ -117,6 +123,10 @@ @implementation DummyClassForPathHack
     GSPLAT_METAL_ADD_KERNEL(compute_cov2d_bounds_kernel);
     GSPLAT_METAL_ADD_KERNEL(map_gaussian_to_intersects_kernel);
     GSPLAT_METAL_ADD_KERNEL(get_tile_bin_edges_kernel);
+    GSPLAT_METAL_ADD_KERNEL(fused_loss_fwd_kernel);
+    GSPLAT_METAL_ADD_KERNEL(fused_loss_reduce_kernel);
+    GSPLAT_METAL_ADD_KERNEL(fused_loss_finalize_kernel);
+    GSPLAT_METAL_ADD_KERNEL(fused_loss_bwd_kernel);
 
     [metal_library release];
 
@@ -182,16 +192,16 @@ static EncodeArg tensor(const torch::Tensor& x) {
     size_t _arrayNumBytes;
     const torch::Tensor* _tensor;
 
-    friend void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize grid_size, MTLSize thread_group_size, std::vector<EncodeArg> args);
+    friend void dispatchKernelEx(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize grid_size, MTLSize thread_group_size, bool by_threadgroups, std::vector<EncodeArg> args);
 };
 
-void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize grid_size, MTLSize thread_group_size, std::vector<EncodeArg> args) {
-    // Get a reference to the command buffer for the MPS stream
-    id<MTLCommandBuffer> command_buffer = torch::mps::get_command_buffer();
-    TORCH_CHECK(command_buffer, "Failed to retrieve command buffer reference");
-
+void dispatchKernelEx(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize grid_size, MTLSize thread_group_size, bool by_threadgroups, std::vector<EncodeArg> args) {
     // Dispatch the kernel
     dispatch_sync(ctx->d_queue, ^(){
+        // Get a reference to the command buffer for the MPS stream
+        id<MTLCommandBuffer> command_buffer = torch::mps::get_command_buffer();
+        TORCH_CHECK(command_buffer, "Failed to retrieve command buffer reference");
+
         // Start a compute pass
         id<MTLComputeCommandEncoder> encoder = [command_buffer computeCommandEncoder];
         TORCH_CHECK(encoder, "Failed to create compute command encoder");
@@ -222,14 +232,22 @@ void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize
         }
 
         // Dispatch the compute command
-        [encoder dispatchThreads:grid_size threadsPerThreadgroup:thread_group_size];
+        if (by_threadgroups) {
+            [encoder dispatchThreadgroups:grid_size threadsPerThreadgroup:thread_group_size];
+        } else {
+            [encoder dispatchThreads:grid_size threadsPerThreadgroup:thread_group_size];
+        }
         [encoder endEncoding];
 
-        // Commit the work
-        torch::mps::synchronize();
+        // Submit the work without blocking
+        torch::mps::commit();
     });
 }
 
+void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize grid_size, MTLSize thread_group_size, std::vector<EncodeArg> args) {
+    dispatchKernelEx(ctx, cpso, grid_size, thread_group_size, false, args);
+}
+
 std::tuple<
     torch::Tensor, // output conics
     torch::Tensor> // output radii
@@ -574,6 +592,48 @@ void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize
     const int img_width = std::get<0>(img_size);
     const int img_height = std::get<1>(img_size);
 
+    uint32_t img_size_dim3[4] = {(uint32_t)std::get<0>(img_size), (uint32_t)std::get<1>(img_size), (uint32_t)std::get<2>(img_size), 0xDEAD};
+    uint32_t tile_bounds_arr[4] = {
+        (uint32_t)std::get<0>(tile_bounds),
+        (uint32_t)std::get<1>(tile_bounds),
+        (uint32_t)std::get<2>(tile_bounds),
+        0xDEAD
+    };
+    int32_t block_size_dim2[2] = {std::get<0>(block), std::get<1>(block)};
+
+    MetalContext* ctx = get_global_context();
+
+    if (channels == 3 && block_size_dim2[0] == BLOCK_X && block_size_dim2[1] == BLOCK_Y) {
+        torch::Tensor out_img = torch::empty(
+            {img_height, img_width, channels}, xys.options().dtype(torch::kFloat32)
+        );
+        torch::Tensor final_Ts = torch::empty(
+            {img_height, img_width}, xys.options().dtype(torch::kFloat32)
+        );
+        torch::Tensor final_idx = torch::empty(
+            {img_height, img_width}, xys.options().dtype(torch::kInt32)
+        );
+
+        MTLSize groups = MTLSizeMake(tile_bounds_arr[0], tile_bounds_arr[1], 1);
+        MTLSize thread_group_size = MTLSizeMake(BLOCK_X, BLOCK_Y, 1);
+        dispatchKernelEx(ctx, ctx->rasterize_forward_kernel_cpso, groups, thread_group_size, true, {
+            EncodeArg::array(tile_bounds_arr, sizeof(tile_bounds_arr)),
+            EncodeArg::array(img_size_dim3, sizeof(img_size_dim3)),
+            EncodeArg::tensor(gaussian_ids_sorted),
+            EncodeArg::tensor(tile_bins),
+            EncodeArg::tensor(xys),
+            EncodeArg::tensor(conics),
+            EncodeArg::tensor(colors),
+            EncodeArg::tensor(opacities),
+            EncodeArg::tensor(final_Ts),
+            EncodeArg::tensor(final_idx),
+            EncodeArg::tensor(out_img),
+            EncodeArg::tensor(background)
+        });
+
+        return std::make_tuple(out_img, final_Ts, final_idx);
+    }
+
     torch::Tensor out_img = torch::zeros(
         {img_height, img_width, channels}, xys.options().dtype(torch::kFloat32)
     );
@@ -584,16 +644,6 @@ void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize
         {img_height, img_width}, xys.options().dtype(torch::kInt32)
     );
 
-    uint32_t img_size_dim3[4] = {(uint32_t)std::get<0>(img_size), (uint32_t)std::get<1>(img_size), (uint32_t)std::get<2>(img_size), 0xDEAD};
-    uint32_t tile_bounds_arr[4] = {
-        (uint32_t)std::get<0>(tile_bounds), 
-        (uint32_t)std::get<1>(tile_bounds), 
-        (uint32_t)std::get<2>(tile_bounds), 
-        0xDEAD
-    };
-    int32_t block_size_dim2[2] = {std::get<0>(block), std::get<1>(block)};
-
-    MetalContext* ctx = get_global_context();
     MTLSize grid_size = MTLSizeMake(img_width, img_height, 1);
     MTLSize thread_group_size = MTLSizeMake(block_size_dim2[0], block_size_dim2[1], 1);
     dispatchKernel(ctx, ctx->nd_rasterize_forward_kernel_cpso, grid_size, thread_group_size, {
@@ -748,9 +798,9 @@ void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize
     };
 
     MetalContext* ctx = get_global_context();
-    MTLSize grid_size = MTLSizeMake(img_width, img_height, 1);
+    MTLSize groups = MTLSizeMake(tile_bounds_arr[0], tile_bounds_arr[1], 1);
     MTLSize thread_group_size = MTLSizeMake(BLOCK_X, BLOCK_Y, 1);
-    dispatchKernel(ctx, ctx->nd_rasterize_backward_kernel_cpso, grid_size, thread_group_size, {
+    dispatchKernelEx(ctx, ctx->nd_rasterize_backward_kernel_cpso, groups, thread_group_size, true, {
         EncodeArg::array(tile_bounds_arr, sizeof(tile_bounds_arr)),
         EncodeArg::array(img_size, sizeof(img_size)),
         EncodeArg::scalar(channels),
@@ -795,7 +845,11 @@ void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize
         const torch::Tensor &final_Ts,
         const torch::Tensor &final_idx,
         const torch::Tensor &v_output, // dL_dout_color
-        const torch::Tensor &v_output_alpha
+        const torch::Tensor &v_output_alpha,
+        const torch::Tensor &error_map,
+        const torch::Tensor &edge_map,
+        const torch::Tensor &densification_info,
+        const torch::Tensor &v_xy_abs
     ) {
     CHECK_INPUT(gaussians_ids_sorted);
     CHECK_INPUT(tile_bins);
@@ -818,6 +872,15 @@ void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize
         torch::zeros({num_points, channels}, xys.options());
     torch::Tensor v_opacity = torch::zeros({num_points, 1}, xys.options());
 
+    const bool has_dinfo = densification_info.numel() > 0;
+    torch::Tensor dummy = torch::zeros({1}, xys.options());
+    torch::Tensor errorMapBuf = error_map.numel() > 0 ? error_map : dummy;
+    torch::Tensor edgeMapBuf = edge_map.numel() > 0 ? edge_map : dummy;
+    const int has_edge = edge_map.numel() > 0 ? 1 : 0;
+    torch::Tensor dinfoBuf = has_dinfo ? densification_info : dummy;
+    const int has_xy_abs = v_xy_abs.numel() > 0 ? 1 : 0;
+    torch::Tensor xyAbsBuf = has_xy_abs ? v_xy_abs : dummy;
+
     // Get a reference to the command buffer for the MPS stream
     id<MTLCommandBuffer> command_buffer = torch::mps::get_command_buffer();
     TORCH_CHECK(command_buffer, "Failed to retrieve command buffer reference");
@@ -831,9 +894,9 @@ void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize
     };
 
     MetalContext* ctx = get_global_context();
-    MTLSize grid_size = MTLSizeMake(img_width, img_height, 1);
+    MTLSize groups = MTLSizeMake(tile_bounds_arr[0], tile_bounds_arr[1], 1);
     MTLSize thread_group_size = MTLSizeMake(BLOCK_X, BLOCK_Y, 1);
-    dispatchKernel(ctx, ctx->rasterize_backward_kernel_cpso, grid_size, thread_group_size, {
+    dispatchKernelEx(ctx, ctx->rasterize_backward_kernel_cpso, groups, thread_group_size, true, {
         EncodeArg::array(tile_bounds_arr, sizeof(tile_bounds_arr)),
         EncodeArg::array(img_size, sizeof(img_size)),
         EncodeArg::tensor(gaussians_ids_sorted),
@@ -850,8 +913,129 @@ void dispatchKernel(MetalContext* ctx, id<MTLComputePipelineState> cpso, MTLSize
         EncodeArg::tensor(v_xy),
         EncodeArg::tensor(v_conic),
         EncodeArg::tensor(v_colors),
-        EncodeArg::tensor(v_opacity)
+        EncodeArg::tensor(v_opacity),
+        EncodeArg::scalar((int32_t)num_points),
+        EncodeArg::scalar((int32_t)(has_dinfo ? 1 : 0)),
+        EncodeArg::scalar((int32_t)has_edge),
+        EncodeArg::scalar((int32_t)has_xy_abs),
+        EncodeArg::tensor(errorMapBuf),
+        EncodeArg::tensor(edgeMapBuf),
+        EncodeArg::tensor(dinfoBuf),
+        EncodeArg::tensor(xyAbsBuf)
     });
 
     return std::make_tuple(v_xy, v_conic, v_colors, v_opacity);
-}
\ No newline at end of file
+}
+std::tuple<torch::Tensor, torch::Tensor> fused_loss_forward_tensor(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask,
+    const float ssim_weight,
+    const bool valid_padding,
+    const bool want_grad
+){
+    CHECK_INPUT(rendered);
+    CHECK_INPUT(gt);
+    const int H = rendered.size(0);
+    const int W = rendered.size(1);
+    const int C = rendered.size(2);
+    const bool hasMask = mask.defined() && mask.numel() > 0;
+    if (hasMask){ CHECK_INPUT(mask); }
+
+    auto opts = rendered.options();
+    torch::Tensor ssimMap = torch::empty({H, W}, opts);
+    torch::Tensor partials = want_grad
+        ? torch::empty({3, static_cast<long long>(H) * W * C}, opts.dtype(torch::kHalf))
+        : torch::empty({1}, opts.dtype(torch::kHalf)); // dummy buffer
+    torch::Tensor maskBuf = hasMask ? mask : torch::empty({1}, opts); // dummy buffer
+    const long long planeSize = static_cast<long long>(H) * W * C;
+    (void)planeSize;
+    torch::Tensor pMu = want_grad ? partials[0] : partials;
+    torch::Tensor pS1 = want_grad ? partials[1] : partials;
+    torch::Tensor pS12 = want_grad ? partials[2] : partials;
+
+    MetalContext* ctx = get_global_context();
+    MTLSize tileGrid = MTLSizeMake((W + 15) / 16, (H + 15) / 16, 1);
+    MTLSize tileGroup = MTLSizeMake(16, 16, 1);
+    dispatchKernelEx(ctx, ctx->fused_loss_fwd_kernel_cpso, tileGrid, tileGroup, true, {
+        EncodeArg::scalar((int32_t)H),
+        EncodeArg::scalar((int32_t)W),
+        EncodeArg::scalar((int32_t)C),
+        EncodeArg::scalar((int32_t)(want_grad ? 1 : 0)),
+        EncodeArg::tensor(rendered),
+        EncodeArg::tensor(gt),
+        EncodeArg::tensor(ssimMap),
+        EncodeArg::tensor(pMu),
+        EncodeArg::tensor(pS1),
+        EncodeArg::tensor(pS12)
+    });
+
+    torch::Tensor stats = torch::zeros({2}, opts);
+    const int numPix = H * W;
+    const int reduceThreads = std::min(numPix, 256 * 1024);
+    MTLSize reduceGrid = MTLSizeMake((reduceThreads + 255) / 256, 1, 1);
+    MTLSize reduceGroup = MTLSizeMake(256, 1, 1);
+    dispatchKernelEx(ctx, ctx->fused_loss_reduce_kernel_cpso, reduceGrid, reduceGroup, true, {
+        EncodeArg::scalar((int32_t)H),
+        EncodeArg::scalar((int32_t)W),
+        EncodeArg::scalar((int32_t)C),
+        EncodeArg::scalar((int32_t)(hasMask ? 1 : 0)),
+        EncodeArg::scalar(ssim_weight),
+        EncodeArg::scalar((int32_t)(valid_padding ? 1 : 0)),
+        EncodeArg::tensor(rendered),
+        EncodeArg::tensor(gt),
+        EncodeArg::tensor(ssimMap),
+        EncodeArg::tensor(maskBuf),
+        EncodeArg::tensor(stats)
+    });
+    dispatchKernel(ctx, ctx->fused_loss_finalize_kernel_cpso, MTLSizeMake(1, 1, 1), MTLSizeMake(1, 1, 1), {
+        EncodeArg::scalar((int32_t)C),
+        EncodeArg::tensor(stats)
+    });
+
+    return std::make_tuple(stats, want_grad ? partials : torch::empty({0}, opts.dtype(torch::kHalf)));
+}
+
+torch::Tensor fused_loss_backward_tensor(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask,
+    const torch::Tensor &partials,
+    const torch::Tensor &stats,
+    const torch::Tensor &v_loss,
+    const float ssim_weight,
+    const bool valid_padding
+){
+    const int H = rendered.size(0);
+    const int W = rendered.size(1);
+    const int C = rendered.size(2);
+    const bool hasMask = mask.defined() && mask.numel() > 0;
+
+    torch::Tensor vRendered = torch::empty_like(rendered);
+    torch::Tensor maskBuf = hasMask ? mask : torch::empty({1}, rendered.options());
+    torch::Tensor pMu = partials[0];
+    torch::Tensor pS1 = partials[1];
+    torch::Tensor pS12 = partials[2];
+
+    MetalContext* ctx = get_global_context();
+    MTLSize tileGrid = MTLSizeMake((W + 15) / 16, (H + 15) / 16, 1);
+    MTLSize tileGroup = MTLSizeMake(16, 16, 1);
+    dispatchKernelEx(ctx, ctx->fused_loss_bwd_kernel_cpso, tileGrid, tileGroup, true, {
+        EncodeArg::scalar((int32_t)H),
+        EncodeArg::scalar((int32_t)W),
+        EncodeArg::scalar((int32_t)C),
+        EncodeArg::scalar((int32_t)(hasMask ? 1 : 0)),
+        EncodeArg::scalar(ssim_weight),
+        EncodeArg::scalar((int32_t)(valid_padding ? 1 : 0)),
+        EncodeArg::tensor(rendered),
+        EncodeArg::tensor(gt),
+        EncodeArg::tensor(maskBuf),
+        EncodeArg::tensor(pMu),
+        EncodeArg::tensor(pS1),
+        EncodeArg::tensor(pS12),
+        EncodeArg::tensor(stats),
+        EncodeArg::tensor(v_loss),
+        EncodeArg::tensor(vRendered)
+    });
+    return vRendered;
+}
diff --git a/rasterizer/gsplat/backward.cu b/rasterizer/gsplat/backward.cu
index c338dba2..baf42e05 100644
--- a/rasterizer/gsplat/backward.cu
+++ b/rasterizer/gsplat/backward.cu
@@ -1,4 +1,4 @@
-#include "backward.cuh"
+#include "backward.cuh"
 #include "helpers.cuh"
 
 #ifdef USE_HIP
@@ -85,7 +85,7 @@ __global__ void nd_rasterize_backward_kernel(
         }
         const float opac = opacities[g];
         const float vis = __expf(-sigma);
-        const float alpha = min(0.99f, opac * vis);
+        const float alpha = min(0.999f, opac * vis);
         if (alpha < 1.f / 255.f) {
             continue;
         }
@@ -161,6 +161,7 @@ inline __device__ void warpSum(float& val, cg::thread_block_tile<32>& tile){
 __global__ void rasterize_backward_kernel(
     const dim3 tile_bounds,
     const dim3 img_size,
+    const unsigned num_points,
     const int32_t* __restrict__ gaussian_ids_sorted,
     const int2* __restrict__ tile_bins,
     const float2* __restrict__ xys,
@@ -172,10 +173,14 @@ __global__ void rasterize_backward_kernel(
     const int* __restrict__ final_index,
     const float3* __restrict__ v_output,
     const float* __restrict__ v_output_alpha,
+    const float* __restrict__ error_map,
+    const float* __restrict__ edge_map,
     float2* __restrict__ v_xy,
+    float2* __restrict__ v_xy_abs,
     float3* __restrict__ v_conic,
     float3* __restrict__ v_rgb,
-    float* __restrict__ v_opacity
+    float* __restrict__ v_opacity,
+    float* __restrict__ densification_info // [4,N]: sum w, sum w*err, sum w*edge, count(err>0.5)
 ) {
     auto block = cg::this_thread_block();
     int32_t tile_id =
@@ -215,6 +220,9 @@ __global__ void rasterize_backward_kernel(
     // df/d_out for this pixel
     const float3 v_out = v_output[pix_id];
     const float v_out_alpha = v_output_alpha[pix_id];
+    const bool has_dinfo = densification_info != nullptr;
+    const float pix_err = (has_dinfo && error_map != nullptr && inside) ? error_map[pix_id] : 1.0f;
+    const float pix_edge = (has_dinfo && edge_map != nullptr && inside) ? edge_map[pix_id] : 0.0f;
 
     // collect and process batches of gaussians
     // each thread loads one gaussian at a time before rasterizing
@@ -269,7 +277,7 @@ __global__ void rasterize_backward_kernel(
                                             conic.z * delta.y * delta.y) +
                                     conic.y * delta.x * delta.y;
                 vis = __expf(-sigma);
-                alpha = min(0.99f, opac * vis);
+                alpha = min(0.999f, opac * vis);
                 if (sigma < 0.f || alpha < 1.f / 255.f) {
                     valid = 0;
                 }
@@ -294,6 +302,11 @@ __global__ void rasterize_backward_kernel(
             float3 v_conic_local = {0.f, 0.f, 0.f};
             float2 v_xy_local = {0.f, 0.f};
             float v_opacity_local = 0.f;
+            float dens_w_local = 0.f;
+            float dens_e_local = 0.f;
+            float dens_g_local = 0.f;
+            float dens_c_local = 0.f;
+            float2 v_xy_abs_local = {0.f, 0.f};
             //initialize everything to 0, only set if the lane is valid
             if(valid){
                 // compute the current T for this gaussian
@@ -301,6 +314,12 @@ __global__ void rasterize_backward_kernel(
                 T *= ra;
                 // update v_rgb for this gaussian
                 const float fac = alpha * T;
+                if (has_dinfo){
+                    dens_w_local = fac;
+                    dens_e_local = fac * pix_err;
+                    dens_g_local = fac * pix_edge;
+                    dens_c_local = pix_err > 0.5f ? 1.0f : 0.0f;
+                }
                 float v_alpha = 0.f;
                 v_rgb_local = {fac * v_out.x, fac * v_out.y, fac * v_out.z};
 
@@ -324,14 +343,26 @@ __global__ void rasterize_backward_kernel(
                 v_conic_local = {0.5f * v_sigma * delta.x * delta.x, 
                                         0.5f * v_sigma * delta.x * delta.y, 
                                         0.5f * v_sigma * delta.y * delta.y};
-                v_xy_local = {v_sigma * (conic.x * delta.x + conic.y * delta.y), 
+                v_xy_local = {v_sigma * (conic.x * delta.x + conic.y * delta.y),
                                     v_sigma * (conic.y * delta.x + conic.z * delta.y)};
+                if (v_xy_abs != nullptr){
+                    v_xy_abs_local = {fabsf(v_xy_local.x), fabsf(v_xy_local.y)};
+                }
                 v_opacity_local = vis * v_alpha;
             }
             warpSum3(v_rgb_local, warp);
             warpSum3(v_conic_local, warp);
             warpSum2(v_xy_local, warp);
             warpSum(v_opacity_local, warp);
+            if (has_dinfo){
+                warpSum(dens_w_local, warp);
+                warpSum(dens_e_local, warp);
+                warpSum(dens_g_local, warp);
+                warpSum(dens_c_local, warp);
+            }
+            if (v_xy_abs != nullptr){
+                warpSum2(v_xy_abs_local, warp);
+            }
             if (warp.thread_rank() == 0) {
                 int32_t g = id_batch[t];
                 float* v_rgb_ptr = (float*)(v_rgb);
@@ -349,6 +380,18 @@ __global__ void rasterize_backward_kernel(
                 atomicAdd(v_xy_ptr + 2*g + 1, v_xy_local.y);
                 
                 atomicAdd(v_opacity + g, v_opacity_local);
+
+                if (has_dinfo){
+                    atomicAdd(densification_info + g, dens_w_local);
+                    atomicAdd(densification_info + num_points + g, dens_e_local);
+                    atomicAdd(densification_info + 2 * num_points + g, dens_g_local);
+                    atomicAdd(densification_info + 3 * num_points + g, dens_c_local);
+                }
+                if (v_xy_abs != nullptr){
+                    float* v_xy_abs_ptr = (float*)(v_xy_abs);
+                    atomicAdd(v_xy_abs_ptr + 2*g + 0, v_xy_abs_local.x);
+                    atomicAdd(v_xy_abs_ptr + 2*g + 1, v_xy_abs_local.y);
+                }
             }
         }
     }
diff --git a/rasterizer/gsplat/backward.cuh b/rasterizer/gsplat/backward.cuh
index bac980f3..6298f407 100644
--- a/rasterizer/gsplat/backward.cuh
+++ b/rasterizer/gsplat/backward.cuh
@@ -59,6 +59,7 @@ __global__ void nd_rasterize_backward_kernel(
 __global__ void rasterize_backward_kernel(
     const dim3 tile_bounds,
     const dim3 img_size,
+    const unsigned num_points,
     const int32_t* __restrict__ gaussian_ids_sorted,
     const int2* __restrict__ tile_bins,
     const float2* __restrict__ xys,
@@ -70,10 +71,14 @@ __global__ void rasterize_backward_kernel(
     const int* __restrict__ final_index,
     const float3* __restrict__ v_output,
     const float* __restrict__ v_output_alpha,
+    const float* __restrict__ error_map,
+    const float* __restrict__ edge_map,
     float2* __restrict__ v_xy,
+    float2* __restrict__ v_xy_abs,
     float3* __restrict__ v_conic,
     float3* __restrict__ v_rgb,
-    float* __restrict__ v_opacity
+    float* __restrict__ v_opacity,
+    float* __restrict__ densification_info
 );
 
 __device__ void project_cov3d_ewa_vjp(
diff --git a/rasterizer/gsplat/bindings.cu b/rasterizer/gsplat/bindings.cu
index 80d581ce..cdb92f81 100644
--- a/rasterizer/gsplat/bindings.cu
+++ b/rasterizer/gsplat/bindings.cu
@@ -8,12 +8,14 @@
 #include <hip/hip_runtime.h>
 #include <hip/hip_runtime_api.h>
 #include <hip/hip_cooperative_groups.h>
+#include <hip/hip_fp16.h>
 #else
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 #include <cuda.h>
 #include <cuda_runtime.h>
 #include <cuda_runtime_api.h>
+#include <cuda_fp16.h>
 #endif
 
 #include <cstdio>
@@ -23,6 +25,391 @@
 
 namespace cg = cooperative_groups;
 
+// Fused L1 + DSSIM loss over [H,W,C] images. One kernel computes the
+// SSIM map and the closed-form partials in a shared-memory tile (two-pass
+// separable 11-tap blur), one reduces to the scalar loss on-device, and one
+// produces dL/dimage directly.
+
+#define LOSS_BX 16
+#define LOSS_BY 16
+#define LOSS_HALO 5
+#define LOSS_SX (LOSS_BX + 2 * LOSS_HALO)
+#define LOSS_SY (LOSS_BY + 2 * LOSS_HALO)
+#define LOSS_C1 0.0001f
+#define LOSS_C2 0.0009f
+
+// 11-tap gaussian (sigma 1.5), matches the SSIM reference window
+__device__ __constant__ float lossGauss[11] = {
+    0.001028380123898387f, 0.0075987582094967365f, 0.036000773310661316f,
+    0.10936068743467331f, 0.21300552785396576f, 0.26601171493530273f,
+    0.21300552785396576f, 0.10936068743467331f, 0.036000773310661316f,
+    0.0075987582094967365f, 0.001028380123898387f};
+
+__device__ __forceinline__ float loss_pix(
+    const float* __restrict__ img, int y, int x, int c, int H, int W, int C
+){
+    if (x < 0 || x >= W || y < 0 || y >= H) return 0.0f;
+    return img[(y * W + x) * C + c];
+}
+
+// A pixel participates in the unmasked loss only away from the blur border
+__device__ __forceinline__ bool loss_valid(int y, int x, int H, int W, bool validPad){
+    if (!validPad || H <= 10 || W <= 10) return true;
+    return x >= LOSS_HALO && x < W - LOSS_HALO && y >= LOSS_HALO && y < H - LOSS_HALO;
+}
+
+__global__ void fused_loss_fwd_kernel(
+    const int H, const int W, const int C,
+    const float* __restrict__ rendered,
+    const float* __restrict__ gt,
+    float* __restrict__ ssimMap, // [H,W] channel mean
+    __half* __restrict__ pMu,    // [H,W,C] each, or nullptr
+    __half* __restrict__ pS1,
+    __half* __restrict__ pS12
+){
+    const int px = blockIdx.x * LOSS_BX + threadIdx.x;
+    const int py = blockIdx.y * LOSS_BY + threadIdx.y;
+    const int tileX = blockIdx.x * LOSS_BX;
+    const int tileY = blockIdx.y * LOSS_BY;
+
+    __shared__ float sTile[LOSS_SY][LOSS_SX][2];
+    __shared__ float sConv[LOSS_SY][LOSS_BX][5];
+
+    float ssimSum = 0.0f;
+    for (int c = 0; c < C; c++){
+        // Load tile + halo
+        {
+            const int tileSize = LOSS_SY * LOSS_SX;
+            const int threads = LOSS_BX * LOSS_BY;
+            const int tRank = threadIdx.y * LOSS_BX + threadIdx.x;
+            for (int tid = tRank; tid < tileSize; tid += threads){
+                const int ly = tid / LOSS_SX;
+                const int lx = tid % LOSS_SX;
+                const int gy = tileY + ly - LOSS_HALO;
+                const int gx = tileX + lx - LOSS_HALO;
+                sTile[ly][lx][0] = loss_pix(rendered, gy, gx, c, H, W, C);
+                sTile[ly][lx][1] = loss_pix(gt, gy, gx, c, H, W, C);
+            }
+        }
+        __syncthreads();
+
+        // Horizontal pass: accumulate moments; each thread covers two rows
+        {
+            const int lx = threadIdx.x + LOSS_HALO;
+            for (int pass = 0; pass < 2; pass++){
+                const int ly = threadIdx.y + pass * LOSS_BY;
+                if (ly >= LOSS_SY) break;
+                float sX = 0.f, sX2 = 0.f, sY = 0.f, sY2 = 0.f, sXY = 0.f;
+                #pragma unroll
+                for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                    const float w = lossGauss[LOSS_HALO + d];
+                    const float x = sTile[ly][lx + d][0];
+                    const float y = sTile[ly][lx + d][1];
+                    sX += x * w;
+                    sX2 += x * x * w;
+                    sY += y * w;
+                    sY2 += y * y * w;
+                    sXY += x * y * w;
+                }
+                sConv[ly][threadIdx.x][0] = sX;
+                sConv[ly][threadIdx.x][1] = sX2;
+                sConv[ly][threadIdx.x][2] = sY;
+                sConv[ly][threadIdx.x][3] = sY2;
+                sConv[ly][threadIdx.x][4] = sXY;
+            }
+        }
+        __syncthreads();
+
+        // Vertical pass + SSIM + partials
+        if (px < W && py < H){
+            const int ly = threadIdx.y + LOSS_HALO;
+            const int lx = threadIdx.x;
+            float m0 = 0.f, m1 = 0.f, m2 = 0.f, m3 = 0.f, m4 = 0.f;
+            #pragma unroll
+            for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                const float w = lossGauss[LOSS_HALO + d];
+                const float* row = sConv[ly + d][lx];
+                m0 += row[0] * w;
+                m1 += row[1] * w;
+                m2 += row[2] * w;
+                m3 += row[3] * w;
+                m4 += row[4] * w;
+            }
+            const float muX = m0;
+            const float muY = m2;
+            const float sigmaX = m1 - muX * muX;
+            const float sigmaY = m3 - muY * muY;
+            const float sigmaXY = m4 - muX * muY;
+
+            const float A = muX * muX + muY * muY + LOSS_C1;
+            const float B = sigmaX + sigmaY + LOSS_C2;
+            const float Cc = 2.f * muX * muY + LOSS_C1;
+            const float Dc = 2.f * sigmaXY + LOSS_C2;
+            const float s = (Cc * Dc) / (A * B);
+            ssimSum += s;
+
+            if (pMu){
+                const int idx = (py * W + px) * C + c;
+                const float dMu = (muY * 2.f * Dc) / (A * B) - (muY * 2.f * Cc) / (A * B)
+                                - (muX * 2.f * Cc * Dc) / (A * A * B) + (muX * 2.f * Cc * Dc) / (A * B * B);
+                pMu[idx] = __float2half(dMu);
+                pS1[idx] = __float2half((-Cc * Dc) / (A * B * B));
+                pS12[idx] = __float2half((2.f * Cc) / (A * B));
+            }
+        }
+        __syncthreads();
+    }
+
+    if (px < W && py < H){
+        ssimMap[py * W + px] = ssimSum / static_cast<float>(C);
+    }
+}
+
+// Reduces the combined loss over pixels: out[0] += sum of gate * ((1-w)*sum_c|d_c| + C*w*(1-ssim)),
+// out[1] += sum of gate. The gate is the mask value or the valid-padding indicator.
+__global__ void fused_loss_reduce_kernel(
+    const int H, const int W, const int C,
+    const float* __restrict__ rendered,
+    const float* __restrict__ gt,
+    const float* __restrict__ ssimMap,
+    const float* __restrict__ mask, // nullptr when unmasked
+    const float ssimWeight,
+    const bool validPad,
+    float* __restrict__ out
+){
+    const int numPix = H * W;
+    float lossSum = 0.0f;
+    float gateSum = 0.0f;
+    for (int p = blockIdx.x * blockDim.x + threadIdx.x; p < numPix; p += blockDim.x * gridDim.x){
+        const int y = p / W;
+        const int x = p % W;
+        float gate;
+        if (mask){
+            gate = mask[p];
+        }else{
+            gate = loss_valid(y, x, H, W, validPad) ? 1.0f : 0.0f;
+        }
+        if (gate != 0.0f){
+            float l1 = 0.0f;
+            for (int c = 0; c < C; c++){
+                l1 += fabsf(rendered[p * C + c] - gt[p * C + c]);
+            }
+            const float contrib = (1.0f - ssimWeight) * l1
+                                + static_cast<float>(C) * ssimWeight * (1.0f - ssimMap[p]);
+            lossSum += gate * contrib;
+            gateSum += gate;
+        }
+    }
+
+    __shared__ float sLoss[256];
+    __shared__ float sGate[256];
+    sLoss[threadIdx.x] = lossSum;
+    sGate[threadIdx.x] = gateSum;
+    __syncthreads();
+    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1){
+        if (threadIdx.x < stride){
+            sLoss[threadIdx.x] += sLoss[threadIdx.x + stride];
+            sGate[threadIdx.x] += sGate[threadIdx.x + stride];
+        }
+        __syncthreads();
+    }
+    if (threadIdx.x == 0){
+        atomicAdd(&out[0], sLoss[0]);
+        atomicAdd(&out[1], sGate[0]);
+    }
+}
+
+// out[0] = loss, out[1] = normalization denominator (gateSum * C)
+__global__ void fused_loss_finalize_kernel(const int C, float* __restrict__ out){
+    const float denom = out[1] * static_cast<float>(C) + 1e-8f;
+    out[0] = out[0] / denom;
+    out[1] = denom;
+}
+
+__global__ void fused_loss_bwd_kernel(
+    const int H, const int W, const int C,
+    const float ssimWeight,
+    const bool validPad,
+    const float* __restrict__ rendered,
+    const float* __restrict__ gt,
+    const float* __restrict__ mask, // nullptr when unmasked
+    const __half* __restrict__ pMu,
+    const __half* __restrict__ pS1,
+    const __half* __restrict__ pS12,
+    const float* __restrict__ stats, // stats[1] = denominator
+    const float* __restrict__ vLoss,
+    float* __restrict__ vRendered
+){
+    const int px = blockIdx.x * LOSS_BX + threadIdx.x;
+    const int py = blockIdx.y * LOSS_BY + threadIdx.y;
+    const int tileX = blockIdx.x * LOSS_BX;
+    const int tileY = blockIdx.y * LOSS_BY;
+    const float chainScale = vLoss[0] / stats[1];
+
+    __shared__ float sData[LOSS_SY][LOSS_SX][3];
+    __shared__ float sConv[LOSS_SY][LOSS_BX][3];
+
+    for (int c = 0; c < C; c++){
+        float p1 = 0.f, p2 = 0.f;
+        if (px < W && py < H){
+            p1 = rendered[(py * W + px) * C + c];
+            p2 = gt[(py * W + px) * C + c];
+        }
+
+        // Load the chain-weighted partials for the tile + halo
+        {
+            const int tileSize = LOSS_SY * LOSS_SX;
+            const int threads = LOSS_BX * LOSS_BY;
+            const int tRank = threadIdx.y * LOSS_BX + threadIdx.x;
+            for (int tid = tRank; tid < tileSize; tid += threads){
+                const int ly = tid / LOSS_SX;
+                const int lx = tid % LOSS_SX;
+                const int gy = tileY + ly - LOSS_HALO;
+                const int gx = tileX + lx - LOSS_HALO;
+                float chain = 0.0f;
+                if (gx >= 0 && gx < W && gy >= 0 && gy < H){
+                    const float gate = mask ? mask[gy * W + gx]
+                                            : (loss_valid(gy, gx, H, W, validPad) ? 1.0f : 0.0f);
+                    chain = -ssimWeight * gate * chainScale;
+                }
+                const int idx = (gy * W + gx) * C + c;
+                const bool inside = gx >= 0 && gx < W && gy >= 0 && gy < H;
+                sData[ly][lx][0] = inside ? __half2float(pMu[idx]) * chain : 0.0f;
+                sData[ly][lx][1] = inside ? __half2float(pS1[idx]) * chain : 0.0f;
+                sData[ly][lx][2] = inside ? __half2float(pS12[idx]) * chain : 0.0f;
+            }
+        }
+        __syncthreads();
+
+        // Horizontal pass
+        {
+            const int lx = threadIdx.x + LOSS_HALO;
+            for (int pass = 0; pass < 2; pass++){
+                const int ly = threadIdx.y + pass * LOSS_BY;
+                if (ly >= LOSS_SY) break;
+                float a0 = 0.f, a1 = 0.f, a2 = 0.f;
+                #pragma unroll
+                for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                    const float w = lossGauss[LOSS_HALO + d];
+                    a0 += sData[ly][lx + d][0] * w;
+                    a1 += sData[ly][lx + d][1] * w;
+                    a2 += sData[ly][lx + d][2] * w;
+                }
+                sConv[ly][threadIdx.x][0] = a0;
+                sConv[ly][threadIdx.x][1] = a1;
+                sConv[ly][threadIdx.x][2] = a2;
+            }
+        }
+        __syncthreads();
+
+        // Vertical pass + L1 term
+        if (px < W && py < H){
+            const int ly = threadIdx.y + LOSS_HALO;
+            const int lx = threadIdx.x;
+            float s0 = 0.f, s1 = 0.f, s2 = 0.f;
+            #pragma unroll
+            for (int d = -LOSS_HALO; d <= LOSS_HALO; d++){
+                const float w = lossGauss[LOSS_HALO + d];
+                const float* row = sConv[ly + d][lx];
+                s0 += row[0] * w;
+                s1 += row[1] * w;
+                s2 += row[2] * w;
+            }
+            const float gradSsim = s0 + 2.f * p1 * s1 + p2 * s2;
+
+            const float gate = mask ? mask[py * W + px]
+                                    : (loss_valid(py, px, H, W, validPad) ? 1.0f : 0.0f);
+            const float sign = (p1 == p2) ? 0.0f : copysignf(1.0f, p1 - p2);
+            const float gradL1 = (1.0f - ssimWeight) * sign * gate * chainScale;
+
+            vRendered[(py * W + px) * C + c] = gradSsim + gradL1;
+        }
+        __syncthreads();
+    }
+}
+
+std::tuple<torch::Tensor, torch::Tensor> fused_loss_forward_tensor(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask,
+    const float ssim_weight,
+    const bool valid_padding,
+    const bool want_grad
+){
+    CHECK_INPUT(rendered);
+    CHECK_INPUT(gt);
+    const int H = rendered.size(0);
+    const int W = rendered.size(1);
+    const int C = rendered.size(2);
+    const bool hasMask = mask.defined() && mask.numel() > 0;
+    if (hasMask){ CHECK_INPUT(mask); }
+
+    auto opts = rendered.options();
+    torch::Tensor ssimMap = torch::empty({H, W}, opts);
+    torch::Tensor partials = want_grad
+        ? torch::empty({3, static_cast<long long>(H) * W * C}, opts.dtype(torch::kHalf))
+        : torch::empty({0}, opts.dtype(torch::kHalf));
+    __half* pBase = want_grad ? reinterpret_cast<__half*>(partials.data_ptr<at::Half>()) : nullptr;
+    const long long planeSize = static_cast<long long>(H) * W * C;
+
+    const dim3 block(LOSS_BX, LOSS_BY);
+    const dim3 grid((W + LOSS_BX - 1) / LOSS_BX, (H + LOSS_BY - 1) / LOSS_BY);
+    fused_loss_fwd_kernel<<<grid, block>>>(
+        H, W, C,
+        rendered.data_ptr<float>(), gt.data_ptr<float>(),
+        ssimMap.data_ptr<float>(),
+        pBase, pBase ? pBase + planeSize : nullptr, pBase ? pBase + 2 * planeSize : nullptr
+    );
+
+    torch::Tensor stats = torch::zeros({2}, opts);
+    const int numPix = H * W;
+    const int reduceBlocks = (std::min)(1024, (numPix + 255) / 256);
+    fused_loss_reduce_kernel<<<reduceBlocks, 256>>>(
+        H, W, C,
+        rendered.data_ptr<float>(), gt.data_ptr<float>(),
+        ssimMap.data_ptr<float>(),
+        hasMask ? mask.data_ptr<float>() : nullptr,
+        ssim_weight, valid_padding,
+        stats.data_ptr<float>()
+    );
+    fused_loss_finalize_kernel<<<1, 1>>>(C, stats.data_ptr<float>());
+
+    return std::make_tuple(stats, partials);
+}
+
+torch::Tensor fused_loss_backward_tensor(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask,
+    const torch::Tensor &partials,
+    const torch::Tensor &stats,
+    const torch::Tensor &v_loss,
+    const float ssim_weight,
+    const bool valid_padding
+){
+    const int H = rendered.size(0);
+    const int W = rendered.size(1);
+    const int C = rendered.size(2);
+    const bool hasMask = mask.defined() && mask.numel() > 0;
+
+    torch::Tensor vRendered = torch::empty_like(rendered);
+    const __half* pBase = reinterpret_cast<const __half*>(partials.data_ptr<at::Half>());
+    const long long planeSize = static_cast<long long>(H) * W * C;
+
+    const dim3 block(LOSS_BX, LOSS_BY);
+    const dim3 grid((W + LOSS_BX - 1) / LOSS_BX, (H + LOSS_BY - 1) / LOSS_BY);
+    fused_loss_bwd_kernel<<<grid, block>>>(
+        H, W, C, ssim_weight, valid_padding,
+        rendered.data_ptr<float>(), gt.data_ptr<float>(),
+        hasMask ? mask.data_ptr<float>() : nullptr,
+        pBase, pBase + planeSize, pBase + 2 * planeSize,
+        stats.data_ptr<float>(),
+        v_loss.data_ptr<float>(),
+        vRendered.data_ptr<float>()
+    );
+    return vRendered;
+}
+
 __global__ void compute_cov2d_bounds_kernel(
     const unsigned num_pts, const float* __restrict__ covs2d, float* __restrict__ conics, float* __restrict__ radii
 ) {
@@ -579,7 +966,11 @@ std::
         const torch::Tensor &final_Ts,
         const torch::Tensor &final_idx,
         const torch::Tensor &v_output, // dL_dout_color
-        const torch::Tensor &v_output_alpha // dL_dout_alpha
+        const torch::Tensor &v_output_alpha, // dL_dout_alpha
+        const torch::Tensor &error_map, // [H,W] or empty
+        const torch::Tensor &edge_map, // [H,W] or empty
+        const torch::Tensor &densification_info, // [4,N] accumulated in place, or empty
+        const torch::Tensor &v_xy_abs // [N,2] accumulated in place, or empty
     ) {
 
     CHECK_INPUT(xys);
@@ -594,6 +985,25 @@ std::
     }
 
     const int num_points = xys.size(0);
+
+    if (densification_info.numel() > 0){
+        CHECK_INPUT(densification_info);
+        if (densification_info.size(0) != 4 || densification_info.size(1) != num_points){
+            AT_ERROR("densification_info must have dimensions (4, num_points)");
+        }
+    }
+    if (v_xy_abs.numel() > 0){
+        CHECK_INPUT(v_xy_abs);
+        if (v_xy_abs.size(0) != num_points || v_xy_abs.size(1) != 2){
+            AT_ERROR("v_xy_abs must have dimensions (num_points, 2)");
+        }
+    }
+    if (error_map.numel() > 0){
+        CHECK_INPUT(error_map);
+        if (error_map.numel() != img_height * img_width){
+            AT_ERROR("error_map must have img_height * img_width elements");
+        }
+    }
     const dim3 tile_bounds = {
         (img_width + BLOCK_X - 1) / BLOCK_X,
         (img_height + BLOCK_Y - 1) / BLOCK_Y,
@@ -612,6 +1022,7 @@ std::
     rasterize_backward_kernel<<<tile_bounds, block>>>(
         tile_bounds,
         img_size,
+        num_points,
         gaussians_ids_sorted.contiguous().data_ptr<int>(),
         (int2 *)tile_bins.contiguous().data_ptr<int>(),
         (float2 *)xys.contiguous().data_ptr<float>(),
@@ -623,10 +1034,14 @@ std::
         final_idx.contiguous().data_ptr<int>(),
         (float3 *)v_output.contiguous().data_ptr<float>(),
         v_output_alpha.contiguous().data_ptr<float>(),
+        error_map.numel() > 0 ? error_map.data_ptr<float>() : nullptr,
+        edge_map.numel() > 0 ? edge_map.data_ptr<float>() : nullptr,
         (float2 *)v_xy.contiguous().data_ptr<float>(),
+        v_xy_abs.numel() > 0 ? (float2 *)v_xy_abs.data_ptr<float>() : nullptr,
         (float3 *)v_conic.contiguous().data_ptr<float>(),
         (float3 *)v_colors.contiguous().data_ptr<float>(),
-        v_opacity.contiguous().data_ptr<float>()
+        v_opacity.contiguous().data_ptr<float>(),
+        densification_info.numel() > 0 ? densification_info.data_ptr<float>() : nullptr
     );
 
     return std::make_tuple(v_xy, v_conic, v_colors, v_opacity);
diff --git a/rasterizer/gsplat/bindings.h b/rasterizer/gsplat/bindings.h
index 441d6f35..a1babb44 100644
--- a/rasterizer/gsplat/bindings.h
+++ b/rasterizer/gsplat/bindings.h
@@ -1,4 +1,4 @@
-#ifdef USE_HIP
+#ifdef USE_HIP
 #include <hip/hip_runtime.h>
 #else
 #include "cuda_runtime.h"
@@ -23,6 +23,28 @@ std::tuple<
     torch::Tensor> // output radii
 compute_cov2d_bounds_tensor(const int num_pts, torch::Tensor &A);
 
+// Fused L1 + DSSIM loss over [H,W,C] float images.
+// Returns {stats, partials}: stats[0] = loss, stats[1] = normalization
+// denominator (both on device); partials holds the SSIM derivative maps
+// needed by the backward pass (empty when want_grad is false).
+std::tuple<torch::Tensor, torch::Tensor> fused_loss_forward_tensor(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask, // [H,W] float or empty
+    const float ssim_weight,
+    const bool valid_padding,
+    const bool want_grad);
+
+torch::Tensor fused_loss_backward_tensor(
+    const torch::Tensor &rendered,
+    const torch::Tensor &gt,
+    const torch::Tensor &mask,
+    const torch::Tensor &partials,
+    const torch::Tensor &stats,
+    const torch::Tensor &v_loss,
+    const float ssim_weight,
+    const bool valid_padding);
+
 torch::Tensor compute_sh_forward_tensor(
     unsigned num_points,
     unsigned degree,
@@ -186,5 +208,9 @@ std::
         const torch::Tensor &final_Ts,
         const torch::Tensor &final_idx,
         const torch::Tensor &v_output, // dL_dout_color
-        const torch::Tensor &v_output_alpha
+        const torch::Tensor &v_output_alpha,
+        const torch::Tensor &error_map, // [H,W] or empty
+        const torch::Tensor &edge_map, // [H,W] or empty
+        const torch::Tensor &densification_info, // [4,N] accumulated in place, or empty
+        const torch::Tensor &v_xy_abs // [N,2] accumulated in place, or empty
     );
\ No newline at end of file
diff --git a/simple_trainer.cpp b/simple_trainer.cpp
index 9912c5d6..45d258c1 100644
--- a/simple_trainer.cpp
+++ b/simple_trainer.cpp
@@ -167,7 +167,7 @@ int main(int argc, char **argv){
                 p[4], // camDepths
                 height,
                 width,
-                background);
+                background)[0];
         }else{
             #if defined(USE_HIP) || defined(USE_CUDA) || defined(USE_MPS)
                 auto p = ProjectGaussians::apply(means, scales, 1, 
@@ -189,7 +189,7 @@ int main(int argc, char **argv){
                     torch::sigmoid(opacities),
                     height,
                     width,
-                    background);
+                    background)[0];
             #else
                 throw std::runtime_error("GPU support not built, use --cpu");
             #endif
diff --git a/ssim.cpp b/ssim.cpp
index f8a152e6..3897a35e 100644
--- a/ssim.cpp
+++ b/ssim.cpp
@@ -1,46 +1,74 @@
-// Ported from https://github.com/Po-Hsun-Su/pytorch-ssim
-// MIT
+// Fused (1-w)*L1 + w*DSSIM photometric loss with a closed-form
+// backward
 
 #include "ssim.hpp"
+#include "gsplat.hpp"
 
-using namespace torch::indexing;
+namespace {
 
-torch::Tensor SSIM::eval(const torch::Tensor& rendered, const torch::Tensor& gt) {
-    torch::Tensor img1 = gt.permute({2, 0, 1}).index({None, "..."});
-    torch::Tensor img2 = rendered.permute({2, 0, 1}).index({None, "..."});
-    
-    if (img1.device() != window.device()){
-        window = window.to(img1.device());
+std::tuple<torch::Tensor, torch::Tensor> fusedLossForwardDispatch(
+    const torch::Tensor &rendered, const torch::Tensor &gt, const torch::Tensor &mask,
+    float ssimWeight, bool validPadding, bool wantGrad){
+#if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_MPS)
+    if (!rendered.is_cpu()){
+        return fused_loss_forward_tensor(rendered, gt, mask, ssimWeight, validPadding, wantGrad);
     }
-    torch::Tensor mu1 = torch::nn::functional::conv2d(img1, window, torch::nn::functional::Conv2dFuncOptions().padding(windowSize / 2).groups(channel));
-    torch::Tensor mu2 = torch::nn::functional::conv2d(img2, window, torch::nn::functional::Conv2dFuncOptions().padding(windowSize / 2).groups(channel));
+#endif
+    return fused_loss_forward_tensor_cpu(rendered, gt, mask, ssimWeight, validPadding, wantGrad);
+}
 
-    torch::Tensor mu1Sq = mu1.pow(2);
-    torch::Tensor mu2Sq = mu2.pow(2);
-    torch::Tensor mu1mu2 = mu1 * mu2;
+torch::Tensor fusedLossBackwardDispatch(
+    const torch::Tensor &rendered, const torch::Tensor &gt, const torch::Tensor &mask,
+    const torch::Tensor &partials, const torch::Tensor &stats, const torch::Tensor &vLoss,
+    float ssimWeight, bool validPadding){
+#if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_MPS)
+    if (!rendered.is_cpu()){
+        return fused_loss_backward_tensor(rendered, gt, mask, partials, stats, vLoss, ssimWeight, validPadding);
+    }
+#endif
+    return fused_loss_backward_tensor_cpu(rendered, gt, mask, partials, stats, vLoss, ssimWeight, validPadding);
+}
 
-    torch::Tensor sigma1Sq = torch::nn::functional::conv2d(img1 * img1, window, torch::nn::functional::Conv2dFuncOptions().padding(windowSize / 2).groups(channel)) - mu1Sq;
-    torch::Tensor sigma2Sq = torch::nn::functional::conv2d(img2 * img2, window, torch::nn::functional::Conv2dFuncOptions().padding(windowSize / 2).groups(channel)) - mu2Sq;
-    torch::Tensor sigma12 = torch::nn::functional::conv2d(img1 * img2, window, torch::nn::functional::Conv2dFuncOptions().padding(windowSize / 2).groups(channel)) - mu1mu2;
-    
-    const float C1 = 0.01 * 0.01;
-    const float C2 = 0.03 * 0.03;
+class FusedL1SsimLossFunction : public torch::autograd::Function<FusedL1SsimLossFunction>{
+public:
+    static torch::Tensor forward(torch::autograd::AutogradContext *ctx,
+                                 torch::Tensor rendered, // [H,W,C]
+                                 torch::Tensor gt,       // [H,W,C]
+                                 torch::Tensor mask,     // [H,W] or empty
+                                 double ssimWeight, bool validPadding){
+        auto r = fusedLossForwardDispatch(rendered.contiguous(), gt, mask,
+                                          static_cast<float>(ssimWeight), validPadding, true);
+        torch::Tensor stats = std::get<0>(r);
+        ctx->save_for_backward({ rendered, gt, mask, std::get<1>(r), stats });
+        ctx->saved_data["ssimWeight"] = ssimWeight;
+        ctx->saved_data["validPadding"] = validPadding;
+        return stats.index({0});
+    }
 
-    torch::Tensor ssimMap = ((2.0f * mu1mu2 + C1) * (2.0f * sigma12 + C2)) / ((mu1Sq + mu2Sq + C1) * (sigma1Sq + sigma2Sq + C2));
+    static torch::autograd::tensor_list backward(torch::autograd::AutogradContext *ctx, torch::autograd::tensor_list gradOutputs){
+        torch::autograd::variable_list saved = ctx->get_saved_variables();
+        torch::Tensor vRendered = fusedLossBackwardDispatch(
+            saved[0], saved[1], saved[2], saved[3], saved[4],
+            gradOutputs[0].contiguous(),
+            static_cast<float>(ctx->saved_data["ssimWeight"].toDouble()),
+            ctx->saved_data["validPadding"].toBool());
+        torch::Tensor none;
+        return { vRendered, none, none, none, none };
+    }
+};
 
-    return ssimMap.mean();
 }
 
-torch::Tensor SSIM::createWindow(){
-    torch::Tensor _1DWindow = gaussian(1.5f).unsqueeze(1);
-    torch::Tensor _2DWindow = _1DWindow.mm(_1DWindow.t()).unsqueeze(0).unsqueeze(0);
-    return _2DWindow.expand({channel, 1, windowSize, windowSize}).contiguous();
+torch::Tensor fusedL1SsimLoss(const torch::Tensor &rendered, const torch::Tensor &gt,
+                              const torch::Tensor &mask, float ssimWeight, bool validPadding){
+    return FusedL1SsimLossFunction::apply(rendered, gt, mask,
+                                          static_cast<double>(ssimWeight), validPadding);
 }
 
-torch::Tensor SSIM::gaussian(float sigma) {
-    torch::Tensor gauss = torch::zeros(windowSize);
-    for (int i = 0; i < windowSize; i++) {
-        gauss[i] = std::exp(-(std::pow(std::floor(static_cast<float>(i - windowSize) / 2.0f), 2.0f)) / (2.0f * sigma * sigma));
-    }
-    return gauss / gauss.sum();
-}
\ No newline at end of file
+torch::Tensor fusedL1SsimLossValue(const torch::Tensor &rendered, const torch::Tensor &gt,
+                                   float ssimWeight){
+    torch::Tensor empty;
+    auto r = fusedLossForwardDispatch(rendered.contiguous(), gt.contiguous(), empty,
+                                      ssimWeight, false, false);
+    return std::get<0>(r).index({0});
+}
diff --git a/ssim.hpp b/ssim.hpp
index b74a86df..e7a99e8a 100644
--- a/ssim.hpp
+++ b/ssim.hpp
@@ -3,24 +3,14 @@
 
 #include <torch/torch.h>
 
-// Ported from https://github.com/Po-Hsun-Su/pytorch-ssim
-// MIT
+// Fused (1-ssimWeight)*L1 + ssimWeight*DSSIM loss with autograd support.
+// mask may be an empty tensor; validPadding crops the blur border from the
+// unmasked loss
+torch::Tensor fusedL1SsimLoss(const torch::Tensor &rendered, const torch::Tensor &gt,
+                              const torch::Tensor &mask, float ssimWeight, bool validPadding);
 
-class SSIM{
-public:
-    SSIM(int windowSize, int channel) : windowSize(windowSize), channel(channel){
-        window = createWindow();
-    };
+// Loss value only
+torch::Tensor fusedL1SsimLossValue(const torch::Tensor &rendered, const torch::Tensor &gt,
+                                   float ssimWeight);
 
-    torch::Tensor eval(const torch::Tensor& rendered, const torch::Tensor& gt);
-private:
-    torch::Tensor createWindow();
-    torch::Tensor gaussian(float sigma);
-
-    int windowSize;
-    int channel;
-    torch::Tensor window;
-};
-
-
-#endif
\ No newline at end of file
+#endif
diff --git a/undistort.cpp b/undistort.cpp
new file mode 100644
index 00000000..0e08c533
--- /dev/null
+++ b/undistort.cpp
@@ -0,0 +1,164 @@
+#include <cmath>
+#include <algorithm>
+#include <limits>
+#include "undistort.hpp"
+
+static void distortPoint(const UndistortParams &p, float x, float y, float &dx, float &dy){
+    float r2 = x * x + y * y;
+    float num = 1.0f + r2 * (p.k1 + r2 * (p.k2 + r2 * p.k3));
+    float den = 1.0f + r2 * (p.k4 + r2 * (p.k5 + r2 * p.k6));
+    float radial = num / den;
+    dx = x * radial + 2.0f * p.p1 * x * y + p.p2 * (r2 + 2.0f * x * x);
+    dy = y * radial + p.p1 * (r2 + 2.0f * y * y) + 2.0f * p.p2 * x * y;
+}
+
+
+// Newton iteration, port of COLMAP IterativeUndistortion
+static bool undistortPoint(const UndistortParams &p, float dx, float dy, float &ux, float &uy){
+    const int kNumIterations = 100;
+    const double kMinStepSquaredNorm = 1e-10;
+    const double kRelStepRadius = 0.1;
+    const double kStepRadius = 0.1;
+
+    const double x0 = dx, y0 = dy;
+    double x = x0, y = y0;
+    bool converged = false;
+
+    for (int it = 0; it < kNumIterations; it++){
+        const double r2 = x * x + y * y;
+        const double num = 1.0 + r2 * (p.k1 + r2 * (p.k2 + r2 * p.k3));
+        const double den = 1.0 + r2 * (p.k4 + r2 * (p.k5 + r2 * p.k6));
+        if (std::fabs(den) < 1e-12) break;
+        const double radial = num / den;
+        // d(radial)/d(r2) by the quotient rule
+        const double dNum = p.k1 + r2 * (2.0 * p.k2 + 3.0 * p.k3 * r2);
+        const double dDen = p.k4 + r2 * (2.0 * p.k5 + 3.0 * p.k6 * r2);
+        const double dRadial = (dNum * den - num * dDen) / (den * den);
+
+        const double fx = x * radial + 2.0 * p.p1 * x * y + p.p2 * (r2 + 2.0 * x * x) - x0;
+        const double fy = y * radial + p.p1 * (r2 + 2.0 * y * y) + 2.0 * p.p2 * x * y - y0;
+        const double j00 = radial + 2.0 * x * x * dRadial + 2.0 * p.p1 * y + 6.0 * p.p2 * x;
+        const double j01 = 2.0 * x * y * dRadial + 2.0 * p.p1 * x + 2.0 * p.p2 * y;
+        const double j10 = j01;
+        const double j11 = radial + 2.0 * y * y * dRadial + 6.0 * p.p1 * y + 2.0 * p.p2 * x;
+        const double det = j00 * j11 - j01 * j10;
+        if (std::fabs(det) < 1e-12) break;
+
+        double sx = (fx * j11 - fy * j01) / det;
+        double sy = (fy * j00 - fx * j10) / det;
+
+        // Trust region: |step| <= max(|x| * kRelStepRadius, kStepRadius)
+        const double radiusSqr = (std::max)(r2 * kRelStepRadius * kRelStepRadius,
+                                            kStepRadius * kStepRadius);
+        const double stepSqr = sx * sx + sy * sy;
+        if (stepSqr > radiusSqr){
+            const double s = std::sqrt(radiusSqr / stepSqr);
+            sx *= s;
+            sy *= s;
+        }
+        x -= sx;
+        y -= sy;
+        if (sx * sx + sy * sy < kMinStepSquaredNorm){
+            converged = true;
+            break;
+        }
+    }
+
+    ux = static_cast<float>(x);
+    uy = static_cast<float>(y);
+    return converged && std::isfinite(x) && std::isfinite(y);
+}
+
+UndistortParams computeUndistortParams(float fx, float fy, float cx, float cy,
+                                       int width, int height,
+                                       float k1, float k2, float k3,
+                                       float k4, float k5, float k6,
+                                       float p1, float p2,
+                                       float blankPixels){
+    UndistortParams p;
+    p.srcFx = fx; p.srcFy = fy; p.srcCx = cx; p.srcCy = cy;
+    p.srcW = width; p.srcH = height;
+    p.dstFx = fx; p.dstFy = fy; p.dstCx = cx; p.dstCy = cy;
+    p.dstW = width; p.dstH = height;
+    p.k1 = k1; p.k2 = k2; p.k3 = k3; p.k4 = k4; p.k5 = k5; p.k6 = k6; p.p1 = p1; p.p2 = p2;
+
+    if (k1 == 0.0f && k2 == 0.0f && k3 == 0.0f && k4 == 0.0f && k5 == 0.0f && k6 == 0.0f &&
+        p1 == 0.0f && p2 == 0.0f) return p;
+
+    const float inf = std::numeric_limits<float>::max();
+    float leftMinX = inf, leftMaxX = -inf, rightMinX = inf, rightMaxX = -inf;
+    float topMinY = inf, topMaxY = -inf, bottomMinY = inf, bottomMaxY = -inf;
+
+    // Undistort a source pixel center (corner-origin convention) and reproject
+    auto trace = [&](float px, float py, float &ox, float &oy){
+        float ux, uy;
+        if (!undistortPoint(p, (px - cx) / fx, (py - cy) / fy, ux, uy)) return false;
+        ox = fx * ux + cx;
+        oy = fy * uy + cy;
+        return true;
+    };
+
+    for (int y = 0; y < height; y++){
+        float ox, oy;
+        if (trace(0.5f, y + 0.5f, ox, oy)){
+            leftMinX = (std::min)(leftMinX, ox);
+            leftMaxX = (std::max)(leftMaxX, ox);
+        }
+        if (trace(width - 0.5f, y + 0.5f, ox, oy)){
+            rightMinX = (std::min)(rightMinX, ox);
+            rightMaxX = (std::max)(rightMaxX, ox);
+        }
+    }
+    for (int x = 0; x < width; x++){
+        float ox, oy;
+        if (trace(x + 0.5f, 0.5f, ox, oy)){
+            topMinY = (std::min)(topMinY, oy);
+            topMaxY = (std::max)(topMaxY, oy);
+        }
+        if (trace(x + 0.5f, height - 0.5f, ox, oy)){
+            bottomMinY = (std::min)(bottomMinY, oy);
+            bottomMaxY = (std::max)(bottomMaxY, oy);
+        }
+    }
+
+    // If a whole border failed to solve there is nothing sane to rescale to
+    if (leftMinX == inf || rightMaxX == -inf || topMinY == inf || bottomMaxY == -inf ||
+        leftMaxX == -inf || rightMinX == inf || topMaxY == -inf || bottomMinY == inf){
+        return p;
+    }
+
+    // Scale such that the undistorted image contains all source pixels (min)
+    // or no blank pixels (max)
+    float minScaleX = (std::min)(cx / (cx - leftMinX), (width - 0.5f - cx) / (rightMaxX - cx));
+    float minScaleY = (std::min)(cy / (cy - topMinY), (height - 0.5f - cy) / (bottomMaxY - cy));
+    float maxScaleX = (std::max)(cx / (cx - leftMaxX), (width - 0.5f - cx) / (rightMinX - cx));
+    float maxScaleY = (std::max)(cy / (cy - topMaxY), (height - 0.5f - cy) / (bottomMinY - cy));
+
+    float scaleX = 1.0f / (minScaleX * blankPixels + maxScaleX * (1.0f - blankPixels));
+    float scaleY = 1.0f / (minScaleY * blankPixels + maxScaleY * (1.0f - blankPixels));
+    scaleX = std::clamp(scaleX, 0.2f, 2.0f);
+    scaleY = std::clamp(scaleY, 0.2f, 2.0f);
+
+    p.dstW = (std::max)(1, static_cast<int>(scaleX * width));
+    p.dstH = (std::max)(1, static_cast<int>(scaleY * height));
+    p.dstCx = cx * static_cast<float>(p.dstW) / static_cast<float>(width);
+    p.dstCy = cy * static_cast<float>(p.dstH) / static_cast<float>(height);
+    return p;
+}
+
+void buildUndistortMaps(const UndistortParams &p, cv::Mat &mapx, cv::Mat &mapy){
+    mapx.create(p.dstH, p.dstW, CV_32FC1);
+    mapy.create(p.dstH, p.dstW, CV_32FC1);
+    for (int oy = 0; oy < p.dstH; oy++){
+        float *rx = mapx.ptr<float>(oy);
+        float *ry = mapy.ptr<float>(oy);
+        for (int ox = 0; ox < p.dstW; ox++){
+            float nx = (ox + 0.5f - p.dstCx) / p.dstFx;
+            float ny = (oy + 0.5f - p.dstCy) / p.dstFy;
+            float dnx, dny;
+            distortPoint(p, nx, ny, dnx, dny);
+            rx[ox] = dnx * p.srcFx + p.srcCx - 0.5f;
+            ry[ox] = dny * p.srcFy + p.srcCy - 0.5f;
+        }
+    }
+}
diff --git a/undistort.hpp b/undistort.hpp
new file mode 100644
index 00000000..5514d8f0
--- /dev/null
+++ b/undistort.hpp
@@ -0,0 +1,22 @@
+#ifndef UNDISTORT_H
+#define UNDISTORT_H
+
+#include <opencv2/core/core.hpp>
+
+struct UndistortParams{
+    float srcFx, srcFy, srcCx, srcCy;
+    int srcW, srcH;
+    float dstFx, dstFy, dstCx, dstCy;
+    int dstW, dstH;
+    float k1, k2, k3, k4, k5, k6, p1, p2;
+};
+
+UndistortParams computeUndistortParams(float fx, float fy, float cx, float cy,
+                                       int width, int height,
+                                       float k1, float k2, float k3,
+                                       float k4, float k5, float k6,
+                                       float p1, float p2,
+                                       float blankPixels = 0.0f);
+void buildUndistortMaps(const UndistortParams &p, cv::Mat &mapx, cv::Mat &mapy);
+
+#endif
diff --git a/utils.cpp b/utils.cpp
deleted file mode 100644
index 0ee624c5..00000000
--- a/utils.cpp
+++ /dev/null
@@ -1 +0,0 @@
-#include "utils.hpp"