diff --git a/benchmarks/Bench/KMeans.hs b/benchmarks/Bench/KMeans.hs
--- a/benchmarks/Bench/KMeans.hs
+++ b/benchmarks/Bench/KMeans.hs
@@ -16,19 +16,30 @@
     g <- createSystemRandom
     fmap fromSeed $ save g
 
-dat :: MU.Matrix Double
-dat = unsafePerformIO $ fmap MU.fromRows $ randVectors 1000 10
+matrix_1000_10 :: MU.Matrix Double
+matrix_1000_10 = unsafePerformIO $ fmap MU.fromRows $ randVectors 1000 10
 
+matrix_30000_50 :: MU.Matrix Double
+matrix_30000_50 = unsafePerformIO $ fmap MU.fromRows $ randVectors 30000 50
+
 benchKMeans :: Benchmark
 benchKMeans = bgroup "KMeans clustering"
     [ bgroup "AI.Clustering.KMeans"
-        [ bench "k-means++ (n = 1000, k = 7)" $
+        [ bench "k-means++ (size = 1000 X 10, k = 7)" $
             whnf ( \x -> membership $ kmeans 7 x defaultKMeansOpts
                 { kmeansMethod = KMeansPP
-                , kmeansSeed = gen } ) dat
-        , bench "forgy (n = 1000, k = 7)" $
+                , kmeansSeed = gen } ) matrix_1000_10
+        , bench "forgy (size = 1000 X 10, k = 7)" $
             whnf ( \x -> membership $ kmeans 7 x defaultKMeansOpts
+                { kmeansMethod = Forgy
+                , kmeansSeed = gen } ) matrix_1000_10
+        , bench "k-means++ (size = 30000 X 50, k = 10)" $
+            whnf ( \x -> membership $ kmeans 10 x defaultKMeansOpts
                 { kmeansMethod = KMeansPP
-                , kmeansSeed = gen } ) dat
+                , kmeansSeed = gen } ) matrix_30000_50
+        , bench "forgy (size = 30000 X 50, k = 10)" $
+            whnf ( \x -> membership $ kmeans 10 x defaultKMeansOpts
+                { kmeansMethod = Forgy
+                , kmeansSeed = gen } ) matrix_30000_50
         ]
     ]
diff --git a/benchmarks/Bench/Utils.hs b/benchmarks/Bench/Utils.hs
--- a/benchmarks/Bench/Utils.hs
+++ b/benchmarks/Bench/Utils.hs
@@ -10,5 +10,5 @@
             -> Int  -- ^ vector length
             -> IO [U.Vector Double]
 randVectors n k = do
-    g <- createSystemRandom
+    g <- create
     replicateM n $ uniformVector g k
diff --git a/clustering.cabal b/clustering.cabal
--- a/clustering.cabal
+++ b/clustering.cabal
@@ -1,5 +1,5 @@
 name:                clustering
-version:             0.3.1
+version:             0.4.0
 synopsis:            High performance clustering algorithms
 description:
   Following clutering methods are included in this library:
@@ -14,7 +14,7 @@
 license-file:        LICENSE
 author:              Kai Zhang
 maintainer:          kai@kzhang.org
-copyright:           (c) 2015 Kai Zhang
+copyright:           (c) 2015-2018 Kai Zhang
 category:            Math
 build-type:          Simple
 cabal-version:       >=1.10
@@ -66,7 +66,7 @@
     , clustering
     , hierarchical-clustering
     , split
-    , Rlang-QQ
+    , inline-r
 
 benchmark bench
   type: exitcode-stdio-1.0
diff --git a/src/AI/Clustering/Hierarchical.hs b/src/AI/Clustering/Hierarchical.hs
--- a/src/AI/Clustering/Hierarchical.hs
+++ b/src/AI/Clustering/Hierarchical.hs
@@ -43,6 +43,7 @@
     , size
     , Linkage(..)
     , hclust
+    , normalize
     , cutAt
     , flatten
     , drawDendrogram
@@ -85,6 +86,16 @@
         Weighted -> weighted
         Ward -> ward
         _ -> error "Not implemented"
+
+-- | Normalize the tree heights so that the highest is 1.
+normalize :: Dendrogram a -> Dendrogram a
+normalize dendro = go dendro
+  where
+    go (Branch n d l r) = Branch n (d / maxHeight) (go l) (go r)
+    go (Leaf x) = Leaf x
+    maxHeight = case dendro of
+        Branch _ x _ _ -> x
+        Leaf _ -> 0
 
 -- | Cut a dendrogram at given height.
 cutAt :: Dendrogram a -> Distance -> [Dendrogram a]
diff --git a/src/AI/Clustering/KMeans.hs b/src/AI/Clustering/KMeans.hs
--- a/src/AI/Clustering/KMeans.hs
+++ b/src/AI/Clustering/KMeans.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE BangPatterns #-}
 
 module AI.Clustering.KMeans
     ( KMeans(..)
@@ -10,6 +11,8 @@
     -- * Initialization methods
     , Method(..)
 
+    , decode
+
     -- * References
     -- $references
     ) where
@@ -17,6 +20,7 @@
 import Control.Monad (forM_)
 import Control.Monad.Primitive (PrimMonad, PrimState)
 import qualified Data.Matrix.Unboxed as MU
+import Data.Matrix.Generic (unsafeTakeRow)
 import qualified Data.Matrix.Unboxed.Mutable as MM
 import Data.Ord (comparing)
 import qualified Data.Vector as V
@@ -36,14 +40,17 @@
        -> MU.Matrix Double   -- ^ Input data stored as rows in a matrix
        -> KMeansOpts
        -> KMeans (U.Vector Double)
-kmeans k mat opts = KMeans member cs grps
+kmeans k mat opts
+    | containNaN = error "Input data contains NaN."
+    | otherwise = KMeans member cs grps sse'
   where
-    (member, cs) = kmeans' initial dat fn
+    containNaN =  U.any isNaN $ MU.flatten mat
+    (member, cs, sse') = kmeans' initial (kmeansMaxIter opts) dat fn
     grps = if kmeansClusters opts
         then Just $ decode member $ MU.toRows mat
         else Nothing
     dat = U.enumFromN 0 $ MU.rows mat
-    fn = MU.takeRow mat
+    fn = unsafeTakeRow mat
     initial = runST $ do
         gen <- initialize $ kmeansSeed opts
         case kmeansMethod opts of
@@ -59,9 +66,12 @@
          -> (a -> U.Vector Double)
          -> KMeansOpts
          -> KMeans a
-kmeansBy k dat fn opts = KMeans member cs grps
+kmeansBy k dat fn opts
+    | containNaN = error "Input data contains NaN."
+    | otherwise = KMeans member cs grps sse'
   where
-    (member, cs) = kmeans' initial dat fn
+    containNaN = G.foldl (\acc x -> acc || U.any isNaN (fn x)) False dat
+    (member, cs, sse') = kmeans' initial (kmeansMaxIter opts) dat fn
     grps = if kmeansClusters opts
         then Just $ decode member $ G.toList dat
         else Nothing
@@ -76,36 +86,40 @@
 -- | K-means algorithm
 kmeans' :: G.Vector v a
         => MU.Matrix Double         -- ^ Initial set of k centroids
+        -> Int                      -- ^ Max inter
         -> v a                      -- ^ Input data
         -> (a -> U.Vector Double)   -- ^ Feature extraction function
-        -> (U.Vector Int, MU.Matrix Double)
-kmeans' initial dat fn
+        -> (U.Vector Int, MU.Matrix Double, Double)
+kmeans' initial maxiter dat fn
     | U.length (fn $ G.head dat) /= d = error "Dimension mismatched."
-    | otherwise = (member, centers)
+    | otherwise = (member, centers, U.sum $ U.imap ( \i x -> sqrt $ sumSquares
+        (fn $ dat G.! i) (centers `MU.takeRow` x) ) member )
   where
-    (member, centers) = loop initial U.empty
-    loop means membership
-        | membership' == membership = (membership, means)
-        | otherwise = loop (update membership') membership'
+    (member, centers) = loop 0 initial U.empty
+    loop !iter means membership
+        | iter >= maxiter || membership' == membership = (membership, means)
+        | otherwise = loop (iter+1) (update membership') membership'
       where
         membership' = assign means
 
     -- Assignment step
     assign means = U.generate n $ \i ->
         let x = fn $ G.unsafeIndex dat i
-        in fst $ minimumBy (comparing snd) $ zip [0..k-1] $ map (sumSquares x) $ MU.toRows means
+            f (!min', !j') j = let d = sumSquares x $ means `unsafeTakeRow` j
+                               in if d < min' then (d, j) else (min', j')
+        in snd $ foldl' f (1/0, -1) [0..k-1]
 
     -- Update step
     update membership = MU.create $ do
         m <- MM.replicate (k,d) 0.0
         count <- UM.replicate k (0 :: Int)
         forM_ [0..n-1] $ \i -> do
-            let x = membership U.! i
-            UM.unsafeRead count x >>= UM.unsafeWrite count x . (+1)
-
-            let vec = fn $ dat G.! i
+            let x = membership `U.unsafeIndex` i
+                vec = fn $ dat `G.unsafeIndex` i
+            UM.unsafeModify count (+1) x
             forM_ [0..d-1] $ \j ->
-                MM.unsafeRead m (x,j) >>= MM.unsafeWrite m (x,j) . (+ (vec U.! j))
+                MM.unsafeRead m (x,j) >>=
+                    MM.unsafeWrite m (x,j) . (+ (vec `U.unsafeIndex` j))
         -- normalize
         forM_ [0..k-1] $ \i -> do
             c <- UM.unsafeRead count i
@@ -127,16 +141,6 @@
   where
     n = U.maximum member + 1
 {-# INLINE decode #-}
-
-{-
--- Compute within-cluster sum of squares
-withinSS :: KMeans -> MU.Matrix Double -> [Double]
-withinSS result mat = zipWith f (decode result [0 .. MU.rows mat-1]) .
-                          MU.toRows . _centers $ result
-  where
-    f c center = foldl' (+) 0 $ map (sumSquares center . MU.takeRow mat) c
-    -}
-
 
 -- $references
 --
diff --git a/src/AI/Clustering/KMeans/Internal.hs b/src/AI/Clustering/KMeans/Internal.hs
--- a/src/AI/Clustering/KMeans/Internal.hs
+++ b/src/AI/Clustering/KMeans/Internal.hs
@@ -59,7 +59,7 @@
 {-# INLINE kmeansPP #-}
 
 sumSquares :: U.Vector Double -> U.Vector Double -> Double
-sumSquares xs = U.sum . U.zipWith (\x y -> (x - y)**2) xs
+sumSquares xs = U.sum . U.zipWith (\x y -> (x - y) * (x - y)) xs
 {-# INLINE sumSquares #-}
 
 -- | Generate N non-duplicated uniformly distributed random variables in a given range.
diff --git a/src/AI/Clustering/KMeans/Types.hs b/src/AI/Clustering/KMeans/Types.hs
--- a/src/AI/Clustering/KMeans/Types.hs
+++ b/src/AI/Clustering/KMeans/Types.hs
@@ -25,13 +25,22 @@
     { kmeansMethod :: Method
     , kmeansSeed :: (U.Vector Word32)   -- ^ Seed for random number generation
     , kmeansClusters :: Bool   -- ^ Wether to return clusters, may use a lot memory
+    , kmeansMaxIter :: Int     -- ^ Maximum iteration
     }
 
+-- | Default options.
+-- > defaultKMeansOpts = KMeansOpts
+-- >     { kmeansMethod = KMeansPP
+-- >     , kmeansSeed = U.fromList [1,2,3,4,5,6,7]
+-- >     , kmeansClusters = True
+-- >     , kmeansMaxIter = 10
+-- >     }
 defaultKMeansOpts :: KMeansOpts
 defaultKMeansOpts = KMeansOpts
     { kmeansMethod = KMeansPP
-    , kmeansSeed = U.fromList [1,2,3,4,5,6,7]
+    , kmeansSeed = U.fromList [2341,2342,3934,425,2345,80006,2343,234491,124,729]
     , kmeansClusters = True
+    , kmeansMaxIter = 10000
     }
 
 -- | Results from running kmeans
@@ -41,6 +50,7 @@
                                     -- point is allocated.
     , centers :: MU.Matrix Double  -- ^ A matrix of cluster centers.
     , clusters :: Maybe [[a]]
+    , sse :: Double                -- ^ the sum of squared error (SSE)
     } deriving (Show)
 
 -- | Different initialization methods
diff --git a/tests/Test/KMeans.hs b/tests/Test/KMeans.hs
--- a/tests/Test/KMeans.hs
+++ b/tests/Test/KMeans.hs
@@ -1,23 +1,31 @@
-{-# LANGUAGE QuasiQuotes #-}
-{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE QuasiQuotes     #-}
+{-# LANGUAGE TemplateHaskell #-}
+
 module Test.KMeans
     ( tests
     ) where
 
-import Control.Monad
-import qualified Data.Matrix.Unboxed as MU
-import qualified Data.Vector.Unboxed as V
-import Data.List
-import RlangQQ
-import System.Random.MWC
-import Test.Tasty
-import Test.Tasty.HUnit
-import Test.Tasty.QuickCheck
+import           Control.Monad
+import           Data.Int                      (Int32)
+import           Data.List
+import qualified Data.Matrix.Unboxed           as MU
+import           Data.Maybe
+import qualified Data.Vector.SEXP              as S
+import qualified Data.Vector.Unboxed           as V
+import qualified Foreign.R                     as R
+import qualified Foreign.R.Type                as R
+import qualified H.Prelude                     as H
+import           Language.R.HExp
+import           Language.R.QQ
+import           System.Random.MWC
+import           Test.Tasty
+import           Test.Tasty.HUnit
+import           Test.Tasty.QuickCheck
 
-import AI.Clustering.KMeans
-import AI.Clustering.KMeans.Internal
+import           AI.Clustering.KMeans
+import           AI.Clustering.KMeans.Internal
 
-import Test.Utils
+import           Test.Utils
 
 tests :: TestTree
 tests = testGroup "KMeans:"
@@ -25,13 +33,14 @@
     ]
 
 rKmeans :: Int -> [Double] -> [Double] -> IO [Int]
-rKmeans n dat center = do
-    o <- [r| x <- matrix(hs_dat, ncol=hs_n,byrow=T);
-             y <- matrix(hs_center, ncol=hs_n,byrow=T);
-             hs_result <- kmeans(x,y,iter.max=1000000,algorithm="Lloyd")$cluster;
-         |]
-    let x = Label :: Label "result"
-    return $ o .!. x
+rKmeans n' dat center = fmap (map (fromIntegral :: Int32 -> Int)) $ H.runRegion $ do
+    xxx <- [r| x <- matrix(dat_hs, ncol=n_hs,byrow=T);
+             y <- matrix(center_hs, ncol=n_hs,byrow=T);
+             kmeans(x,y,iter.max=10000,algorithm="Lloyd")$cluster
+    |]
+    return $ H.fromSEXP $ H.cast R.SInt xxx
+  where
+    n = fromIntegral n' :: Double
 
 testKMeans :: Assertion
 testKMeans = do
@@ -45,13 +54,12 @@
         dat = V.enumFromN 0 $ MU.rows mat
         fn = MU.takeRow mat
 
-    centers <- kmeansPP g k dat fn
+    init_centers <- kmeansPP g k dat fn
 
-    r <- rKmeans d (MU.toList mat) (MU.toList centers)
-    let test = sort $ map sort $ decode result xs
-        result = kmeansWith centers dat fn
-        true = sort $ map sort $ decode result{_clusters=V.fromList $ map (subtract 1) r} xs
-        show' xs = unlines $ map (show . map (unwords . map show . V.toList)) xs
+    result_r <- rKmeans d (MU.toList mat) (MU.toList init_centers)
 
-    assertBool ("Expect: " ++ show' true ++ "\nBut saw: " ++ show' test) $
-        test == true
+    let result = sort $ map sort $ fromJust $ clusters $ kmeans k mat defaultKMeansOpts{kmeansMethod=Centers init_centers}
+        true = sort $ map sort $ decode (V.fromList $ map (subtract 1) result_r) xs
+
+    assertBool ("Expect: " ++ show (map length true) ++ "\nBut saw: " ++ show (map length result)) $
+        result == true